Add river ortho recovery from saved masks

geodata_pipeline/river_erosion.py

@@ -425,8 +425,18 @@ def _apply_water_mask_to_ortho(tile_id: str, mask: np.ndarray, cfg: Config) -> N
     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()
@@ -455,11 +465,6 @@ def _apply_water_mask_to_ortho(tile_id: str, mask: np.ndarray, cfg: Config) -> N
     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:

scripts/rebuild_ortho_jpg_river.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Rebuild river-masked orthophotos from existing tile masks."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+
+import numpy as np
+from osgeo import gdal
+
+from geodata_pipeline.config import Config
+from geodata_pipeline.river_erosion import _apply_water_mask_to_ortho
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Rebuild export_unity/ortho_jpg_river from work/river_masks and export_unity/ortho_jpg."
+    )
+    p.add_argument("--config", default="geodata_config.toml", help="Path to pipeline config TOML.")
+    p.add_argument("--manifest", default="", help="Optional manifest CSV override.")
+    p.add_argument("--mask-dir", default="", help="Optional mask directory override.")
+    p.add_argument("--output-dir", default="", help="Optional output ortho_jpg_river directory override.")
+    p.add_argument("--skip-existing", action="store_true", help="Skip tiles that already exist in output.")
+    p.add_argument("--mask-res", type=int, default=0, help="Fallback mask resolution when a tile mask is missing.")
+    return p.parse_args()
+
+
+def load_tile_ids(manifest_path: Path) -> list[str]:
+    with manifest_path.open("r", encoding="utf-8", newline="") as f:
+        rows = list(csv.DictReader(f))
+    return [row["tile_id"].strip() for row in rows if row.get("tile_id")]
+
+
+def load_mask(mask_path: Path) -> np.ndarray | None:
+    ds = gdal.Open(str(mask_path), gdal.GA_ReadOnly)
+    if ds is None:
+        return None
+    arr = ds.ReadAsArray()
+    ds = None
+    if arr is None:
+        return None
+    if arr.ndim == 3:
+        arr = arr[0]
+    arr = arr.astype(np.float32)
+    if arr.max() > 1.0:
+        arr /= 255.0
+    return np.clip(arr, 0.0, 1.0)
+
+
+def main() -> int:
+    args = parse_args()
+    cfg = Config.load(args.config)
+    gdal.UseExceptions()
+
+    manifest_path = Path(args.manifest or cfg.export.manifest_path)
+    if not manifest_path.exists():
+        raise SystemExit(f"[rebuild_ortho_jpg_river] Missing manifest: {manifest_path}")
+
+    mask_dir = Path(args.mask_dir or Path(cfg.work.work_dir) / "river_masks")
+    output_dir = Path(args.output_dir or cfg.ortho.river_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    fallback_res = int(args.mask_res or cfg.heightmap.out_res)
+    if fallback_res <= 0:
+        fallback_res = 1025
+
+    tile_ids = load_tile_ids(manifest_path)
+    written = 0
+    skipped = 0
+    with_mask = 0
+    without_mask = 0
+    missing_source = 0
+
+    print(f"[rebuild_ortho_jpg_river] Tiles in manifest: {len(tile_ids)}")
+    print(f"[rebuild_ortho_jpg_river] Mask dir: {mask_dir}")
+    print(f"[rebuild_ortho_jpg_river] Output dir: {output_dir}")
+
+    for tile_id in tile_ids:
+        out_path = output_dir / f"{tile_id}.jpg"
+        if args.skip_existing and out_path.exists():
+            skipped += 1
+            continue
+
+        source_path = Path(cfg.export.ortho_dir) / f"{tile_id}.jpg"
+        if not source_path.exists():
+            missing_source += 1
+            print(f"[rebuild_ortho_jpg_river] Missing source ortho for {tile_id}: {source_path}")
+            continue
+
+        mask_path = mask_dir / f"{tile_id}_mask.png"
+        mask = load_mask(mask_path)
+        if mask is None:
+            without_mask += 1
+            mask = np.zeros((fallback_res, fallback_res), dtype=np.float32)
+        else:
+            with_mask += 1
+
+        _apply_water_mask_to_ortho(tile_id, mask, cfg)
+        if out_path.exists():
+            written += 1
+
+    print(
+        "[rebuild_ortho_jpg_river] Summary: "
+        f"written={written}, skipped={skipped}, with_mask={with_mask}, "
+        f"without_mask={without_mask}, missing_source={missing_source}"
+    )
+    return 0 if missing_source == 0 else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())