Skip to content

yolov5_object_detection

YOLOv5ObjectDetection

Bases: ObjectDetectionBaseOnnxRoboflowInferenceModel

Roboflow ONNX Object detection model (Implements an object detection specific infer method).

This class is responsible for performing object detection using the YOLOv5 model with ONNX runtime.

Attributes:

Name Type Description
weights_file str

Path to the ONNX weights file.

Source code in inference/models/yolov5/yolov5_object_detection.py
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
class YOLOv5ObjectDetection(ObjectDetectionBaseOnnxRoboflowInferenceModel):
    """Roboflow ONNX Object detection model (Implements an object detection specific infer method).

    This class is responsible for performing object detection using the YOLOv5 model
    with ONNX runtime.

    Attributes:
        weights_file (str): Path to the ONNX weights file.
    """

    @property
    def weights_file(self) -> str:
        """Gets the weights file for the YOLOv5 model.

        Returns:
            str: Path to the ONNX weights file.
        """
        return "yolov5s_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.
        """
        predictions = self.onnx_session.run(None, {self.input_name: img_in})[0]
        return (predictions,)

weights_file: str property

Gets the weights file for the YOLOv5 model.

Returns:

Name Type Description
str str

Path to the ONNX weights file.

predict(img_in, **kwargs)

Performs object detection on the given image using the ONNX session.

Parameters:

Name Type Description Default
img_in ndarray

Input image as a NumPy array.

required

Returns:

Type Description
Tuple[ndarray]

Tuple[np.ndarray]: NumPy array representing the predictions.

Source code in inference/models/yolov5/yolov5_object_detection.py
29
30
31
32
33
34
35
36
37
38
39
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.
    """
    predictions = self.onnx_session.run(None, {self.input_name: img_in})[0]
    return (predictions,)