Upload modal_inference.py

modal_inference.py

@@ -0,0 +1,71 @@
+import modal
+import os
+
+app = modal.App("surgisight")
+
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .apt_install("libgl1", "libglib2.0-0")
+    .pip_install(
+        "ultralytics",
+        "pillow",
+        "numpy",
+        "opencv-python-headless",
+        "huggingface_hub",
+    )
+)
+
+# Cache the model weights inside the Modal image so it doesn't re-download every call
+with image.imports():
+    from ultralytics import YOLO
+    from PIL import Image as PILImage
+    import numpy as np
+    import cv2
+    import io
+
+
+@app.cls(gpu="T4", image=image, secrets=[modal.Secret.from_name("hf-secret")])
+class SurgiSightDetector:
+
+    @modal.enter()
+    def load_model(self):
+        from huggingface_hub import hf_hub_download
+        model_path = hf_hub_download(
+            repo_id="sugan04/cholec-yolo26n-seg",
+            filename="best.pt",
+            token=os.environ.get("HF_TOKEN")
+        )
+        self.model = YOLO(model_path)
+
+    @modal.method()
+    def run(self, image_bytes: bytes, conf_threshold: float = 0.25):
+        nparr = np.frombuffer(image_bytes, np.uint8)
+        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        results = self.model(frame, task="segment", conf=conf_threshold)
+        annotated = results[0].plot()
+
+        # Encode annotated image back to bytes
+        _, buffer = cv2.imencode(".png", annotated)
+        annotated_bytes = buffer.tobytes()
+
+        # Extract detections
+        boxes = results[0].boxes
+        detections = []
+        if boxes is not None and len(boxes) > 0:
+            for box in boxes:
+                detections.append({
+                    "cls_id": int(box.cls[0]),
+                    "conf": float(box.conf[0])
+                })
+
+        return {"annotated_bytes": annotated_bytes, "detections": detections}
+
+
+# For local testing
+@app.local_entrypoint()
+def main():
+    from PIL import Image as PILImage
+    import io
+    detector = SurgiSightDetector()
+    print("Modal SurgiSight detector ready.")