
# ── Standard library ─────────────────────────────────────────────────────────
import os
from pathlib import Path

# ── Point cloud - LiDAR data ───────────────────────────────────────────────────────────────
import laspy

# ── Geospatial / raster ───────────────────────────────────────────────────────
import rasterio
from rasterio.transform import from_origin
from rasterio.crs import CRS

# ── Numerical ─────────────────────────────────────────────────────────────────
import numpy as np
from scipy.ndimage import distance_transform_edt, binary_dilation, gaussian_filter
from skimage.feature import peak_local_max
from skimage.segmentation import watershed, find_boundaries

# ── Visualisation ─────────────────────────────────────────────────────────────
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.colors import ListedColormap
from matplotlib.ticker import ScalarFormatter

print("✓ All imports successful")

✓ All imports successful


# ── Project configuration ─────────────
DATA_DIR = Path("E:/Py/3D/LAS/Neigh")   # ← change this to your data folder

# Input
LAS_PATH = DATA_DIR / "rgb_675501c.laz"   # RGB-coloured LiDAR point cloud; fetched from (https://kartta.hel.fi/?link=aw57fg#)

# Outputs
OUT_DTM  = DATA_DIR / "dtm_05m.tif"                  # Digital Terrain Model
OUT_NDSM = DATA_DIR / "ndsm_05m.tif"                 # Normalised DSM
OUT_CHM  = DATA_DIR / "chm_veg_05m_clean.tif"        # Vegetation CHM (buildings removed)

# Processing parameters
RESOLUTION           = 0.5   # metres per pixel
BUILDING_DILATION_M  = 1.0   # buffer around building footprints to remove edge contamination
MIN_VEG_HEIGHT_M     = 0.5   # below this, pixels are treated as ground / noise
MIN_TREE_HEIGHT_M    = 5.0   # minimum tree height while detecting a tree-top
SMOOTH_SIGMA         = 1.5   # Gaussian sigma for CHM smoothing (pixels)
MIN_DISTANCE_PX      = 5     # minimum separation between tree tops (pixels)
PROMINENCE_WINDOW    = 12    # local window for prominence filter (pixels)
MIN_PROMINENCE       = 0.8   # minimum height above local mean (metres)
MIN_CROWN_HEIGHT_M   = 2.0   # minimum CHM value to include in crown mask
POINT_SAMPLE_N       = 10    # 1-in-N subsampling for scatter visualisations

# Derived
BUILDING_DILATION_PX = int(BUILDING_DILATION_M / RESOLUTION)

print("✓ Configuration set")
print(f"  Resolution        : {RESOLUTION} m/px")
print(f"  Building buffer   : {BUILDING_DILATION_M} m ({BUILDING_DILATION_PX} px)")
print(f"  Min tree height   : {MIN_TREE_HEIGHT_M} m")
print(f"  Min tree spacing  : {MIN_DISTANCE_PX * RESOLUTION} m")

✓ Configuration set
  Resolution        : 0.5 m/px
  Building buffer   : 1.0 m (2 px)
  Min tree height   : 5.0 m
  Min tree spacing  : 2.5 m


def points_to_grid(x, y, xmin, ymax, res, width, height):
    col = ((x - xmin) / res).astype(np.int32)
    row = ((ymax - y) / res).astype(np.int32)
    valid = (col >= 0) & (col < width) & (row >= 0) & (row < height)
    return row[valid], col[valid], valid


def rasterize_stat(x, y, z, xmin, ymax, res, width, height, stat="max"):
    """
    Rasterise 3-D point data to a 2-D grid by computing a per-pixel statistic.

    Parameters
    ----------
    x, y, z : array-like — point coordinates and values
    stat    : 'max' (DSM, CHM) or 'min' (DTM ground surface)

    Returns
    -------
    np.ndarray, shape (height, width), dtype float32
        Pixels with no points are NaN.
    """
    x, y, z = np.asarray(x), np.asarray(y), np.asarray(z, dtype=np.float32)
    row, col, valid = points_to_grid(x, y, xmin, ymax, res, width, height)
    z_valid = z[valid]

    if len(z_valid) == 0:
        raise ValueError("No valid points fell inside the raster extent.")

    if stat == "max":
        raster = np.full((height, width), -np.inf, dtype=np.float32)
        np.maximum.at(raster, (row, col), z_valid)
        raster[raster == -np.inf] = np.nan
    elif stat == "min":
        raster = np.full((height, width), np.inf, dtype=np.float32)
        np.minimum.at(raster, (row, col), z_valid)
        raster[raster == np.inf] = np.nan
    else:
        raise ValueError("stat must be 'min' or 'max'")

    return raster


def rasterize_binary_mask(x, y, xmin, ymax, res, width, height):
    """
    Rasterise XY point locations to a binary presence/absence grid.

    Returns
    -------
    np.ndarray, shape (height, width), dtype bool
    """
    x, y = np.asarray(x), np.asarray(y)
    row, col, valid = points_to_grid(x, y, xmin, ymax, res, width, height)
    mask = np.zeros((height, width), dtype=bool)
    mask[row, col] = True
    return mask


def fill_nearest(arr):
    """
    Fill NaN pixels by nearest-neighbour propagation.

    Used to produce a continuous DTM surface where ground points are sparse
    (e.g. under dense canopy). Uses scipy distance transform to find the
    nearest valid pixel for each NaN pixel.

    Parameters
    ----------
    arr : np.ndarray — 2-D raster with NaN gaps

    Returns
    -------
    np.ndarray — fully filled raster, same shape
    """
    nan_mask = np.isnan(arr)
    if np.all(nan_mask):
        raise ValueError("Raster is entirely NaN; cannot fill.")
    indices = distance_transform_edt(
        nan_mask, return_distances=False, return_indices=True
    )
    return arr[tuple(indices)]


def save_raster(path, arr, transform, crs):
    """
    Save a 2-D NumPy array as a single-band GeoTIFF.
    """
    with rasterio.open(
        path, "w",
        driver="GTiff",
        height=arr.shape[0], width=arr.shape[1],
        count=1, dtype="float32",
        crs=crs, transform=transform,
        nodata=np.nan
    ) as dst:
        dst.write(arr.astype(np.float32), 1)


def nan_aware_gaussian(arr, sigma):
    """
    Apply Gaussian smoothing while correctly handling NaN pixels.

    Standard scipy gaussian_filter treats NaN as zero, which suppresses
    heights near data gaps. This normalised convolution approach avoids
    that bias by smoothing the numerator and denominator separately.

    Parameters
    ----------
    arr   : np.ndarray — 2-D raster, may contain NaN
    sigma : float      — Gaussian standard deviation in pixels

    Returns
    -------
    np.ndarray — smoothed array; NaN pixels in input remain NaN in output
    """
    nan_mask = np.isnan(arr)
    tmp = arr.copy()
    tmp[nan_mask] = 0.0

    smooth_num = gaussian_filter(tmp,               sigma=sigma)
    smooth_den = gaussian_filter((~nan_mask).astype(float), sigma=sigma)

    result = smooth_num / np.maximum(smooth_den, 1e-6)
    result[nan_mask] = np.nan
    return result


print("✓ Helper functions defined")

✓ Helper functions defined


# ── Load full point cloud ─────────────────────────────────────────────────────
las = laspy.read(str(LAS_PATH))

x   = np.asarray(las.x)
y   = np.asarray(las.y)
z   = np.asarray(las.z)
cls = np.asarray(las.classification)

print(f"Point cloud loaded:")
print(f"  Total points   : {len(x):,}")
print(f"  Classes found  : {np.unique(cls).tolist()}")
print(f"  Z range        : {z.min():.1f} m → {z.max():.1f} m")
print(f"  Easting range  : {x.min():.0f} → {x.max():.0f}")
print(f"  Northing range : {y.min():.0f} → {y.max():.0f}")

# ── Subsample for visualisation only ─────────────────────────────────────────
x_vis  = x[::POINT_SAMPLE_N]
y_vis  = y[::POINT_SAMPLE_N]
cls_vis = cls[::POINT_SAMPLE_N]

# ── RGB colours for visualisation ────────────────────────────────────────────
r = las.red[::POINT_SAMPLE_N].astype(np.float32)
g = las.green[::POINT_SAMPLE_N].astype(np.float32)
b = las.blue[::POINT_SAMPLE_N].astype(np.float32)
rgb_vis = np.stack([r, g, b], axis=1)
rgb_vis /= rgb_vis.max()   # normalise to [0, 1]

print(f"\n  Visualisation subsample: 1-in-{POINT_SAMPLE_N} → {len(x_vis):,} points")

Point cloud loaded:
  Total points   : 16,530,887
  Classes found  : [1, 2, 3, 4, 5, 6, 8, 9, 17]
  Z range        : 0.1 m → 60.5 m
  Easting range  : 25501500 → 25502000
  Northing range : 6675000 → 6675500

  Visualisation subsample: 1-in-10 → 1,653,089 points


# ── Classification map — dark background style for LiDAR visualisation ────────
CLASS_MAP = {
    1:  ("#aaaaaa", "Unclassified"),
    2:  ("#8B6914", "Ground"),
    3:  ("#90EE90", "Low vegetation"),
    4:  ("#228B22", "Medium vegetation"),
    5:  ("#006400", "High vegetation"),
    6:  ("#FF4500", "Building"),
    7:  ("#FF69B4", "Noise"),
    8:  ("#9B59B6", "Keypoints"),
    9:  ("#1E90FF", "Water"),
    10: ("#7F8C8D", "Rail"),
    17: ("#FFD700", "Bridge"),
}

fig, ax = plt.subplots(figsize=(14, 11))
fig.patch.set_facecolor("#0f0f1a")
ax.set_facecolor("#0f0f1a")

for c in np.unique(cls_vis):
    pt_mask = cls_vis == c
    color, label = CLASS_MAP.get(int(c), ("#ffffff", f"Class {c}"))
    ax.scatter(
        x_vis[pt_mask], y_vis[pt_mask],
        c=color, label=f"{int(c)} — {label}",
        s=0.4, alpha=0.7, linewidths=0
    )

ax.set_aspect("equal")
ax.set_title("LiDAR Point Cloud — ASPRS Classification",
             color="white", fontsize=14, pad=10)
ax.tick_params(colors="#666", labelsize=8)
ax.set_xlabel("Easting (m)",  color="#666", fontsize=9)
ax.set_ylabel("Northing (m)", color="#666", fontsize=9)
for sp in ax.spines.values():
    sp.set_edgecolor("#333")

leg = ax.legend(
    loc="upper right", fontsize=8, markerscale=8,
    facecolor="#1a1a2e", edgecolor="#444", labelcolor="white",
    title="Classification", title_fontsize=9
)
leg.get_title().set_color("white")

plt.tight_layout()
plt.show()

print("\nObservation: Some building rooftop points are misclassified as vegetation.")
print("We address this in Step 3 with a dilated building mask applied to the CHM.")

Observation: Some building rooftop points are misclassified as vegetation.
We address this in Step 3 with a dilated building mask applied to the CHM.


# ── Define raster grid extent from point cloud bounding box ──────────────────
xmin, xmax = x.min(), x.max()
ymin, ymax = y.min(), y.max()

grid_width  = int(np.ceil((xmax - xmin) / RESOLUTION))
grid_height = int(np.ceil((ymax - ymin) / RESOLUTION))
transform   = from_origin(xmin, ymax, RESOLUTION, RESOLUTION)

print(f"Raster grid:")
print(f"  Size       : {grid_width} cols × {grid_height} rows")
print(f"  Resolution : {RESOLUTION} m")
print(f"  Extent     : ({xmin:.0f}, {ymin:.0f}) → ({xmax:.0f}, {ymax:.0f})")

Raster grid:
  Size       : 1000 cols × 1000 rows
  Resolution : 0.5 m
  Extent     : (25501500, 6675000) → (25502000, 6675500)


# ── DTM — ground points only, minimum z per pixel ─────────────────────────────
# Using min-z picks the lowest LiDAR return, which best represents bare earth.
# Gaps (no ground return) are filled by nearest-neighbour interpolation.

ground_mask  = (cls == 2)
xg, yg, zg  = x[ground_mask], y[ground_mask], z[ground_mask]
print(f"Ground points  : {len(xg):,}")

dtm        = rasterize_stat(xg, yg, zg, xmin, ymax, RESOLUTION, grid_width, grid_height, stat="min")
dtm_filled = fill_nearest(dtm)   # interpolate gaps under dense canopy

# ── DSM — all surface classes, maximum z per pixel ────────────────────────────
surface_mask = np.isin(cls, [2, 3, 4, 5, 6, 17])
xs, ys, zs  = x[surface_mask], y[surface_mask], z[surface_mask]
print(f"Surface points : {len(xs):,}")

dsm        = rasterize_stat(xs, ys, zs, xmin, ymax, RESOLUTION, grid_width, grid_height, stat="max")
dsm_filled = fill_nearest(dsm)

# ── nDSM — Normalised DSM — height of all features above ground ────────────────────────────────
ndsm = dsm_filled - dtm_filled
ndsm[ndsm < 0] = 0   # clip negatives from ground return noise

print(f"\nDerived surface models:")
print(f"  DTM  : {np.nanmin(dtm_filled):.1f} m → {np.nanmax(dtm_filled):.1f} m")
print(f"  DSM  : {np.nanmin(dsm_filled):.1f} m → {np.nanmax(dsm_filled):.1f} m")
print(f"  nDSM : {ndsm.min():.1f} m → {ndsm.max():.1f} m (height above ground)")

Ground points  : 8,096,019
Surface points : 16,428,304

Derived surface models:
  DTM  : 0.8 m → 17.3 m
  DSM  : 0.4 m → 60.5 m
  nDSM : 0.0 m → 46.2 m (height above ground)


# ── Building mask — with dilation buffer ─────────────────────────────────────
# LiDAR classification occasionally assigns rooftop returns to vegetation
# classes. We rasterise class-6 points, dilate by BUILDING_DILATION_PX pixels,
# and use this mask to suppress false tree tops on or near buildings.

building_pts   = (cls == 6)
xb, yb         = x[building_pts], y[building_pts]
print(f"Building points : {len(xb):,}")

building_mask_raw = rasterize_binary_mask(
    xb, yb, xmin, ymax, RESOLUTION, grid_width, grid_height
)
building_mask_dilated = binary_dilation(
    building_mask_raw, iterations=BUILDING_DILATION_PX
)

print(f"  Building pixels before dilation : {building_mask_raw.sum():,}")
print(f"  Building pixels after dilation  : {building_mask_dilated.sum():,}")

# ── CHM — vegetation points only ─────────────────────────────────────────────
# Canopy Height Model = max vegetation z − DTM, with buildings removed.

veg_mask   = np.isin(cls, [3, 4, 5])
xv, yv, zv = x[veg_mask], y[veg_mask], z[veg_mask]
print(f"Vegetation points : {len(xv):,}")

veg_top = rasterize_stat(
    xv, yv, zv, xmin, ymax, RESOLUTION, grid_width, grid_height, stat="max"
)

chm = veg_top - dtm_filled
chm[np.isnan(veg_top)] = np.nan               # preserve NoData where no veg points exist
chm[chm < 0] = 0                              # clip ground noise
chm[chm < MIN_VEG_HEIGHT_M] = np.nan          # remove very low returns (grass, ground clutter)

# Apply building mask with dilation buffer
chm_clean = chm.copy()
chm_clean[building_mask_dilated] = np.nan

print(f"\nCHM:")
print(f"  Valid pixels (raw)    : {np.sum(~np.isnan(chm)):,}")
print(f"  Valid pixels (cleaned) : {np.sum(~np.isnan(chm_clean)):,}")
print(f"  Height range           : {np.nanmin(chm_clean):.1f} m → {np.nanmax(chm_clean):.1f} m")

Building points : 3,424,562
  Building pixels before dilation : 246,755
  Building pixels after dilation  : 289,945
Vegetation points : 4,905,488

CHM:
  Valid pixels (raw)    : 340,199
  Valid pixels (cleaned) : 238,797
  Height range           : 0.5 m → 40.1 m


# ── Save all derived rasters ──────────────────────────────────────────────────
# 
crs = CRS.from_epsg(3879) #Helsinki CRS
print(crs)

save_raster(OUT_DTM,  dtm_filled, transform, crs)
save_raster(OUT_NDSM, ndsm,       transform, crs)
save_raster(OUT_CHM,  chm_clean,  transform, crs)

print(f"Saved:")
print(f"  DTM      → {OUT_DTM}")
print(f"  nDSM     → {OUT_NDSM}")
print(f"  CHM      → {OUT_CHM}")

EPSG:3879
Saved:
  DTM      → E:\Py\3D\LAS\Neigh\dtm_05m.tif
  nDSM     → E:\Py\3D\LAS\Neigh\ndsm_05m.tif
  CHM      → E:\Py\3D\LAS\Neigh\chm_veg_05m_clean.tif


# ── Five-panel overview of derived surface products ───────────────────────────
products = [
    (dtm_filled,         "terrain",  "DTM — Bare Earth (m)"),
    (ndsm,               "viridis",  "nDSM — Height Above Ground (m)"),
    (building_mask_dilated.astype(float), "gray", "Building Mask"),
    (chm,                "YlGn",     "CHM — Vegetation Height, raw (m)"),
    (chm_clean,          "YlGn",     "CHM — Vegetation Height, cleaned (m)"),
]

fig, axes = plt.subplots(1, 5, figsize=(26, 5))

for ax, (data, cmap, title) in zip(axes, products):
    im = ax.imshow(data, cmap=cmap)
    ax.set_title(title, fontsize=10, fontweight="bold")
    ax.axis("off")
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

plt.suptitle("Derived Surface Models from LiDAR Point Cloud",
             fontsize=13, fontweight="bold", y=1.01)
plt.tight_layout()
plt.show()


# ── Classify vegetation CHM into height tiers ─────────────────────────────────
canopy_classes = np.zeros(chm_clean.shape, dtype=np.uint8)

canopy_classes[(chm_clean >= 0)  & (chm_clean < 2)]  = 1   # low shrub
canopy_classes[(chm_clean >= 2)  & (chm_clean < 5)]  = 2   # high shrub
canopy_classes[(chm_clean >= 5)  & (chm_clean < 10)] = 3   # sub-canopy
canopy_classes[(chm_clean >= 10)]                     = 4   # canopy
canopy_classes[np.isnan(chm_clean)]                   = 0   # background / no data

# ── Building height tiers ─────────────────────────────────────────────────────
# Compute building nDSM from class-6 points using maximum z per pixel
zb          = z[building_pts]
building_top = rasterize_stat(
    xb, yb, zb, xmin, ymax, RESOLUTION, grid_width, grid_height, stat="max"
)
building_ndsm = building_top - dtm_filled
building_ndsm[np.isnan(building_top)] = np.nan
building_ndsm[building_ndsm < 0]      = 0

# ── Combined raster: vegetation + buildings ───────────────────────────────────
combined = canopy_classes.copy()
# Re-derive binary building mask (1-px dilation only for clean boundaries)
bld_mask_1px = binary_dilation(building_mask_raw, iterations=1)
combined[bld_mask_1px] = 5   # overwrite with building class

# Count pixels per class
labels = {
    0: "Background", 1: "Low shrub (0–2 m)", 2: "High shrub (2–5 m)",
    3: "Sub-canopy (5–10 m)", 4: "Canopy (>10 m)", 5: "Buildings"
}
print("Pixel counts per class:")
for c, label in labels.items():
    n = (combined == c).sum()
    print(f"  {c} {label:25s}: {n:>8,}")

Pixel counts per class:
  0 Background               :  491,995
  1 Low shrub (0–2 m)        :   61,778
  2 High shrub (2–5 m)       :   69,945
  3 Sub-canopy (5–10 m)      :   86,529
  4 Canopy (>10 m)           :   20,545
  5 Buildings                :  269,208


# ── Colour scheme for vegetation tiers ────────────────────────────────────────
VEGE_COLORS = ["white", "#b8e186", "#7fbc41", "#4d9221", "#00441b", "#bdbdbd"]
cmap_struct  = ListedColormap(VEGE_COLORS)

legend_patches = [
    mpatches.Patch(color="#b8e186", label="Low shrub (0–2 m)"),
    mpatches.Patch(color="#7fbc41", label="High shrub (2–5 m)"),
    mpatches.Patch(color="#4d9221", label="Sub-canopy (5–10 m)"),
    mpatches.Patch(color="#00441b", label="Canopy (>10 m)"),
    mpatches.Patch(color="#bdbdbd", label="Buildings"),
]

# Get extent for labelled axes
extent = [xmin, xmax, ymin, ymax]

# ── Three-panel: vegetation only | combined | with building heights ────────────
fig, axes = plt.subplots(1, 2, figsize=(18, 8))

# Panel 1: Vegetation structure only
axes[0].imshow(canopy_classes, cmap=ListedColormap(VEGE_COLORS[:-1]),
               extent=extent, origin="upper")
axes[0].set_title("Vegetation Structure Only", fontsize=12, fontweight="bold")
axes[0].legend(handles=legend_patches[:-1], loc="upper left", fontsize=8)

# Panel 2: Combined vegetation + buildings
axes[1].imshow(combined, cmap=cmap_struct, vmin=0, vmax=5,
               extent=extent, origin="upper")
axes[1].set_title("Vertical Structure + Buildings", fontsize=12, fontweight="bold")
axes[1].legend(handles=legend_patches, loc="upper left", fontsize=8)

for ax in axes:
    ax.set_xlabel("Easting (m)", fontsize=9)
    ax.set_ylabel("Northing (m)", fontsize=9)
    ax.ticklabel_format(style="plain")
    ax.xaxis.set_major_formatter(ScalarFormatter(useOffset=False))
    ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))

plt.suptitle("Urban 3D Vertical Structure — Helsinki",
             fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()


# More informed: Vegetation structure + Building height overlay
fig, ax = plt.subplots(figsize=(12, 14))

# ── Vegetation structure base layer ──────────────────────────────────────────
ax.imshow(
    canopy_classes,
    cmap=ListedColormap(VEGE_COLORS[:-1]),  # exclude building grey
    extent=extent,
    origin="upper",
    vmin=0, vmax=4   # pin colour scale to 0–4 so classes map correctly
)

# ── Building height overlay ───────────────────────────────────────────────────
im = ax.imshow(
    building_ndsm,
    cmap="YlOrBr",
    alpha=0.65,
    extent=extent,
    origin="upper"
)

# ── Axes labels and formatting ────────────────────────────────────────────────
ax.set_title("Vegetation Structure + Building Height (m)",
             fontsize=14, fontweight="bold", pad=12)
ax.set_xlabel("Easting (m)",  fontsize=10)
ax.set_ylabel("Northing (m)", fontsize=10)
ax.set_aspect("equal")

# Remove scientific notation on coordinate axes
ax.ticklabel_format(style="plain")
ax.xaxis.set_major_formatter(ScalarFormatter(useOffset=False))
ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))

# ── Horizontal colorbar (manual placement, bottom of plot area) ───────────────
# [left, bottom, width, height] — all in figure-fraction coordinates
cax  = fig.add_axes([0.125, 0.155, 0.595, 0.018])
cbar = fig.colorbar(im, cax=cax, orientation="horizontal")
cbar.set_label("Building height (m)", fontsize=10)

# ── Vegetation legend below the colorbar ─────────────────────────────────────
ax.legend(
    handles=legend_patches[:-1],   # exclude "Buildings" patch
    loc="lower left",
    bbox_to_anchor=(-0.01, -0.12),
    ncol=4,
    frameon=True,
    columnspacing=1.2,
    handlelength=1.5,
    fontsize=9
)

# Leave space on the right so the tall axes doesn't clip
fig.subplots_adjust(right=0.72)

plt.show()


# ── Load saved CHM and apply NaN-aware smoothing ─────────────────────────────
# (Re-reading from file so this step is independent and reproducible)

with rasterio.open(str(OUT_CHM)) as src:
    chm_loaded = src.read(1)
    raster_transform = src.transform
    raster_bounds    = src.bounds
    res_x, res_y     = src.res

print(f"CHM loaded from disk:")
print(f"  Resolution : {res_x} m × {res_y} m")
print(f"  Shape      : {chm_loaded.shape}")
print(f"  Height range : {np.nanmin(chm_loaded):.1f} m → {np.nanmax(chm_loaded):.1f} m")

# NaN-aware Gaussian smoothing
chm_smooth = nan_aware_gaussian(chm_loaded, sigma=SMOOTH_SIGMA)

# Detection image: set NaN and sub-threshold to 0
detect_img = chm_smooth.copy()
detect_img[np.isnan(detect_img)] = 0
detect_img[detect_img < MIN_TREE_HEIGHT_M] = 0

print(f"\n  Smoothed with NaN-aware Gaussian (sigma={SMOOTH_SIGMA} px)")
print(f"  Detection threshold: {MIN_TREE_HEIGHT_M} m")

CHM loaded from disk:
  Resolution : 0.5 m × 0.5 m
  Shape      : (1000, 1000)
  Height range : 0.5 m → 40.1 m

  Smoothed with NaN-aware Gaussian (sigma=1.5 px)
  Detection threshold: 5.0 m


# ── Stage 1: peak local maximum ───────────────────────────────────────────────
# Detects all local maxima separated by at least MIN_DISTANCE_PX pixels.

tree_tops_initial = peak_local_max(
    detect_img,
    min_distance  = MIN_DISTANCE_PX,
    threshold_abs = MIN_TREE_HEIGHT_M,
    exclude_border= False
)
print(f"Stage 1 — peak_local_max:  {len(tree_tops_initial):,} candidates")

Stage 1 — peak_local_max:  913 candidates


def filter_peaks_by_prominence(chm_img, peaks, window_size, min_prominence):
    """
    Filter tree top candidates by local prominence.

    A peak is kept only if its height exceeds the mean of all valid
    (finite, positive) pixels within a local neighbourhood by at least
    `min_prominence` metres. This removes false peaks on gently sloping
    or flat canopy surfaces.

    Parameters
    ----------
    chm_img       : np.ndarray — smoothed CHM, may contain NaN
    peaks         : np.ndarray, shape (N, 2) — (row, col) indices from peak_local_max
    window_size   : int   — side length of the local neighbourhood (pixels)
    min_prominence: float — minimum height above local mean (metres)

    Returns
    -------
    np.ndarray, shape (M, 2) — filtered peaks, M ≤ N
    """
    keep = []
    half = window_size // 2
    nrows, ncols = chm_img.shape

    for r, c in peaks:
        peak_height = chm_img[r, c]
        if np.isnan(peak_height) or peak_height <= 0:
            continue

        # Extract local window, handling image boundaries
        r0, r1 = max(0, r - half), min(nrows, r + half + 1)
        c0, c1 = max(0, c - half), min(ncols, c + half + 1)
        local   = chm_img[r0:r1, c0:c1]

        # Compute mean of valid pixels in the window
        local_valid = local[np.isfinite(local) & (local > 0)]
        if len(local_valid) == 0:
            continue

        if (peak_height - np.mean(local_valid)) >= min_prominence:
            keep.append([r, c])

    return np.array(keep, dtype=int) if keep else np.empty((0, 2), dtype=int)


# ── Stage 2: prominence filter ────────────────────────────────────────────────
tree_tops_final = filter_peaks_by_prominence(
    chm_smooth,
    tree_tops_initial,
    window_size    = PROMINENCE_WINDOW,
    min_prominence = MIN_PROMINENCE
)

print(f"Stage 2 — prominence filter: {len(tree_tops_final):,} trees retained")
print(f"  Removed : {len(tree_tops_initial) - len(tree_tops_final):,} false peaks")
print(f"  Retention rate: {len(tree_tops_final)/len(tree_tops_initial)*100:.1f}%")

# Convert pixel indices to map coordinates
rows_f, cols_f = tree_tops_final[:, 0], tree_tops_final[:, 1]
xs_final, ys_final = rasterio.transform.xy(raster_transform, rows_f, cols_f, offset="center")
xs_final = np.array(xs_final)
ys_final = np.array(ys_final)

rows_i, cols_i = tree_tops_initial[:, 0], tree_tops_initial[:, 1]
xs_initial, ys_initial = rasterio.transform.xy(raster_transform, rows_i, cols_i, offset="center")
xs_initial = np.array(xs_initial)
ys_initial = np.array(ys_initial)

# Heights at detected tree tops
top_heights = chm_loaded[rows_f, cols_f]
print(f"\n  Tree height stats:")
print(f"    Min    : {top_heights.min():.1f} m")
print(f"    Median : {np.median(top_heights):.1f} m")
print(f"    Max    : {top_heights.max():.1f} m")

Stage 2 — prominence filter: 872 trees retained
  Removed : 41 false peaks
  Retention rate: 95.5%

  Tree height stats:
    Min    : 4.8 m
    Median : 9.2 m
    Max    : 39.8 m


# ── Side-by-side: before and after prominence filtering ───────────────────────
fig, axes = plt.subplots(1, 2, figsize=(18, 8), sharex=True, sharey=True)

titles = [
    f"Before prominence filter — {len(tree_tops_initial):,} candidates",
    f"After prominence filter — {len(tree_tops_final):,} trees",
]
xs_list = [xs_initial, xs_final]
ys_list = [ys_initial, ys_final]
colors  = ["dodgerblue", "crimson"]

for ax, title, xs_plot, ys_plot, color in zip(
    axes, titles, xs_list, ys_list, colors
):
    # RGB point cloud base
    ax.scatter(x_vis, y_vis, c=rgb_vis, s=0.3, rasterized=True)
    # Tree top markers
    ax.scatter(xs_plot, ys_plot, s=12, c=color, marker="x",
               linewidths=1.2, label=f"Tree tops ({len(xs_plot):,})")
    ax.set_title(title, fontsize=12, fontweight="bold")
    ax.set_xlabel("Easting (m)", fontsize=9)
    ax.set_ylabel("Northing (m)", fontsize=9)
    ax.legend(loc="upper right", fontsize=9)
    ax.set_aspect("equal")
    ax.ticklabel_format(style="plain")

plt.suptitle("Tree Top Detection — Effect of Prominence Filtering",
             fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()


# ── Tree height distribution ──────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(9, 4))

ax.hist(top_heights, bins=30, color="#4d9221", edgecolor="white",
        linewidth=0.5, alpha=0.85)
ax.axvline(np.median(top_heights), color="crimson", linewidth=2,
           linestyle="--", label=f"Median: {np.median(top_heights):.1f} m")
ax.set_xlabel("Tree height at detected top (m)", fontsize=11)
ax.set_ylabel("Count", fontsize=11)
ax.set_title(f"Height Distribution of {len(top_heights):,} Detected Trees",
             fontsize=12, fontweight="bold")
ax.legend(fontsize=10)
ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()


# ── Prepare smoothed CHM for watershed ───────────────────────────────────────
# Use a fresh smooth on the clean CHM (sigma=1 for crown delineation — 
# less aggressive than for peak detection to preserve crown boundaries)

chm_seg    = chm_clean.copy().astype(float)
chm_smooth_crown = nan_aware_gaussian(chm_seg, sigma=1.0)

# Crown mask: only pixels above MIN_CROWN_HEIGHT_M are eligible
crown_mask = chm_smooth_crown > MIN_CROWN_HEIGHT_M

print(f"Crown mask coverage: {crown_mask.sum():,} pixels "
      f"({crown_mask.mean()*100:.1f}% of raster)")

# ── Build marker image from tree top coordinates ──────────────────────────────
# Each tree top gets a unique integer label. Only tops that fall within
# the crown mask (i.e. above MIN_CROWN_HEIGHT_M) are used as seeds.

markers = np.zeros(chm_smooth_crown.shape, dtype=np.int32)
n_seeds = 0

for i, (r, c) in enumerate(tree_tops_final, start=1):
    if crown_mask[r, c]:
        markers[r, c] = i
        n_seeds += 1

print(f"Watershed seeds: {n_seeds} of {len(tree_tops_final)} tree tops "
      f"fall within crown mask")

Crown mask coverage: 182,509 pixels (18.3% of raster)
Watershed seeds: 872 of 872 tree tops fall within crown mask


# ── Run watershed ─────────────────────────────────────────────────────────────
# Invert CHM: peaks → basins. watershed fills from low values upward,
# which grows crowns downward from each tree top marker.

elevation = -np.nan_to_num(chm_smooth_crown, nan=0)

crowns = watershed(
    elevation,
    markers = markers,
    mask    = crown_mask     # restrict growth to vegetation pixels only
)

n_crowns = len(np.unique(crowns)) - 1   # subtract background (label 0)
print(f"Watershed complete: {n_crowns} individual crowns delineated")

# Crown size statistics
crown_sizes = [
    (crowns == c).sum() * RESOLUTION**2
    for c in np.unique(crowns) if c > 0
]
print(f"  Crown area range  : {min(crown_sizes):.1f} m² → {max(crown_sizes):.1f} m²")
print(f"  Median crown area : {np.median(crown_sizes):.1f} m²")
print(f"  (equivalent crown diameter: ~{2*(np.median(crown_sizes)/3.14)**0.5:.1f} m)")

Watershed complete: 872 individual crowns delineated
  Crown area range  : 0.2 m² → 334.2 m²
  Median crown area : 38.6 m²
  (equivalent crown diameter: ~7.0 m)


import geopandas as gpd
from rasterio.features import shapes
from shapely.geometry import shape

# ── Vectorise raster crown segments → polygon GeoDataFrame ───────────────────
# rasterio.features.shapes yields (geometry, value) pairs for each connected region
crown_polygons = []

for geom, value in shapes(
    crowns.astype(np.int32),
    mask=(crowns > 0).astype(np.uint8),   # exclude background
    transform=raster_transform
):
    tree_id = int(value)
    crown_poly = shape(geom)

    # Look up height at the tree top pixel for this tree
    r, c = tree_tops_final[tree_id - 1]  # tree IDs are 1-indexed
    height = float(chm_loaded[r, c]) if 0 <= tree_id - 1 < len(tree_tops_final) else np.nan

    crown_polygons.append({
        "tree_id":    tree_id,
        "height_m":   round(height, 2),
        "crown_m2":   round(crown_poly.area, 2),
        "geometry":   crown_poly
    })

# Build GeoDataFrame with the correct CRS
crowns_gdf = gpd.GeoDataFrame(crown_polygons, crs="EPSG:3879")

# Sort by height descending for easy inspection
crowns_gdf = crowns_gdf.sort_values("height_m", ascending=False).reset_index(drop=True)

print(f"Crown GeoDataFrame: {len(crowns_gdf)} polygons")
print(crowns_gdf[["tree_id", "height_m", "crown_m2"]].head(10).to_string(index=False))

# ── Save to GeoPackage (single file, no sidecar files unlike shapefile) ───────
out_gpkg = DATA_DIR / "tree_crowns.gpkg"
crowns_gdf.to_file(out_gpkg, driver="GPKG", layer="tree_crowns")
print(f"\n✓ Saved → {out_gpkg}")
print(f"  Open in QGIS or ArcGIS — each polygon is one tree")

# ── Quick summary statistics ───────────────────────────────────────────────────
print(f"\nTree inventory summary:")
print(f"  Total trees detected  : {len(crowns_gdf):,}")
print(f"  Height range          : {crowns_gdf.height_m.min():.1f} m → {crowns_gdf.height_m.max():.1f} m")
print(f"  Median tree height    : {crowns_gdf.height_m.median():.1f} m")
print(f"  Total canopy area     : {crowns_gdf.crown_m2.sum()/10000:.2f} ha")
print(f"  Median crown area     : {crowns_gdf.crown_m2.median():.1f} m²")

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[18], line 10
      5 # ── Vectorise raster crown segments → polygon GeoDataFrame ───────────────────
      6 # rasterio.features.shapes yields (geometry, value) pairs for each connected region
      7 crown_polygons = []
      9 for geom, value in shapes(
---> 10     crowns.astype(np.int32),
     11     mask=(crowns > 0).astype(np.uint8),   # exclude background
     12     transform=raster_transform
     13 ):
     14     tree_id = int(value)
     15     crown_poly = shape(geom)

NameError: name 'crowns' is not defined


# ── Crown map and boundaries ──────────────────────────────────────────────────
boundaries = find_boundaries(crowns, mode="outer")

fig, axes = plt.subplots(1, 2, figsize=(18, 8))

# Panel 1: crown segments coloured by tree ID
axes[0].imshow(crowns, cmap="nipy_spectral", origin="upper")
axes[0].set_title(f"Individual Tree Crowns ({n_crowns} trees)",
                  fontsize=12, fontweight="bold")
axes[0].axis("off")

# Panel 2: CHM with crown boundaries overlaid
axes[1].imshow(chm_clean, cmap="YlGn", origin="upper")
boundary_overlay = np.where(boundaries, 1.0, np.nan)
axes[1].imshow(boundary_overlay, cmap="autumn", alpha=0.9, origin="upper")
axes[1].set_title("CHM with Individual Crown Boundaries",
                  fontsize=12, fontweight="bold")
axes[1].axis("off")
# Proper ScalarMappable with explicit norm so the colorbar scale is correct
chm_valid = chm_clean[~np.isnan(chm_clean)]
sm = plt.cm.ScalarMappable(
    cmap="YlGn",
    norm=plt.Normalize(vmin=np.nanmin(chm_clean), vmax=np.nanmax(chm_clean))
)
sm.set_array([])
plt.colorbar(sm, ax=axes[1], label="Canopy height (m)", shrink=0.8)

plt.suptitle("Crown Delineation — Helsinki Urban Trees",
             fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()


# ── Convert crown boundary pixels to map coordinates ──────────────────────────
rows_b, cols_b = np.where(boundaries)
xb_map = xmin + (cols_b + 0.5) * RESOLUTION
yb_map = ymax - (rows_b + 0.5) * RESOLUTION

# ── Final composite map ───────────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(12, 10))

# RGB point cloud base (subsampled)
ax.scatter(x_vis, y_vis, c=rgb_vis, s=0.3, rasterized=True, label="_nolegend_")

# Crown boundaries
ax.scatter(xb_map, yb_map, s=0.05, c="cornflowerblue",
           marker=".", label="Crown boundaries")

# Detected tree tops
ax.scatter(xs_final, ys_final, s=8, c="crimson",
           marker="x", linewidths=1.2, label=f"Tree tops ({len(xs_final):,})")

ax.set_title(
    f"Urban Tree Inventory — {len(xs_final):,} Trees Detected\n"
    f"Helsinki, Finland | LiDAR {RESOLUTION} m resolution",
    fontsize=13, fontweight="bold"
)
ax.set_xlabel("Easting (m)",  fontsize=10)
ax.set_ylabel("Northing (m)", fontsize=10)
ax.set_aspect("equal")
ax.ticklabel_format(style="plain")
ax.xaxis.set_major_formatter(ScalarFormatter(useOffset=False))
ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))
ax.legend(loc="upper right", fontsize=10, markerscale=4,
          framealpha=0.85, edgecolor="#ccc")

plt.tight_layout()
plt.show()

Stage	Description
1. Point cloud exploration	Load and visualise RGB-coloured LiDAR data with classification labels
2. Surface model derivation	Rasterise point cloud → DTM, DSM, and nDSM at 0.5 m resolution
3. Vegetation CHM	Isolate vegetation points, compute Canopy Height Model, mask buildings for visulisation
4. Vertical structure classification	Classify canopy into four height so we know what height of vegetation is where: like shrub, sub-canopy and canopy + building layer
5. Tree top detection	Peak local maximum with NaN-aware Gaussian smoothing + prominence filtering
6. Crown delineation	Marker-controlled watershed segmentation using CHM as elevation surface
7. Final maps	RGB point cloud overlaid with detected tree tops and crown boundaries

Item	Detail
Input	RGB-coloured airborne LiDAR `.laz`, Helsinki urban neighbourhood
Resolution	0.5 m
Surface models	DTM, DSM, nDSM, vegetation CHM
Building mask	Class-6 points + 1 m dilation to remove rooftop contamination
CHM smoothing	NaN-aware normalised Gaussian convolution
Tree detection	`peak_local_max` + prominence filter
Crown delineation	Marker-controlled watershed on inverted CHM
Output	Per-tree location, height, and crown extent

🌳 Urban Tree Inventory from LiDAR : CHM, Tree Tops & Crown Delineation¶

Project Overview¶

What this notebook includes¶

Key design decisions¶

Data Source¶

0. Imports and Configuration¶

Step 1 — Helper Functions¶

Step 2 — Load and Explore the LiDAR Point Cloud¶

Step 3 — Rasterise to Surface Models (DTM, DSM, nDSM, CHM)¶

Step 4 — Vertical Structure Classification¶

Step 5 — Tree Top Detection¶

Step 6 — Crown Delineation via Marker-Controlled Watershed¶

Crown Export — GeoPackage / Shapefile¶

Step 7 — Final Map: RGB Point Cloud + Tree Tops + Crown Boundaries¶

Summary¶

What this pipeline delivers for clients¶

Further extensions¶

Class	Label
2	Ground
3	Low vegetation
4	Medium vegetation
5	High vegetation
6	Building
9	Water
17	Bridge

Class	Height range	Ecological meaning
1	0–2 m	Low shrub / herbaceous
2	2–5 m	High shrub / young trees
3	5–10 m	Sub-canopy trees
4	> 10 m	Full canopy trees
5	—	Buildings