Update handler.py

handler.py

@@ -8,13 +8,24 @@ class EndpointHandler:
     """
     Hugging Face Inference Endpoint handler for LocustGuard YOLO model.
 
-    Expected JSON:
+    Accepts BOTH formats:
+
+    Direct HTTP / Spaces:
     {
-      "image": "<base64 encoded image>",
+      "image": "<base64>",
       "conf": 0.25,
       "iou": 0.45
     }
 
+    Playground / Hosted API:
+    {
+      "inputs": {
+        "image": "<base64>",
+        "conf": 0.25,
+        "iou": 0.45
+      }
+    }
+
     Returns:
     {
       "detections": [
@@ -32,9 +43,12 @@ class EndpointHandler:
         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))
+        # ✅ Compatibility layer (does NOT break anything)
+        payload = data.get("inputs", data)
+
+        image_b64 = payload["image"]
+        conf = float(payload.get("conf", 0.25))
+        iou = float(payload.get("iou", 0.45))
 
         image_bytes = base64.b64decode(image_b64)
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")