Skip to content

YOLOv5 object detection

Aditya Lohia edited this page Jun 9, 2022 · 11 revisions

This guide explains how to run CVU for YOLOv5 object detection.

YOLOv5

YOLOv5 is one of the state-of-the-art Object Detection models. Please check out more about it and train your own custom models through it's official repo. CVU also supports custom weights for all the backends.

Checkout following backends for more specific information


from cvu.detector import Detector
detector = Detector(classes="coco", backend = "torch",
                    weight = "yolov5s", device = "auto"))

Detector Arguments

  • classes (Union[str, List[str]]): name of classes to be detected. It can be set to individual classes like 'coco', 'person', 'cat' etc. Alternatively, it also accepts list of classes such as ['person', 'cat']. For default models/weights, 'classes' is used to filter out objects according to provided argument from coco class. For custom models, all classes should be provided in original order as list of strings.

  • backend (str, optional): name of the backend to be used for inference purposes. Defaults to "torch".

  • weight (str, optional): path to weight files (according to selected backend). Alternatively, it also accepts identifiers (such as yolvo5s, yolov5m, etc.) to load pretrained models. Defaults to "yolov5s".

  • device (str, optional): name of the device to be used. Valid devices can be "cpu", "gpu", "cuda", "tpu", "auto". Defaults to "auto" which tries to use the device best suited for selected backend and the hardware avaibility.


Every object detector expects BGR formatted single image (batching is not supported yet), and returns a Predictions object which represents a group/list of detected objects. You can access individual detections using indexing or looping through Predictions. A single detection is represented by Prediction.

import cv2
from cvu.detector import Detector

# read image, initiate detector
img = cv2.imread("example.jpg")
detector = Detector(classes="coco")

# inference
predictions = detector(img)

# predictions info
print(predictions)

# draw on frame (inplace + returns img)
predictions.draw(img)

# class-wise counter object
print(predictions.count())

# loop through
for prediction in predictions:
    # print info
    print(prediction)

    # access specific things
    print(prediction.bbox, prediction.confidence)
    print(prediction.class_id, prediction.class_name)

    # draw specific prediction (only inplace, doesn't return img)
    prediction.draw(img)

# save img
cv2.imwrite("output.jpg", img)

These wrappers around detections provides various functionalities for drawing boxes, accessing detection info (individually as well as a group). And they implement CVU's common predictions interface for consistency.


Every Object detector is implemented in cvu.detector, following a common interface.

As of our first alpha release, we only support Yolov5s models.


Precission Accuracy (Yolov5)

Index

For a successful deployment, precission is also important (i.e. closness to results of native-framework of trained weights). We will add more numerical results soon, for now we provide image comparison.

Torch


ONNX


TensorFlow


TFLite

Clone this wiki locally