Upload inference.py

src/inference.py

@@ -0,0 +1,144 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from PIL import ImageDraw, ImageFont
+
+from src.backbone import Backbone, backbone_transform
+from src.rpn import RPN_Head, RegionProposalNetwork
+from src.roi import RoIPool
+from src.detection_net import DetectionHead, DetectionNet
+from src.dataset import VOC_CLASSES
+
+from huggingface_hub import hf_hub_download
+
+# One fixed, distinguishable color per VOC class, keyed by VOC_CLASSES order.
+_PALETTE = [
+    "#e6194b", "#3cb44b", "#ffe119", "#4363d8", "#f58231",
+    "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe",
+    "#008080", "#e6beff", "#9a6324", "#fffac8", "#800000",
+    "#aaffc3", "#808000", "#ffd8b1", "#000075", "#808080",
+]
+_CLASS_COLORS = dict(zip(VOC_CLASSES, _PALETTE))
+
+
+@dataclass
+class Pipeline:
+    backbone: Backbone
+    rpn_network: RegionProposalNetwork
+    roi_pool: RoIPool
+    detection_network: DetectionNet
+    device: torch.device
+
+
+_CHECKPOINT_NAME = "faster_rcnn_final.bin"
+
+
+def load_pipeline(checkpoint_dir, device):
+    checkpoint_dir = Path(checkpoint_dir)
+    if not (checkpoint_dir / _CHECKPOINT_NAME).exists():
+        print(f"Missing checkpoint file in '{checkpoint_dir}': {_CHECKPOINT_NAME}. ")
+
+        checkpoint_path = hf_hub_download(repo_id="0Curious0/faster_rcnn_resnet50", filename="checkpoints/faster_rcnn_final.bin")
+    else:
+        checkpoint_path = checkpoint_dir / _CHECKPOINT_NAME
+
+    backbone = Backbone().to(device)
+    rpn_head = RPN_Head(in_channels=1024, mid_channels=512)
+    detection_head = DetectionHead()
+
+    unified_ckpt = torch.load(checkpoint_path, map_location=device)
+
+    backbone.load_state_dict(unified_ckpt["backbone_state_dict"])
+    rpn_head.load_state_dict(unified_ckpt["rpn_state_dict"])
+    detection_head.load_state_dict(unified_ckpt["detection_state_dict"])
+
+    rpn_network = RegionProposalNetwork(rpn_head=rpn_head).to(device)
+    roi_pool = RoIPool(output_size=(7, 7), pooling_mode="adaptive").to(device)
+    detection_network = DetectionNet(detection_head=detection_head).to(device)
+
+    for module in (backbone, rpn_network, roi_pool, detection_network):
+        module.eval()
+        for param in module.parameters():
+            param.requires_grad = False
+
+    return Pipeline(
+        backbone=backbone,
+        rpn_network=rpn_network,
+        roi_pool=roi_pool,
+        detection_network=detection_network,
+        device=device,
+    )
+
+
+def predict(pipeline, pil_image, score_thresh, nms_iou_thresh):
+    # pil_image: original, un-resized PIL image (RGB). backbone_transform only reads target["size"] (and doesn't need it to hold anything), so an empty dict is enough
+    img_tensor, _ = backbone_transform(pil_image, {"size": {}})
+    batch_imgs = img_tensor.unsqueeze(0).to(pipeline.device)
+
+    tensor_height, tensor_width = batch_imgs.shape[2], batch_imgs.shape[3]
+    img_sizes_before_pad = [(tensor_height, tensor_width)]
+
+    pipeline.detection_network.score_thresh = score_thresh
+    pipeline.detection_network.nms_iou_thresh = nms_iou_thresh
+
+    with torch.inference_mode():
+        feature_map = pipeline.backbone(batch_imgs)
+        _, proposals = pipeline.rpn_network(
+            feature_map,
+            batch_img_height=tensor_height,
+            batch_img_width=tensor_width,
+            img_sizes_before_pad=img_sizes_before_pad,
+            pre_nms_top_n=6000,
+            post_nms_top_n=2000,
+        )
+        pooled = pipeline.roi_pool(feature_map, proposals, tensor_height, tensor_width)
+        labels_list, scores_list, boxes_list = pipeline.detection_network(
+            proposals, pooled, img_sizes_before_pad
+        )
+
+    labels, scores, boxes = labels_list[0], scores_list[0], boxes_list[0]
+
+    # Rescaling per-axis by the tensor's own dims maps back onto the original image regardless of that swap, since it's the exact inverse of whatever TF.resize did.
+    orig_width, orig_height = pil_image.size
+    scale_x = orig_width / tensor_width
+    scale_y = orig_height / tensor_height
+
+    detections = []
+    for box, label, score in zip(boxes, labels, scores):
+        x1, y1, x2, y2 = box.tolist()
+        rescaled_box = (x1 * scale_x, y1 * scale_y, x2 * scale_x, y2 * scale_y)
+        detections.append((rescaled_box, VOC_CLASSES[label.item()], score.item()))
+
+    return detections
+
+
+def draw_boxes(pil_image, detections):
+    annotated = pil_image.convert("RGB").copy()
+    draw = ImageDraw.Draw(annotated)
+
+    # Scale font to image size so labels stay legible on both small thumbnails and
+    font_size = max(16, round(min(annotated.size) / 40))            # divide by 40 to make font size 2.5% of the smaller image dimension
+    font = ImageFont.load_default(size=font_size)
+
+    for (x1, y1, x2, y2), label, score in detections:
+        color = _CLASS_COLORS[label]
+        draw.rectangle((x1, y1, x2, y2), outline=color, width=3)
+
+        text = f"{label} {score:.2f}"
+        text_bbox = draw.textbbox((0, 0), text, font=font)
+        text_width = text_bbox[2] - text_bbox[0]
+        text_height = text_bbox[3] - text_bbox[1]
+        pad = 2
+
+        # Label goes above the box unless that would run off the top of the image,
+        # in which case it's drawn just inside the box instead.
+        label_top = y1 - text_height - 2 * pad
+        if label_top < 0:
+            label_top = y1
+        label_bg = (x1, label_top, x1 + text_width + 2 * pad, label_top + text_height + 2 * pad)
+
+        draw.rectangle(label_bg, fill=color)
+        draw.text((x1 + pad, label_top + pad - text_bbox[1]), text, fill="white", font=font)
+
+    return annotated