#!/usr/bin/env python3 """Refine official footprint boxes into auditable image-visible SAM roof boxes.""" from __future__ import annotations import argparse import gc import hashlib import json import math import os import shutil import sys from pathlib import Path SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) from training_release_manifest import ( # noqa: E402 TrainingReleaseError, assert_yolo_summary_bound_to_training_release, file_sha256, training_release_paths, ) def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def iou(left: tuple[float, float, float, float], right: tuple[float, float, float, float]) -> float: ix = max(0.0, min(left[2], right[2]) - max(left[0], right[0])) iy = max(0.0, min(left[3], right[3]) - max(left[1], right[1])) intersection = ix * iy union = (left[2] - left[0]) * (left[3] - left[1]) + (right[2] - right[0]) * ( right[3] - right[1] ) - intersection return intersection / union if union > 0 else 0.0 def plausible_refinement( source: tuple[float, float, float, float], refined: tuple[float, float, float, float], *, min_iou: float, min_area_ratio: float, max_area_ratio: float, max_center_shift_ratio: float, min_dimension_ratio: float, max_dimension_ratio: float, ) -> bool: source_width, source_height = source[2] - source[0], source[3] - source[1] refined_width, refined_height = refined[2] - refined[0], refined[3] - refined[1] source_area = source_width * source_height refined_area = refined_width * refined_height ratio = refined_area / source_area if source_area > 0 else 0.0 if source_width <= 0 or source_height <= 0: return False width_ratio = refined_width / source_width height_ratio = refined_height / source_height source_centre = ((source[0] + source[2]) / 2, (source[1] + source[3]) / 2) refined_centre = ((refined[0] + refined[2]) / 2, (refined[1] + refined[3]) / 2) centre_shift_ratio = math.dist(source_centre, refined_centre) / math.hypot(source_width, source_height) return ( min_area_ratio <= ratio <= max_area_ratio and min_dimension_ratio <= width_ratio <= max_dimension_ratio and min_dimension_ratio <= height_ratio <= max_dimension_ratio and centre_shift_ratio <= max_center_shift_ratio and iou(source, refined) >= min_iou ) def read_boxes(path: Path, width: int, height: int) -> list[tuple[float, float, float, float]]: boxes = [] for line in path.read_text(encoding="utf-8").splitlines() if path.is_file() else []: class_id, cx, cy, bw, bh = map(float, line.split()) if class_id != 0: raise ValueError(f"Unexpected class in {path}: {class_id}") boxes.append(((cx - bw / 2) * width, (cy - bh / 2) * height, (cx + bw / 2) * width, (cy + bh / 2) * height)) return boxes def yolo_line(box: tuple[float, float, float, float], width: int, height: int) -> str: x1, y1, x2, y2 = box return f"0 {(x1+x2)/(2*width):.8f} {(y1+y2)/(2*height):.8f} {(x2-x1)/width:.8f} {(y2-y1)/height:.8f}" def expanded(box: tuple[float, float, float, float], factor: float, width: int, height: int) -> list[float]: x1, y1, x2, y2 = box cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 hw, hh = (x2 - x1) * factor / 2, (y2 - y1) * factor / 2 return [max(0, cx - hw), max(0, cy - hh), min(width - 1, cx + hw), min(height - 1, cy + hh)] def link_or_copy(source: Path, target: Path) -> None: target.parent.mkdir(parents=True, exist_ok=True) try: os.link(source, target) except OSError: shutil.copy2(source, target) def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--summary", type=Path, required=True) parser.add_argument("--train-yaml", type=Path, required=True) parser.add_argument("--corpus-manifest", type=Path, required=True) parser.add_argument("--model", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--device", default="cuda:0") parser.add_argument("--prompt-scale", type=float, default=1.5) parser.add_argument("--min-source-iou", type=float, default=0.15) parser.add_argument("--min-area-ratio", type=float, default=0.25) parser.add_argument("--max-area-ratio", type=float, default=4.0) parser.add_argument("--max-center-shift-ratio", type=float, default=0.75) parser.add_argument("--min-dimension-ratio", type=float, default=0.5) parser.add_argument("--max-dimension-ratio", type=float, default=2.0) parser.add_argument("--max-prompts-per-pass", type=int, default=96) parser.add_argument("--fallback-policy", choices=("retain", "drop"), default="retain") parser.add_argument("--force", action="store_true") parser.add_argument( "--fixture-mode", action="store_true", help="Accept only an explicitly fixture-only source release; refined labels remain non-trainable.", ) args = parser.parse_args() try: source_release = assert_yolo_summary_bound_to_training_release( summary_path=args.summary, train_yaml=args.train_yaml, corpus_manifest=args.corpus_manifest, fixture_mode=args.fixture_mode, ) except TrainingReleaseError as exc: raise SystemExit(str(exc)) from exc if args.output_dir.exists(): if not args.force: raise SystemExit(f"Output exists: {args.output_dir}") shutil.rmtree(args.output_dir) from PIL import Image import torch from ultralytics import SAM summary = json.loads(args.summary.read_text(encoding="utf-8")) model = SAM(str(args.model)) refined_count = fallback_count = 0 reason_counts: dict[str, int] = {} output_tiles = [] for index, tile in enumerate(summary["tiles"], start=1): source_image = Path(tile["image_path"]) source_label = Path(tile["label_path"]) split = tile["split"] target_image = args.output_dir / "images" / split / source_image.name target_label = args.output_dir / "labels" / split / source_label.name link_or_copy(source_image, target_image) target_label.parent.mkdir(parents=True, exist_ok=True) with Image.open(source_image) as image: width, height = image.size source_boxes = read_boxes(source_label, width, height) output_boxes: list[tuple[float, float, float, float] | None] = list(source_boxes) if source_boxes: for start in range(0, len(source_boxes), args.max_prompts_per_pass): source_chunk = source_boxes[start : start + args.max_prompts_per_pass] prompts = [expanded(box, args.prompt_scale, width, height) for box in source_chunk] result = model.predict(str(source_image), bboxes=prompts, device=args.device, verbose=False)[0] masks = result.masks.data.cpu().numpy() if result.masks is not None else [] candidates = [] for mask in masks: ys, xs = mask.nonzero() if len(xs): candidates.append((float(xs.min()), float(ys.min()), float(xs.max() + 1), float(ys.max() + 1))) unmatched = set(range(len(candidates))) for local_index, source_box in enumerate(source_chunk): ranked = sorted(((iou(source_box, candidates[item]), item) for item in unmatched), reverse=True) overlap, candidate_index = ranked[0] if ranked else (0.0, -1) if candidate_index < 0 or overlap < args.min_source_iou: reason_counts["unmatched_mask"] = reason_counts.get("unmatched_mask", 0) + 1 fallback_count += 1 if args.fallback_policy == "drop": output_boxes[start + local_index] = None continue candidate = candidates[candidate_index] if plausible_refinement( source_box, candidate, min_iou=args.min_source_iou, min_area_ratio=args.min_area_ratio, max_area_ratio=args.max_area_ratio, max_center_shift_ratio=args.max_center_shift_ratio, min_dimension_ratio=args.min_dimension_ratio, max_dimension_ratio=args.max_dimension_ratio, ): output_boxes[start + local_index] = candidate unmatched.remove(candidate_index) refined_count += 1 else: reason_counts["geometry_gate"] = reason_counts.get("geometry_gate", 0) + 1 fallback_count += 1 if args.fallback_policy == "drop": output_boxes[start + local_index] = None del result, masks if model.predictor is not None: model.predictor.reset_image() gc.collect() torch.cuda.empty_cache() retained_boxes = [box for box in output_boxes if box is not None] target_label.write_text("\n".join(yolo_line(box, width, height) for box in retained_boxes) + ("\n" if retained_boxes else ""), encoding="utf-8") output_tile = dict(tile) output_tile.update({ "image_path": str(target_image), "label_path": str(target_label), "label_count": len(retained_boxes), "is_negative": not retained_boxes, }) output_tiles.append(output_tile) print(f"{index}/{len(summary['tiles'])} {tile['sample_slug']}: {len(source_boxes)}", flush=True) output_summary = dict(summary) # SAM changes label bytes and semantics. It therefore cannot inherit the # source release; keep the parent evidence separately and require a new # governed corpus/review/release before any training consumer can use it. for field_name in ( "training_release_manifest", "training_release_manifest_sha256", "training_release_freeze", "training_asset_manifest", ): output_summary.pop(field_name, None) output_summary.update( { "output_dir": str(args.output_dir), "dataset_yaml": str(args.output_dir / "dataset.yaml"), "tiles": output_tiles, "label_semantics": ( "sam_visible_roof_only" if args.fallback_policy == "drop" else "sam_visible_roof_with_official_footprint_fallback" ), "label_count": refined_count if args.fallback_policy == "drop" else summary["label_count"], "positive_tile_count": sum(not tile["is_negative"] for tile in output_tiles), "negative_tile_count": sum(tile["is_negative"] for tile in output_tiles), "training_eligible": False, "training_eligibility_reason": ( "SAM-refined labels require a new governed corpus, validation report, human review and immutable training release." ), } ) summary_path = args.output_dir / "yolo_tile_dataset_summary.json" summary_path.write_text(json.dumps(output_summary, indent=2), encoding="utf-8") (args.output_dir / "dataset.yaml").write_text( f"path: {args.output_dir}\ntrain: images/train\nval: images/val\nnames:\n 0: building\n", encoding="utf-8", ) evidence = { "schema_version": 1, "status": "ok", "source_summary": str(args.summary), "source_summary_sha256": sha256(args.summary), "source_training_release": str(training_release_paths(args.train_yaml)["release_manifest"]), "source_training_release_sha256": file_sha256( training_release_paths(args.train_yaml)["release_manifest"] ), "source_corpus_manifest_sha256": source_release["corpus"]["manifest_sha256"], "sam_model": str(args.model), "sam_model_sha256": sha256(args.model), "device": args.device, "prompt_scale": args.prompt_scale, "min_source_iou": args.min_source_iou, "min_area_ratio": args.min_area_ratio, "max_area_ratio": args.max_area_ratio, "max_center_shift_ratio": args.max_center_shift_ratio, "min_dimension_ratio": args.min_dimension_ratio, "max_dimension_ratio": args.max_dimension_ratio, "max_prompts_per_pass": args.max_prompts_per_pass, "fallback_policy": args.fallback_policy, "refined_label_count": refined_count, "fallback_label_count": fallback_count, "dropped_fallback_label_count": fallback_count if args.fallback_policy == "drop" else 0, "fallback_reason_counts": reason_counts, "training_eligible": False, "fixture_mode": bool(args.fixture_mode), } (args.output_dir / "sam-refinement.json").write_text(json.dumps(evidence, indent=2), encoding="utf-8") print(json.dumps(evidence, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())