"""Class map + nDSM -> discrete 3D objects: trees, buildings, water. The seg head says *what* each pixel is and the height heads say *how tall*. A renderer wants neither: it wants "a 9 m tree here, a 12 m building with this footprint there". This module is that step. It is pure numpy / scipy / scikit-image and runs on the CPU after the GPU reservation has been handed back. trees tree pixels >= `min_tree_height_m`, split into crowns by a watershed on the smoothed canopy height, seeded at its peaks buildings building components; height = median of the footprint's core (edges eroded away, where the head blends roof into ground) water water components, outline only Output contract (`objects.json`, `version` 1) --------------------------------------------- Positions are in **pixels of the height grid**, origin at the top-left *corner* of the grid: pixel (row, col) covers x in [col, col + 1), y in [row, row + 1), so its centre is (col + 0.5, row + 0.5). Multiply by `grid.gsd_m` for metres. Heights, radii and areas are already metric. trees[] {x, y, h, r} crown centroid, top height (m), crown radius (m) buildings[] {poly, h, h_max, area_m2} water[] {poly, area_m2} `poly` is the outer ring only, as [[x, y], ...], first point not repeated, clockwise on screen (x right, y down). Every list is sorted largest-first and capped, and `truncated` says whether a cap was hit. """ from __future__ import annotations from dataclasses import asdict, dataclass import numpy as np from config import CLASS_NAMES SCHEMA_VERSION = 1 TREE_ID = CLASS_NAMES.index("tree") BUILDING_ID = CLASS_NAMES.index("building") WATER_ID = CLASS_NAMES.index("water") @dataclass class ObjectParams: # trees min_tree_height_m: float = 2.0 # below this it is shrub, drawn as ground min_crown_sep_m: float = 2.5 # peaks closer than this are one crown canopy_smooth_m: float = 0.75 # Gaussian sigma on the canopy before peak finding min_tree_area_m2: float = 3.0 # watershed slivers smaller than this are dropped max_trees: int = 20_000 # buildings min_building_area_m2: float = 20.0 min_building_height_m: float = 2.0 # a "building" lower than this is a misread slab building_edge_m: float = 1.0 # eroded off the footprint before taking the median building_hole_m2: float = 25.0 # courtyards smaller than this are filled max_buildings: int = 5_000 # water min_water_area_m2: float = 50.0 max_water: int = 500 # shared open_m: float = 1.0 # morphological opening: drops features thinner than this outline_tolerance_m: float = 0.75 # Douglas-Peucker tolerance for footprints def extract_objects(seg: np.ndarray, height: np.ndarray, gsd_m: float, params: ObjectParams | None = None) -> dict: """`seg` (H, W) class ids, `height` (H, W) metres, same grid -> objects dict.""" if seg.shape != height.shape: raise ValueError(f"class map {seg.shape} and height map {height.shape} differ") if not gsd_m or gsd_m <= 0: raise ValueError(f"gsd_m must be positive, got {gsd_m}") p = params or ObjectParams() seg = np.asarray(seg) height = np.where(np.isfinite(height), height, 0.0).astype(np.float32) trees, trees_cut = _trees(seg == TREE_ID, height, gsd_m, p) buildings, buildings_cut = _buildings(seg == BUILDING_ID, height, gsd_m, p) water, water_cut = _water(seg == WATER_ID, gsd_m, p) H, W = seg.shape return { "version": SCHEMA_VERSION, "grid": {"width": int(W), "height": int(H), "gsd_m": float(gsd_m), "origin": "top-left corner; pixel centre = (col + 0.5, row + 0.5)"}, "units": {"x": "px", "y": "px", "poly": "px", "h": "m", "h_max": "m", "r": "m", "area_m2": "m^2"}, "trees": trees, "buildings": buildings, "water": water, "counts": {"trees": len(trees), "buildings": len(buildings), "water": len(water)}, "truncated": {"trees": trees_cut, "buildings": buildings_cut, "water": water_cut}, "params": asdict(p), } # --------------------------------------------------------------------- # masks # --------------------------------------------------------------------- def _drop_small(mask: np.ndarray, min_px: int) -> np.ndarray: from scipy import ndimage as ndi if min_px <= 1 or not mask.any(): return mask lab, _ = ndi.label(mask) sizes = np.bincount(lab.ravel()) keep = sizes >= min_px keep[0] = False return keep[lab] def _clean(mask: np.ndarray, gsd_m: float, min_area_m2: float, open_m: float, hole_m2: float = 0.0) -> np.ndarray: """Opening, then small-blob removal, then (optionally) small-hole filling.""" from scipy import ndimage as ndi radius = int(round(open_m / (2.0 * gsd_m))) if radius > 0 and mask.any(): # a square, not a disk: a disk's opening clips every footprint corner size = 2 * radius + 1 mask = ndi.binary_opening(mask, structure=np.ones((size, size), bool)) mask = _drop_small(mask, int(np.ceil(min_area_m2 / gsd_m**2))) if hole_m2 > 0: # a hole is a background component that is not the big outside one mask = mask | ~_drop_small(~mask, int(np.ceil(hole_m2 / gsd_m**2))) return mask # --------------------------------------------------------------------- # trees # --------------------------------------------------------------------- def _trees(mask: np.ndarray, height: np.ndarray, gsd_m: float, p: ObjectParams) -> tuple[list[dict], bool]: from scipy import ndimage as ndi from skimage.feature import peak_local_max from skimage.segmentation import watershed mask = _clean(mask, gsd_m, p.min_tree_area_m2, p.open_m) mask &= height >= p.min_tree_height_m mask = _drop_small(mask, int(np.ceil(p.min_tree_area_m2 / gsd_m**2))) if not mask.any(): return [], False canopy = np.where(mask, height, 0.0).astype(np.float32) sigma = p.canopy_smooth_m / gsd_m smooth = ndi.gaussian_filter(canopy, sigma) if sigma >= 0.3 else canopy # Tie-break plateaus towards the blob interior (1 mm per pixel from the # edge): a flat-topped canopy otherwise yields a grid of equal "peaks", # i.e. one tree becomes four. Far below any real crown relief. smooth = smooth + 1e-3 * ndi.distance_transform_edt(mask).astype(np.float32) comps, n_comps = ndi.label(mask) min_dist = max(1, int(round(p.min_crown_sep_m / gsd_m))) peaks = peak_local_max(smooth, min_distance=min_dist, labels=comps, exclude_border=False) markers = np.zeros(mask.shape, np.int32) if len(peaks): markers[peaks[:, 0], peaks[:, 1]] = np.arange(1, len(peaks) + 1) # a flat-topped blob can yield no peak at all; it is still one tree seeded = np.zeros(n_comps + 1, bool) seeded[comps[markers > 0]] = True missing = np.flatnonzero(~seeded[1:]) + 1 if len(missing): pos = ndi.maximum_position(smooth, comps, missing) start = int(markers.max()) + 1 for k, (r, c) in enumerate(pos): markers[int(r), int(c)] = start + k crowns = watershed(-smooth, markers, mask=mask) # per-crown statistics, vectorised: tens of thousands of crowns is normal rr, cc = np.nonzero(crowns) lab = crowns[rr, cc] n_lab = int(crowns.max()) + 1 count = np.bincount(lab, minlength=n_lab).astype(np.float64) ok = count > 0 cy = np.bincount(lab, rr, n_lab)[ok] / count[ok] cx = np.bincount(lab, cc, n_lab)[ok] / count[ok] # raw maximum: the height head is already smooth at pixel scale (decoded at # H/2 and resampled from the canonical GSD), and any filter here shaves the # apex off a conical crown — a 3x3 median read an 18 m tree as 17.1 m h = np.zeros(n_lab, np.float64) np.maximum.at(h, lab, height[rr, cc]) h = h[ok] area_m2 = count[ok] * gsd_m**2 keep = (area_m2 >= p.min_tree_area_m2) & (h >= p.min_tree_height_m) cy, cx, h, area_m2 = cy[keep], cx[keep], h[keep], area_m2[keep] order = np.argsort(-h, kind="stable") truncated = len(order) > p.max_trees order = order[:p.max_trees] trees = [ {"x": round(float(cx[i]) + 0.5, 2), "y": round(float(cy[i]) + 0.5, 2), "h": round(float(h[i]), 2), "r": round(float(np.sqrt(area_m2[i] / np.pi)), 2)} for i in order ] return trees, truncated # --------------------------------------------------------------------- # footprints # --------------------------------------------------------------------- def _outline(component: np.ndarray, origin: tuple[int, int], tolerance_px: float) -> list[list[float]] | None: """Outer ring of a single boolean component, in corner-origin pixels.""" from skimage.measure import approximate_polygon, find_contours padded = np.pad(component, 1).astype(np.uint8) contours = find_contours(padded, 0.5) if not contours: return None ring = max(contours, key=len) if tolerance_px > 0: ring = approximate_polygon(ring, tolerance=tolerance_px) if len(ring) > 1 and np.allclose(ring[0], ring[-1]): ring = ring[:-1] if len(ring) < 3: return None # (row, col) index space of the padded crop -> corner-origin (x, y) on the grid y = ring[:, 0] - 1 + origin[0] + 0.5 x = ring[:, 1] - 1 + origin[1] + 0.5 # shoelace > 0 with y pointing down is clockwise on screen if np.sum(x * np.roll(y, -1) - np.roll(x, -1) * y) < 0: x, y = x[::-1], y[::-1] return [[round(float(a), 2), round(float(b), 2)] for a, b in zip(x, y)] def _components(mask: np.ndarray): from scipy import ndimage as ndi lab, _ = ndi.label(mask) for i, sl in enumerate(ndi.find_objects(lab), start=1): if sl is not None: yield sl, lab[sl] == i def _buildings(mask: np.ndarray, height: np.ndarray, gsd_m: float, p: ObjectParams) -> tuple[list[dict], bool]: from scipy import ndimage as ndi mask = _clean(mask, gsd_m, p.min_building_area_m2, p.open_m, p.building_hole_m2) erode = int(round(p.building_edge_m / gsd_m)) tol = p.outline_tolerance_m / gsd_m out = [] for sl, comp in _components(mask): core = ndi.binary_erosion(comp, iterations=erode) if erode else comp # thin wings erode to nothing; fall back to the whole footprint if core.sum() < max(4, 0.1 * comp.sum()): core = comp hv = height[sl][core] h = float(np.median(hv)) if h < p.min_building_height_m: continue poly = _outline(comp, (sl[0].start, sl[1].start), tol) if poly is None: continue out.append({"poly": poly, "h": round(h, 2), "h_max": round(float(np.percentile(hv, 95)), 2), "area_m2": round(float(comp.sum()) * gsd_m**2, 1)}) out.sort(key=lambda b: -b["area_m2"]) return out[:p.max_buildings], len(out) > p.max_buildings def _water(mask: np.ndarray, gsd_m: float, p: ObjectParams) -> tuple[list[dict], bool]: mask = _clean(mask, gsd_m, p.min_water_area_m2, p.open_m, p.min_water_area_m2) tol = p.outline_tolerance_m / gsd_m out = [] for sl, comp in _components(mask): poly = _outline(comp, (sl[0].start, sl[1].start), tol) if poly is not None: out.append({"poly": poly, "area_m2": round(float(comp.sum()) * gsd_m**2, 1)}) out.sort(key=lambda w: -w["area_m2"]) return out[:p.max_water], len(out) > p.max_water def summarise(objects: dict) -> str: """One status line: counts plus the mean tree / building height.""" t, b, w = objects["trees"], objects["buildings"], objects["water"] parts = [f"{len(t)} trees"] if t: parts[-1] += f" (mean {np.mean([o['h'] for o in t]):.1f} m)" parts.append(f"{len(b)} buildings") if b: parts[-1] += f" (median {np.median([o['h'] for o in b]):.1f} m)" parts.append(f"{len(w)} water bodies") if any(objects["truncated"].values()): parts.append("capped — see objects.json `truncated`") return " · ".join(parts)