|
from typing import Tuple |
|
|
|
import numpy as np |
|
|
|
from inference.core.models.object_detection_base import ( |
|
ObjectDetectionBaseOnnxRoboflowInferenceModel, |
|
) |
|
|
|
|
|
class YOLONASObjectDetection(ObjectDetectionBaseOnnxRoboflowInferenceModel): |
|
box_format = "xyxy" |
|
|
|
@property |
|
def weights_file(self) -> str: |
|
"""Gets the weights file for the YOLO-NAS 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}) |
|
boxes = predictions[0] |
|
class_confs = predictions[1] |
|
confs = np.expand_dims(np.max(class_confs, axis=2), axis=2) |
|
predictions = np.concatenate([boxes, confs, class_confs], axis=2) |
|
return (predictions,) |
|
|
