Ok, try this

astro-20250209-seg.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c71465d7cd15546cdc30165ef23248b3cd8598b9d744c2383c643d3b79a1355a
-size 45236271

config.json

@@ -1 +1,8 @@
-{"model_type": "YOLOv11-Segmentation", "framework": "pytorch", "architecture": "YOLO", "image_size": 640, "num_classes": 2}
+{
+    "architectures": [
+        "YOLOTransformersModel"
+    ],
+    "model_type": "yolo-segmentation",
+    "onnx_model": "astro-yolo11m-seg.onnx",
+    "task": "image-segmentation"
+}

inference.py

@@ -1,17 +0,0 @@
-from ultralytics import YOLO
-from PIL import Image
-import torch
-
-# Load model from the same directory as this script
-MODEL_PATH = "pytorch_model.bin"
-
-class YOLOSegmentation:
-    def __init__(self):
-        self.model = YOLO(MODEL_PATH)  # Load YOLO model
-    
-    def __call__(self, image: Image.Image):
-        results = self.model(image)  # Run inference
-        return results[0].tojson()  # Convert output to JSON format
-
-# Hugging Face Inference API expects a `model` variable
-model = YOLOSegmentation()

model-index.json

@@ -1,10 +0,0 @@
-{
-  "model-index": [
-    {
-      "name": "YOLOv11 Segmentation Model",
-      "library_name": "ultralytics",
-      "pipeline_tag": "image-segmentation",
-      "tags": ["object-detection", "segmentation", "YOLO"]
-    }
-  ]
-}

modelling_yolo.py

@@ -0,0 +1,22 @@
+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

pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c71465d7cd15546cdc30165ef23248b3cd8598b9d744c2383c643d3b79a1355a
-size 45236271

requirements.txt

@@ -1,3 +1,6 @@
-ultralytics>=8.0.0
-torch>=2.0.0
-pillow
+ultralytics
+transformers
+onnxruntime
+torch
+numpy
+Pillow