detection.py

from __future__ import annotations

from functools import lru_cache
from typing import Any

from PIL import Image


@lru_cache(maxsize=1)
def _load_detr(model_name: str = "facebook/detr-resnet-50", revision: str = "no_timm"):
    try:
        from transformers import DetrForObjectDetection, DetrImageProcessor
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "Missing dependency. Install requirements with: "
            "pip install -r object_detection/requirements.txt"
        ) from exc

    processor = DetrImageProcessor.from_pretrained(model_name, revision=revision)
    model = DetrForObjectDetection.from_pretrained(model_name, revision=revision)
    model.eval()
    return processor, model


def detect_objects(
    image: Image.Image,
    *,
    threshold: float = 0.9,
    model_name: str = "facebook/detr-resnet-50",
    revision: str = "no_timm",
    device: str | None = None,
) -> list[dict[str, Any]]:
    try:
        import torch
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "PyTorch is required for detection. Install it (and other deps) with: "
            "pip install -r object_detection/requirements.txt"
        ) from exc

    def _is_no_kernel_image_error(exc: Exception) -> bool:
        msg = str(exc).lower()
        return "no kernel image is available for execution on the device" in msg or (
            "cudaerrornokernelimagefordevice" in msg
        )

    def _run_on_device(run_device: str) -> list[dict[str, Any]]:
        processor, model = _load_detr(model_name=model_name, revision=revision)
        model = model.to(run_device)

        inputs = processor(images=image, return_tensors="pt")
        inputs = {k: v.to(run_device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)

        target_sizes = torch.tensor([image.size[::-1]], device=run_device)
        results = processor.post_process_object_detection(
            outputs, target_sizes=target_sizes, threshold=threshold
        )[0]

        rows: list[dict[str, Any]] = []
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            box_vals = [round(float(i), 2) for i in box.tolist()]
            rows.append(
                {
                    "label": model.config.id2label[int(label.item())],
                    "confidence": round(float(score.item()), 3),
                    "xmin": box_vals[0],
                    "ymin": box_vals[1],
                    "xmax": box_vals[2],
                    "ymax": box_vals[3],
                }
            )
        return rows

    if device is None:
        preferred = "cuda" if torch.cuda.is_available() else "cpu"
        if preferred == "cpu":
            return _run_on_device("cpu")

        try:
            return _run_on_device("cuda")
        except RuntimeError as exc:
            if _is_no_kernel_image_error(exc):
                # Common on GPUs whose compute capability isn't supported by the installed PyTorch CUDA build.
                # Fall back to CPU for usability.
                return _run_on_device("cpu")
            raise

    if device not in {"cpu", "cuda"}:
        raise ValueError("device must be 'cpu', 'cuda', or None (auto).")

    try:
        return _run_on_device(device)
    except RuntimeError as exc:
        if device == "cuda" and _is_no_kernel_image_error(exc):
            raise RuntimeError(
                "Your CUDA-enabled PyTorch build does not support this GPU (no kernel image available). "
                "Pick CPU in the UI, or install a PyTorch build compatible with your GPU/CUDA/driver."
            ) from exc
        raise


if __name__ == "__main__":
    image = Image.open("test.jpg").convert("RGB")
    for row in detect_objects(image, threshold=0.9):
        print(
            f"Detected {row['label']} with confidence {row['confidence']} "
            f"at location {[row['xmin'], row['ymin'], row['xmax'], row['ymax']]}"
        )