diff --git a/geodata_pipeline/river_erosion.py b/geodata_pipeline/river_erosion.py index a861e82..bde8a96 100644 --- a/geodata_pipeline/river_erosion.py +++ b/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: diff --git a/scripts/rebuild_ortho_jpg_river.py b/scripts/rebuild_ortho_jpg_river.py new file mode 100644 index 0000000..8fa5da5 --- /dev/null +++ b/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())