864 lines
30 KiB
Python
864 lines
30 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import glob
|
|
import json
|
|
import math
|
|
import os
|
|
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 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_depth_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")
|
|
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)
|
|
|
|
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)
|
|
|
|
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
|
|
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:
|
|
sigma_px = profile.smooth_sigma_m / max(_pixel_size_m(gt), 1e-6)
|
|
if sigma_px > 0.0:
|
|
depth = ndimage.gaussian_filter(depth, sigma=sigma_px, mode="nearest")
|
|
|
|
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
|
|
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 _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 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 = 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"])
|
|
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
|
|
|
|
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,
|
|
"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
|