#!/usr/bin/env python3 """Split global master masks back into tile masks. Optional filtering modes: - keep only source/sink tiles that contain non-black pixels - keep only water tiles that differ from a reference river mask directory """ from __future__ import annotations import argparse import json import re from pathlib import Path from typing import Dict import numpy as np from PIL import Image _TILE_KEY_RE = re.compile(r"^(dgm\d+_\d+_\d+_\d+)") # Master masks can exceed PIL's safety threshold; these files are trusted local data. Image.MAX_IMAGE_PIXELS = None def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="Split master masks into per-tile water/source/sink masks.") p.add_argument("--master-dir", default="work/mask_master", help="Directory containing master_meta.json and master images.") p.add_argument("--out-water-dir", default="raw/water_masks", help="Output directory for water tile masks.") p.add_argument("--out-source-dir", default="raw/water_source_masks", help="Output directory for source tile masks.") p.add_argument("--out-sink-dir", default="raw/water_sink_masks", help="Output directory for sink tile masks.") p.add_argument( "--ref-water-dir", default="work/river_masks", help="Reference water mask dir used to keep only changed water tiles.", ) p.add_argument( "--keep-informative-only", action="store_true", help="Write only informative masks (non-black source/sink; water changed vs ref water dir).", ) p.add_argument( "--prune-existing", action="store_true", help="Remove existing PNG/WLD outputs in out dirs before writing.", ) p.add_argument("--write-wld", action="store_true", help="Write world files for output masks.") return p.parse_args() def write_worldfile(path: Path, px: float, py: float, minx: float, maxy: float) -> None: c = minx + (px / 2.0) f = maxy + (py / 2.0) text = "\n".join( [ f"{px:.12f}", "0.0", "0.0", f"{py:.12f}", f"{c:.12f}", f"{f:.12f}", ] ) path.write_text(text + "\n", encoding="utf-8") def ensure_dirs(*dirs: Path) -> None: for d in dirs: d.mkdir(parents=True, exist_ok=True) def tile_key_from_name(name: str) -> str: stem = Path(name).stem m = _TILE_KEY_RE.match(stem) return m.group(1) if m else "" def remove_existing_outputs(out_dir: Path) -> None: if not out_dir.exists(): return for p in out_dir.glob("*.png"): p.unlink(missing_ok=True) for p in out_dir.glob("*.wld"): p.unlink(missing_ok=True) def has_non_black(arr: np.ndarray) -> bool: return bool(np.any(arr != 0)) def index_reference_masks(ref_dir: Path) -> Dict[str, Path]: out: Dict[str, Path] = {} if not ref_dir.exists(): return out for p in sorted(ref_dir.glob("*.png")): key = tile_key_from_name(p.name) if not key or key in out: continue out[key] = p return out def resize_nearest(arr: np.ndarray, width: int, height: int) -> np.ndarray: if arr.shape[0] == height and arr.shape[1] == width: return arr img = Image.fromarray(arr, mode="RGB") resized = img.resize((width, height), resample=Image.Resampling.NEAREST) return np.array(resized, dtype=np.uint8) def water_differs_from_reference(water_arr: np.ndarray, key: str, ref_index: Dict[str, Path]) -> bool: ref_path = ref_index.get(key) if ref_path is None: # No baseline tile: treat as informative so we do not lose manual work. return True ref_arr = np.array(Image.open(ref_path).convert("RGB"), dtype=np.uint8) ref_arr = resize_nearest(ref_arr, water_arr.shape[1], water_arr.shape[0]) return not np.array_equal(water_arr, ref_arr) def source_name_for_tile(tile: Dict[str, object]) -> str: source_file = str(tile.get("source_file") or "").strip() if source_file: return source_file stem = str(tile.get("water_stem") or "").strip() if stem: return f"{stem}_source_mask.png" return f"{tile['key']}_source_mask.png" def sink_name_for_tile(tile: Dict[str, object]) -> str: sink_file = str(tile.get("sink_file") or "").strip() if sink_file: return sink_file stem = str(tile.get("water_stem") or "").strip() if stem: return f"{stem}_sink_mask.png" return f"{tile['key']}_sink_mask.png" def crop_tile(master: np.ndarray, xoff: int, yoff: int, width: int, height: int) -> np.ndarray: return master[yoff : yoff + height, xoff : xoff + width].copy() def main() -> int: args = parse_args() master_dir = Path(args.master_dir) meta_path = master_dir / "master_meta.json" if not meta_path.exists(): raise SystemExit(f"[mask_split_master] Missing meta: {meta_path}") meta = json.loads(meta_path.read_text(encoding="utf-8")) m = meta["master"] tiles = meta["tiles"] px = float(m["px"]) py = float(m["py"]) water_master = np.array(Image.open(master_dir / m["water_master"]).convert("RGB"), dtype=np.uint8) source_master = np.array(Image.open(master_dir / m["source_master"]).convert("RGB"), dtype=np.uint8) sink_master = np.array(Image.open(master_dir / m["sink_master"]).convert("RGB"), dtype=np.uint8) ensure_dirs(Path(args.out_water_dir), Path(args.out_source_dir), Path(args.out_sink_dir)) out_water = Path(args.out_water_dir) out_source = Path(args.out_source_dir) out_sink = Path(args.out_sink_dir) if args.prune_existing: remove_existing_outputs(out_water) remove_existing_outputs(out_source) remove_existing_outputs(out_sink) ref_index = index_reference_masks(Path(args.ref_water_dir)) written_water = 0 written_source = 0 written_sink = 0 skipped_water = 0 skipped_source = 0 skipped_sink = 0 for tile in tiles: xoff = int(tile["xoff"]) yoff = int(tile["yoff"]) width = int(tile["width"]) height = int(tile["height"]) water_arr = crop_tile(water_master, xoff, yoff, width, height) source_arr = crop_tile(source_master, xoff, yoff, width, height) sink_arr = crop_tile(sink_master, xoff, yoff, width, height) water_name = str(tile.get("water_file") or f"{tile['key']}_mask_viz.png") source_name = source_name_for_tile(tile) sink_name = sink_name_for_tile(tile) water_path = out_water / water_name source_path = out_source / source_name sink_path = out_sink / sink_name write_water = True write_source = True write_sink = True if args.keep_informative_only: write_source = has_non_black(source_arr) write_sink = has_non_black(sink_arr) write_water = water_differs_from_reference(water_arr, str(tile["key"]), ref_index) if write_water: Image.fromarray(water_arr, mode="RGB").save(water_path) if args.write_wld: minx = float(tile["minx"]) maxy = float(tile["maxy"]) write_worldfile(water_path.with_suffix(".wld"), px, py, minx, maxy) written_water += 1 else: skipped_water += 1 if write_source: Image.fromarray(source_arr, mode="RGB").save(source_path) if args.write_wld: minx = float(tile["minx"]) maxy = float(tile["maxy"]) write_worldfile(source_path.with_suffix(".wld"), px, py, minx, maxy) written_source += 1 else: skipped_source += 1 if write_sink: Image.fromarray(sink_arr, mode="RGB").save(sink_path) if args.write_wld: minx = float(tile["minx"]) maxy = float(tile["maxy"]) write_worldfile(sink_path.with_suffix(".wld"), px, py, minx, maxy) written_sink += 1 else: skipped_sink += 1 total_tiles = len(tiles) print(f"[mask_split_master] Tiles processed: {total_tiles}") print(f"[mask_split_master] Water written: {written_water}, skipped: {skipped_water}") print(f"[mask_split_master] Source written: {written_source}, skipped: {skipped_source}") print(f"[mask_split_master] Sink written: {written_sink}, skipped: {skipped_sink}") print(f"[mask_split_master] Water dir: {out_water}") print(f"[mask_split_master] Source dir: {out_source}") print(f"[mask_split_master] Sink dir: {out_sink}") return 0 if __name__ == "__main__": raise SystemExit(main())