OMG-InstantID / inference /models /yolov8 /yolov8_instance_segmentation.py
Fucius's picture
Upload 422 files
2eafbc4 verified
from typing import List, Tuple
import numpy as np
from inference.core.models.instance_segmentation_base import (
InstanceSegmentationBaseOnnxRoboflowInferenceModel,
)
class YOLOv8InstanceSegmentation(InstanceSegmentationBaseOnnxRoboflowInferenceModel):
"""YOLOv8 Instance Segmentation ONNX Inference Model.
This class is responsible for performing instance segmentation using the YOLOv8 model
with ONNX runtime.
Attributes:
weights_file (str): Path to the ONNX weights file.
Methods:
predict: Performs inference on the given image using the ONNX session.
"""
@property
def weights_file(self) -> str:
"""Gets the weights file for the YOLOv8 model.
Returns:
str: Path to the ONNX weights file.
"""
return "weights.onnx"
def predict(self, img_in: np.ndarray, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
"""Performs inference on the given image using the ONNX session.
Args:
img_in (np.ndarray): Input image as a NumPy array.
Returns:
Tuple[np.ndarray, np.ndarray]: Tuple containing two NumPy arrays representing the predictions and protos. The predictions include boxes, confidence scores, class confidence scores, and masks.
"""
predictions = self.onnx_session.run(None, {self.input_name: img_in})
protos = predictions[1]
predictions = predictions[0]
predictions = predictions.transpose(0, 2, 1)
boxes = predictions[:, :, :4]
class_confs = predictions[:, :, 4:-32]
confs = np.expand_dims(np.max(class_confs, axis=2), axis=2)
masks = predictions[:, :, -32:]
predictions = np.concatenate([boxes, confs, class_confs, masks], axis=2)
return predictions, protos