infer/objects.py

"""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)