Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,275 Bytes
2eafbc4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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)
|