fix(river): Fix ValueError in lake centroid iteration
This commit is contained in:
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user