YOLOv5DS

Multi-task yolov5 with detection and segmentation based on yolov5(branch v6.0)

  • decoupled head
  • anchor free
  • segmentation head

README中文

Ablation experiment

All experiments is trained on a small dataset with 47 classes ,2.6k+ images for training and 1.5k+ images for validation:

model P R [email protected] [email protected]:95
yolov5s 0.536 0.368 0.374 0.206
yolov5s+train scrach 0.452 0.314 0.306 0.152
yolov5s+decoupled head 0.555 0.375 0.387 0.214
yolov5s + decoupled head+class balance weights 0.541 0.392 0.396 0.217
yolov5s + decoupled head+class balance weights 0.574 0.386 0.403 0.22
yolov5s + decoupled head+seghead 0.533 0.383 0.396 0.212

The baseline model is yolov5s. and decoupled head, add class balance weights all helps to improve MAP.

Adding a segmentation head can still get equivalent MAP as single detection model.

Training Method

python trainds.py

As VOC dataset do not offer the box labels and mask labels, so we forward this model with a detection batch and a segmention batch , and accumulate the gradient , than update the whole model parameters.

MAP

To compare with the SSD512, we use VOC07+12 training set as the detection training set, VOC07 test data as detection test data, for segmentation ,we use VOC12 segmentation datset as training and test set.

the input size is 512(letter box).

model VOC2007 test
SSD512 79.8
yolov5s+seghead(512) 79.2

The above results only trained less than 200 epoch, weights

demo

see detectds.py.

Train custom data

  1. Use labelme to label box and mask on your dataset;

    the box label format is voc, you can use voc2yolo.py to convert to yolo format,

    the mask label is json files , you should convert to mask .png image labels,like VOC2012 segmentation labels.

  2. see how to arrange your detection dataset with yolov5 , then arrange your segmentation dataset same as yolo files , see data/voc.yaml:

    
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: .  # dataset root dir
train: VOC/det/images/train  # train images (relative to 'path') 118287 images
val: VOC/det/images/test  # train images (relative to 'path') 5000 images
road_seg_train: VOC/seg/images/train   # road segmentation data
road_seg_val: VOC/seg/images/val

# Classes
nc: 20  # number of classes
segnc: 20

names: ['aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat', 'chair',
           'cow', 'diningtable', 'dog', 'horse',
           'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']  # class names

segnames: ['aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat', 'chair',
           'cow', 'diningtable', 'dog', 'horse',
           'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']
    1. change the config in trainds.py and :
    python trainds.py

    1. test image folder with :

      python detectds.py

Reference

  2. yolov5

GitHub

View Github