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