114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
#!/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())
|