Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,193 Bytes
df6c67d |
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 |
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,)
|