import onnxruntime as ort
import numpy as np
from PIL import Image

class YOLOSegmentationModel:
    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