Uisng the model.py approach

config.json

@@ -4,5 +4,11 @@
     ],
     "model_type": "yolo-segmentation",
     "onnx_model": "astro-yolo11m-seg.onnx",
-    "task": "image-segmentation"
+    "task": "image-segmentation",
+    "widgets": [
+        {
+            "src": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/image-segmentation-input.jpg"
+        }
+    ],
+    "library_name": "onnxruntime"
 }

model.py

@@ -0,0 +1,29 @@
+import numpy as np
+import onnxruntime as ort
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+
+class YOLOSegmentationModel:
+    def __init__(self):
+        # Download and load the ONNX model from Hugging Face Hub
+        model_path = hf_hub_download(repo_id="rayh/astro-seg", filename="astro-yolo11m-seg.onnx")
+        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+
+    def preprocess(self, image: Image.Image):
+        # Convert image to RGB and preprocess for ONNX model
+        input_array = np.array(image.convert("RGB")).astype(np.float32)
+        input_array = np.expand_dims(input_array, axis=0)  # Add batch dimension
+        return input_array
+
+    def predict(self, image: Image.Image):
+        input_tensor = self.preprocess(image)
+        outputs = self.session.run(None, {"images": input_tensor})
+        return outputs  # Modify if needed to return bounding boxes/masks
+
+model = YOLOSegmentationModel()
+
+# HF Inference API expects a `predict` function
+def predict(image: Image.Image):
+    return model.predict(image)

modelling_yolo.py

@@ -1,22 +0,0 @@
-import torch
-import onnxruntime as ort
-from transformers import PreTrainedModel, PretrainedConfig
-from PIL import Image
-import numpy as np
-
-class YOLOConfig(PretrainedConfig):
-    model_type = "yolo-segmentation"
-
-class YOLOTransformersModel(PreTrainedModel):
-    config_class = YOLOConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.session = ort.InferenceSession(config.onnx_model, providers=["CPUExecutionProvider"])
-
-    def forward(self, images):
-        input_array = np.array(images.convert("RGB")).astype(np.float32)
-        input_array = np.expand_dims(input_array, axis=0)  # Add batch dimension
-
-        outputs = self.session.run(None, {"images": input_array})
-        return outputs  # Modify as needed to match Transformers format