Skip to content

salihmarangoz/yolo_segmentation_ros

Repository files navigation

yolo_segmentation_ros

Sweet pepper detection+segmentation+tracking via YOLOv8. Also an alternative easy replacement to agrobot_mrcnn_ros and mask_rcnn_ros packages. Watch a quick experiment on YouTube (yolov8l-seg-V1):

Watch the video

Also for yolov8n-seg-V2: https://youtu.be/UmDr89TXKYQ

Also for yolov8l-seg-V1 but with agrobot_mrcnn_ros visualization style: https://youtu.be/kM_W8cpG5QE

Installation

Choose a or b for installing the Python dependencies.

a) User Python Env

Easiest way:

$ pip install -r requirements.txt

b) Virtual Python Env

If you don't want to mess up with your current environment.

$ mkdir "$HOME/.virtualenvs"
$ cd "$HOME/.virtualenvs"
$ python3 -m venv yolo_segmentation_ros
$ source yolo_segmentation_ros/bin/activate
$ cd ~/catkin_ws/src/yolo_segmentation_ros
$ pip install -r requirements.txt

If you choose this option, don't forget to update yolo_segmentation_ros_python param in the launch/start.launch file.

Download Weights

Read First! This table is copied from here. mAPval values are for single-model single-scale on COCO val2017 dataset. Speed averaged over COCO val images using an Amazon EC2 P4d instance.

Model mAPbox 50-95 mAPmask 50-95 Speed CPU ONNX (ms) Speed A100 TensorRT (ms) params (M) FLOPs (B)
YOLOv8n-seg 36.7 30.5 96.1 1.21 3.4 12.6
YOLOv8s-seg 44.6 36.8 155.7 1.47 11.8 42.6
YOLOv8m-seg 49.9 40.8 317.0 2.18 27.3 110.2
YOLOv8l-seg 52.3 42.6 572.4 2.79 46.0 220.5
YOLOv8x-seg 53.4 43.4 712.1 4.02 71.8 344.1

Download Our Pretrained Models: Models are named in yolov8[SIZE]-seg-[VERSION]. Always prefer latest version of that specific model size. Sizes are nano (n), small (s), medium (m), large (l), x-large (x). DOWNLOAD the weights/best.pt and place it into the model folder. Small models run faster but worse, large models run slower but better. Currently recommended best large model for GPU is yolov8l-seg-V1, best nano model for CPU is yolov8n-seg-V2.

Tested on;

  • Ubuntu 20.04 + ROS Noetic
  • Ryzen 5900HX + RTX3080M 16GB
  • Nvidia Driver 525.105.17
  • CUDA 11.7
  • ultralytics 8.0.114

Training

You can skip this step if you would like to use pre-trained models for sweet pepper detection.

This dataset also may be available here. If you can download the dataset you can skip to the training part.

$ cd yolo_segmentation_ros/train

Prepare Dataset: Follow the steps below.

  1. (Mandatory) Prepare BUP20
    1. Download BUP20 dataset and extract BUP20.tar.gz into a folder. (it should have the name CKA_sweet_pepper_2020_summer)
    2. Modify DATASET_PATH at the top of the notebook then run the notebook train/bup20_to_yolov8.ipynb to create yolov8 dataset.
  2. (Optional) Prepare Sweet Pepper and Peduncle Segmentation Dataset from Kaggle:
    1. Download Sweet Pepper and Peduncle Segmentation Dataset from Kaggle and extract into a folder. (it should have the name dataset_620_red_yellow_cart_only)
    2. Modify DATASET_PATH at the top of the notebook then run the notebook train/dataset_620_to_yolov8.ipynb to create yolov8 dataset.
  3. (Mandatory) Split train/validation using split_train_val.ipynb

Training: After these steps, the dataset should be ready. For training the model:

  1. Run the script train/train.py for training. This should take about 15-20 mins.
  2. Copy the train/runs/segment/train/weights/best.pt to model/best.pt

Prediction

$ roscd yolo_segmentation_ros/predict

# Predict a region of your screen:
$ python predict_screen.py

# Predict using your web-cam:
$ python predict_webcam.py

# Predict using a video:
$ python predict_video.py

# Predict using images:
$ python predict_images.py

ROS Node

ROS Dependencies

These dependencies are optional. Outputs will be enabled if the related dependency is installed.

Demo

The ROS node script is located here: scripts/ros_node.py

$ sudo apt install ros-noetic-video_stream_opencv # Install dependency
$ roslaunch yolo_segmentation_ros demo.launch # DEMO (with User Python Env)
$ roslaunch yolo_segmentation_ros demo.launch yolo_segmentation_ros_python:="$HOME/.virtualenvs/yolo_segmentation_ros/bin/python3" # DEMO (with Virtual Python Env)

FAQ

  • AGPL-3.0 License!?

  • Good resource for datasets?

    • Kaggle is a source for datasets. But datasets may need manual work for adaptation to YOLOv8 format.
    • Roboflow is very easy to use and supports exporting datasets in YOLOv8 format. But I can personally say that annotation quality is usually not good.
  • How to convert other annotation types to YOLOv8 format?

    • JSON2YOLO may be useful. If you are planning to do manual conversion, take into account that each instance in YOLOv8 format includes a single segment. So, there is a small trick to include instances with multiple semantic segments.
  • Any labelling tools or helpers?

    • ultralytics has automatic annotation feature. Check autoannotate folder for scripts. For more information visit this page.

      Original Auto-annotated
      2 2_vis