Update handler.py

handler.py

@@ -1,65 +1,72 @@
-from typing import Dict, Any
-import base64
-import io
-from PIL import Image
-import torch
-
-class EndpointHandler:
-    """
-    Hugging Face Inference Endpoint handler for LocustGuard YOLO model.
-
-    Expected JSON:
-    {
-      "image": "<base64 encoded image>",
-      "conf": 0.25,
-      "iou": 0.45
-    }
-
-    Returns:
-    {
-      "detections": [
-        {
-          "label": str,
-          "confidence": float,
-          "coordinates": [xmin, ymin, xmax, ymax]
-        }
-      ]
-    }
-    """
-
-    def __init__(self, path: str = "."):
-        self.model = torch.hub.load(
-            "ultralytics/yolov5",
-            "custom",
-            path=f"{path}/best.pt",
-            force_reload=False
-        )
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        image_b64 = data["image"]
-        conf = float(data.get("conf", 0.25))
-        iou = float(data.get("iou", 0.45))
-
-        image_bytes = base64.b64decode(image_b64)
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-        self.model.conf = conf
-        self.model.iou = iou
-
-        results = self.model(image, size=640)
-        r = results.pandas().xyxy[0]
-
-        detections = []
-        for _, row in r.iterrows():
-            detections.append({
-                "label": row["name"],
-                "confidence": float(row["confidence"]),
-                "coordinates": [
-                    round(float(row["xmin"]), 2),
-                    round(float(row["ymin"]), 2),
-                    round(float(row["xmax"]), 2),
-                    round(float(row["ymax"]), 2),
-                ]
-            })
-
-        return {"detections": detections}
+from typing import Dict, Any
+import base64
+import io
+from PIL import Image
+from ultralytics import YOLO
+
+class EndpointHandler:
+    """
+    Hugging Face Inference Endpoint handler for LocustGuard YOLO model.
+
+    Expected JSON:
+    {
+      "image": "<base64 encoded image>",
+      "conf": 0.25,
+      "iou": 0.45
+    }
+
+    Returns:
+    {
+      "detections": [
+        {
+          "label": str,
+          "confidence": float,
+          "coordinates": [xmin, ymin, xmax, ymax]
+        }
+      ]
+    }
+    """
+
+    def __init__(self, path: str = "."):
+        # ✅ Native ultralytics loader (NO torch.hub)
+        self.model = YOLO(f"{path}/best.pt")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        image_b64 = data["image"]
+        conf = float(data.get("conf", 0.25))
+        iou = float(data.get("iou", 0.45))
+
+        image_bytes = base64.b64decode(image_b64)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+        # Run inference
+        results = self.model(
+            image,
+            conf=conf,
+            iou=iou,
+            imgsz=640,
+            verbose=False
+        )
+
+        r = results[0]
+        detections = []
+
+        if r.boxes is not None:
+            for box in r.boxes:
+                x1, y1, x2, y2 = box.xyxy[0].tolist()
+                cls_id = int(box.cls[0])
+                conf_score = float(box.conf[0])
+                label = self.model.names[cls_id]
+
+                detections.append({
+                    "label": label,
+                    "confidence": round(conf_score, 3),
+                    "coordinates": [
+                        round(float(x1), 2),
+                        round(float(y1), 2),
+                        round(float(x2), 2),
+                        round(float(y2), 2),
+                    ]
+                })
+
+        return {"detections": detections}