from ultralytics import YOLO
import numpy as np
from PIL import Image

class YOLO11SegmentationModel:
    def __init__(self, model_path: str):
        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