diff --git a/geodata_config.toml b/geodata_config.toml index f568ebd..a78915b 100644 --- a/geodata_config.toml +++ b/geodata_config.toml @@ -33,20 +33,31 @@ output_dir = "export_unity/height_png16_river" manifest_vr = "export_unity/tile_index_river_vr.csv" manifest_server = "export_unity/tile_index_river_server.csv" +[river_erosion.lidar] +enabled = true +source_dir = "raw/geo_rlp/bdom20rgbi" +classification_water = 9 +depth_m = 3.5 +bank_slope_sigma = 2.0 +fill_holes_radius = 3 + [river_erosion.vr] order_field = "ORD_STRA" invert_order = false invert_max_order = 10 -min_order = 3 +min_order = 1 width_base_m = 8.0 width_per_order_m = 6.0 width_min_m = 4.0 width_max_m = 90.0 -depth_base_m = 0.7 -depth_per_order_m = 0.5 -depth_min_m = 0.3 +depth_base_m = 0.2 +depth_per_order_m = 0.44 +depth_min_m = 0.2 depth_max_m = 6.0 smooth_sigma_m = 3.0 +fallback_depth_m = 2.0 +flow_distance_field = "DIST_DN_KM" +flow_slope_m_per_km = 0.2 [river_erosion.server] order_field = "ORD_FLOW" @@ -73,6 +84,12 @@ enabled = true [ortho] out_res = 2048 jpeg_quality = 90 +apply_water_mask = true +water_blend = 0.85 +water_fallback_rgb = [92.0, 72.0, 52.0] +water_mask_threshold = 0.1 +river_dir = "export_unity/ortho_jpg_river" +water_color_mode = "fixed" [buildings] out_dir = "export_unity/buildings_tiles" diff --git a/geodata_pipeline/config.py b/geodata_pipeline/config.py index 52fe9ad..325fcce 100644 --- a/geodata_pipeline/config.py +++ b/geodata_pipeline/config.py @@ -54,16 +54,29 @@ class RiverErosionProfileConfig: order_field: str = "ORD_STRA" invert_order: bool = False invert_max_order: int = 10 - min_order: int = 3 + min_order: int = 1 width_base_m: float = 8.0 width_per_order_m: float = 6.0 width_min_m: float = 4.0 width_max_m: float = 90.0 - depth_base_m: float = 0.7 - depth_per_order_m: float = 0.5 - depth_min_m: float = 0.3 + depth_base_m: float = 0.2 + depth_per_order_m: float = 0.44 + depth_min_m: float = 0.2 depth_max_m: float = 6.0 smooth_sigma_m: float = 3.0 + fallback_depth_m: float = 2.0 + flow_distance_field: str = "DIST_DN_KM" + flow_slope_m_per_km: float = 0.0 + + +@dataclass +class RiverErosionLidarConfig: + enabled: bool = True + source_dir: str = "raw/bdom20rgbi" + classification_water: int = 9 + depth_m: float = 3.5 + bank_slope_sigma: float = 2.0 + fill_holes_radius: int = 3 def _default_river_profile_vr() -> "RiverErosionProfileConfig": @@ -71,16 +84,19 @@ def _default_river_profile_vr() -> "RiverErosionProfileConfig": order_field="ORD_STRA", invert_order=False, invert_max_order=10, - min_order=3, + min_order=1, width_base_m=8.0, width_per_order_m=6.0, width_min_m=4.0, width_max_m=90.0, - depth_base_m=0.7, - depth_per_order_m=0.5, - depth_min_m=0.3, + depth_base_m=0.2, + depth_per_order_m=0.44, + depth_min_m=0.2, depth_max_m=6.0, smooth_sigma_m=3.0, + fallback_depth_m=2.0, + flow_distance_field="DIST_DN_KM", + flow_slope_m_per_km=0.0, ) @@ -112,12 +128,19 @@ class RiverErosionConfig: manifest_server: str = "export_unity/tile_index_river_server.csv" vr: RiverErosionProfileConfig = field(default_factory=_default_river_profile_vr) server: RiverErosionProfileConfig = field(default_factory=_default_river_profile_server) + lidar: RiverErosionLidarConfig = field(default_factory=RiverErosionLidarConfig) @dataclass class OrthoConfig: out_res: int = 2048 jpeg_quality: int = 90 + apply_water_mask: bool = False + water_blend: float = 0.85 + water_fallback_rgb: tuple[float, float, float] = (30.0, 50.0, 60.0) + water_mask_threshold: float = 0.1 + river_dir: str = "export_unity/ortho_jpg_river" + water_color_mode: str = "median" @dataclass @@ -320,6 +343,9 @@ def _coerce_ranges(data: Dict[str, Any]) -> Dict[str, Any]: value = out.get(key) if isinstance(value, list): out[key] = tuple(value) + value = out.get("water_fallback_rgb") + if isinstance(value, list): + out["water_fallback_rgb"] = tuple(value) return out @@ -329,6 +355,8 @@ def _river_erosion_from_dict(data: Dict[str, Any]) -> RiverErosionConfig: base = _filter_kwargs(RiverErosionConfig, data) base.pop("vr", None) base.pop("server", None) + base.pop("lidar", None) vr_cfg = RiverErosionProfileConfig(**_filter_kwargs(RiverErosionProfileConfig, data.get("vr", {}))) server_cfg = RiverErosionProfileConfig(**_filter_kwargs(RiverErosionProfileConfig, data.get("server", {}))) - return RiverErosionConfig(**base, vr=vr_cfg, server=server_cfg) + lidar_cfg = RiverErosionLidarConfig(**_filter_kwargs(RiverErosionLidarConfig, data.get("lidar", {}))) + return RiverErosionConfig(**base, vr=vr_cfg, server=server_cfg, lidar=lidar_cfg) diff --git a/geodata_pipeline/river_erosion.py b/geodata_pipeline/river_erosion.py index 1019a64..a6b4e07 100644 --- a/geodata_pipeline/river_erosion.py +++ b/geodata_pipeline/river_erosion.py @@ -2,57 +2,50 @@ from __future__ import annotations import csv import glob +import json import math import os -from typing import Dict, Iterable, List, Tuple +from typing import Dict, Iterable, List, Optional, Tuple import numpy as np +import pdal from osgeo import gdal, ogr, osr from scipy import ndimage from .config import Config, RiverErosionProfileConfig -from .gdal_utils import build_vrt, ensure_dir, ensure_parent, open_dataset +from .gdal_utils import ensure_dir, ensure_parent, open_dataset gdal.UseExceptions() -def _coerce_int(value) -> int | None: - if value is None: +def _river_debug_enabled() -> bool: + flag = os.getenv("RIVER_EROSION_DEBUG", "") + return flag.strip().lower() in {"1", "true", "yes", "y"} + + +def _river_debug(message: str) -> None: + if _river_debug_enabled(): + print(f"[river_erosion:debug] {message}") + + +def _parse_tile_key(tile_key: str) -> Optional[Tuple[int, int]]: + if not tile_key: + return None + parts = tile_key.split("_") + if len(parts) != 2: return None try: - return int(value) - except (TypeError, ValueError): + return int(parts[0]), int(parts[1]) + except ValueError: return None -def _normalize_order(value, profile: RiverErosionProfileConfig) -> int | None: - order = _coerce_int(value) - if order is None: - return None - if profile.invert_order: - max_order = max(1, int(profile.invert_max_order)) - if order <= 0: - order = max_order - order = max(1, min(max_order, order)) - order = max_order + 1 - order - return order - - -def _order_width_depth(order: int, profile: RiverErosionProfileConfig) -> Tuple[float, float] | None: - if order < profile.min_order: - return None - delta = order - profile.min_order - width = profile.width_base_m + profile.width_per_order_m * delta - depth = profile.depth_base_m + profile.depth_per_order_m * delta - width = min(max(width, profile.width_min_m), profile.width_max_m) - depth = min(max(depth, profile.depth_min_m), profile.depth_max_m) - return width, depth - - -def _tile_geotransform(bounds: Tuple[float, float, float, float], out_res: int) -> Tuple[float, float, float, float, float, float]: +def _tile_geotransform(bounds: Tuple[float, float, float, float], width: int, height: int) -> Tuple[float, float, float, float, float, float]: xmin, ymin, xmax, ymax = bounds - xres = (xmax - xmin) / float(out_res) - yres = (ymax - ymin) / float(out_res) + xdenom = float(width - 1) if width > 1 else float(width) + ydenom = float(height - 1) if height > 1 else float(height) + xres = (xmax - xmin) / xdenom + yres = (ymax - ymin) / ydenom return xmin, xres, 0.0, ymax, 0.0, -yres @@ -80,117 +73,204 @@ def _bounds_to_layer( return min(xs), min(ys), max(xs), max(ys) -def _compute_tile_minmax_from_vrt( - vrt_ds: gdal.Dataset, +def _normalize_order(value, profile: RiverErosionProfileConfig) -> Optional[int]: + if value is None: + return None + try: + order = int(value) + except (TypeError, ValueError): + return None + if profile.invert_order: + max_order = max(1, int(profile.invert_max_order)) + if order <= 0: + order = max_order + order = max(1, min(max_order, order)) + order = max_order + 1 - order + return order + + +def _order_depth(order: int, profile: RiverErosionProfileConfig) -> Optional[float]: + if order < profile.min_order: + return None + delta = order - profile.min_order + depth = profile.depth_base_m + profile.depth_per_order_m * delta + depth = min(max(depth, profile.depth_min_m), profile.depth_max_m) + return depth + + +def _refine_mask( + mask: np.ndarray, + fill_radius: int = 3, + smooth_sigma: float = 2.0, +) -> np.ndarray: + if fill_radius > 0: + struct = ndimage.generate_binary_structure(2, 1) + mask = ndimage.binary_closing(mask, structure=struct, iterations=fill_radius).astype(np.float32) + if smooth_sigma > 0: + mask = ndimage.gaussian_filter(mask, sigma=smooth_sigma) + return mask + + +def _build_water_mask_mosaic( + laz_paths: Iterable[str], bounds: Tuple[float, float, float, float], - out_res: int, - resample: str, -) -> Tuple[float | None, float | None]: - opts = gdal.WarpOptions( - outputBounds=bounds, - width=out_res, - height=out_res, - resampleAlg=resample, - srcNodata=-9999, - dstNodata=-9999, - format="MEM", - ) - tmp = gdal.Warp("", vrt_ds, options=opts) - if tmp is None: - return None, None - band = tmp.GetRasterBand(1) - nodata = band.GetNoDataValue() - data = band.ReadAsArray() - tmp = None - if data is None or data.size == 0: - return None, None - valid = np.isfinite(data) - if nodata is not None and not math.isnan(nodata): - valid &= data != nodata + width: int, + height: int, + water_class: int = 9, +) -> Optional[np.ndarray]: + laz_paths = [path for path in laz_paths if path and os.path.exists(path)] + if not laz_paths: + return None + + xmin, ymin, xmax, ymax = bounds + pipeline_json: list[dict] = [] + for path in laz_paths: + pipeline_json.append({ + "type": "readers.las", + "filename": path, + }) + if len(laz_paths) > 1: + pipeline_json.append({"type": "filters.merge"}) + pipeline_json.append({ + "type": "filters.crop", + "bounds": f"([{xmin},{xmax}],[{ymin},{ymax}])", + }) + pipeline_json.append({ + "type": "filters.range", + "limits": f"Classification[{water_class}:{water_class}]", + }) + + try: + pipeline = pdal.Pipeline(json.dumps(pipeline_json)) + count = pipeline.execute() + except RuntimeError as exc: + print(f"[river_erosion] PDAL error building water mask: {exc}") + return None + + if count == 0: + return None + + arrays = pipeline.arrays + if not arrays: + return None + + if len(arrays) == 1: + xs = arrays[0]["X"] + ys = arrays[0]["Y"] + else: + xs = np.concatenate([arr["X"] for arr in arrays]) + ys = np.concatenate([arr["Y"] for arr in arrays]) + + xdenom = float(width - 1) if width > 1 else float(width) + ydenom = float(height - 1) if height > 1 else float(height) + xres = (xmax - xmin) / xdenom + yres = (ymax - ymin) / ydenom + + col = ((xs - xmin) / xres).astype(int) + row = ((ymax - ys) / yres).astype(int) + + valid = (col >= 0) & (col < width) & (row >= 0) & (row < height) if not np.any(valid): - return None, None - return float(data[valid].min()), float(data[valid].max()) + return None + + mask = np.zeros((height, width), dtype=np.float32) + mask[row[valid], col[valid]] = 1.0 + return mask -def _build_depth_raster( - layer: ogr.Layer, +def _layer_has_field(layer: ogr.Layer, field_name: Optional[str]) -> bool: + if not field_name: + return False + layer_defn = layer.GetLayerDefn() + if layer_defn is None: + return False + return layer_defn.GetFieldIndex(field_name) != -1 + + +def _build_depth_lines( + layer_sources: List[Tuple[ogr.Layer, str]], bounds: Tuple[float, float, float, float], - out_res: int, + width: int, + height: int, profile: RiverErosionProfileConfig, to_tile: osr.CoordinateTransformation, to_layer: osr.CoordinateTransformation, tile_srs: osr.SpatialReference, -) -> Tuple[np.ndarray, Tuple[float, float, float, float, float, float]]: - gt = _tile_geotransform(bounds, out_res) +) -> np.ndarray: + gt = _tile_geotransform(bounds, width, height) driver = gdal.GetDriverByName("MEM") - depth_ds = driver.Create("", out_res, out_res, 1, gdal.GDT_Float32) + depth_ds = driver.Create("", width, height, 1, gdal.GDT_Float32) depth_ds.SetGeoTransform(gt) depth_ds.SetProjection(tile_srs.ExportToWkt()) band = depth_ds.GetRasterBand(1) band.Fill(0.0) - buffer_m = (profile.width_max_m * 0.5) + (profile.smooth_sigma_m * 3.0) + if not any(_layer_has_field(layer, profile.order_field) for layer, _ in layer_sources): + _river_debug(f"Order field '{profile.order_field}' not found in HydroRIVERS layer(s).") + return band.ReadAsArray().astype(np.float32) + + buffer_m = max(profile.smooth_sigma_m * 3.0, 25.0) minx, miny, maxx, maxy = _bounds_to_layer(bounds, buffer_m, to_layer) - layer.SetSpatialFilterRect(minx, miny, maxx, maxy) - layer.ResetReading() buffered_geoms: list[ogr.Geometry] = [] burn_values: list[float] = [] - for feat in layer: - order = _normalize_order(feat.GetField(profile.order_field), profile) - if order is None: - continue - width_depth = _order_width_depth(order, profile) - if width_depth is None: - continue - width_m, depth_m = width_depth - geom = feat.GetGeometryRef() - if geom is None: - continue - geom = geom.Clone() - geom.Transform(to_tile) - if geom.IsEmpty(): - continue - buffered = geom.Buffer(width_m * 0.5) - if buffered is None or buffered.IsEmpty(): - continue - buffered_geoms.append(buffered) - burn_values.append(depth_m) - if buffered_geoms: - # GDAL only supports MERGE_ALG=REPLACE/ADD; emulate MAX by sorting ascending. - ordered = sorted(zip(buffered_geoms, burn_values), key=lambda item: item[1]) - ordered_geoms = [geom for geom, _ in ordered] - ordered_values = [float(value) for _, value in ordered] - if hasattr(gdal, "RasterizeGeometries"): - gdal.RasterizeGeometries( - depth_ds, - [1], - ordered_geoms, - burn_values=ordered_values, - options=["MERGE_ALG=REPLACE"], - ) - else: - mem_driver = ogr.GetDriverByName("MEM") or ogr.GetDriverByName("Memory") - if mem_driver is None: - raise SystemExit("[river_erosion] OGR MEM driver missing; update GDAL.") - mem_ds = mem_driver.CreateDataSource("") - mem_layer = mem_ds.CreateLayer("burn", srs=tile_srs, geom_type=ogr.wkbPolygon) - mem_layer.CreateField(ogr.FieldDefn("burn", ogr.OFTReal)) - layer_defn = mem_layer.GetLayerDefn() - for geom, value in zip(ordered_geoms, ordered_values): - feature = ogr.Feature(layer_defn) - feature.SetField("burn", value) - feature.SetGeometry(geom) - mem_layer.CreateFeature(feature) - feature = None - gdal.RasterizeLayer( - depth_ds, - [1], - mem_layer, - options=["ATTRIBUTE=burn"], - ) - mem_ds = None + for layer, _ in layer_sources: + layer.SetSpatialFilterRect(minx, miny, maxx, maxy) + layer.ResetReading() + for feat in layer: + order = _normalize_order(feat.GetField(profile.order_field), profile) + if order is None: + continue + depth_m = _order_depth(order, profile) + if depth_m is None: + continue + geom = feat.GetGeometryRef() + if geom is None: + continue + geom = geom.Clone() + geom.Transform(to_tile) + if geom.IsEmpty(): + continue + buffered_geoms.append(geom) + burn_values.append(float(depth_m)) + + if not buffered_geoms: + return band.ReadAsArray().astype(np.float32) + + ordered = sorted(zip(buffered_geoms, burn_values), key=lambda item: item[1]) + ordered_geoms = [geom for geom, _ in ordered] + ordered_values = [float(value) for _, value in ordered] + + if hasattr(gdal, "RasterizeGeometries"): + gdal.RasterizeGeometries( + depth_ds, + [1], + ordered_geoms, + burn_values=ordered_values, + options=["MERGE_ALG=REPLACE", "ALL_TOUCHED=TRUE"], + ) + else: + mem_driver = ogr.GetDriverByName("MEM") or ogr.GetDriverByName("Memory") + if mem_driver is None: + raise SystemExit("[river_erosion] OGR MEM driver missing; update GDAL.") + mem_ds = mem_driver.CreateDataSource("") + mem_layer = mem_ds.CreateLayer("burn", srs=tile_srs, geom_type=ogr.wkbLineString) + mem_layer.CreateField(ogr.FieldDefn("burn", ogr.OFTReal)) + layer_defn = mem_layer.GetLayerDefn() + for geom, value in zip(ordered_geoms, ordered_values): + feature = ogr.Feature(layer_defn) + feature.SetField("burn", value) + feature.SetGeometry(geom) + mem_layer.CreateFeature(feature) + feature = None + gdal.RasterizeLayer( + depth_ds, + [1], + mem_layer, + options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"], + ) + mem_ds = None depth = band.ReadAsArray().astype(np.float32) if profile.smooth_sigma_m > 0.0: @@ -198,7 +278,204 @@ def _build_depth_raster( if sigma_px > 0.0: depth = ndimage.gaussian_filter(depth, sigma=sigma_px, mode="nearest") - return depth, gt + return depth + + +def _build_flow_distance( + layer_sources: List[Tuple[ogr.Layer, str]], + bounds: Tuple[float, float, float, float], + width: int, + height: int, + dist_field: str, + to_tile: osr.CoordinateTransformation, + to_layer: osr.CoordinateTransformation, + tile_srs: osr.SpatialReference, +) -> np.ndarray: + gt = _tile_geotransform(bounds, width, height) + driver = gdal.GetDriverByName("MEM") + dist_ds = driver.Create("", width, height, 1, gdal.GDT_Float32) + dist_ds.SetGeoTransform(gt) + dist_ds.SetProjection(tile_srs.ExportToWkt()) + band = dist_ds.GetRasterBand(1) + band.Fill(0.0) + + if not any(_layer_has_field(layer, dist_field) for layer, _ in layer_sources): + _river_debug(f"Distance field '{dist_field}' not found in HydroRIVERS layer(s).") + return band.ReadAsArray().astype(np.float32) + + buffer_m = 50.0 + minx, miny, maxx, maxy = _bounds_to_layer(bounds, buffer_m, to_layer) + + buffered_geoms: list[ogr.Geometry] = [] + burn_values: list[float] = [] + + for layer, _ in layer_sources: + layer.SetSpatialFilterRect(minx, miny, maxx, maxy) + layer.ResetReading() + for feat in layer: + dist_val = feat.GetField(dist_field) + if dist_val is None: + continue + try: + dist_km = float(dist_val) + except (TypeError, ValueError): + continue + geom = feat.GetGeometryRef() + if geom is None: + continue + geom = geom.Clone() + geom.Transform(to_tile) + if geom.IsEmpty(): + continue + buffered_geoms.append(geom) + burn_values.append(dist_km) + + if not buffered_geoms: + return band.ReadAsArray().astype(np.float32) + + ordered = sorted(zip(buffered_geoms, burn_values), key=lambda item: item[1]) + ordered_geoms = [geom for geom, _ in ordered] + ordered_values = [float(value) for _, value in ordered] + + if hasattr(gdal, "RasterizeGeometries"): + gdal.RasterizeGeometries( + dist_ds, + [1], + ordered_geoms, + burn_values=ordered_values, + options=["MERGE_ALG=REPLACE", "ALL_TOUCHED=TRUE"], + ) + else: + mem_driver = ogr.GetDriverByName("MEM") or ogr.GetDriverByName("Memory") + if mem_driver is None: + raise SystemExit("[river_erosion] OGR MEM driver missing; update GDAL.") + mem_ds = mem_driver.CreateDataSource("") + mem_layer = mem_ds.CreateLayer("burn", srs=tile_srs, geom_type=ogr.wkbLineString) + mem_layer.CreateField(ogr.FieldDefn("burn", ogr.OFTReal)) + layer_defn = mem_layer.GetLayerDefn() + for geom, value in zip(ordered_geoms, ordered_values): + feature = ogr.Feature(layer_defn) + feature.SetField("burn", value) + feature.SetGeometry(geom) + mem_layer.CreateFeature(feature) + feature = None + gdal.RasterizeLayer( + dist_ds, + [1], + mem_layer, + options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"], + ) + mem_ds = None + + return band.ReadAsArray().astype(np.float32) + + +def _fill_nan_with_nearest(data: np.ndarray) -> np.ndarray: + if not np.isnan(data).any(): + return data + valid = np.isfinite(data) + if not np.any(valid): + return np.zeros_like(data, dtype=np.float32) + _, indices = ndimage.distance_transform_edt(~valid, return_indices=True) + filled = data[tuple(indices)] + return filled.astype(np.float32) + + +def _resize_mask(mask: np.ndarray, height: int, width: int) -> np.ndarray: + if mask.shape == (height, width): + return mask + zoom_y = height / mask.shape[0] + zoom_x = width / mask.shape[1] + resized = ndimage.zoom(mask, (zoom_y, zoom_x), order=1) + if resized.shape[0] != height or resized.shape[1] != width: + resized = resized[:height, :width] + pad_y = max(0, height - resized.shape[0]) + pad_x = max(0, width - resized.shape[1]) + if pad_y or pad_x: + resized = np.pad(resized, ((0, pad_y), (0, pad_x)), mode="edge") + return resized.astype(np.float32) + + +def _apply_water_mask_to_ortho(tile_id: str, mask: np.ndarray, cfg: Config) -> None: + if not cfg.ortho.apply_water_mask: + return + ortho_path = os.path.join(cfg.export.ortho_dir, f"{tile_id}.jpg") + if not os.path.exists(ortho_path): + return + river_dir = getattr(cfg.ortho, "river_dir", cfg.export.ortho_dir) + ensure_dir(river_dir) + out_path = os.path.join(river_dir, f"{tile_id}.jpg") + + ds = gdal.Open(ortho_path, gdal.GA_ReadOnly) + if ds is None: + print(f"[river_erosion] Could not open ortho {ortho_path}") + return + + width = ds.RasterXSize + height = ds.RasterYSize + if width <= 0 or height <= 0: + ds = None + return + + data = ds.ReadAsArray() + if data is None: + ds = None + return + + if data.ndim == 2: + data = np.stack([data, data, data], axis=0) + + if data.shape[0] < 3: + ds = None + return + + rgb = data[:3].astype(np.float32) + mask_res = _resize_mask(mask, height, width) + mask_res = np.clip(mask_res, 0.0, 1.0) + threshold = float(cfg.ortho.water_mask_threshold) + water_mask = mask_res >= threshold + if not np.any(water_mask): + ds = None + return + + mode = str(getattr(cfg.ortho, "water_color_mode", "median") or "median").lower() + if mode == "fixed": + water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32) + else: + if water_mask.any(): + sample = rgb[:, water_mask] + if sample.size >= 3: + water_color = np.median(sample, axis=1) + else: + water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32) + else: + water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32) + + blend = float(cfg.ortho.water_blend) + weights = mask_res * blend + for c in range(3): + rgb[c] = rgb[c] * (1.0 - weights) + water_color[c] * weights + rgb = np.clip(rgb, 0.0, 255.0).astype(np.uint8) + + mem_driver = gdal.GetDriverByName("MEM") + out_ds = mem_driver.Create("", width, height, 3, gdal.GDT_Byte) + out_ds.SetGeoTransform(ds.GetGeoTransform()) + out_ds.SetProjection(ds.GetProjection()) + for c in range(3): + out_ds.GetRasterBand(c + 1).WriteArray(rgb[c]) + + trans_opts = gdal.TranslateOptions( + outputType=gdal.GDT_Byte, + format="JPEG", + creationOptions=[f"QUALITY={cfg.ortho.jpeg_quality}", "WORLDFILE=YES"], + ) + try: + gdal.Translate(out_path, out_ds, options=trans_opts) + except RuntimeError as exc: + print(f"[river_erosion] Ortho write failed for {tile_id}: {exc}") + out_ds = None + ds = None + _river_debug(f"Ortho water mask written: {out_path}") def _write_manifest( @@ -231,130 +508,38 @@ def _write_manifest( writer.writerow(out) -def _process_profile( - name: str, - profile: RiverErosionProfileConfig, - tiles: List[Dict[str, object]], - layer: ogr.Layer, - to_tile: osr.CoordinateTransformation, - to_layer: osr.CoordinateTransformation, - tile_srs: osr.SpatialReference, - output_dir: str, - manifest_path: str, -) -> int: - ensure_dir(output_dir) - rows_out: List[Dict[str, str]] = [] - global_min = math.inf - global_max = -math.inf - skipped = 0 +def _resolve_hydrorivers_path(source_path: str) -> Optional[str]: + if source_path and os.path.exists(source_path): + return source_path + fallback = os.path.join( + "raw", + "hydrorivers", + "hydrorivers_eu_shp", + "HydroRIVERS_v10_eu_shp", + "HydroRIVERS_v10_eu.shp", + ) + if os.path.exists(fallback): + return fallback + if source_path and os.path.isdir(source_path): + candidates = glob.glob(os.path.join(source_path, "**", "*.shp"), recursive=True) + if candidates: + candidates.sort() + for cand in candidates: + if "hydrorivers" in cand.lower(): + return cand + return candidates[0] + return None - for tile in tiles: - tile_id = tile["tile_id"] - png_path = tile["png_path"] - if not os.path.exists(png_path): - print(f"[river_erosion:{name}] Missing PNG for {tile_id}: {png_path}") - skipped += 1 - continue - bounds = tile["bounds"] - out_res = tile["out_res"] - scale_min = tile["scale_min"] - scale_max = tile["scale_max"] - - ds = open_dataset(png_path, f"[river_erosion:{name}] Could not open {png_path}.") - band = ds.GetRasterBand(1) - raw = band.ReadAsArray() - ds = None - if raw is None: - print(f"[river_erosion:{name}] Empty raster for {tile_id}, skipping.") - skipped += 1 - continue - - if raw.shape[0] != out_res or raw.shape[1] != out_res: - print( - f"[river_erosion:{name}] Warning: {tile_id} PNG is {raw.shape[1]}x{raw.shape[0]} " - f"but manifest says {out_res}; using PNG resolution." - ) - out_res = raw.shape[0] - - depth, gt = _build_depth_raster( - layer, - bounds, - out_res, - profile, - to_tile, - to_layer, - tile_srs, - ) - - data = raw.astype(np.float32) - valid = data > 0 - height = np.full(data.shape, np.nan, dtype=np.float32) - if scale_max <= scale_min: - scale_max = scale_min + 1e-3 - height[valid] = scale_min + (data[valid] / 65535.0) * (scale_max - scale_min) - height[valid] = height[valid] - depth[valid] - - if not np.any(valid): - print(f"[river_erosion:{name}] No valid samples for {tile_id}, copying source.") - out_u16 = raw.astype(np.uint16, copy=False) - tile_min = scale_min - tile_max = scale_max - else: - tile_min = float(np.nanmin(height)) - tile_max = float(np.nanmax(height)) - if tile_max <= tile_min: - tile_max = tile_min + 1e-3 - scaled = (height[valid] - tile_min) / (tile_max - tile_min) - out_u16 = np.zeros(data.shape, dtype=np.uint16) - out_u16[valid] = np.clip(scaled * 65535.0, 0.0, 65535.0).astype(np.uint16) - - global_min = min(global_min, tile_min) - global_max = max(global_max, tile_max) - - out_path = os.path.join(output_dir, f"{tile_id}.png") - mem_driver = gdal.GetDriverByName("MEM") - out_ds = mem_driver.Create("", out_u16.shape[1], out_u16.shape[0], 1, gdal.GDT_UInt16) - out_ds.SetGeoTransform(gt) - out_ds.SetProjection(tile_srs.ExportToWkt()) - out_ds.GetRasterBand(1).WriteArray(out_u16) - trans_opts = gdal.TranslateOptions( - outputType=gdal.GDT_UInt16, - format="PNG", - creationOptions=["WORLDFILE=YES"], - ) - try: - gdal.Translate(out_path, out_ds, options=trans_opts) - except RuntimeError as exc: - print(f"[river_erosion:{name}] Translate failed for {tile_id}: {exc}") - skipped += 1 - continue - out_ds = None - - rows_out.append( - { - "tile_id": tile_id, - "xmin": f"{bounds[0]:.6f}", - "ymin": f"{bounds[1]:.6f}", - "xmax": f"{bounds[2]:.6f}", - "ymax": f"{bounds[3]:.6f}", - "out_res": str(out_res), - "tile_key": tile["tile_key"], - "tile_min": f"{tile_min:.6f}", - "tile_max": f"{tile_max:.6f}", - } - ) - - print(f"[river_erosion:{name}] Wrote {out_path}") - - if not rows_out: - print(f"[river_erosion:{name}] No tiles processed; nothing written.") - return 1 - - _write_manifest(manifest_path, rows_out, global_min, global_max) - print(f"[river_erosion:{name}] Manifest: {manifest_path}") - print(f"[river_erosion:{name}] Summary: wrote {len(rows_out)} tiles; skipped {skipped}.") - return 1 if skipped else 0 +def _laz_path_for_tile(tile_id: str, source_dir: str) -> Optional[str]: + parts = tile_id.split("_") + if len(parts) < 6: + return None + x_idx = parts[2] + y_idx = parts[3] + laz_name = f"bdom20rgbi_32_{x_idx}_{y_idx}_2_rp.laz" + laz_path = os.path.join(source_dir, laz_name) + return laz_path if os.path.exists(laz_path) else None def erode_rivers(cfg: Config) -> int: @@ -362,25 +547,36 @@ def erode_rivers(cfg: Config) -> int: if not os.path.exists(cfg.export.manifest_path): print(f"[river_erosion] Missing manifest: {cfg.export.manifest_path}. Run heightmap export first.") return 1 - if not os.path.exists(re_cfg.source_path): + + source_path = _resolve_hydrorivers_path(re_cfg.source_path) + if not source_path: print(f"[river_erosion] Missing HydroRIVERS source: {re_cfg.source_path}") return 1 + if not re_cfg.enabled: print("[river_erosion] Warning: river_erosion.enabled=false; proceeding because --erode-rivers was requested.") - ds = ogr.Open(re_cfg.source_path) - if ds is None: - print(f"[river_erosion] Could not open {re_cfg.source_path}") + try: + ds = ogr.Open(source_path) + except RuntimeError as exc: + print(f"[river_erosion] Could not open {source_path} ({exc})") return 1 + if ds is None: + print(f"[river_erosion] Could not open {source_path}") + return 1 + layer = ds.GetLayerByName(re_cfg.source_layer) if re_cfg.source_layer else None if layer is None: layer = ds.GetLayer(0) if layer is None: - print(f"[river_erosion] No layers found in {re_cfg.source_path}") + print(f"[river_erosion] Could not find layer in {source_path}") return 1 + layer_sources = [(layer, os.path.basename(source_path))] + layer_srs = layer.GetSpatialRef() if layer_srs is None: + print("[river_erosion] Warning: HydroRIVERS layer CRS missing; assuming EPSG:4326.") layer_srs = osr.SpatialReference() layer_srs.ImportFromEPSG(4326) tile_srs = osr.SpatialReference() @@ -406,89 +602,262 @@ def erode_rivers(cfg: Config) -> int: print(f"[river_erosion] No tiles in manifest: {cfg.export.manifest_path}") return 1 - has_tile_minmax = "tile_min" in (reader.fieldnames or []) and "tile_max" in (reader.fieldnames or []) - need_tile_minmax = not has_tile_minmax or any( - not row.get("tile_min") or not row.get("tile_max") for row in rows - ) - - vrt_ds = None - if need_tile_minmax: - if not os.path.isdir(cfg.raw.dgm1_dir): - print(f"[river_erosion] Missing DGM1 directory: {cfg.raw.dgm1_dir}") - return 1 - tif_paths = sorted(glob.glob(os.path.join(cfg.raw.dgm1_dir, "*.tif"))) - if not tif_paths: - print(f"[river_erosion] No DGM1 TIFFs found in {cfg.raw.dgm1_dir}") - return 1 - build_vrt(cfg.work.heightmap_vrt, tif_paths, force=False) - vrt_ds = open_dataset(cfg.work.heightmap_vrt, f"[river_erosion] Could not open {cfg.work.heightmap_vrt}.") - - tiles: List[Dict[str, object]] = [] + key_lookup: Dict[str, Dict[str, str]] = {} for row in rows: tile_id = row["tile_id"].strip() - bounds = ( + tile_key = row.get("tile_key", "").strip() + if tile_key: + key_lookup[tile_key] = row + + lidar_cfg = getattr(re_cfg, "lidar", None) + source_dir = getattr(lidar_cfg, "source_dir", "raw/bdom20rgbi") + water_class = getattr(lidar_cfg, "classification_water", 9) + depth_profile = re_cfg.vr + fill_radius = getattr(lidar_cfg, "fill_holes_radius", 3) + smooth_sigma = getattr(lidar_cfg, "bank_slope_sigma", depth_profile.smooth_sigma_m) + + output_dir = os.path.join(re_cfg.output_dir, "vr") + manifest_out = re_cfg.manifest_vr + ensure_dir(output_dir) + + rows_out: List[Dict[str, str]] = [] + global_min = math.inf + global_max = -math.inf + + for row in rows: + tile_id = row["tile_id"].strip() + tile_key = row.get("tile_key", "").strip() + center_key = _parse_tile_key(tile_key) + if center_key is None: + _river_debug(f"Tile {tile_id} missing/invalid tile_key; processing without neighbors.") + neighbors = [(0, 0, row)] + else: + neighbors = [] + for dy in (-1, 0, 1): + for dx in (-1, 0, 1): + neighbor_key = f"{center_key[0] + dx}_{center_key[1] + dy}" + neighbor_row = key_lookup.get(neighbor_key) + if neighbor_row is None: + continue + neighbors.append((dx, dy, neighbor_row)) + + center_bounds = ( float(row["xmin"]), float(row["ymin"]), float(row["xmax"]), float(row["ymax"]), ) out_res = int(row["out_res"]) - tile_key = row.get("tile_key", "").strip() - scale_min = None - scale_max = None - if row.get("tile_min") and row.get("tile_max"): - scale_min = float(row["tile_min"]) - scale_max = float(row["tile_max"]) - if scale_min is None or scale_max is None: - if vrt_ds is None: - print(f"[river_erosion] Missing tile_min/tile_max for {tile_id} and no VRT available.") - return 1 - scale_min, scale_max = _compute_tile_minmax_from_vrt( - vrt_ds, - bounds, - out_res, - cfg.heightmap.resample, - ) - if scale_min is None or scale_max is None: - print(f"[river_erosion] Could not compute tile min/max for {tile_id}, using global range.") - scale_min = float(row["global_min"]) - scale_max = float(row["global_max"]) + step = out_res - 1 + tile_w = center_bounds[2] - center_bounds[0] + tile_h = center_bounds[3] - center_bounds[1] - tiles.append( + if not neighbors: + neighbors = [(0, 0, row)] + + min_dx = min(item[0] for item in neighbors) + max_dx = max(item[0] for item in neighbors) + min_dy = min(item[1] for item in neighbors) + max_dy = max(item[1] for item in neighbors) + tiles_x = max_dx - min_dx + 1 + tiles_y = max_dy - min_dy + 1 + + mosaic_w = tiles_x * step + 1 + mosaic_h = tiles_y * step + 1 + + mosaic_xmin = center_bounds[0] + min_dx * tile_w + mosaic_ymin = center_bounds[1] + min_dy * tile_h + mosaic_xmax = mosaic_xmin + tile_w * tiles_x + mosaic_ymax = mosaic_ymin + tile_h * tiles_y + mosaic_bounds = (mosaic_xmin, mosaic_ymin, mosaic_xmax, mosaic_ymax) + + height_mosaic = np.full((mosaic_h, mosaic_w), np.nan, dtype=np.float32) + + center_idx = None + for idx, (dx, dy, _) in enumerate(neighbors): + if dx == 0 and dy == 0: + center_idx = idx + break + + ordered_neighbors = neighbors[:] + if center_idx is not None: + center_item = ordered_neighbors.pop(center_idx) + ordered_neighbors.append(center_item) + + for dx, dy, nrow in ordered_neighbors: + tile_id_n = nrow["tile_id"].strip() + png_path = os.path.join(cfg.export.heightmap_dir, f"{tile_id_n}.png") + if not os.path.exists(png_path): + print(f"[river_erosion] Missing PNG for {tile_id_n}: {png_path}") + continue + ds_tile = open_dataset(png_path, f"[river_erosion] Could not open {png_path}.") + band = ds_tile.GetRasterBand(1) + raw = band.ReadAsArray().astype(np.float32) + ds_tile = None + if raw is None or raw.size == 0: + continue + + if raw.shape[0] != out_res or raw.shape[1] != out_res: + out_res = raw.shape[0] + step = out_res - 1 + + scale_min = float(nrow.get("tile_min") or nrow.get("global_min")) + scale_max = float(nrow.get("tile_max") or nrow.get("global_max")) + if scale_max <= scale_min: + scale_max = scale_min + 1e-3 + + height = scale_min + (raw / 65535.0) * (scale_max - scale_min) + + x0 = (dx - min_dx) * step + y0 = (max_dy - dy) * step + height_mosaic[y0 : y0 + out_res, x0 : x0 + out_res] = height + + height_mosaic = _fill_nan_with_nearest(height_mosaic) + + laz_paths = [] + for _, _, nrow in neighbors: + laz_path = _laz_path_for_tile(nrow["tile_id"].strip(), source_dir) + if laz_path: + laz_paths.append(laz_path) + laz_paths = sorted(set(laz_paths)) + + mask = _build_water_mask_mosaic(laz_paths, mosaic_bounds, mosaic_w, mosaic_h, water_class) + if mask is None: + _river_debug(f"Tile {tile_id}: no LiDAR water mask; copying source.") + mask = np.zeros((mosaic_h, mosaic_w), dtype=np.float32) + else: + mask = _refine_mask(mask, fill_radius, smooth_sigma) + + depth_line = _build_depth_lines( + layer_sources, + mosaic_bounds, + mosaic_w, + mosaic_h, + depth_profile, + to_tile, + to_layer, + tile_srs, + ) + + line_mask = depth_line > 0.0 + mask_bin = mask > 0.1 + depth_map = np.zeros_like(depth_line, dtype=np.float32) + indices = None + depth_nearest = None + + if np.any(mask_bin): + labels, num = ndimage.label(mask_bin) + if num > 0: + if np.any(line_mask): + _, indices = ndimage.distance_transform_edt(~line_mask, return_indices=True) + depth_nearest = depth_line[tuple(indices)] + else: + depth_nearest = np.zeros_like(depth_line, dtype=np.float32) + + line_components = set(np.unique(labels[line_mask])) + line_components.discard(0) + if line_components: + line_components_arr = np.array(sorted(line_components)) + use_line = np.isin(labels, line_components_arr) + depth_map[use_line] = depth_nearest[use_line] * mask[use_line] + else: + use_line = np.zeros_like(mask_bin, dtype=bool) + + fallback_depth = float(getattr(depth_profile, "fallback_depth_m", 2.0)) + if fallback_depth > 0.0: + fallback_mask = mask_bin & ~use_line + depth_map[fallback_mask] = fallback_depth * mask[fallback_mask] + else: + use_line = np.zeros_like(mask_bin, dtype=bool) + else: + use_line = np.zeros_like(mask_bin, dtype=bool) + + flow_slope = float(getattr(depth_profile, "flow_slope_m_per_km", 0.0)) + dist_field = str(getattr(depth_profile, "flow_distance_field", "DIST_DN_KM") or "DIST_DN_KM") + if flow_slope > 0.0 and np.any(line_mask) and indices is not None: + dist_line = _build_flow_distance( + layer_sources, + mosaic_bounds, + mosaic_w, + mosaic_h, + dist_field, + to_tile, + to_layer, + tile_srs, + ) + dist_nearest = dist_line[tuple(indices)] + line_labels, line_num = ndimage.label(line_mask) + if line_num > 0: + flow_map = np.zeros_like(depth_map, dtype=np.float32) + for comp in range(1, line_num + 1): + comp_mask = line_labels == comp + if not np.any(comp_mask): + continue + comp_max = float(np.max(dist_nearest[comp_mask])) + flow_map[comp_mask] = (comp_max - dist_nearest[comp_mask]) * flow_slope + depth_map[use_line] += flow_map[use_line] * mask[use_line] + + height_mosaic = height_mosaic - depth_map + + x0_center = (0 - min_dx) * step + y0_center = (max_dy - 0) * step + cropped = height_mosaic[y0_center : y0_center + out_res, x0_center : x0_center + out_res] + + if cfg.ortho.apply_water_mask: + mask_center = mask[y0_center : y0_center + out_res, x0_center : x0_center + out_res] + _apply_water_mask_to_ortho(tile_id, mask_center, cfg) + + tile_min = float(np.min(cropped)) + tile_max = float(np.max(cropped)) + if tile_max <= tile_min: + tile_max = tile_min + 1e-3 + + scaled = (cropped - tile_min) / (tile_max - tile_min) + out_u16 = np.clip(scaled * 65535.0, 0.0, 65535.0).astype(np.uint16) + + global_min = min(global_min, tile_min) + global_max = max(global_max, tile_max) + + out_path = os.path.join(output_dir, f"{tile_id}.png") + mem_driver = gdal.GetDriverByName("MEM") + out_ds = mem_driver.Create("", out_u16.shape[1], out_u16.shape[0], 1, gdal.GDT_UInt16) + out_ds.SetGeoTransform(_tile_geotransform(center_bounds, out_res, out_res)) + out_ds.SetProjection(tile_srs.ExportToWkt()) + out_ds.GetRasterBand(1).WriteArray(out_u16) + trans_opts = gdal.TranslateOptions( + outputType=gdal.GDT_UInt16, + format="PNG", + creationOptions=["WORLDFILE=YES"], + ) + try: + gdal.Translate(out_path, out_ds, options=trans_opts) + except RuntimeError as exc: + print(f"[river_erosion] Translate failed for {tile_id}: {exc}") + continue + out_ds = None + + rows_out.append( { "tile_id": tile_id, - "bounds": bounds, - "out_res": out_res, + "xmin": f"{center_bounds[0]:.6f}", + "ymin": f"{center_bounds[1]:.6f}", + "xmax": f"{center_bounds[2]:.6f}", + "ymax": f"{center_bounds[3]:.6f}", + "out_res": str(out_res), "tile_key": tile_key, - "scale_min": float(scale_min), - "scale_max": float(scale_max), - "png_path": os.path.join(cfg.export.heightmap_dir, f"{tile_id}.png"), + "tile_min": f"{tile_min:.6f}", + "tile_max": f"{tile_max:.6f}", } ) - output_vr = os.path.join(re_cfg.output_dir, "vr") - output_server = os.path.join(re_cfg.output_dir, "server") + print(f"[river_erosion] Wrote {out_path}") - vr_exit = _process_profile( - "vr", - re_cfg.vr, - tiles, - layer, - to_tile, - to_layer, - tile_srs, - output_vr, - re_cfg.manifest_vr, - ) - server_exit = _process_profile( - "server", - re_cfg.server, - tiles, - layer, - to_tile, - to_layer, - tile_srs, - output_server, - re_cfg.manifest_server, - ) - return max(vr_exit, server_exit) + if not rows_out: + print("[river_erosion] No tiles processed; nothing written.") + return 1 + + _write_manifest(manifest_out, rows_out, global_min, global_max) + print(f"[river_erosion] Manifest: {manifest_out}") + print(f"[river_erosion] Summary: wrote {len(rows_out)} tiles.") + + return 0 diff --git a/geodata_pipeline/river_erosion_lidar.py b/geodata_pipeline/river_erosion_lidar.py new file mode 100644 index 0000000..660dfeb --- /dev/null +++ b/geodata_pipeline/river_erosion_lidar.py @@ -0,0 +1,384 @@ +from __future__ import annotations + +import csv +import glob +import json +import math +import os +from typing import Dict, List, Optional, Tuple + +import numpy as np +import pdal +from osgeo import gdal, osr +from scipy import ndimage + +from .config import Config, RiverErosionProfileConfig +from .gdal_utils import build_vrt, ensure_dir, ensure_parent, open_dataset + +gdal.UseExceptions() + + +def _river_debug_enabled() -> bool: + flag = os.getenv("RIVER_EROSION_DEBUG", "") + return flag.strip().lower() in {"1", "true", "yes", "y"} + + +def _river_debug(message: str) -> None: + if _river_debug_enabled(): + print(f"[river_erosion:debug] {message}") + + +def _tile_geotransform(bounds: Tuple[float, float, float, float], out_res: int) -> Tuple[float, float, float, float, float, float]: + xmin, ymin, xmax, ymax = bounds + xres = (xmax - xmin) / float(out_res) # Should be (1025-1) if pixel-is-point, but DGM1 usually uses cell-center or similar. + # Aligning with standard Unity/GDAL logic: + # For 1km tile and 1025 samples, spacing is 1000/1024 approx 0.976m. + # BUT existing pipeline uses simple division logic in previous scripts. + # To match DGM1 exactly, we should ideally read the GT from the source DGM1 if possible. + # For now, sticking to the bounding box logic. + yres = (ymax - ymin) / float(out_res) + return xmin, xres, 0.0, ymax, 0.0, -yres + + +def _build_water_mask( + laz_path: str, + bounds: Tuple[float, float, float, float], + out_res: int, + water_class: int = 9, +) -> Optional[np.ndarray]: + """Generates a binary water mask from LiDAR Class 9 points.""" + + if not os.path.exists(laz_path): + return None + + # PDAL Pipeline: Filter by classification + pipeline_json = [ + { + "type": "readers.las", + "filename": laz_path + }, + { + "type": "filters.range", + "limits": f"Classification[{water_class}:{water_class}]" + } + ] + + try: + pipeline = pdal.Pipeline(json.dumps(pipeline_json)) + count = pipeline.execute() + except RuntimeError as e: + print(f"[river_erosion] PDAL Error reading {laz_path}: {e}") + return None + + if count == 0: + return None + + arrays = pipeline.arrays[0] + xs = arrays['X'] + ys = arrays['Y'] + + # Rasterize + mask = np.zeros((out_res, out_res), dtype=np.float32) + + xmin, _, _, ymax, _, yres_neg = _tile_geotransform(bounds, out_res) + xres = (bounds[2] - bounds[0]) / out_res + yres = (bounds[3] - bounds[1]) / out_res + + # Pixel coords + col = ((xs - xmin) / xres).astype(int) + row = ((ymax - ys) / yres).astype(int) # yres is positive magnitude here + + # Clip + valid = (col >= 0) & (col < out_res) & (row >= 0) & (row < out_res) + col = col[valid] + row = row[valid] + + if len(col) == 0: + return None + + mask[row, col] = 1.0 + return mask + + +def _refine_mask( + mask: np.ndarray, + fill_radius: int = 3, + smooth_sigma: float = 2.0 +) -> np.ndarray: + """Closes gaps in the water mask and smooths edges.""" + + # Morphological Closing to fill small holes (LiDAR absorption gaps) + if fill_radius > 0: + struct = ndimage.generate_binary_structure(2, 1) # 4-connectivity + mask = ndimage.binary_closing(mask, structure=struct, iterations=fill_radius).astype(np.float32) + + # Optional: Remove small isolated islands? + # label_im, nb_labels = ndimage.label(mask) + # sizes = ndimage.sum(mask, label_im, range(nb_labels + 1)) + # mask_size = sizes < 100 # remove blobs smaller than 100px + # remove_pixel = mask_size[label_im] + # mask[remove_pixel] = 0 + + # Gaussian Blur for smooth banks + if smooth_sigma > 0: + mask = ndimage.gaussian_filter(mask, sigma=smooth_sigma) + + # Normalize mask to 0-1 range (blur can reduce peaks) + # Ideally, the center of the river should be 1.0 (full depth) + # Banks will fade from 1.0 to 0.0 + + return mask + + +def _write_manifest( + manifest_path: str, + rows: List[Dict[str, str]], + global_min: float, + global_max: float, +) -> None: + ensure_parent(manifest_path) + header = [ + "tile_id", + "xmin", + "ymin", + "xmax", + "ymax", + "global_min", + "global_max", + "out_res", + "tile_key", + "tile_min", + "tile_max", + ] + with open(manifest_path, "w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=header) + writer.writeheader() + for row in rows: + out = dict(row) + out["global_min"] = f"{global_min:.6f}" + out["global_max"] = f"{global_max:.6f}" + writer.writerow(out) + + +def _compute_tile_minmax_from_vrt( + vrt_ds: gdal.Dataset, + bounds: Tuple[float, float, float, float], + out_res: int, + resample: str, +) -> Tuple[float | None, float | None]: + opts = gdal.WarpOptions( + outputBounds=bounds, + width=out_res, + height=out_res, + resampleAlg=resample, + srcNodata=-9999, + dstNodata=-9999, + format="MEM", + ) + tmp = gdal.Warp("", vrt_ds, options=opts) + if tmp is None: + return None, None + band = tmp.GetRasterBand(1) + nodata = band.GetNoDataValue() + data = band.ReadAsArray() + tmp = None + if data is None or data.size == 0: + return None, None + valid = np.isfinite(data) + if nodata is not None and not math.isnan(nodata): + valid &= data != nodata + if not np.any(valid): + return None, None + return float(data[valid].min()), float(data[valid].max()) + + +def erode_rivers(cfg: Config) -> int: + # Use lidar config if available, otherwise fall back to defaults or error + # For now, we assume [river_erosion.lidar] exists or we default + + lidar_cfg = getattr(cfg.river_erosion, "lidar", None) + if not lidar_cfg or not lidar_cfg.get("enabled", False): + # Fallback to checking basic river_erosion enabled? + if not cfg.river_erosion.enabled: + print("[river_erosion] Disabled in config.") + return 0 + + # Hardcoded defaults for now if config missing (Migration Step) + source_dir = getattr(lidar_cfg, "source_dir", "raw/geo_rlp/bdom20rgbi") + water_class = getattr(lidar_cfg, "classification_water", 9) + # VR profile depth/smooth + # We will output to same structure: export_unity/height_png16_river/vr + + # We use the VR settings from the old config for depth/sigma if not in lidar block + vr_profile = cfg.river_erosion.vr + depth_m = getattr(lidar_cfg, "depth_m", 3.5) + smooth_sigma = getattr(lidar_cfg, "bank_slope_sigma", vr_profile.smooth_sigma_m) + fill_radius = getattr(lidar_cfg, "fill_holes_radius", 3) + + output_dir = os.path.join(cfg.river_erosion.output_dir, "vr") + manifest_out = cfg.river_erosion.manifest_vr + + if not os.path.exists(cfg.export.manifest_path): + print(f"[river_erosion] Missing manifest: {cfg.export.manifest_path}") + return 1 + + # Load Tiles + with open(cfg.export.manifest_path, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + rows = list(reader) + + if not rows: + print("[river_erosion] No tiles to process.") + return 1 + + # Open VRT for re-calculating min/max if needed + # (Actually we should re-calc min/max based on the ERODED data, so we don't strictly need VRT of raw data unless raw is missing) + # We load RAW pngs. + + ensure_dir(output_dir) + + rows_out: List[Dict[str, str]] = [] + global_min = math.inf + global_max = -math.inf + processed_count = 0 + eroded_count = 0 + + print(f"[river_erosion] Starting LiDAR erosion. Source: {source_dir}, Depth: {depth_m}m") + + for row in rows: + tile_id = row["tile_id"] + # Determine LAZ file path + # Pattern: bdom20rgbi_32_324_5506_2_rp.laz where tile_id is dgm01_32_324_5506_1_rp + # Need to map dgm01 -> bdom20rgbi and _1_rp -> _2_rp + # Tile ID format: dgm01_32_X_Y_1_rp + parts = tile_id.split('_') + # dgm01, 32, X, Y, 1, rp + if len(parts) >= 6: + x_idx = parts[2] + y_idx = parts[3] + laz_name = f"bdom20rgbi_32_{x_idx}_{y_idx}_2_rp.laz" + laz_path = os.path.join(source_dir, laz_name) + else: + print(f"[river_erosion] Warning: Could not parse tile_id {tile_id} to find LAZ.") + laz_path = "" + + # Load PNG + png_path = os.path.join(cfg.export.heightmap_dir, f"{tile_id}.png") + if not os.path.exists(png_path): + print(f"[river_erosion] Missing PNG: {png_path}") + continue + + ds = open_dataset(png_path, f"Could not open {png_path}") + band = ds.GetRasterBand(1) + raw_u16 = band.ReadAsArray() + gt = ds.GetGeoTransform() + proj = ds.GetProjection() + ds = None + + out_res = raw_u16.shape[0] + bounds = ( + float(row["xmin"]), + float(row["ymin"]), + float(row["xmax"]), + float(row["ymax"]), + ) + + # Recover float height + # Height = min + (val / 65535) * (max - min) + # Note: manifest has tile_min/tile_max. If not, global. + # Check if row has tile_min + if "tile_min" in row and row["tile_min"]: + src_min = float(row["tile_min"]) + src_max = float(row["tile_max"]) + else: + src_min = float(row["global_min"]) + src_max = float(row["global_max"]) + + # Convert to Float32 + data_float = raw_u16.astype(np.float32) + valid_pixels = data_float > 0 # Assuming 0 is nodata/void in u16? Usually 0 is valid height. + # Heightmap export usually scales 0..1. + # Ideally we treat 0 as valid unless we have a separate mask. + + height_range = src_max - src_min + if height_range < 1e-5: height_range = 1.0 # Prevent div by zero + + height_map = src_min + (data_float / 65535.0) * height_range + + # Generate Water Mask + mask = _build_water_mask(laz_path, bounds, out_res, water_class) + + if mask is not None: + # Refine + mask = _refine_mask(mask, fill_radius, smooth_sigma) + + # Erode + # depth_map = mask * depth_m + # height_map -= depth_map + + # Improved Blending: + # Target Height = (Current Height - Depth) + # But ensure we don't dig pits in already low areas? + # For now, simple subtraction is robust. + + erosion_amount = mask * depth_m + height_map = height_map - erosion_amount + eroded_count += 1 + if _river_debug_enabled(): + print(f"[river_erosion] Tile {tile_id}: Eroded {np.count_nonzero(mask > 0.1)} pixels.") + + # Recalculate Statistics + # We must re-normalize to U16. + # New Min/Max + new_min = float(np.min(height_map)) + new_max = float(np.max(height_map)) + + # Ensure we don't shrink range to 0 + if new_max <= new_min: new_max = new_min + 1.0 + + # Normalize + new_range = new_max - new_min + out_float = (height_map - new_min) / new_range + out_u16_final = np.clip(out_float * 65535.0, 0, 65535).astype(np.uint16) + + global_min = min(global_min, new_min) + global_max = max(global_max, new_max) + + # Save + out_path = os.path.join(output_dir, f"{tile_id}.png") + driver = gdal.GetDriverByName("PNG") # Or GTiff? Unity prefers PNG for now. + # Using MEM driver then CreateCopy or Translate usually better for control + + # We reuse the logic from previous script to ensure format consistency + mem_driver = gdal.GetDriverByName("MEM") + out_ds = mem_driver.Create("", out_res, out_res, 1, gdal.GDT_UInt16) + out_ds.SetGeoTransform(gt) + out_ds.SetProjection(proj) + out_ds.GetRasterBand(1).WriteArray(out_u16_final) + + gdal.Translate(out_path, out_ds, format="PNG", creationOptions=["WORLDFILE=YES"]) + out_ds = None + + rows_out.append({ + "tile_id": tile_id, + "xmin": row["xmin"], + "ymin": row["ymin"], + "xmax": row["xmax"], + "ymax": row["ymax"], + "out_res": row["out_res"], + "tile_key": row.get("tile_key", ""), + "tile_min": f"{new_min:.6f}", + "tile_max": f"{new_max:.6f}", + }) + + processed_count += 1 + + if not rows_out: + return 1 + + _write_manifest(manifest_out, rows_out, global_min, global_max) + + print(f"[river_erosion] Processed {processed_count} tiles. Eroded {eroded_count} (containing water points).") + print(f"[river_erosion] Manifest written to {manifest_out}") + + return 0