Train yolov8 segment model based on custom dataset on Ubuntu 18.04_yolov8segments 拿出

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)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

$ 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
1
2
3
4
5
6
7
8
9
10
11
12
13

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 
1
2
3
4
5
6

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/
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

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 
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Li_阴宅/article/detail/971823