#!/usr/bin/env python3 """Mine detector proposals into a leak-free binary crop dataset.""" from __future__ import annotations import argparse import hashlib import json import sys from collections import Counter from pathlib import Path from typing import Any from PIL import Image SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) from evaluate_belgium_building_candidate import iou, read_references PROTECTED_SPLITS = {"calibration", "test", "background-test"} 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 classify_proposals( predictions: list[tuple[tuple[float, float, float, float], float]], references: list[tuple[float, float, float, float]], match_iou: float, ) -> list[tuple[str, tuple[float, float, float, float], float]]: unmatched = set(range(len(references))) classified = [] for box, score in sorted(predictions, key=lambda item: -item[1]): candidates = [(iou(box, references[index]), index) for index in unmatched] overlap, index = max(candidates, default=(0.0, -1)) label = "positive" if overlap >= match_iou else "negative" if label == "positive": unmatched.remove(index) classified.append((label, box, score)) return classified def eligible_tiles( summary: dict[str, Any], manifest: dict[str, Any], region: str | None ) -> list[dict[str, Any]]: samples = {item["sample_slug"]: item for item in manifest["samples"]} selected = [] for tile in summary["tiles"]: sample = samples.get(tile["sample_slug"]) if sample is None: raise ValueError(f"unknown sample: {tile['sample_slug']}") tile_split = str(tile.get("split") or "") manifest_split = str(sample.get("split") or "") if tile_split in PROTECTED_SPLITS or manifest_split in PROTECTED_SPLITS: raise ValueError(f"protected split in proposal source: {tile['sample_slug']}") if tile_split != manifest_split: raise ValueError(f"tile/manifest split mismatch: {tile['sample_slug']}") if tile.get("kept", True) and tile_split in {"train", "val"} and ( region is None or sample.get("region") == region ): selected.append(tile) if not selected or not any(tile["split"] == "val" for tile in selected): raise ValueError("proposal dataset requires eligible train and validation tiles") return selected def crop_square( source: Image.Image, box: tuple[float, float, float, float], scale: float ) -> Image.Image: x1, y1, x2, y2 = box cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 side = max(x2 - x1, y2 - y1, 8.0) * scale left = max(0, min(source.width - 1, int(cx - side / 2))) top = max(0, min(source.height - 1, int(cy - side / 2))) right = max(left + 1, min(source.width, int(cx + side / 2))) bottom = max(top + 1, min(source.height, int(cy + side / 2))) return source.crop((left, top, right, bottom)).convert("RGB") def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--model", type=Path, required=True) parser.add_argument("--summary", type=Path, required=True) parser.add_argument("--corpus-manifest", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--region") parser.add_argument("--confidence", type=float, default=0.05) parser.add_argument("--match-iou", type=float, default=0.25) parser.add_argument("--crop-scale", type=float, default=1.4) parser.add_argument("--max-positive-per-tile", type=int, default=24) parser.add_argument("--max-negative-per-tile", type=int, default=24) parser.add_argument("--device", default="cuda:0") parser.add_argument("--imgsz", type=int, default=640) args = parser.parse_args() if args.output_dir.exists(): parser.error(f"output already exists: {args.output_dir}") from ultralytics import YOLO summary = json.loads(args.summary.read_text(encoding="utf-8")) manifest = json.loads(args.corpus_manifest.read_text(encoding="utf-8")) tiles = eligible_tiles(summary, manifest, args.region) args.output_dir.mkdir(parents=True) counts: Counter[str] = Counter() sample_counts: Counter[str] = Counter() model = YOLO(str(args.model)) for start in range(0, len(tiles), 16): batch_tiles = tiles[start : start + 16] results = model.predict( [tile["image_path"] for tile in batch_tiles], conf=args.confidence, device=args.device, imgsz=args.imgsz, max_det=1000, iou=0.7, verbose=False, ) for tile, result in zip(batch_tiles, results, strict=True): with Image.open(tile["image_path"]) as opened: source = opened.convert("RGB") references = read_references(Path(tile["label_path"]), source.width, source.height) proposals = [ (tuple(map(float, box)), float(score)) for box, score in zip( result.boxes.xyxy.cpu().tolist(), result.boxes.conf.cpu().tolist(), strict=True ) ] classified = classify_proposals(proposals, references, args.match_iou) limits = {"positive": args.max_positive_per_tile, "negative": args.max_negative_per_tile} per_label: Counter[str] = Counter() for proposal_index, (label, box, score) in enumerate(classified): if per_label[label] >= limits[label]: continue per_label[label] += 1 split = tile["split"] target_dir = args.output_dir / split / label target_dir.mkdir(parents=True, exist_ok=True) name = f"{tile['sample_slug']}__{Path(tile['image_path']).stem}__{proposal_index:04d}.jpg" crop_square(source, box, args.crop_scale).save(target_dir / name, quality=92) counts[f"{split}/{label}"] += 1 sample_counts[tile["sample_slug"]] += 1 for split in ("train", "val"): for label in ("positive", "negative"): if counts[f"{split}/{label}"] == 0: raise RuntimeError(f"empty proposal class: {split}/{label}") evidence = { "schema_version": 1, "status": "ok", "model": str(args.model), "model_sha256": sha256(args.model), "summary": str(args.summary), "summary_sha256": sha256(args.summary), "corpus_manifest": str(args.corpus_manifest), "corpus_manifest_sha256": sha256(args.corpus_manifest), "region": args.region, "confidence": args.confidence, "match_iou": args.match_iou, "crop_scale": args.crop_scale, "counts": dict(sorted(counts.items())), "sample_counts": dict(sorted(sample_counts.items())), "protected_samples_in_training": [], "tile_count": len(tiles), } (args.output_dir / "proposal-dataset-evidence.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())