| """Canopy-height cue for tree_removal from Meta CHMv2 (DINOv3 ViT-L/16 SAT-493M + DPT head, DINOv3 License). |
| |
| heights(pre) and heights(post) are canopy-height maps in metres; a pixel is a tree-removal cue when the |
| pre canopy was tall and the post surface is near the ground. |
| """ |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
| from transformers import AutoImageProcessor, AutoModelForDepthEstimation |
|
|
|
|
| class CHM: |
| def __init__(self, path, device, flip_tta=True, scale=1.0): |
| self.proc = AutoImageProcessor.from_pretrained(path) |
| self.m = AutoModelForDepthEstimation.from_pretrained(path).to(device).eval() |
| self.device, self.flip_tta, self.scale = device, flip_tta, scale |
|
|
| @torch.no_grad() |
| def _run(self, imgs): |
| h, w = imgs[0].shape[:2] |
| if self.scale != 1.0: |
| sz = (int(round(w * self.scale / 16)) * 16, int(round(h * self.scale / 16)) * 16) |
| imgs = [np.asarray(Image.fromarray(im).resize(sz, Image.BICUBIC)) for im in imgs] |
| inp = {k: v.to(self.device) for k, v in self.proc(images=list(imgs), return_tensors="pt").items()} |
| out = self.m(**inp) |
| res = self.proc.post_process_depth_estimation(out, target_sizes=[(h, w)] * len(imgs)) |
| return torch.stack([r["predicted_depth"].float().reshape(h, w) for r in res]) |
|
|
| @torch.no_grad() |
| def heights(self, imgs): |
| """list of (H,W,3) uint8 -> (N,H,W) float32 tensor (metres).""" |
| h = self._run(imgs) |
| if self.flip_tta: |
| h = (h + self._run([im[:, ::-1].copy() for im in imgs]).flip(-1)) / 2 |
| return h |
|
|
|
|
| def _smooth(h, k): |
| return F.avg_pool2d(h[:, None], k, 1, k // 2, count_include_pad=False)[:, 0] |
|
|
|
|
| def cue_abs(hp, hq, k=7, lo=3.0, span=4.0, gmax=4.0, gspan=3.0): |
| """tall before (>= lo+span m -> 1) and near ground after (<= gmax-gspan m -> 1).""" |
| hp, hq = _smooth(hp, k), _smooth(hq, k) |
| return ((hp - lo) / span).clamp(0, 1) * ((gmax - hq) / gspan).clamp(0, 1) |
|
|
|
|
| def cue_diff(hp, hq, k=7, lo=2.0, span=4.0): |
| """height drop >= lo+span m -> 1.""" |
| return ((_smooth(hp, k) - _smooth(hq, k) - lo) / span).clamp(0, 1) |
|
|
