From 6c434882b85287d57f93d3cb9e94fade9f79cb33 Mon Sep 17 00:00:00 2001 From: "s0wlz (Matthias Puchstein)" Date: Thu, 5 Feb 2026 00:31:22 +0100 Subject: [PATCH] fix(river): Fix ValueError in lake centroid iteration --- geodata_pipeline/river_erosion.py | 245 +++++++++++++++++++++++++++++- 1 file changed, 238 insertions(+), 7 deletions(-) diff --git a/geodata_pipeline/river_erosion.py b/geodata_pipeline/river_erosion.py index 717d572..a861e82 100644 --- a/geodata_pipeline/river_erosion.py +++ b/geodata_pipeline/river_erosion.py @@ -5,14 +5,14 @@ import glob import json import math import os -from typing import Dict, Iterable, List, Optional, Tuple +from typing import Any, 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 .config import Config, LakeConfig, RiverErosionProfileConfig from .gdal_utils import ensure_dir, ensure_parent, open_dataset gdal.UseExceptions() @@ -209,7 +209,7 @@ def _build_order_lines( _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) + buffer_m = max(profile.smooth_sigma_m * 3.0, 250.0) minx, miny, maxx, maxy = _bounds_to_layer(bounds, buffer_m, to_layer) buffered_geoms: list[ogr.Geometry] = [] @@ -561,6 +561,161 @@ def _laz_path_for_tile(tile_id: str, source_dir: str) -> Optional[str]: return laz_path if os.path.exists(laz_path) else None +def _find_lake_depths( + manual_geoms: List[ogr.Geometry], + manual_srs: osr.SpatialReference, + lake_cfg: LakeConfig, + hydrolakes_path: str, +) -> List[Tuple[float, Dict[str, Any]]]: + results = [] + default_depth = lake_cfg.default_depth_m + tolerance = lake_cfg.match_tolerance_m + + if not manual_geoms: + return [] + + if not os.path.exists(hydrolakes_path): + _river_debug(f"HydroLAKES path not found: {hydrolakes_path}") + return [(default_depth, {"error": "HydroLAKES missing"}) for _ in manual_geoms] + + ds = ogr.Open(hydrolakes_path) + if ds is None: + return [(default_depth, {"error": "Could not open HydroLAKES"}) for _ in manual_geoms] + + layer = ds.GetLayer(0) + layer_srs = layer.GetSpatialRef() + + # Transformation from Manual (Tile) to HydroLAKES + transform = None + if not manual_srs.IsSame(layer_srs): + transform = osr.CoordinateTransformation(manual_srs, layer_srs) + + # Reverse Transform (HydroLAKES -> Manual) for distance check + rev_transform = None + if not manual_srs.IsSame(layer_srs): + rev_transform = osr.CoordinateTransformation(layer_srs, manual_srs) + + for geom in manual_geoms: + # Working geometry in HydroLAKES CRS for spatial filter + search_geom = geom.Clone() + if transform: + search_geom.Transform(transform) + + centroid = search_geom.Centroid() + pt = centroid.GetPoint() # (lon, lat) + + # Simple spatial filter around centroid (approx 0.1 deg ~ 10km) + search_buffer = 0.1 + layer.SetSpatialFilterRect(pt[0] - search_buffer, pt[1] - search_buffer, pt[0] + search_buffer, pt[1] + search_buffer) + + best_dist = float("inf") + best_depth = default_depth + best_id = -1 + matched = False + + layer.ResetReading() + for feat in layer: + feat_geom = feat.GetGeometryRef() + if feat_geom is None: + continue + + # Clone and transform feature to Manual CRS (Meters) for accurate distance + feat_geom_m = feat_geom.Clone() + if rev_transform: + feat_geom_m.Transform(rev_transform) + + dist = feat_geom_m.Distance(geom) + + if dist < best_dist: + best_dist = dist + best_id = feat.GetFID() + # Get depth + d = feat.GetField("Depth_avg") + if d is not None and float(d) > 0: + best_depth = float(d) + else: + best_depth = default_depth + + if best_dist <= tolerance: + matched = True + results.append((best_depth, {"id": best_id, "distance": best_dist, "matched": True})) + else: + # If closest is too far, use default + results.append((default_depth, {"id": best_id, "distance": best_dist, "matched": False, "note": "Outside tolerance"})) + + return results + + +def _load_manual_mask( + tile_id: str, + search_dir: str = "raw/water_masks" +) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: + if not os.path.exists(search_dir): + return None, None, None + + # Try exact match first + # Expected pattern: {tile_id}_mask_viz.png or similar + # User file: dgm01_32_324_5507_rp_mask_viz.png + # Tile ID: dgm01_32_324_5507_1_rp + + # Heuristic: try to construct the filename from tile_id parts + # If tile_id has "_1_rp", try removing "_1" + + candidates = [ + f"{tile_id}.png", + f"{tile_id}_mask.png", + f"{tile_id}_mask_viz.png", + f"{tile_id}_viz.png" + ] + + # Try stripping "_1" + base_id = tile_id.replace("_1_rp", "_rp") + candidates.extend([ + f"{base_id}.png", + f"{base_id}_mask.png", + f"{base_id}_mask_viz.png", + f"{base_id}_viz.png" + ]) + + match_path = None + for cand in candidates: + path = os.path.join(search_dir, cand) + if os.path.exists(path): + match_path = path + break + + if not match_path: + return None, None, None + + ds = gdal.Open(match_path) + if ds is None: + return None, None, None + + # Expecting RGB + if ds.RasterCount < 3: + return None, None, None + + red = ds.GetRasterBand(1).ReadAsArray() + green = ds.GetRasterBand(2).ReadAsArray() + blue = ds.GetRasterBand(3).ReadAsArray() + + # Thresholds (assuming UInt8 0-255 or UInt16 0-65535) + # If UInt16, max is 65535. + # Check max value to guess? Or just > 0. + + # Semantic: + # Blue = Lake (Water) + # Green = Island (No Erosion) + # Red = Bridge (No Erosion + Fix) + + # Masks are boolean (1 where color is present) + mask_r = (red > 0) + mask_g = (green > 0) + mask_b = (blue > 0) + + return mask_b, mask_g, mask_r + + def erode_rivers(cfg: Config) -> int: re_cfg = cfg.river_erosion if not os.path.exists(cfg.export.manifest_path): @@ -684,10 +839,20 @@ def erode_rivers(cfg: Config) -> int: 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 + if center_key and cfg.tile_key.enabled: + grid_x = center_key[0] * cfg.tile_key.tile_size_x + grid_y = center_key[1] * cfg.tile_key.tile_size_y + + mosaic_xmin = grid_x + min_dx * cfg.tile_key.tile_size_x + mosaic_ymin = grid_y + min_dy * cfg.tile_key.tile_size_y + mosaic_xmax = mosaic_xmin + tiles_x * cfg.tile_key.tile_size_x + mosaic_ymax = mosaic_ymin + tiles_y * cfg.tile_key.tile_size_y + else: + 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) @@ -794,6 +959,72 @@ def erode_rivers(cfg: Config) -> int: if depth_val and depth_val > 0.0: depth_map[comp_mask] = depth_val * mask[comp_mask] + # --- Manual Mask Integration --- + # Load manual mask for the center tile (if available) + man_b, man_g, man_r = _load_manual_mask(tile_id) + + if man_b is not None or man_g is not None or man_r is not None: + # Determine slice for center tile in the mosaic + x0_c = (0 - min_dx) * step + y0_c = (max_dy - 0) * step + sl_y = slice(y0_c, y0_c + out_res) + sl_x = slice(x0_c, x0_c + out_res) + + # 1. Apply Erasers (Green=Island, Red=Bridge) + if man_g is not None: + depth_map[sl_y, sl_x][man_g] = 0.0 + mask[sl_y, sl_x][man_g] = 0.0 + + if man_r is not None: + depth_map[sl_y, sl_x][man_r] = 0.0 + mask[sl_y, sl_x][man_r] = 0.0 + + # 2. Apply Lakes (Blue) + if man_b is not None: + # Force water mask + mask[sl_y, sl_x][man_b] = 1.0 + + # Calculate depths + lbl_lakes, num_lakes = ndimage.label(man_b) + if num_lakes > 0: + gt_c = _tile_geotransform(center_bounds, out_res, out_res) + lake_geoms = [] + # Get centers of mass + # center_of_mass returns list of tuples if index is a sequence + centers = ndimage.center_of_mass(man_b, lbl_lakes, range(1, num_lakes + 1)) + + for (r, c) in centers: + # Pixel to World + # x = xmin + c * xres + # y = ymax + r * yres (yres is negative) + wx = gt_c[0] + c * gt_c[1] + wy = gt_c[3] + r * gt_c[5] + pt = ogr.Geometry(ogr.wkbPoint) + pt.AddPoint(wx, wy) + lake_geoms.append(pt) + + # Lookup depths + lake_depths_meta = _find_lake_depths( + lake_geoms, + tile_srs, + re_cfg.lakes, + re_cfg.lakes.source_path + ) + + # Apply depths + for idx, (depth, meta) in enumerate(lake_depths_meta): + comp_idx = idx + 1 + lake_mask = (lbl_lakes == comp_idx) + # Apply to mosaic depth map + depth_map[sl_y, sl_x][lake_mask] = depth + + match_status = "MATCH" if meta.get("matched") else "DEFAULT" + dist_info = f"{meta.get('distance', -1):.1f}m" + _river_debug( + f"Tile {tile_id} Manual Lake {idx+1}: " + f"Depth={depth:.2f}m ({match_status}, {dist_info})" + ) + height_mosaic = height_mosaic - depth_map flow_slope = float(getattr(depth_profile, "flow_slope_m_per_km", 0.0))