Pytorch版yolov5训练自定义数据集

目录

1.下载源码

2.制作数据集

3.生成文件

4.两个yaml文件

5.训练

6.测试

7.python test.py

---

相比yolov3版本，整理了生成train、val txt文档的方式，可以将数据集放在别的目录下了

1.下载源码

https://github.com/ultralytics/yolov5，下载源码并解压

https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data，官方教程

2.制作数据集

仍由VOC格式为例

JPEGImages改名为images，目录下只含有两个文件夹即可：

datasetname
    -images
    -Annotations

3.生成文件

代码1:split_train_val.py

import os  
import random  
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', type=str, help='input xml label path')
parser.add_argument('--txt_path', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.8  
xmlfilepath = opt.xml_path  
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)  
if not os.path.exists(txtsavepath):
	os.makedirs(txtsavepath)
  
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)  
  
ftrainval = open(txtsavepath + '/trainval.txt', 'w')  
ftest = open(txtsavepath + '/test.txt', 'w')  
ftrain = open(txtsavepath + '/train.txt', 'w')  
fval = open(txtsavepath + '/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()

把这个文件随便放在哪，执行python split_train_val.py --xml_path xx --txt_path xx

指定xml的label存放位置和希望输出的训练验证集划分txt位置，代码会从xml_path读取label文件，然后进行划分，将文件名保存在txt_path的txt文件中，txt_path 一般为 .../ImageSets/Main

代码2：voc_label.py

# -*- 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 = ['一次性快餐盒', '书籍纸张', '充电宝', '剩饭剩菜', '包', '垃圾桶', '塑料器皿', '塑料玩具', '塑料衣架', '大骨头', '干电池', '快递纸袋', 
'插头电线', '旧衣服', '易拉罐', '枕头', '果皮果肉', '毛绒玩具', '污损塑料', '污损用纸', '洗护用品', '烟蒂', '牙签', '玻璃器皿', '砧板', '筷子', 
'纸盒纸箱', '花盆', '茶叶渣', '菜帮菜叶', '蛋壳', '调料瓶', '软膏', '过期药物', '酒瓶', '金属厨具', '金属器皿', '金属食品罐', '锅', '陶瓷器皿', '鞋', '食用油桶', '饮料瓶', '鱼骨']
     #修改为自己的类别
#classes = ["eye", "nose"]
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('Annotations/%s.xml'%( image_id))  
    out_file = open('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('labels/'):
        os.makedirs('labels/')
    image_ids = open('ImageSets/Main/%s.txt'%(image_set)).read().strip().split()
    list_file = open('%s.txt'%(image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path+'/images/%s.jpg\n'%(image_id))
        convert_annotation(image_id)
    list_file.close()   
#os.system("cat 2007_train.txt 2007_val.txt > train.txt")     

把这个文件放在数据集根目录下，修改其中的classes列表，然后运行 python voc_label.py

代码会根据上一步划分的txt文件读取xml格式的label文件，转换成txt格式的label保存在labels目录下，并在当前目录下生成train、val、test.txt文件，里面保存的是完整的文件名，是绝对路径

两个代码运行结束后目录如下：

datasetname
    -Annotations
        -xxx.xml
    -images
        -xxx.jpg
    -ImageSets
        -Main
            -train.txt
            -val.txt
    -labels
        -xxx.txt

4.两个yaml文件

创建./data/xx.yaml

train: ./data/train.txt  # 128 images
val: ./data/val.txt  # 128 images

# number of classes
nc: 44

# class names
names: ['一次性快餐盒', '书籍纸张', '充电宝', '剩饭剩菜', '包', '垃圾桶', '塑料器皿', '塑料玩具', '塑料衣架', '大骨头', '干电池', '快递纸袋', 
'插头电线', '旧衣服', '易拉罐', '枕头', '果皮果肉', '毛绒玩具', '污损塑料', '污损用纸']

其中train和val是voc_label.py生成的txt文件，nc是类别数量, names中的内容与顺序要和voc_label.py文件中保持一致

修改./models/xx.yaml

在./models里选一个模型进行修改，只修改其中的nc为自己的类别数量即可。

5.训练

cd 到yolov5目录下

python train.py --img 640 --batch 8 --epoch 300 --data ./data/HuaWeigarbage.yaml --cfg ./models/yolov5s.yaml --weights weights/yolov5s.pt --logdir runs/temp --device '1'

 日志保存目录默认保存在./runs/exp0下，想要修改目录可以指定--logdir参数

6.测试

python detect.py --weights runs/exp0/weights/best.pt --source data/val.txt --device 0

对原始代码进行了修改，可以指定存放待检测图片路径的txt文件

7.python test.py

python test.py --weights runs/mask/exp2/weights/best.pt --data data/mask.yaml --batch-size 16 --conf-thres 0.4 --device '0'

一些说明：

默认保存模型的标准：

实际应用中可以根据需要修改这个比重。

保存模型的内容：中间过程保存的是模型、优化器、epoch等信息，所以会比单一模型占用空间大很多。

with open(results_file, 'r') as f:  # create checkpoint
    ckpt = {'epoch': epoch,
        'best_fitness': best_fitness,
        'training_results': f.read(),
        'model': ema.ema,
        'optimizer': None if final_epoch else optimizer.state_dict()}

# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
    torch.save(ckpt, best)

在训练结束后进行了strip optimizer操作，只保存模型，所以占用空间会小很多：

def strip_optimizer(f='weights/best.pt', s=''):  # from utils.general import *; strip_optimizer()
    # Strip optimizer from 'f' to finalize training, optionally save as 's'
    x = torch.load(f, map_location=torch.device('cpu'))
    x['optimizer'] = None
    x['training_results'] = None
    x['epoch'] = -1
    x['model'].half()  # to FP16
    for p in x['model'].parameters():
        p.requires_grad = False
    torch.save(x, s or f)
    mb = os.path.getsize(s or f) / 1E6  # filesize
    print('Optimizer stripped from %s,%s %.1fMB' % (f, (' saved as %s,' % s) if s else '', mb))