YOLOv5DS
Multi-task yolov5 with detection and segmentation based on yolov5(branch v6.0)
- decoupled head
- anchor free
- segmentation head
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
-
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.
-
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']- change the config in trainds.py and :
python trainds.py-
test image folder with :
python detectds.py
1 Oct 15, 2021
12 Nov 11, 2022
719 Jan 2, 2023
12 Dec 5, 2022
73 Dec 16, 2022
27 Oct 14, 2022
27 Dec 20, 2022
487 Dec 27, 2022
23 Dec 19, 2022
2 Jan 20, 2022
109 Dec 31, 2022
43 Dec 21, 2022
37 Dec 21, 2022
1 Dec 31, 2021
46 Dec 28, 2022
92 Dec 30, 2022
262 Dec 27, 2022
45 Jan 01, 2023
5.1k Dec 30, 2022
42 Dec 23, 2022
53 Jul 17, 2022
31 Dec 05, 2022
1 Jan 17, 2022
2 Oct 25, 2022
102 Dec 25, 2022
18 May 29, 2022
12 Jan 01, 2023
5 Nov 24, 2021
30 Dec 01, 2022