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

# ── Geospatial ────────────────────────────────────────────────────────────────
import requests
import rasterio
from rasterio.transform import from_bounds

# ── GeoAI — deep learning canopy height model ─────────────────────────────────
import geoai
from geoai.canopy import CanopyHeightEstimation

# ── Numerical ─────────────────────────────────────────────────────────────────
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

# ── Visualisation ─────────────────────────────────────────────────────────────
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.ticker import ScalarFormatter

print("✓ All imports successful")

✓ All imports successful

Import derived CHM clean we had from LiDAR already

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

# Input
LAS_PATH = DATA_DIR / "rgb_683495c.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   : 26,225,998
  Classes found  : [1, 2, 3, 4, 5, 8, 9]
  Z range        : 21.8 m → 67.1 m
  Easting range  : 25495500 → 25496000
  Northing range : 6683000 → 6683500

  Visualisation subsample: 1-in-10 → 2,622,600 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     : (25495500, 6683000) → (25496000, 6683500)

# ── 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  : 2,878,540
Surface points : 26,013,553

Derived surface models:
  DTM  : 21.8 m → 37.8 m
  DSM  : 22.2 m → 67.1 m
  nDSM : 0.0 m → 37.9 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 : 0
  Building pixels before dilation : 0
  Building pixels after dilation  : 0
Vegetation points : 23,135,013

CHM:
  Valid pixels (raw)    : 934,203
  Valid pixels (cleaned) : 934,203
  Height range           : 0.5 m → 37.9 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\CHM\dtm_05m.tif
  nDSM     → E:\Py\3D\CHM\ndsm_05m.tif
  CHM      → E:\Py\3D\CHM\chm_veg_05m_clean.tif

# =========================
# READ LiDAR derived CHM
# =========================
chm_path = r"E:\Py\3D\CHM\chm_veg_05m_clean.tif"

with rasterio.open(chm_path) as src:
    chm = src.read(1)
    transform = src.transform
    bounds = src.bounds
    res_x, res_y = src.res
    width = src.width
    height = src.height
    crs = src.crs

print("Resolution:", res_x, res_y)
print("Width, Height:", width, height)
print("Bounds:", bounds)
print("CRS:", crs)

Resolution: 0.5 0.5
Width, Height: 1000 1000
Bounds: BoundingBox(left=25495500.0, bottom=6682999.999, right=25496000.0, top=6683499.999)
CRS: EPSG:3879

# Then lets fetch orthophoto for similar area from WMS
wms_url = "https://kartta.hel.fi/ws/geoserver/avoindata/wms"
layer = "avoindata:Ortoilmakuva_2021_5cm"

out_tif = r"E:/Py/3D/CHM/orthophoto_2021.tif"

params = {
    "SERVICE": "WMS",
    "VERSION": "1.1.1",
    "REQUEST": "GetMap",
    "LAYERS": layer,
    "STYLES": "",
    "SRS": "EPSG:3879",
    "BBOX": f"{bounds.left},{bounds.bottom},{bounds.right},{bounds.top}",
    "WIDTH": width,
    "HEIGHT": height,
    "FORMAT": "image/tiff"
}

r = requests.get(wms_url, params=params)
r.raise_for_status()

with open(out_tif, "wb") as f:
    f.write(r.content)

print("Saved:", out_tif)

Saved: E:/Py/3D/CHM/orthophoto_2021.tif

out_tif = r"E:/Py/3D/CHM/orthophoto_2021.tif"
with rasterio.open(out_tif) as src:
    img = src.read()

rgb = np.moveaxis(img[:3],0,2)

plt.figure(figsize=(6,6))
plt.imshow(rgb)
plt.title("Orthophoto 2021")
plt.axis("off")

C:\Users\Asus\anaconda3\envs\GEOAIcourse\Lib\site-packages\rasterio\__init__.py:356: NotGeoreferencedWarning: Dataset has no geotransform, gcps, or rpcs. The identity matrix will be returned.
  dataset = DatasetReader(path, driver=driver, sharing=sharing, **kwargs)

(np.float64(-0.5), np.float64(999.5), np.float64(999.5), np.float64(-0.5))

ortho_path = r"E:/Py/3D/CHM/orthophoto_2021.tif"

with rasterio.open(chm_path) as chm_src:
    print("CHM")
    print(" shape :", chm_src.height, chm_src.width)
    print(" crs   :", chm_src.crs)
    print(" res   :", chm_src.res)
    print(" bounds:", chm_src.bounds)

with rasterio.open(ortho_path) as ortho_src:
    print("\nORTHO")
    print(" shape :", ortho_src.height, ortho_src.width)
    print(" crs   :", ortho_src.crs)
    print(" res   :", ortho_src.res)
    print(" bounds:", ortho_src.bounds)

CHM
 shape : 1000 1000
 crs   : EPSG:3879
 res   : (0.5, 0.5)
 bounds: BoundingBox(left=25495500.0, bottom=6682999.999, right=25496000.0, top=6683499.999)

ORTHO
 shape : 1000 1000
 crs   : None
 res   : (1.0, 1.0)
 bounds: BoundingBox(left=0.0, bottom=1000.0, right=1000.0, top=0.0)

import rasterio
from rasterio.transform import from_bounds

ortho_path = r"E:/Py/3D/CHM/orthophoto_2021.tif"
ortho_fixed = r"E:/Py/3D/CHM/orthophoto_2021_geo.tif"

# read CHM reference
with rasterio.open(chm_path) as chm_src:
    bounds = chm_src.bounds
    crs = chm_src.crs
    width = chm_src.width
    height = chm_src.height

# read orthophoto pixels
with rasterio.open(ortho_path) as src:
    data = src.read()
    profile = src.profile

# create proper transform
transform = from_bounds(
    bounds.left,
    bounds.bottom,
    bounds.right,
    bounds.top,
    width,
    height
)

profile.update({
    "crs": crs,
    "transform": transform,
    "width": width,
    "height": height
})

with rasterio.open(ortho_fixed, "w", **profile) as dst:
    dst.write(data)

print("Fixed orthophoto:", ortho_fixed)

Fixed orthophoto: E:/Py/3D/CHM/orthophoto_2021_geo.tif

with rasterio.open(ortho_fixed) as src:
    print(src.crs)
    print(src.bounds)
    print(src.res)

EPSG:3879
BoundingBox(left=25495500.0, bottom=6682999.999, right=25496000.0, top=6683499.999)
(0.5, 0.5)

pred_path = r"E:/Py/3D/CHM/chm_geoai_predict.tif"

estimator = CanopyHeightEstimation(
    model_name="compressed_SSLhuge_aerial"
)

pred = estimator.predict(
    ortho_fixed,
    output_path=pred_path,
    batch_size=4
)

print("GeoAI CHM saved:", pred_path)

C:\Users\Asus\anaconda3\envs\GEOAIcourse\Lib\site-packages\torch\_utils.py:445: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  device=storage.device,

GeoAI CHM saved: E:/Py/3D/CHM/chm_geoai_predict.tif

pred_path = r"E:/Py/3D/CHM/chm_geoai_predict.tif"
with rasterio.open(chm_path) as src:
    chm_lidar = src.read(1)

with rasterio.open(pred_path) as src:
    chm_geoai = src.read(1)

fig, ax = plt.subplots(1,2, figsize=(12,5))

ax[0].imshow(chm_lidar, cmap="viridis", vmin=0, vmax=40)
ax[0].set_title("LiDAR CHM")

ax[1].imshow(chm_geoai, cmap="viridis", vmin=0, vmax=40)
ax[1].set_title("GeoAI CHM")

plt.show()

To make this visually good.. we can visuialise veg only area in GeoAI predicted too

geoai_masked_path = r"E:/Py/3D/CHM/chm_geoai_predict_vegmasked.tif"

out_profile = pred_profile.copy()
out_profile.update(dtype="float32", nodata=np.nan)

with rasterio.open(geoai_masked_path, "w", **out_profile) as dst:
    dst.write(geoai_masked, 1)

print("Saved:", geoai_masked_path)

# ── Build vegetation mask from LiDAR CHM ─────────────────────────────────────
# Any pixel where LiDAR detected canopy above 1 cm is considered vegetation.
veg_mask = chm_lidar > 0.01

print(f"Vegetation mask:")
print(f"  Vegetated pixels : {veg_mask.sum():,} ({veg_mask.mean()*100:.1f}% of raster)")
print(f"  Non-vegetated    : {(~veg_mask).sum():,}")

# Apply mask to both CHMs — non-vegetated pixels become NaN
lidar_masked = chm_lidar.astype("float32").copy()
lidar_masked[~veg_mask] = np.nan

geoai_masked = chm_geoai.astype("float32").copy()
geoai_masked[~veg_mask] = np.nan

print(f"\nMasked value ranges:")
print(f"  LiDAR : {np.nanmin(lidar_masked):.2f} m → {np.nanmax(lidar_masked):.2f} m")
print(f"  GeoAI : {np.nanmin(geoai_masked):.2f} m → {np.nanmax(geoai_masked):.2f} m")

# Save masked GeoAI CHM for use in other analyses
out_profile = chm_profile.copy()
out_profile.update(dtype="float32", nodata=np.nan)

with rasterio.open(str(geoai_masked_path), "w", **out_profile) as dst:
    dst.write(geoai_masked, 1)

print(f"\n✓ Masked GeoAI CHM saved → {geoai_masked_path}")

Saved: E:/Py/3D/CHM/chm_geoai_predict_vegmasked.tif
Vegetation mask:
  Vegetated pixels : 934,203 (93.4% of raster)
  Non-vegetated    : 65,797

Masked value ranges:
  LiDAR : 0.50 m → 37.87 m
  GeoAI : 0.00 m → 25.61 m

✓ Masked GeoAI CHM saved → E:/Py/3D/CHM/chm_geoai_predict_vegmasked.tif

veg_mask = chm_lidar > 0.01

plt.figure(figsize=(6,6))
plt.imshow(veg_mask, cmap="gray")
plt.title("Vegetation mask from LiDAR CHM")
plt.axis("off")
plt.show()

geoai_masked = chm_geoai.astype("float32").copy()
geoai_masked[~veg_mask] = np.nan

lidar_masked = chm_lidar.astype("float32").copy()
lidar_masked[~veg_mask] = np.nan

# ── Shared colour scale anchored to LiDAR distribution ───────────────────────
lidar_vals = lidar_masked[np.isfinite(lidar_masked)]
vmin = 0
vmax = np.percentile(lidar_vals, 99)   # exclude extreme outliers from colour scale

fig, axes = plt.subplots(1, 2, figsize=(16, 7), constrained_layout=True)

im1 = axes[0].imshow(lidar_masked, cmap="YlGn", vmin=vmin, vmax=vmax)
axes[0].set_title("LiDAR-derived CHM\n(ground truth)",
                  fontsize=13, fontweight="bold")

im2 = axes[1].imshow(geoai_masked, cmap="YlGn", vmin=vmin, vmax=vmax)
axes[1].set_title("GeoAI Predicted CHM\n(compressed_SSLhuge_aerial)",
                  fontsize=13, fontweight="bold")

for ax in axes:
    ax.axis("on")
    ax.ticklabel_format(style="plain")
    ax.xaxis.set_major_formatter(ScalarFormatter(useOffset=False))
    ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))
    ax.set_xlabel("Column (pixels)", fontsize=9)
    ax.set_ylabel("Row (pixels)", fontsize=9)

cbar = fig.colorbar(im1, ax=axes, location="right",
                    shrink=0.85, pad=0.02)
cbar.set_label("Canopy height (m)", fontsize=11)

plt.suptitle("Canopy Height Comparison — Vegetation Pixels Only",
             fontsize=14, fontweight="bold")
plt.show()

print("LiDAR min/max:", np.nanmin(lidar_masked), np.nanmax(lidar_masked))
print("GeoAI min/max:", np.nanmin(geoai_masked), np.nanmax(geoai_masked))

LiDAR min/max: 0.5 37.869
GeoAI min/max: 0.0 25.614182

# ── Extract valid pixel pairs for statistical comparison ──────────────────────
valid = np.isfinite(lidar_masked) & np.isfinite(geoai_masked)
y_true = lidar_masked[valid]   # LiDAR = ground truth
y_pred = geoai_masked[valid]   # GeoAI = prediction

print(f"Valid pixel pairs for comparison: {valid.sum():,}")

# ── Fit linear regression ─────────────────────────────────────────────────────
reg   = LinearRegression().fit(y_true.reshape(-1, 1), y_pred)
y_fit = reg.predict(y_true.reshape(-1, 1))

# ── Compute metrics ───────────────────────────────────────────────────────────
r2        = reg.score(y_true.reshape(-1, 1), y_pred)
rmse      = np.sqrt(mean_squared_error(y_true, y_pred))
mean_bias = np.mean(y_pred - y_true)   # positive = GeoAI overestimates
slope     = reg.coef_[0]
intercept = reg.intercept_

print("\nValidation metrics (vegetation pixels only):")
print(f"  R²         : {r2:.3f}")
print(f"  RMSE       : {rmse:.2f} m")
print(f"  Mean bias  : {mean_bias:+.2f} m  "
      f"({'underestimate' if mean_bias < 0 else 'overestimate'})")
print(f"  Regression : GeoAI = {slope:.2f} × LiDAR + {intercept:.2f}")

Valid pixel pairs for comparison: 934,203

Validation metrics (vegetation pixels only):
  R²         : 0.183
  RMSE       : 10.95 m
  Mean bias  : -8.14 m  (underestimate)
  Regression : GeoAI = 0.28 × LiDAR + 6.55

# ── Height distribution comparison ───────────────────────────────────────────
fig, ax = plt.subplots(figsize=(9, 4))

ax.hist(y_true, bins=60, alpha=0.65, color="#2ca02c",
        label=f"LiDAR  (mean={y_true.mean():.1f} m)", density=True)
ax.hist(y_pred, bins=60, alpha=0.65, color="#1f77b4",
        label=f"GeoAI  (mean={y_pred.mean():.1f} m)", density=True)

ax.axvline(y_true.mean(), color="#2ca02c", linestyle="--", linewidth=1.5)
ax.axvline(y_pred.mean(), color="#1f77b4", linestyle="--", linewidth=1.5)

ax.set_xlabel("Canopy height (m)", fontsize=11)
ax.set_ylabel("Density", fontsize=11)
ax.set_title("Height Distribution — LiDAR vs GeoAI (vegetation pixels)",
             fontsize=12, fontweight="bold")
ax.legend(fontsize=10)
ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()

# ── Scatter plot: GeoAI vs LiDAR ─────────────────────────────────────────────
# Subsample for plotting speed (scatter with millions of points is slow)
n_plot = min(50_000, len(y_true))
idx    = np.random.choice(len(y_true), n_plot, replace=False)

fig, ax = plt.subplots(figsize=(7, 7))

ax.scatter(y_true[idx], y_pred[idx],
           s=1, alpha=0.12, color="steelblue", rasterized=True)

# 1:1 perfect agreement line
max_val = max(y_true.max(), y_pred.max())
ax.plot([0, max_val], [0, max_val], "r--",
        linewidth=1.5, label="Perfect agreement (1:1)")

# Fitted regression line
x_line = np.linspace(0, max_val, 100)
y_line = slope * x_line + intercept
ax.plot(x_line, y_line, "k-",
        linewidth=2, label=f"Regression  (slope={slope:.2f})")

# Annotate metrics
ax.text(0.05, 0.95,
        f"R²   = {r2:.3f}\nRMSE = {rmse:.2f} m\nBias = {mean_bias:+.2f} m",
        transform=ax.transAxes,
        fontsize=11, verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.8))

ax.set_xlabel("LiDAR height (m)  [ground truth]", fontsize=11)
ax.set_ylabel("GeoAI height (m)  [prediction]",   fontsize=11)
ax.set_title("GeoAI vs LiDAR Canopy Height — Pixel-Level Validation",
             fontsize=12, fontweight="bold")
ax.legend(fontsize=10)
ax.set_xlim(0, max_val)
ax.set_ylim(0, max_val)
ax.set_aspect("equal")
ax.grid(alpha=0.25)

plt.tight_layout()
plt.show()

# ── Compute residual: GeoAI − LiDAR ──────────────────────────────────────────
# Positive = GeoAI overestimates; Negative = GeoAI underestimates
residual = geoai_masked - lidar_masked   # NaN outside vegetation mask

# Symmetric colour scale centred on zero
res_vals  = residual[np.isfinite(residual)]
clim      = np.percentile(np.abs(res_vals), 97)   # clip extreme outliers

print(f"Residual statistics:")
print(f"  Mean bias  : {res_vals.mean():+.2f} m")
print(f"  Std dev    : {res_vals.std():.2f} m")
print(f"  P5  / P95  : {np.percentile(res_vals, 5):+.1f} m / {np.percentile(res_vals, 95):+.1f} m")
print(f"  Colour scale: ±{clim:.1f} m")

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

# ── Left: residual map ────────────────────────────────────────────────────────
im_res = axes[0].imshow(
    residual,
    cmap="RdBu",        # red = underestimate, blue = overestimate
    vmin=-clim, vmax=clim
)
axes[0].set_title("Residual Map  (GeoAI − LiDAR)",
                  fontsize=13, fontweight="bold")
axes[0].axis("off")
cbar_res = plt.colorbar(im_res, ax=axes[0], shrink=0.8)
cbar_res.set_label("Error (m)  [+ = overestimate, − = underestimate]",
                   fontsize=10)

# ── Right: absolute error map ─────────────────────────────────────────────────
abs_error = np.abs(residual)
im_abs = axes[1].imshow(
    abs_error,
    cmap="YlOrRd",
    vmin=0, vmax=clim
)
axes[1].set_title("Absolute Error Map  |GeoAI − LiDAR|",
                  fontsize=13, fontweight="bold")
axes[1].axis("off")
cbar_abs = plt.colorbar(im_abs, ax=axes[1], shrink=0.8)
cbar_abs.set_label("Absolute error (m)", fontsize=10)

plt.suptitle("Spatial Error Distribution — GeoAI vs LiDAR CHM",
             fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()

Residual statistics:
  Mean bias  : -8.14 m
  Std dev    : 7.32 m
  P5  / P95  : -19.6 m / +5.8 m
  Colour scale: ±22.0 m