Upload folder using huggingface_hub

.gitignore

@@ -208,4 +208,6 @@ __marimo__/
 
 *.jpg`
 runs/
-data/
+data/
+
+*yolov11-segmentation_earth-worm/

handler.py

@@ -0,0 +1,100 @@
+# handler.py
+# Hugging Face Inference Endpoints - Custom Handler for Ultralytics YOLOv11-seg
+# Returns: {"instances":[{"label":str,"score":float,"polygon":[[x,y],...]},...],
+#           "width": int, "height": int}
+
+import io
+import base64
+from typing import Any, Dict, List, Union
+
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+
+class EndpointHandler:
+    def __init__(self, path: str = "."):
+        """
+        Called once on container startup.
+        `path` points to the repo root mounted in the endpoint.
+        """
+        # Resolve weights using Hub API to get the raw binary (handles LFS/private).
+        self.repo_id = "dashingzombie/yolov11-segmentation_earth-worm"
+        self.filename = "best.pt"  # change if you prefer last.pt
+
+        weights_path = hf_hub_download(
+            repo_id=self.repo_id,
+            filename=self.filename,
+            repo_type="model"
+        )
+        self.model = YOLO(weights_path)
+
+        # If class names were not baked into the checkpoint, you can force them:
+        if not getattr(self.model, "names", None):
+            self.model.names = {0: "body_mask"}  # single-class fallback
+
+    def _to_image(self, payload: Dict[str, Any]) -> Image.Image:
+        """
+        Accepts either:
+          - {"inputs": {"image": <base64-string>}}  (Serverless-style)
+          - {"inputs": <base64-string>}
+          - {"image_bytes": <raw-bytes>}            (Toolkit raw)
+        """
+        if "image_bytes" in payload:
+            return Image.open(io.BytesIO(payload["image_bytes"])).convert("RGB")
+
+        inputs = payload.get("inputs", payload.get("instances", None))
+        if isinstance(inputs, dict):
+            img_b64 = inputs.get("image") or inputs.get("img") or inputs.get("data")
+        else:
+            img_b64 = inputs
+
+        if isinstance(img_b64, str):
+            # strip possible 'data:image/jpeg;base64,' prefix
+            if "," in img_b64:
+                img_b64 = img_b64.split(",", 1)[1]
+            data = base64.b64decode(img_b64)
+            return Image.open(io.BytesIO(data)).convert("RGB")
+
+        raise ValueError("No image provided. Expected 'image_bytes' or base64 string under 'inputs'.")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Runs per request. `data` is the incoming JSON/body parsed by the Toolkit.
+        Returns JSON-serializable dict.
+        """
+        image = self._to_image(data)
+        W, H = image.size
+
+        # confidence threshold can be overridden via params
+        params = data.get("parameters", {}) or data.get("options", {})
+        conf = float(params.get("conf", 0.25))
+
+        results = self.model(image, conf=conf, verbose=False)[0]
+        names = results.names
+
+        instances: List[Dict[str, Any]] = []
+        if results.masks is not None:
+            # polygons per instance: results.masks.xy (list of Nx2 arrays)
+            for i, poly in enumerate(results.masks.xy):
+                cls_id = int(results.boxes.cls[i].item())
+                score = float(results.boxes.conf[i].item())
+                polygon = [[float(x), float(y)] for x, y in poly]
+                instances.append({
+                    "label": names[cls_id],
+                    "score": score,
+                    "polygon": polygon
+                })
+        else:
+            # Fallback to boxes if masks missing (rare for -seg)
+            for i, b in enumerate(results.boxes.xyxy.tolist()):
+                x1, y1, x2, y2 = [float(v) for v in b]
+                cls_id = int(results.boxes.cls[i].item())
+                score = float(results.boxes.conf[i].item())
+                instances.append({
+                    "label": names[cls_id],
+                    "score": score,
+                    "bbox_xyxy": [x1, y1, x2, y2]
+                })
+
+        return {"instances": instances, "width": W, "height": H}

requiremnts.txt

@@ -1,4 +1,6 @@
 ultralytics>=8.3
 torch
 torchvision
-pillow
+pillow
+huggingface_hub
+fastapi