diff --git a/scripts/build_building_proposal_classifier_dataset.py b/scripts/build_building_proposal_classifier_dataset.py new file mode 100644 index 00000000..eb5fa222 --- /dev/null +++ b/scripts/build_building_proposal_classifier_dataset.py @@ -0,0 +1,166 @@ +#!/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()) diff --git a/scripts/train_building_proposal_classifier.py b/scripts/train_building_proposal_classifier.py new file mode 100644 index 00000000..e9ff9986 --- /dev/null +++ b/scripts/train_building_proposal_classifier.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Train and export the binary building-proposal filter.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def binary_metrics(scores: list[float], labels: list[int], threshold: float = 0.5) -> dict[str, float | int]: + tp = sum(score >= threshold and label == 1 for score, label in zip(scores, labels, strict=True)) + fp = sum(score >= threshold and label == 0 for score, label in zip(scores, labels, strict=True)) + fn = sum(score < threshold and label == 1 for score, label in zip(scores, labels, strict=True)) + precision = tp / (tp + fp) if tp + fp else 1.0 + recall = tp / (tp + fn) if tp + fn else 1.0 + return {"tp": tp, "fp": fp, "fn": fn, "precision": precision, "recall": recall, + "f1": 2 * precision * recall / (precision + recall) if precision + recall else 0.0} + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--epochs", type=int, default=12) + parser.add_argument("--batch", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + if args.output_dir.exists(): + parser.error(f"output already exists: {args.output_dir}") + + import torch + from torch import nn + from torch.utils.data import DataLoader + from torchvision.datasets import ImageFolder + from torchvision.models import ResNet18_Weights, resnet18 + + weights = ResNet18_Weights.DEFAULT + transform = weights.transforms() + train_ds = ImageFolder(args.dataset_dir / "train", transform=transform) + val_ds = ImageFolder(args.dataset_dir / "val", transform=transform) + if train_ds.class_to_idx != {"negative": 0, "positive": 1}: + raise RuntimeError(f"unexpected class order: {train_ds.class_to_idx}") + train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=0) + val_loader = DataLoader(val_ds, batch_size=args.batch, shuffle=False, num_workers=0) + device = torch.device(args.device) + model = resnet18(weights=weights) + model.fc = nn.Linear(model.fc.in_features, 1) + model.to(device) + positives = sum(label == 1 for _path, label in train_ds.samples) + negatives = len(train_ds) - positives + loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([negatives / positives], device=device)) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) + args.output_dir.mkdir(parents=True) + history = [] + best_f1 = -1.0 + for epoch in range(1, args.epochs + 1): + model.train() + train_loss = 0.0 + for images, labels in train_loader: + images, labels = images.to(device), labels.float().to(device) + optimizer.zero_grad(set_to_none=True) + logits = model(images).flatten() + loss = loss_fn(logits, labels) + loss.backward() + optimizer.step() + train_loss += float(loss) * len(images) + model.eval() + scores: list[float] = [] + labels_out: list[int] = [] + with torch.inference_mode(): + for images, labels in val_loader: + scores.extend(torch.sigmoid(model(images.to(device)).flatten()).cpu().tolist()) + labels_out.extend(labels.tolist()) + metric = binary_metrics(scores, labels_out) + row = {"epoch": epoch, "train_loss": train_loss / len(train_ds), **metric} + history.append(row) + print(json.dumps(row), flush=True) + if float(metric["f1"]) > best_f1: + best_f1 = float(metric["f1"]) + torch.save(model.state_dict(), args.output_dir / "best-state.pt") + model.load_state_dict(torch.load(args.output_dir / "best-state.pt", map_location=device)) + model.eval() + scripted = torch.jit.script(model) + scripted.save(str(args.output_dir / "proposal-classifier.torchscript.pt")) + report = {"schema_version": 1, "status": "ok", "classes": train_ds.class_to_idx, + "train_count": len(train_ds), "validation_count": len(val_ds), + "best_validation_f1": best_f1, "history": history, + "model": str(args.output_dir / "proposal-classifier.torchscript.pt")} + (args.output_dir / "training-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_building_proposal_classifier.py b/tests/test_building_proposal_classifier.py new file mode 100644 index 00000000..3eb9206b --- /dev/null +++ b/tests/test_building_proposal_classifier.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] + + +def load(name: str): + path = ROOT / "scripts" / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +miner = load("build_building_proposal_classifier_dataset") +trainer = load("train_building_proposal_classifier") + + +def test_classify_proposals_consumes_reference_once() -> None: + reference = [(0.0, 0.0, 10.0, 10.0)] + proposals = [((0.0, 0.0, 10.0, 10.0), 0.9), ((0.0, 0.0, 10.0, 10.0), 0.8)] + assert [item[0] for item in miner.classify_proposals(proposals, reference, 0.25)] == ["positive", "negative"] + + +def test_eligible_tiles_rejects_protected_manifest_split() -> None: + manifest = {"samples": [{"sample_slug": "x", "region": "flanders", "split": "test"}]} + summary = {"tiles": [{"sample_slug": "x", "split": "train", "image_path": "x.png"}]} + with pytest.raises(ValueError, match="protected"): + miner.eligible_tiles(summary, manifest, "flanders") + + +def test_binary_metrics() -> None: + result = trainer.binary_metrics([0.9, 0.8, 0.2, 0.1], [1, 0, 1, 0]) + assert result == {"tp": 1, "fp": 1, "fn": 1, "precision": 0.5, "recall": 0.5, "f1": 0.5}