File size: 2,162 Bytes
6d4d684 4ed7709 6d4d684 4ed7709 6d4d684 4ed7709 6d4d684 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """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)
|