File size: 455 Bytes
2b50ca3 543627a 2b50ca3 543627a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
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
|