Update handler.py

handler.py

@@ -4,27 +4,32 @@ import io
 from PIL import Image
 from ultralytics import YOLO
 
+
 class EndpointHandler:
     """
     Hugging Face Inference Endpoint handler for LocustGuard YOLO model.
 
-    Accepts BOTH formats:
+    Accepts:
+    1) Direct HTTP / Spaces:
+       {
+         "image": "<base64>",
+         "conf": 0.25,
+         "iou": 0.45
+       }
 
-    Direct HTTP / Spaces:
-    {
-      "image": "<base64>",
-      "conf": 0.25,
-      "iou": 0.45
-    }
+    2) Playground / Hosted API:
+       {
+         "inputs": {
+           "image": "<base64>",
+           "conf": 0.25,
+           "iou": 0.45
+         }
+       }
 
-    Playground / Hosted API:
-    {
-      "inputs": {
-        "image": "<base64>",
-        "conf": 0.25,
-        "iou": 0.45
-      }
-    }
+    3) HF standard:
+       {
+         "inputs": "<base64>"
+       }
 
     Returns:
     {
@@ -39,27 +44,38 @@ class EndpointHandler:
     """
 
     def __init__(self, path: str = "."):
-        # ✅ Native ultralytics loader (NO torch.hub)
+        # ✅ Native ultralytics loader
         self.model = YOLO(f"{path}/best.pt")
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        # ✅ Compatibility layer (does NOT break anything)
+        # ---------------- NORMALIZE INPUT ----------------
         payload = data.get("inputs", data)
 
-        # ✅ Robust guard: prevents KeyError and hard crashes
-        if not isinstance(payload, dict) or "image" not in payload:
+        # Case 1: HF standard sends raw base64 string
+        if isinstance(payload, str):
+            image_b64 = payload
+            conf = 0.25
+            iou = 0.45
+
+        # Case 2: Dict-based payload
+        elif isinstance(payload, dict) and "image" in payload:
+            image_b64 = payload["image"]
+            conf = float(payload.get("conf", 0.25))
+            iou = float(payload.get("iou", 0.45))
+
+        else:
             return {
-                "error": "Missing required key: 'image'. Expected base64-encoded image."
+                "error": "Invalid input. Expected base64 image under key 'image' or 'inputs'."
             }
 
-        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")
+        # ---------------- DECODE IMAGE ----------------
+        try:
+            image_bytes = base64.b64decode(image_b64)
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to decode image: {str(e)}"}
 
-        # Run inference
+        # ---------------- RUN INFERENCE ----------------
         results = self.model(
             image,
             conf=conf,