Train yolov8 segment model based on custom dataset on Ubuntu 18.04

1 创建yolov8 segment格式数据集

这里创建了 tpadseg数据集，我是基于Label Studio进行图片分割标注，导出COCO格式数据集后，采用以下脚本进行格式转换的。

# -*- encoding:utf-8 -*-
# https://github.com/ultralytics/yolov5/issues/10161

import os
import json


def coco2cocoseg(main_dir, js_path, out_labels_path, b_debug=False):
    js_path = os.path.join(main_dir, js_path)
    out_labels_path = os.path.join(main_dir, out_labels_path)

    if not os.path.exists(out_labels_path):
        os.makedirs(out_labels_path)

    if b_debug:
        print(f'main_dir:{main_dir}')
        print(f'json_file:{js_path}')
        print(f'out_labels_foder:{out_labels_path}')

    js = json.load(open(js_path, 'r'))

    all_images = list()
    annotations_by_id = {}
    annotations_by_image_id = {}
    for ann in js['annotations']:
        ann_id = ann['id']
        assert (ann_id not in annotations_by_id), 'error ,annotation id already in annotations_by_id'
        annotations_by_id[ann_id] = ann
        image_id = ann['image_id']
        if image_id not in annotations_by_image_id:
            annotations_by_image_id[image_id] = list()
        annotations_by_image_id[image_id].append(ann_id)

    for im in js['images']:
        im_id = im['id']
        fname = os.path.basename(im['file_name']).split('.')[0]
        all_images.append(im['file_name'])
        w = im['width']
        h = im['height']
        label_name = fname + '.txt'
        with open(os.path.join(out_labels_path, label_name), 'w') as out_labels:
            if im_id in annotations_by_image_id:
                for ann_id in annotations_by_image_id[im_id]:
                    ann = annotations_by_id[ann_id]
                    # category_id = ann['category_id'] - 1  # coco category id start with 1, and yolo category id start with 0
                    category_id = ann['category_id']
                    out_seg = []
                    for seg in ann['segmentation']:
                        for i in range(int(len(seg)/2)):
                            x = seg[i*2] / w
                            y = seg[i*2 + 1] / h
                            out_seg.append(x)
                            out_seg.append(y)
                    line = str(category_id)
                    for coord in out_seg:
                        line += ' ' + str(coord)
                    out_labels.write(line)
                    out_labels.write('\n')

    with open(os.path.join(main_dir, js_path.split('.')[0] + '.txt'), 'w') as f:
        for im_name in all_images:
            f.write(im_name)
            f.write('\n')


if __name__=='__main__':
    js_path = 'result.json'
    out_labels_path = 'labels/'
    main_dir = '/home/epbox/AI/dataset/tpadseg/project-4-at-2023-01-13-01-55-4684c84f/'
    coco2cocoseg(main_dir, js_path, out_labels_path, True)

$ cd /home/epbox/AI/dataset/tpadseg/v0.0.1/
$ tree . -h --filelimit=10 --dirsfirst
.
├── [4.0K]  images
│   ├── [ 12K]  test [135 entries exceeds filelimit, not opening dir]
│   ├── [ 64K]  train [941 entries exceeds filelimit, not opening dir]
│   └── [ 20K]  valid [268 entries exceeds filelimit, not opening dir]
└── [4.0K]  labels
    ├── [ 12K]  test [135 entries exceeds filelimit, not opening dir]
    ├── [ 64K]  train [941 entries exceeds filelimit, not opening dir]
    └── [ 20K]  valid [268 entries exceeds filelimit, not opening dir]

8 directories, 0 files

2 创建yolov8虚拟环境

# 1. 创建虚拟环境
$ cd ~/Github/yolov8
$ mkvirtualenv yolov8 -p /usr/bin/python3
(yolov8) $ pip install ultralytics -i https://pypi.douban.com/simple
(yolov8) $ pip install albumentations -i https://pypi.douban.com/simple
(yolov8) $ pip install --upgrade pip -i https://pypi.douban.com/simple 

3 训练基于yolov8s-seg模型

# 1. 创建data/tpadseg.yaml文件
(yolov8) $ mkdir data
(yolov8) $ cat data/tpadseg.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, ..]
train: /home/epbox/AI/dataset/tpadseg/v0.0.1/images/valid/  # train images (relative to 'path') 128 images
val:   /home/epbox/AI/dataset/tpadseg/v0.0.1/images/valid/  # val images (relative to 'path') 128 images
test:  /home/epbox/AI/dataset/tpadseg/v0.0.1/images/test/ # test images (optional)

# Classes
names:
  0: blur
  1: phone
  2: reflectLight
  3: reflection

# 2. 训练基于yolov8s-seg的模型 https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt 
(yolov8) $ yolo segment train epochs=100 imgsz=640 device=0 cache=True \
data=data/tpadseg.yaml \
model=/home/epbox/AI/pre_weights/yolov8/yolov8s-seg.pt \
project=runs/segment/tpadseg \
name=yolov8s/exp1/train/

4 测试训练好的模型

# 1. 测试新训练的模型
(yolov8) $ yolo segment val batch=1 device=0 \
data=data/tpadseg.yaml \
model=runs/segment/tpadseg/yolov8s/exp1/train/weights/best.pt \
project=runs/segment/tpadseg \
name=yolov8s/exp1/val/

# 2. 单张图片测试 https://docs.ultralytics.com/cfg/
(yolov8) $ yolo segment predict device=0 save_crop=True \
model=runs/segment/tpadseg/yolov8s/exp1/train/weights/best.pt \
source=1e8b179c-P92975500159211090001_0.jpg \
project=runs/segment/tpadseg/yolov8s/exp1/ \
name=test

Ultralytics YOLOv8.0.29 🚀 Python-3.8.10 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3060, 12047MiB)
YOLOv8s-seg summary (fused): 195 layers, 11781148 parameters, 0 gradients, 42.4 GFLOPs

/home/epbox/.virtualenvs/yolov8/lib/python3.8/site-packages/ultralytics/yolo/utils/plotting.py:153: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  xyxy = torch.tensor(xyxy).view(-1, 4)
image 1/1 /home/epbox/Github/yolov8/1e8b179c-P92975500159211090001_0.jpg: 640x480 1 phone, 1 reflectLight, 1 reflection, 13.2ms
Speed: 1.5ms pre-process, 13.2ms inference, 1.4ms postprocess per image at shape (1, 3, 640, 640)
(yolov8) $ tree runs/segment/tpadseg/yolov8s/exp1/test/
runs/segment/tpadseg/yolov8s/exp1/test/
└── crops
    ├── phone
    │   └── 1e8b179c-P92975500159211090001_0.jpg
    ├── reflection
    │   └── 1e8b179c-P92975500159211090001_0.jpg
    └── reflectLight
        └── 1e8b179c-P92975500159211090001_0.jpg

4 directories, 3 files

1e8b179c-P92975500159211090001_0.jpg	phone	reflectLight	reflection