removed deprecated river erosion
This commit is contained in:
@@ -1,384 +0,0 @@
|
||||
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
|
||||
Reference in New Issue
Block a user