File size: 505 Bytes
5957194 543627a 2b50ca3 83c31a9 543627a 5957194 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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
|