Add hard-negative proposal classifier pipeline

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())

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())