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)