Add learned building proposal filtering
This commit is contained in:
@@ -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,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user