routes/DL_FasterRCNN_ObjectDetection.py

from io import BytesIO
from typing import Any

from fastapi import APIRouter, File, HTTPException, UploadFile

router = APIRouter(tags=["Machine Learning"])


# Lazy-loaded model state, same shape as the image classifier: pretrained
# Faster R-CNN (ResNet-50 FPN backbone, COCO) weights download to the torch hub
# cache on first request. The ResNet-50 backbone is heavier/slower on CPU than
# the MobileNetV3 variant but noticeably sharper (better localization + fewer
# misclassifications), which we prefer for the demo.
MODEL_STATE: dict[str, Any] = {
    "model": None,
    "labels": None,
    "transforms": None,
    "error": None,
}

MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10 MB
SCORE_THRESHOLD = 0.5
MAX_DETECTIONS = 30


def _ensure_model_loaded() -> None:
    if MODEL_STATE["model"] is not None:
        return

    try:
        from torchvision.models.detection import (
            FasterRCNN_ResNet50_FPN_Weights,
            fasterrcnn_resnet50_fpn,
        )

        weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
        model = fasterrcnn_resnet50_fpn(weights=weights)
        model.eval()

        MODEL_STATE["model"] = model
        MODEL_STATE["labels"] = list(weights.meta["categories"])
        MODEL_STATE["transforms"] = weights.transforms()
        MODEL_STATE["error"] = None
    except Exception as e:
        MODEL_STATE["error"] = str(e)
        raise


@router.post("/models/object-detect", summary="Detect objects with Faster R-CNN (COCO)")
async def detect_objects(file: UploadFile = File(...)):
    content_type = (file.content_type or "").lower()
    if not content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    raw = await file.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Uploaded image is empty.")
    if len(raw) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=400, detail="Image exceeds the 10 MB size limit.")

    try:
        _ensure_model_loaded()
    except Exception:
        detail = "Model not loaded."
        if MODEL_STATE["error"]:
            detail = f"Model not loaded: {MODEL_STATE['error']}"
        return {"error": detail, "status": 500}

    import torch
    from PIL import Image, UnidentifiedImageError

    try:
        image = Image.open(BytesIO(raw)).convert("RGB")
    except UnidentifiedImageError:
        raise HTTPException(status_code=400, detail="Could not decode the uploaded image.")

    model = MODEL_STATE["model"]
    labels = MODEL_STATE["labels"]
    transforms = MODEL_STATE["transforms"]
    if model is None or labels is None or transforms is None:
        return {"error": "Model not loaded.", "status": 500}

    width, height = image.size
    input_tensor = transforms(image)
    with torch.no_grad():
        output = model([input_tensor])[0]

    detections = []
    for box, score, label_idx in zip(
        output["boxes"].tolist(),
        output["scores"].tolist(),
        output["labels"].tolist(),
    ):
        if score < SCORE_THRESHOLD:
            continue  # scores are sorted descending, so once below threshold we can stop
        x1, y1, x2, y2 = box
        detections.append(
            {
                "label": labels[int(label_idx)],
                "score": float(score),
                "box": [x1, y1, x2, y2],
            }
        )
        if len(detections) >= MAX_DETECTIONS:
            break

    # Box coordinates are in original-image pixels; width/height let the client
    # scale them to whatever size the image is displayed at.
    return {"detections": detections, "width": width, "height": height}