Add SWE boundary mask pipeline and mask tooling
This commit is contained in:
247
scripts/mask_split_master.py
Normal file
247
scripts/mask_split_master.py
Normal file
@@ -0,0 +1,247 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user