|
from typing import Tuple |
|
|
|
import numpy as np |
|
|
|
from inference.core.exceptions import ModelArtefactError |
|
from inference.core.models.keypoints_detection_base import ( |
|
KeypointsDetectionBaseOnnxRoboflowInferenceModel, |
|
) |
|
from inference.core.models.utils.keypoints import superset_keypoints_count |
|
|
|
|
|
class YOLOv8KeypointsDetection(KeypointsDetectionBaseOnnxRoboflowInferenceModel): |
|
"""Roboflow ONNX keypoints detection model (Implements an object detection specific infer method). |
|
|
|
This class is responsible for performing keypoints detection using the YOLOv8 model |
|
with ONNX runtime. |
|
|
|
Attributes: |
|
weights_file (str): Path to the ONNX weights file. |
|
|
|
Methods: |
|
predict: Performs object detection 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, ...]: |
|
"""Performs object detection on the given image using the ONNX session. |
|
|
|
Args: |
|
img_in (np.ndarray): Input image as a NumPy array. |
|
|
|
Returns: |
|
Tuple[np.ndarray]: NumPy array representing the predictions, including boxes, confidence scores, and class confidence scores. |
|
""" |
|
predictions = self.onnx_session.run(None, {self.input_name: img_in})[0] |
|
predictions = predictions.transpose(0, 2, 1) |
|
boxes = predictions[:, :, :4] |
|
number_of_classes = len(self.get_class_names) |
|
class_confs = predictions[:, :, 4 : 4 + number_of_classes] |
|
keypoints_detections = predictions[:, :, 4 + number_of_classes :] |
|
confs = np.expand_dims(np.max(class_confs, axis=2), axis=2) |
|
bboxes_predictions = np.concatenate( |
|
[boxes, confs, class_confs, keypoints_detections], axis=2 |
|
) |
|
return (bboxes_predictions,) |
|
|
|
def keypoints_count(self) -> int: |
|
"""Returns the number of keypoints in the model.""" |
|
if self.keypoints_metadata is None: |
|
raise ModelArtefactError("Keypoints metadata not available.") |
|
return superset_keypoints_count(self.keypoints_metadata) |
|
|
