Files
GeoData/geodata_pipeline/river_erosion.py

1139 lines
40 KiB
Python

from __future__ import annotations
import csv
import glob
import json
import math
import os
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, LakeConfig, RiverErosionProfileConfig
from .gdal_utils import 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 _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(parts[0]), int(parts[1])
except ValueError:
return None
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
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
def _pixel_size_m(gt: Tuple[float, ...]) -> float:
return (abs(gt[1]) + abs(gt[5])) / 2.0
def _bounds_to_layer(
bounds: Tuple[float, float, float, float],
buffer_m: float,
to_layer: osr.CoordinateTransformation,
) -> Tuple[float, float, float, float]:
xmin, ymin, xmax, ymax = bounds
xmin -= buffer_m
ymin -= buffer_m
xmax += buffer_m
ymax += buffer_m
corners = [(xmin, ymin), (xmin, ymax), (xmax, ymin), (xmax, ymax)]
xs: list[float] = []
ys: list[float] = []
for x, y in corners:
lon, lat, _ = to_layer.TransformPoint(x, y)
xs.append(lon)
ys.append(lat)
return min(xs), min(ys), max(xs), max(ys)
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],
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
mask = np.zeros((height, width), dtype=np.float32)
mask[row[valid], col[valid]] = 1.0
return mask
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_order_lines(
layer_sources: List[Tuple[ogr.Layer, str]],
bounds: Tuple[float, float, float, float],
width: int,
height: int,
profile: RiverErosionProfileConfig,
to_tile: osr.CoordinateTransformation,
to_layer: osr.CoordinateTransformation,
tile_srs: osr.SpatialReference,
) -> np.ndarray:
gt = _tile_geotransform(bounds, width, height)
driver = gdal.GetDriverByName("MEM")
order_ds = driver.Create("", width, height, 1, gdal.GDT_Float32)
order_ds.SetGeoTransform(gt)
order_ds.SetProjection(tile_srs.ExportToWkt())
band = order_ds.GetRasterBand(1)
band.Fill(0.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, 250.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:
order = _normalize_order(feat.GetField(profile.order_field), profile)
if order 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(order))
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(
order_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(
order_ds,
[1],
mem_layer,
options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"],
)
mem_ds = None
return band.ReadAsArray().astype(np.float32)
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
trans_opts = gdal.TranslateOptions(
outputType=gdal.GDT_Byte,
format="JPEG",
creationOptions=[f"QUALITY={cfg.ortho.jpeg_quality}", "WORLDFILE=YES"],
)
if not np.any(water_mask):
try:
gdal.Translate(out_path, ds, options=trans_opts)
except RuntimeError as exc:
print(f"[river_erosion] Ortho write failed for {tile_id}: {exc}")
ds = None
print(f"[river_erosion] Wrote ortho {out_path} (copied; empty mask)")
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])
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
print(f"[river_erosion] Wrote ortho {out_path}")
def _write_manifest(
manifest_path: str,
rows: Iterable[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 _write_mask_png(
out_path: str,
mask: np.ndarray,
bounds: Tuple[float, float, float, float],
tile_srs: osr.SpatialReference,
*,
bands: int = 1,
) -> None:
ensure_parent(out_path)
height, width = mask.shape[-2], mask.shape[-1]
gt = _tile_geotransform(bounds, width, height)
driver = gdal.GetDriverByName("MEM")
ds = driver.Create("", width, height, bands, gdal.GDT_Byte)
ds.SetGeoTransform(gt)
ds.SetProjection(tile_srs.ExportToWkt())
if bands == 1:
ds.GetRasterBand(1).WriteArray(mask.astype(np.uint8))
else:
for idx in range(bands):
ds.GetRasterBand(idx + 1).WriteArray(mask[idx].astype(np.uint8))
options = gdal.TranslateOptions(format="PNG", creationOptions=["WORLDFILE=YES"])
try:
gdal.Translate(out_path, ds, options=options)
except RuntimeError as exc:
print(f"[river_erosion] Mask write failed for {out_path}: {exc}")
ds = None
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
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 _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):
print(f"[river_erosion] Missing manifest: {cfg.export.manifest_path}. Run heightmap export first.")
return 1
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.")
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] 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()
tile_srs.ImportFromEPSG(25832)
layer_srs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
tile_srs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
to_tile = osr.CoordinateTransformation(layer_srs, tile_srs)
to_layer = osr.CoordinateTransformation(tile_srs, layer_srs)
with open(cfg.export.manifest_path, newline="", encoding="utf-8") as f:
reader = csv.DictReader(f)
if not reader.fieldnames:
print(f"[river_erosion] Empty manifest: {cfg.export.manifest_path}")
return 1
required = {"tile_id", "xmin", "ymin", "xmax", "ymax", "global_min", "global_max", "out_res"}
missing = required - set(reader.fieldnames)
if missing:
print(f"[river_erosion] Manifest missing columns: {', '.join(sorted(missing))}")
return 1
rows = list(reader)
if not rows:
print(f"[river_erosion] No tiles in manifest: {cfg.export.manifest_path}")
return 1
key_lookup: Dict[str, Dict[str, str]] = {}
for row in rows:
tile_id = row["tile_id"].strip()
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 = re_cfg.output_dir
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"])
step = out_res - 1
tile_w = center_bounds[2] - center_bounds[0]
tile_h = center_bounds[3] - center_bounds[1]
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
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)
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)
order_line = _build_order_lines(
layer_sources,
mosaic_bounds,
mosaic_w,
mosaic_h,
depth_profile,
to_tile,
to_layer,
tile_srs,
)
line_mask = order_line > 0.0
mask_bin = mask > 0.1
depth_map = np.zeros_like(order_line, dtype=np.float32)
labels = None
num = 0
indices = None
order_nearest = None
dist_nearest = None
if np.any(line_mask):
_, indices = ndimage.distance_transform_edt(~line_mask, return_indices=True)
order_nearest = order_line[tuple(indices)]
if np.any(mask_bin):
labels, num = ndimage.label(mask_bin)
fallback_depth = float(getattr(depth_profile, "fallback_depth_m", 2.0))
if num > 0:
for comp in range(1, num + 1):
comp_mask = labels == comp
if order_nearest is not None:
comp_orders = order_nearest[comp_mask]
comp_orders = comp_orders[comp_orders > 0.0]
if comp_orders.size > 0:
order_val = int(round(float(np.median(comp_orders))))
depth_val = _order_depth(order_val, depth_profile)
else:
depth_val = None
else:
depth_val = None
if depth_val is None or depth_val <= 0.0:
depth_val = fallback_depth
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))
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 and labels is not None and num > 0:
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)]
for comp in range(1, num + 1):
comp_mask = labels == comp
if not np.any(comp_mask):
continue
dist_vals = dist_nearest[comp_mask]
heights = height_mosaic[comp_mask]
if dist_vals.size == 0 or heights.size == 0:
continue
order = np.argsort(-dist_vals)
dist_sorted = dist_vals[order]
heights_sorted = heights[order]
adjusted = heights_sorted + flow_slope * dist_sorted
enforced = np.minimum.accumulate(adjusted)
new_heights_sorted = enforced - flow_slope * dist_sorted
new_heights = heights.copy()
new_heights[order] = new_heights_sorted
height_mosaic[comp_mask] = np.minimum(heights, new_heights)
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)
mask_center = mask[y0_center : y0_center + out_res, x0_center : x0_center + out_res]
threshold = float(cfg.ortho.water_mask_threshold)
water_bin = mask_center >= threshold
mask_dir = os.path.join(cfg.work.work_dir, "river_masks")
ensure_dir(mask_dir)
mask_path = os.path.join(mask_dir, f"{tile_id}_mask.png")
_write_mask_png(mask_path, (water_bin.astype(np.uint8) * 255), center_bounds, tile_srs)
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,
"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,
"tile_min": f"{tile_min:.6f}",
"tile_max": f"{tile_max:.6f}",
}
)
print(f"[river_erosion] Wrote {out_path}")
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