grassjelly / yolov5_ros

yolov5_ros

Colconized yolov5 Python library for easier inference deployment on a ROS2 system.

Installation

pip3 install torch torchvision
cd <your_colcon_ws>
rosdep install --from-paths src -i -r -y
colcon build

Usage

Import inference from yolov5_ros and call the predict function.

from yolov5_ros.inference import Yolov5

model = Yolov5(
    'test.pt',
    'test.yaml',
    conf_thres=0.4,
    iou_thres=0.5,
    device=0
)

image = image.imread('test.jpg')
classes, bounding_boxes = model.predict(image)