Add learned building proposal filtering

backend/tests/test_building_proposal_filter.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+
+SCRIPT = Path(__file__).parents[2] / "scripts" / "train_building_proposal_filter.py"
+SPEC = importlib.util.spec_from_file_location("building_proposal_filter", SCRIPT)
+assert SPEC and SPEC.loader
+MODULE = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(MODULE)
+
+
+def test_evenly_limited_is_deterministic() -> None:
+    assert MODULE.evenly_limited(list(range(10)), 3) == [0, 3, 6]
+
+
+def test_validation_threshold_is_selected_without_test_data() -> None:
+    threshold, metrics = MODULE.choose_threshold([0.9, 0.8, 0.2, 0.1], [1, 1, 0, 0])
+    assert 0.2 < threshold <= 0.8
+    assert metrics["f1"] == 1.0
+
+
+def test_negative_match_iou() -> None:
+    assert MODULE.iou((0, 0, 10, 10), (0, 0, 10, 10)) == 1.0
+    assert MODULE.iou((0, 0, 10, 10), (20, 20, 30, 30)) == 0.0

scripts/evaluate_belgium_building_candidate.py

@@ -97,9 +97,20 @@ def main() -> int:
         help="Maximum detections retained per tile; dense Belgian urban tiles exceed YOLO's default 300.",
     )
     parser.add_argument("--box-scale", type=float, default=1.0)
+    parser.add_argument("--proposal-classifier", type=Path)
+    parser.add_argument("--proposal-classifier-threshold", type=float, default=0.5)
+    parser.add_argument("--proposal-crop-scale", type=float, default=1.4)
     args = parser.parse_args()
 
     from ultralytics import YOLO
+    proposal_classifier = None
+    proposal_transform = None
+    if args.proposal_classifier:
+        import torch
+        from torchvision.models import ResNet18_Weights
+
+        proposal_classifier = torch.jit.load(str(args.proposal_classifier), map_location=args.device).eval()
+        proposal_transform = ResNet18_Weights.DEFAULT.transforms()
 
     summary = json.loads(args.summary.read_text(encoding="utf-8"))
     manifest = json.loads(args.corpus_manifest.read_text(encoding="utf-8"))
@@ -131,6 +142,27 @@ def main() -> int:
             (scale_box(tuple(map(float, box)), args.box_scale), float(score))
             for box, score in zip(result.boxes.xyxy.cpu().tolist(), result.boxes.conf.cpu().tolist(), strict=True)
         ]
+        if proposal_classifier is not None:
+            from PIL import Image
+            import torch
+
+            with Image.open(tile["image_path"]) as opened:
+                source = opened.convert("RGB")
+            crops = []
+            for box, _score in predictions:
+                x1, y1, x2, y2 = box
+                cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+                side = max(x2 - x1, y2 - y1, 8.0) * args.proposal_crop_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)))
+                crop = source.crop((left, top, right, bottom))
+                crops.append(proposal_transform(crop.convert("RGB")))
+            if crops:
+                with torch.inference_mode():
+                    probabilities = torch.sigmoid(proposal_classifier(torch.stack(crops).to(args.device)).flatten()).cpu().tolist()
+                predictions = [item for item, probability in zip(predictions, probabilities, strict=True) if probability >= args.proposal_classifier_threshold]
         observations.append(
             {
                 "sample_slug": tile["sample_slug"],
@@ -173,6 +205,9 @@ def main() -> int:
         "inference_imgsz": args.imgsz,
         "max_detections_per_tile": args.max_det,
         "box_scale": args.box_scale,
+        "proposal_classifier": str(args.proposal_classifier) if args.proposal_classifier else None,
+        "proposal_classifier_threshold": args.proposal_classifier_threshold if args.proposal_classifier else None,
+        "proposal_crop_scale": args.proposal_crop_scale if args.proposal_classifier else None,
         "tile_count": len(tiles),
         "sweeps": sweeps,
     }

scripts/train_building_proposal_filter.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""Train a leak-free PyTorch building/background filter on YOLO proposals."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import random
+import shutil
+from pathlib import Path
+from typing import Any, TypeVar
+
+
+T = TypeVar("T")
+
+
+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 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())
+        boxes.append(((cx - bw / 2) * width, (cy - bh / 2) * height, (cx + bw / 2) * width, (cy + bh / 2) * height))
+    return boxes
+
+
+def expanded_crop(image: Any, box: tuple[float, float, float, float], factor: float) -> Any:
+    x1, y1, x2, y2 = box
+    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+    side = max(x2 - x1, y2 - y1, 8.0) * factor
+    left, top = max(0, int(cx - side / 2)), max(0, int(cy - side / 2))
+    right, bottom = min(image.width, int(cx + side / 2)), min(image.height, int(cy + side / 2))
+    return image.crop((left, top, max(left + 1, right), max(top + 1, bottom))).convert("RGB")
+
+
+def evenly_limited(items: list[T], limit: int) -> list[T]:
+    if len(items) <= limit:
+        return items
+    return [items[index * len(items) // limit] for index in range(limit)]
+
+
+def materialize_split(
+    *,
+    tiles: list[dict[str, Any]],
+    results: list[Any],
+    output_dir: Path,
+    split: str,
+    crop_scale: float,
+    max_positives_per_tile: int,
+    max_negatives_per_tile: int,
+    negative_match_iou: float,
+) -> dict[str, int]:
+    from PIL import Image
+
+    counts = {"positive": 0, "negative": 0}
+    for tile, result in zip(tiles, results, strict=True):
+        with Image.open(tile["image_path"]) as source:
+            image = source.convert("RGB")
+        references = read_boxes(Path(tile["label_path"]), image.width, image.height)
+        positives = evenly_limited(references, max_positives_per_tile)
+        predictions = [tuple(map(float, box)) for box in result.boxes.xyxy.cpu().tolist()]
+        negatives = [
+            box
+            for box in predictions
+            if max((iou(box, reference) for reference in references), default=0.0) < negative_match_iou
+        ][:max_negatives_per_tile]
+        for label, boxes in (("positive", positives), ("negative", negatives)):
+            directory = output_dir / split / label
+            directory.mkdir(parents=True, exist_ok=True)
+            for box in boxes:
+                path = directory / f"{tile['sample_slug']}_{tile['tile_index']:04d}_{counts[label]:07d}.png"
+                expanded_crop(image, box, crop_scale).save(path)
+                counts[label] += 1
+    return counts
+
+
+def choose_threshold(probabilities: list[float], labels: list[int]) -> tuple[float, dict[str, float | int]]:
+    best: tuple[float, float, dict[str, float | int]] | None = None
+    for threshold in [value / 100 for value in range(10, 96, 5)]:
+        tp = sum(p >= threshold and y == 1 for p, y in zip(probabilities, labels, strict=True))
+        fp = sum(p >= threshold and y == 0 for p, y in zip(probabilities, labels, strict=True))
+        fn = sum(p < threshold and y == 1 for p, y in zip(probabilities, labels, strict=True))
+        precision = tp / (tp + fp) if tp + fp else 1.0
+        recall = tp / (tp + fn) if tp + fn else 1.0
+        f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
+        metrics = {"true_positive": tp, "false_positive": fp, "false_negative": fn, "precision": precision, "recall": recall, "f1": f1}
+        candidate = (f1, threshold, metrics)
+        if best is None or candidate[:2] > best[:2]:
+            best = candidate
+    assert best
+    return best[1], best[2]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--summary", type=Path, required=True)
+    parser.add_argument("--proposal-model", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--proposal-confidence", type=float, default=0.01)
+    parser.add_argument("--crop-scale", type=float, default=1.4)
+    parser.add_argument("--max-positives-per-tile", type=int, default=40)
+    parser.add_argument("--max-negatives-per-tile", type=int, default=40)
+    parser.add_argument("--negative-match-iou", type=float, default=0.05)
+    parser.add_argument("--epochs", type=int, default=8)
+    parser.add_argument("--batch", type=int, default=64)
+    parser.add_argument("--seed", type=int, default=20260727)
+    parser.add_argument("--force", action="store_true")
+    args = parser.parse_args()
+    if args.output_dir.exists():
+        if not args.force:
+            raise SystemExit(f"Output exists: {args.output_dir}")
+        shutil.rmtree(args.output_dir)
+    random.seed(args.seed)
+    import torch
+    from torch import nn
+    from torch.utils.data import DataLoader
+    from torchvision import datasets, transforms
+    from torchvision.models import ResNet18_Weights, resnet18
+    from ultralytics import YOLO
+
+    torch.manual_seed(args.seed)
+    summary = json.loads(args.summary.read_text(encoding="utf-8"))
+    split_tiles = {
+        split: [tile for tile in summary["tiles"] if tile.get("kept", True) and tile["split"] == split]
+        for split in ("train", "val")
+    }
+    proposal_model = YOLO(str(args.proposal_model))
+    counts = {}
+    for split, tiles in split_tiles.items():
+        results = proposal_model.predict(
+            [tile["image_path"] for tile in tiles], conf=args.proposal_confidence, imgsz=640,
+            max_det=1000, device=args.device, verbose=False,
+        )
+        counts[split] = materialize_split(
+            tiles=tiles, results=results, output_dir=args.output_dir / "crops", split=split,
+            crop_scale=args.crop_scale, max_positives_per_tile=args.max_positives_per_tile,
+            max_negatives_per_tile=args.max_negatives_per_tile, negative_match_iou=args.negative_match_iou,
+        )
+        del results
+
+    del proposal_model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    weights = ResNet18_Weights.DEFAULT
+    transform = weights.transforms()
+    datasets_by_split = {
+        split: datasets.ImageFolder(args.output_dir / "crops" / split, transform=transform)
+        for split in ("train", "val")
+    }
+    loaders = {
+        split: DataLoader(data, batch_size=args.batch, shuffle=split == "train", num_workers=0)
+        for split, data in datasets_by_split.items()
+    }
+    model = resnet18(weights=weights)
+    model.fc = nn.Linear(model.fc.in_features, 1)
+    model.to(args.device)
+    positives = counts["train"]["positive"]
+    negatives = counts["train"]["negative"]
+    criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([negatives / max(positives, 1)], device=args.device))
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
+    history = []
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        total = 0.0
+        for images, labels in loaders["train"]:
+            images, labels = images.to(args.device), labels.float().to(args.device)
+            optimizer.zero_grad(set_to_none=True)
+            loss = criterion(model(images).flatten(), labels)
+            loss.backward()
+            optimizer.step()
+            total += float(loss.detach()) * len(images)
+        history.append({"epoch": epoch, "train_loss": total / len(datasets_by_split["train"])})
+        print(history[-1], flush=True)
+    model.eval()
+    probabilities: list[float] = []
+    labels: list[int] = []
+    with torch.inference_mode():
+        for images, batch_labels in loaders["val"]:
+            probabilities.extend(torch.sigmoid(model(images.to(args.device)).flatten()).cpu().tolist())
+            labels.extend(batch_labels.tolist())
+    threshold, validation_metrics = choose_threshold(probabilities, labels)
+    scripted = torch.jit.script(model.cpu())
+    model_path = args.output_dir / "proposal-filter.torchscript.pt"
+    scripted.save(str(model_path))
+    evidence = {
+        "schema_version": 1, "status": "ok", "architecture": "resnet18_binary_proposal_filter",
+        "summary": str(args.summary), "summary_sha256": sha256(args.summary),
+        "proposal_model": str(args.proposal_model), "proposal_model_sha256": sha256(args.proposal_model),
+        "device": args.device, "proposal_confidence": args.proposal_confidence, "crop_scale": args.crop_scale,
+        "negative_match_iou": args.negative_match_iou, "counts": counts, "epochs": args.epochs,
+        "history": history, "selected_threshold_source": "validation_only",
+        "selected_threshold": threshold, "validation_metrics": validation_metrics,
+        "model": str(model_path), "model_sha256": sha256(model_path),
+    }
+    (args.output_dir / "proposal-filter.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())