add inference endpoint handler and requirements

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

handler.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import base64
+import io
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from models import register_ultralytics_modules
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        register_ultralytics_modules()
+
+        from ultralytics import YOLO
+
+        weights = Path(path) / "weights" / "symbolic_capsule_network_segmentation.pt"
+        self.model = YOLO(str(weights))
+
+    def __call__(self, data: dict) -> list[dict]:
+        """
+        Args:
+            data: {"inputs": <PIL Image | bytes | str path>}
+
+        Returns:
+            List of dicts compatible with HF image-segmentation pipeline:
+            [{"score": float, "label": str, "mask": "<base64 PNG>"}]
+        """
+        image = data.get("inputs")
+        if isinstance(image, bytes):
+            image = Image.open(io.BytesIO(image)).convert("RGB")
+
+        results = self.model.predict(image, imgsz=640, conf=0.25, verbose=False)
+        r = results[0]
+
+        if r.boxes is None or r.masks is None:
+            return []
+
+        h, w = r.orig_shape
+        output = []
+        for box, mask_tensor in zip(r.boxes, r.masks.data):
+            # Resize binary mask back to original image size
+            mask_np = (mask_tensor.cpu().numpy() * 255).astype(np.uint8)
+            mask_img = Image.fromarray(mask_np).resize((w, h), Image.NEAREST)
+
+            buf = io.BytesIO()
+            mask_img.save(buf, format="PNG")
+            mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+
+            output.append({
+                "score": round(float(box.conf), 4),
+                "label": self.model.names[int(box.cls)],
+                "mask": mask_b64,
+            })
+
+        return output

requirements.txt

@@ -0,0 +1,3 @@
+ultralytics>=8.4.9
+Pillow
+numpy