Create detection_service.py

services/detection_service.py

@@ -0,0 +1,32 @@
+from transformers import DetrImageProcessor, DetrForObjectDetection
+import torch
+from PIL import Image
+
+class DetectionService:
+    def __init__(self, model_name="facebook/detr-resnet-50"):
+        self.processor = DetrImageProcessor.from_pretrained(model_name, revision="no_timm")
+        self.model = DetrForObjectDetection.from_pretrained(model_name, revision="no_timm")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        self.model.eval()
+
+    def detect_objects(self, image, confidence_threshold=0.9):
+        """Detect objects in an image."""
+        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        target_sizes = torch.tensor([image.size[::-1]]).to(self.device)
+        results = self.processor.post_process_object_detection(
+            outputs, target_sizes=target_sizes, threshold=confidence_threshold
+        )[0]
+        detections = []
+        for score, label, box in zip(
+            results["scores"], results["labels"], results["boxes"]
+        ):
+            box = box.cpu().numpy().astype(int)
+            detections.append({
+                "score": score.item(),
+                "label": self.model.config.id2label[label.item()],
+                "box": {"xmin": box[0], "ymin": box[1], "xmax": box[2], "ymax": box[3]}
+            })
+        return detections