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 |
|---|---|---|---|
 |
 |
 |
 |