Upload model.py with huggingface_hub

model.py

@@ -1,13 +1,16 @@
-import onnxruntime as ort
+from ultralytics import YOLO
 import numpy as np
 from PIL import Image
 
-class YOLOSegmentationModel:
+class YOLOv8SegmentationModel:
     def __init__(self, model_path: str):
-        self.session = ort.InferenceSession(model_path)
-
-    def predict(self, image: Image):
-        input_data = np.array(image).astype(np.float32)
-        input_data = np.expand_dims(input_data, axis=0)  # Add batch dimension
-        outputs = self.session.run(None, {"images": input_data})
-        return outputs
+        self.model = YOLO(model_path)
+        
+    def predict(self, image: Image.Image):
+        # Convert PIL Image to numpy array if needed
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        
+        # Run inference
+        results = self.model(image, task='segment')
+        return results[0]  # Return first batch result