YOLOv5-7.0-seg+YOLOv8-seg自定义数据集训练

YOLOv5 

下载源码   https://github.com/ultralytics/yolov5.git
参考链接   yolov5-实例分割 1.如何使用yolov5实现实例分割，并训练自己的数据集_哔哩哔哩_bilibili

• 目录:

- datasets
    - JPEImages #存放图片和标注后的json文件以及转换后的txt文件
    - classes-4 #存放切分好的数据集
        - images
        - labels
    labelme2seg.py
    split_dataset.py
    visual_coco128.py
- segment
    train.py
    predict.py

• (1)-Labelme标注图片，生成json文件
  
  (52条消息) Labelme 安装以及使用_期待686的博客-CSDN博客

• (2)-labelme2seg.py   #讲json文件转换成yolo可以识别的txt文件

        修改类别 classes = ['clothing',…]
        图片和json文件路径 base_path = 'JPEGImages' 

​
import os
import cv2
import json
import numpy as np

# classes = ['clothing','face_shield','boot']
# classes = ['clothing','face_shield','boot','clothing_f']
classes = ['boot_f','clothing','face_shield','clothing_f']
base_path = '22'
path_list = [i.split('.')[0] for i in os.listdir(base_path)]
for path in path_list:
    image = cv2.imread(f'{base_path}/{path}.jpg')
    h,w,c = image.shape
    try:
        with open(f'{base_path}/{path}.json') as f:
            masks = json.load(f)['shapes']
        with open(f'{base_path}/{path}.txt','w+') as f:
            for idx,mask_data in enumerate(masks):
                mask_label = mask_data['label']
                mask = np.array([np.array(i) for i in mask_data['points']],dtype=np.float)
                mask[:,0] /= w
                mask[:,1] /= h
                mask = mask.reshape((-1))
                if idx != 0:
                    f.write('\n')
                f.write(f'{classes.index(mask_label)} {" ".join(list(map(lambda x:f"{x:.6f}",mask)))}')
    except FileNotFoundError as f:
        print(f)

​

• (3)-split_dataset.py   #切分数据集（需要shutil.copy()尽量绝对路径以防出错）

        图片和txt文件路径 base_path = r'F:\wzy\AI\yolov5-master-23-1-6\datasets\JPEGImages';保存路径dataset_path = r'F:\wzy\AI\yolov5-master-23-1-6\datasets\classes-4'

import os
import shutil
import random
import numpy as np

base_path = r'F:\wzy\AI\yolov5-master-23-1-6\datasets\22'
dataset_path = r'F:\wzy\AI\yolov5-master-23-1-6\datasets\classes-4-2'
val_size,test_size = 0.1,0.2

os.makedirs(dataset_path,exist_ok=True)
os.makedirs(f'{dataset_path}/images',exist_ok=True)
os.makedirs(f'{dataset_path}/images/train',exist_ok=True)
os.makedirs(f'{dataset_path}/images/val',exist_ok=True)
os.makedirs(f'{dataset_path}/images/test',exist_ok=True)
os.makedirs(f'{dataset_path}/labels/train',exist_ok=True)
os.makedirs(f'{dataset_path}/labels/val',exist_ok=True)
os.makedirs(f'{dataset_path}/labels/test',exist_ok=True)

path_list = np.array([i.split('.')[0] for i in os.listdir(base_path) if 'txt' in i])
random.shuffle(path_list)
train_id = path_list[:int(len(path_list) * (1 - val_size - test_size))]
val_id = path_list[int(len(path_list) * (1 - val_size - test_size)):int(len(path_list) * (1 - test_size))]
test_id = path_list[int(len(path_list) * (1 - test_size)):]
print(train_id,val_id,test_id)

for i in train_id:
    shutil.copy(f'{base_path}/{i}.jpg', f'./classes-4-2/images/train/{i}.jpg')
    shutil.copy(f'{base_path}/{i}.txt', f'./classes-4-2/labels/train/{i}.txt')
for i in val_id:
    shutil.copy(f'{base_path}/{i}.jpg', f'./classes-4-2/images/val/{i}.jpg')
    shutil.copy(f'{base_path}/{i}.txt', f'./classes-4-2/labels/val/{i}.txt')
for i in test_id:
    shutil.copy(f'{base_path}/{i}.jpg', f'./classes-4-2/images/test/{i}.jpg')
    shutil.copy(f'{base_path}/{i}.txt', f'./classes-4-2/labels/test/{i}.txt')

• (4)-visual_coco128.py   #验证数据集是否正确

        图片路径 img_base_path = 'classes-4/images/train';txt类别路径lab_base_path = 'classes-4/labels/train'

import os
import cv2
import numpy as np

img_base_path = 'classes-4-2/images/train'
lab_base_path = 'classes-4-2/labels/train'

label_path_list = [i.split('.')[0] for i in os.listdir(img_base_path)]
for path in label_path_list:
    image = cv2.imread(f'{img_base_path}/{path}.jpg')
    h,w,c = image.shape
    label = np.zeros((h,w),dtype=np.uint8)
    with open(f'{lab_base_path}/{path}.txt') as f:
        mask = np.array(list(map(lambda x:np.array(x.strip().split()),f.readlines())))
    for i in mask:
        i = np.array(i,dtype=np.float32)[1:].reshape((-1,2))
        i[:,0] *= w
        i[:,1] *= h
        label = cv2.fillPoly(label,[np.array(i,dtype=np.int32)],color=255)
    image = cv2.bitwise_and(image,image,mask=label)
    cv2.imshow('Pic',image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

• (5)-config_me.yaml   #训练的参数：路径/类别(顺序和labelme2seg.py的classes一样)…

train: F:\wzy\AI\yolov5-master-23-1-6\datasets\classes-4-2\images\train  # train images (relative to 'path') 128 images
val: F:\wzy\AI\yolov5-master-23-1-6\datasets\classes-4-2\images\val  # val images (relative to 'path') 128 images
test:  # test images (optional)

names:
  0: boot_f
  1: clothing
  2: face_shield
  3: clothing_f

• (6)-train.py   #训练模型

        run'''python train.py --weights ../weights/yolov5s-seg.pt --cfg ../models/segment/yolov5s-seg.yaml --data ../data/config_me.yaml --epochs 100 --batch-size 16 --device 0'''

• (7)-predict.py   #调用模型

        run '''python predict.py --weights ../runs/train-seg/expm-e100-bs8-epochs-1/weights/best.pt --source ../datasets/images/test --data ../data/config_me.yaml'''

---

YOLOv8

数据处理承接至上文步骤(1)~(5)

• 下载安装包ultralytics可以直接在终端使用yolo命令：

pip install ultralytics

• 训练：

yolo task=segment mode=train model="weights/yolov8n-seg.pt" data="data/config_me.yaml" epochs=100 batch=16 device=0

• 测试：

yolo task=segment mode=predict model="runs/segment/train4/weights/best.pt" source="input"

python3.7版本  cuda11.1  torch1.8.0   安装包  requirements.txt

absl-py==1.3.0
antlr4-python3-runtime==4.9.3
attr==0.3.2
attrs @ file:///F:/su/Programs/conda/envs/cv/Tools/attrs-22.2.0-py3-none-any.whl
backcall==0.2.0
cachetools==5.2.1
certifi @ file:///C:/b/abs_85o_6fm0se/croot/certifi_1671487778835/work/certifi
charset-normalizer==2.1.1
cloudpickle==2.2.0
colorama==0.4.6
cycler==0.11.0
Cython==0.29.32
decorator==5.1.1
exceptiongroup==1.1.0
fairscale==0.4.6
filelock==3.9.0
fonttools==4.38.0
fvcore==0.1.5.post20221221
gitdb @ file:///F:/su/Programs/conda/envs/cv/Tools/gitdb-4.0.10-py3-none-any.whl
GitPython @ file:///F:/su/Programs/conda/envs/cv/Tools/GitPython-3.1.30-py3-none-any.whl
google-auth==2.15.0
google-auth-oauthlib==0.4.6
grpcio==1.51.1
huggingface-hub==0.11.1
hydra-core==1.3.1
idna==3.4
importlib-metadata==6.0.0
importlib-resources==5.10.2
iniconfig==2.0.0
iopath==0.1.10
ipython==7.34.0
jedi==0.18.2
kiwisolver==1.4.4
Markdown==3.4.1
MarkupSafe==2.1.1
matplotlib==3.5.3
matplotlib-inline==0.1.6
numpy==1.21.6
oauthlib==3.2.2
omegaconf==2.3.0
opencv-python @ file:///D:/ana/envs/wzy/Lib/site-packages/opencv_python-4.5.5-cp37-cp37m-win_amd64.whl
packaging==22.0
pandas==1.3.5
parso==0.8.3
pickleshare==0.7.5
Pillow==9.4.0
pluggy==1.0.0
portalocker==2.6.0
prompt-toolkit==3.0.36
protobuf==3.20.3
psutil==5.9.4
pyasn1==0.4.8
pyasn1-modules==0.2.8
pycocotools-windows @ file:///D:/ana/Tools/pycocotools/pycocotools_windows-2.0.0.1-cp37-cp37m-win_amd64.whl
Pygments==2.14.0
pyparsing==3.0.9
pytest==7.2.0
python-dateutil==2.8.2
pytz @ file:///F:/su/Programs/conda/envs/cv/Tools/pytz-2022.7-py2.py3-none-any.whl
pywin32==305
PyYAML==6.0
requests==2.28.1
requests-oauthlib==1.3.1
rsa==4.9
scipy==1.7.3
seaborn==0.12.2
six==1.16.0
smmap @ file:///F:/su/Programs/conda/envs/cv/Tools/smmap-5.0.0-py3-none-any.whl
tabulate==0.9.0
tensorboard==2.11.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
termcolor==2.1.1
thop==0.1.1.post2209072238
timm==0.6.12
tomli==2.0.1
torch @ file:///D:/ana/envs/wzy/Tools/torch-1.8.0%2Bcu111-cp37-cp37m-win_amd64.whl
torchvision @ file:///D:/ana/envs/wzy/Tools/torchvision-0.9.0%2Bcu111-cp37-cp37m-win_amd64.whl
tqdm==4.64.1
traitlets==5.8.0
typing_extensions==4.4.0
ultralytics @ file:///F:/su/Programs/conda/envs/cv/Tools/ultralytics-8.0.5-py3-none-any.whl
urllib3==1.26.13
wcwidth==0.2.5
Werkzeug==2.2.2
wincertstore==0.2
yacs==0.1.8
zipp==3.11.0