Upload 3 files

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:295dee8c572bf2d3e530f8bdb09537ecc4806894fb20fd5eacc034a00a387956
+size 6248810

handler.py

@@ -0,0 +1,65 @@
+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}

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.1.2
+torchvision==0.16.2
+ultralytics==8.1.0
+pillow==10.3.0
+pandas==2.2.2
+numpy==1.26.4