Skip to content

Latest commit

 

History

History
101 lines (79 loc) · 4.16 KB

object-detection-2d-yolov5.md

File metadata and controls

101 lines (79 loc) · 4.16 KB
Raw

YOLOv5DetectorLearner module

The yolov5 module contains the YOLOv5DetectorLearner class, which inherits from the abstract class Learner.

Class YOLOv5DetectorLearner

Bases: engine.learners.Learner

The YOLOv5DetectorLearner class is a wrapper of the YOLO detector[1] Ultralytics implementation based on its availability in the Pytorch Hub. It can be used to perform object detection on images (inference only).

The YOLOv5DetectorLearner class has the following public methods:

YOLOv5DetectorLearner constructor

YOLOv5DetectorLearner(self, model_name, path, device)

Constructor parameters:

  • model_name: str
    Specifies the name of the model to be used. Available models:
    • 'yolov5n' (46.0% mAP, 1.9M parameters)
    • 'yolov5s' (56.0% mAP, 7.2M parameters)
    • 'yolov5m' (63.9% mAP, 21.2M parameters)
    • 'yolov5l' (67.2% mAP, 46.5M parameters)
    • 'yolov5x' (68.9% mAP, 86.7M parameters)
    • 'yolov5n6' (50.7% mAP, 3.2M parameters)
    • 'yolov5s6' (63.0% mAP, 16.8M parameters)
    • 'yolov5m6' (69.0% mAP, 35.7 parameters)
    • 'yolov5l6' (71.6% mAP, 76.8M parameters)
    • 'custom' (for custom models, the path parameter must be set to point to the location of the weights file.) Note that mAP (0.5) is reported on the COCO val2017 dataset.
  • path: str, default=None
    For custom-trained models, specifies the path to the weights to be loaded.
  • device: {'cuda', 'cpu'}, default='cuda' Specifies the device used for inference.
  • temp_path: str, default='.'
    Specifies the path to where the weights will be downloaded when using pretrained models.
  • force_reload: bool, default=False
    Sets the force_reload parameter of the pytorch hub load method. This fixes issues with caching when set to True.

YOLOv5DetectorLearner.infer

The infer method:

YOLOv5DetectorLearner.infer(self, img)

Performs inference on a single image.

Parameters:

  • img: object
    Object of type engine.data.Image or OpenCV.
  • size: int, default=640
    Size of image for inference. The image is resized to this in both sides before being fed to the model.

YOLOv5DetectorLearner.download

YOLOv5DetectorLearner.download(self, path, mode, verbose, url, model_name, img_name)

Downloads the pretrained weights of a YOLOv5s model fine-tuned for truck detection, along with sample truck images for inference, stored in .pt and image files respectively.

Parameters:

  • path: str, default=None
    Specifies the folder where data will be downloaded. If None, the self.temp_path directory is used instead.
  • mode: {'pretrained', 'images'}, default='pretrained'
    If 'pretrained', downloads a pretrained detector model. If 'images', downloads an image to perform inference on.
  • verbose: bool default=False
    If True, enables maximum verbosity.
  • url: str, default=OpenDR FTP URL
    URL of the FTP server.
  • model_name: name of model ftp server, currently only supports yolov5s_finetuned_in_trucks.pt, default = 'yolov5s_finetuned_in_trucks.pt'.\
  • image_name: name of image in ftp server, available files are truckX.jpg for X=1 to 10, default = 'truck7.jpg'.\

Examples

  • Inference and result drawing example on a test .jpg image using OpenCV: 
    import torch
from opendr.engine.data import Image
from opendr.perception.object_detection_2d import YOLOv5DetectorLearner
from opendr.perception.object_detection_2d import draw_bounding_boxes

yolo = YOLOv5DetectorLearner(model_name='yolov5s', device='cpu')

torch.hub.download_url_to_file('https://ultralytics.com/images/zidane.jpg', 'zidane.jpg')  # download image
im1 = Image.open('zidane.jpg')  # OpenDR image

results = yolo.infer(im1)
draw_bounding_boxes(im1.opencv(), results, yolo.classes, show=True, line_thickness=3)

References

[1] YOLOv5: The friendliest AI architecture you'll ever use.