Add learned building proposal filtering

This commit is contained in:
Jens
2026-07-27 07:35:50 +02:00
parent e5a642a7f7
commit 0acde4aba0
3 changed files with 281 additions and 0 deletions
@@ -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,
}