Create handler.py

handler.py

@@ -0,0 +1,71 @@
+from typing import Dict, List, Any
+from ultralytics import YOLO
+import base64
+from io import BytesIO
+from PIL import Image
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the YOLO model
+        self.model = YOLO(f"{path}/FFDNet-L.pt")
+        self.id_to_cls = {0: "TextBox", 1: "ChoiceButton", 2: "Signature"}
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Args:
+            data: A dictionary containing:
+                - "inputs": base64 encoded image or image URL
+                - "parameters": optional dict with confidence, iou, imgsz
+        Returns:
+            List of predictions with bounding boxes and classes
+        """
+        # Extract image from request
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        
+        # Handle image input (base64 or URL)
+        if isinstance(inputs, str):
+            if inputs.startswith("http"):
+                image = inputs
+            else:
+                # Decode base64
+                image_data = base64.b64decode(inputs)
+                image = Image.open(BytesIO(image_data))
+        else:
+            image = inputs
+        
+        # Get parameters with defaults
+        confidence = parameters.get("conf", 0.3)
+        iou = parameters.get("iou", 0.1)
+        imgsz = parameters.get("imgsz", 1600)
+        augment = parameters.get("augment", True)
+        
+        # Run inference
+        results = self.model.predict(
+            image,
+            conf=confidence,
+            iou=iou,
+            imgsz=imgsz,
+            augment=augment
+        )
+        
+        # Format results
+        predictions = []
+        for result in results:
+            if result.boxes is not None:
+                for box in result.boxes.cpu().numpy():
+                    x, y, w, h = box.xywhn[0]
+                    cls_id = int(box.cls.item())
+                    
+                    predictions.append({
+                        "widget_type": self.id_to_cls[cls_id],
+                        "confidence": float(box.conf[0]),
+                        "bounding_box": {
+                            "cx": float(x),
+                            "cy": float(y),
+                            "w": float(w),
+                            "h": float(h)
+                        }
+                    })
+        
+        return predictions