detection.py

"""Object detection using facebook/detr-resnet-50.

The model is loaded once at import time so the (slow) cold start happens at boot
rather than on the first request. For the MVP we use the high-level
`object-detection` pipeline; a later version can switch to DetrImageProcessor /
DetrForObjectDetection for finer control.
"""

import torch
from PIL import Image, ImageOps
from transformers import pipeline

MODEL_NAME = "facebook/detr-resnet-50"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Downscale very large uploads before inference. Keeps CPU latency and memory
# bounded; boxes are reported in this (resized) space and the frontend scales by
# image_size, so the visualization stays correct.
MAX_SIDE = 1000
DEFAULT_THRESHOLD = 0.7

# Loaded once at module import.
_detector = pipeline(
    "object-detection",
    model=MODEL_NAME,
    device=0 if DEVICE == "cuda" else -1,
)


def _prepare(image: Image.Image) -> Image.Image:
    """Normalize orientation, ensure RGB, and bound the largest side."""
    image = ImageOps.exif_transpose(image)  # respect camera rotation
    image = image.convert("RGB")
    w, h = image.size
    longest = max(w, h)
    if longest > MAX_SIDE:
        scale = MAX_SIDE / longest
        image = image.resize((round(w * scale), round(h * scale)))
    return image


def detect(image: Image.Image, threshold: float = DEFAULT_THRESHOLD) -> dict:
    """Run object detection on a PIL image.

    Returns a dict shaped for the frontend:
        {
          "image_size": {"width": int, "height": int},
          "objects": [
            {"label": str, "score": float, "box": {xmin, ymin, xmax, ymax}},
            ...
          ]
        }
    Boxes are integer pixel coordinates in the (possibly resized) image space,
    matching image_size so the frontend can scale them onto its rendered image.
    """
    image = _prepare(image)
    width, height = image.size

    raw = _detector(image)

    objects = []
    for item in raw:
        score = float(item["score"])
        if score < threshold:
            continue
        box = item["box"]
        # Clamp to image bounds so overlay boxes never spill past the edges.
        xmin = max(0, min(int(round(box["xmin"])), width))
        ymin = max(0, min(int(round(box["ymin"])), height))
        xmax = max(0, min(int(round(box["xmax"])), width))
        ymax = max(0, min(int(round(box["ymax"])), height))
        if xmax <= xmin or ymax <= ymin:
            continue
        objects.append(
            {
                "label": item["label"],
                "score": round(score, 4),
                "box": {"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax},
            }
        )

    # Highest-confidence objects first.
    objects.sort(key=lambda o: o["score"], reverse=True)

    return {
        "image_size": {"width": width, "height": height},
        "objects": objects,
    }