YOLOv5训练自定义的烟火数据集和COCO2017数据集
一. 训练自定义数据集
1、准备数据集
(1)数据集具体结构内容如下：
yolov5/data/images存放训练的jpg图片
yolov5/data/annotations存放对应图片标注的xml文件
yolov5/data/ImageSets新建Main文件夹存放划分的训练集、测试集等文件
(2)运行voc2yolo5.py生成划分的训练集、测试集等文件
import os
import random
xmlfilepath=r'annotations'
saveBasePath=r"ImageSets/Main/"
trainval_percent=0.66
train_percent=0.5
temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
if xml.endswith(".xml"):
total_xml.append(xml)
num=len(total_xml)
list=range(num)
tv=int(num*trainval_percent)
tr=int(tv*train_percent)
trainval= random.sample(list,tv)
train=random.sample(trainval,tr)
print("train and val size",tv)
print("traub suze",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
for i in list:
name=total_xml[i][:-4]+'\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest .close()
(3)运行voc_label.py生成对应图片的labels标签文件
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'val', 'test']
classes = ['smoke','fire'] #修改训练的标签类别
abs_path = os.getcwd()
def convert(size, box):
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def convert_annotation(image_id):
in_file = open('data/annotations/%s.xml' % (image_id))
out_file = open('data/labels/%s.txt' % (image_id), 'w')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
for image_set in sets:
if not os.path.exists('data/labels/'):
os.makedirs('data/labels/')
image_ids = open('data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
list_file = open('data/%s.txt' % (image_set), 'w')
for image_id in image_ids:
list_file.write('data/images/%s.jpg\n' % (image_id))
convert_annotation(image_id)
list_file.close()
(4)完成以上步骤后整体数据集结构如下：
data
-annotations
-xxx.xml
-images
-xxx.jpg
-ImageSets
-Main
-train.txt
-test.txt
-val.txt
-trainval.txt
-labels
-xxx.txt
2、yaml文件的创建和修改
(1)创建/data/smoke_fire.yaml
./data/train.txt
./data/val.txt
# number of classes
nc: 2 # 修改训练的类别数
# class names
names: ['smoke','fire'] # 修改训练的标签类别
(2)修改/models/yolov5x.yaml（此处使用yolov5x）
# parameters
nc: 4 # number of classes 修改为训练的类别数
depth_multiple: 1.33 # model depth multiple
width_multiple: 1.25 # layer channel multiple
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, BottleneckCSP, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, BottleneckCSP, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, BottleneckCSP, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, BottleneckCSP, [1024, False]], # 9
]
# YOLOv5 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, BottleneckCSP, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
3、修改train.py文件
if __name__ == '__main__':
parser = argparse.ArgumentParser()
# 修改为加载的预训练模型，此处使用yolov5x.pt
parser.add_argument('--weights', type=str, default='weights/yolov5x.pt', help='initial weights path')
# 修改为上述修改的第二个yaml文件
parser.add_argument('--cfg', type=str, default='models/yolov5x.yaml', help='model.yaml path')
# 修改为上述新建的第一个yaml文件
parser.add_argument('--data', type=str, default='data/smoke_fire.yaml', help='data.yaml path')
parser.add_argument('--hyp', type=str, default='', help='hyperparameters path, i.e. data/hyp.scratch.yaml')
# 设置训练轮次
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
parser.add_argument('--rect', action='store_true', help='rectangular training')
# 修改default=True,可以接着上次训练中断的结果继续训练
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--notest', action='store_true', help='only test final epoch')
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
parser.add_argument('--logdir', type=str, default='runs/', help='logging directory')
parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
opt = parser.parse_args()
接下来就可以开始训练啦★,°:.☆(￣▽￣)/$:.°★ 。
4、修改test.py文件
if __name__ == '__main__':
parser = argparse.ArgumentParser(prog='test.py')
# 修改为上面训练好的模型文件
parser.add_argument('--weights', nargs='+', type=str, default='smoke_fire.pt', help='model.pt path(s)')
# 修改为上述新建的第一个yaml文件
parser.add_argument('--data', type=str, default='data/smoke_fire.yaml', help='*.data path')
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
# 修改default='val'/ 'test'，可查看验证集和测试集的PR曲线
parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
二. 训练COCO2017数据集
1、COCO2017数据集
| name | description | download url |
|---|---|---|
| train2017.zip | 19G, 118k images | http://images.cocodataset.org/zips/train2017.zip |
| val2017.zip | 1G, 5k images | http://images.cocodataset.org/zips/val2017.zip |
| test2017.zip | 7G, 41k images | http://images.cocodataset.org/zips/test2017.zip |
(1)下载以上数据集后，将图片解压到yolov5/data/datasets/coco/images/
(2)下载标注文件annotations_trainval2017.zip，
将标注文件解压到yolov5/data/datasets/coco/annotations/
(3)下载coco2017labels.zip
将coco2017labels/coco/labels解压到yolov5/data/datasets/coco/labels/；
再将coco2017labels/coco/train2017.txt、coco2017labels/coco/val2017.txt和coco2017labels/coco/test-dev2017.txt解压到yolov5/data/datasets/coco
数据集目录为
yolov5
-data
-datasets
-coco
-images
-xxxxxxxxxxxx.jpg
-annotations
-instance_train2017.json
-instance_val2017.json
-labels
-train2017
-xxxxxxxxxxxx.txt
-val2017
-xxxxxxxxxxxx.txt
-train2017.txt
-val2017.txt
-test-dev2017.txt
2、yaml文件的修改
data/coco.yaml
注意检查train val test的路径
# YOLOv5 by Ultralytics, GPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
# └── coco ← downloads here (20.1 GB)
# 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, ..]
# path: public/DL_DATA/COCO2017 # dataset root dir
train: ./data/datasets/coco/train2017.txt # train images (relative to 'path') 118287 images
val: ./data/datasets/coco/val2017.txt # val images (relative to 'path') 5000 images
test: ./data/datasets/coco/test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
# Classes
nc: 80 # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush'] # class names
3、训练
python train.py
