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