YOLOv3训练自定义数据集

1. 下载YOLOv3和权重文件

YOLOv3:git clone https://github.com/pjreddie/darknet
yolov3.weights：wget https://pjreddie.com/media/files/yolov3.weights
darknet53.conv.74：wget https://pjreddie.com/media/files/darknet53.conv.74

2. 数据集处理

（1）在scripts文件夹下按如下目录创建VOCdevkit 文件夹，放自己的训练数据。

VOCdevkit
--VOC2007
----Annotations  #（XML标签文件）
----ImageSets
------Main
----JPEGImages   # (原始图片)

（2）运行voc2yolo5.py 生成划分的训练集、测试集等文件

import os
import random
import sys
root_path = './scripts/VOCdevkit/VOC2007'
xmlfilepath = root_path + '/Annotations'
txtsavepath = root_path + '/ImageSets/Main'
if not os.path.exists(root_path):
    print("cannot find such directory: " + root_path)
    exit()
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)
trainval_percent = 0.8
train_percent = 0.2
total_xml = os.listdir(xmlfilepath)
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("train size:", 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()

（3）修改scripts/voc_label.py并运行,生成对应图片的labels标签文件

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
classes = ["smoke","fire","xxx1",,"xxx2",,"xxx3",,"xxx4"] #按自己的类别修改，顺序要和data/voc.names保持一致

以上2步运行后目录如下所示：

VOCdevkit
--VOC2007
----Annotations  #（XML标签文件）
----ImageSets
------Main
------------test.txt
------------train.txt
------------trainval.txt
------------val.txt
----JPEGImages   # (原始图片)
----labels
------------00001.txt
------------00002.txt
------------00003.txt
------------......
------------......
------------xxxxx.txt

3. 修改配置文件

3.1 修改data/voc.names

smoke
fire
xxx1
xxx2
xxx3
xxx4

3.2 修改cfg/voc.data

classes=6  # 类别的数量
# 数据集处理中的voc_label.py生成
train=./scripts/2007_train.txt  # 训练过程中训练数据的txt文件
valid=./scripts/2007_val.txt  # 训练过程中验证数据的txt文件
names=data/voc.names  # 类别标签名称
backup=backup/  # 存放权重的路径

3.3 修改cfg/yolov3.cfg

[yolo]
mask = 3,4,5
anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401
classes=6  # 修改为自己的类别

修改filters个数：（classes + 5）*3；修改类别数classes
（共mask = 0,1,2 ; mask = 3,4,5 ; mask = 6,7,8 三处的class和上方的三处filters）

[convolutional]
size=1
stride=1
pad=1
filters=33  # filters = (classes + 5）*3
activation=linear


[yolo]
mask = 0,1,2
anchors = 10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326
classes=6  # 自己的类别
num=9
jitter=.3
ignore_thresh = .7

4.修改Makefile文件

4.1修改

cd darknet
vi Makefile

GPU=1 #使用GPU设置为1，CPU设置为0
CUDNN=1  #使用CUDNN设置为1，否则为0
OPENCV=0 #调用摄像头，还需要设置OPENCV为1，否则为0
OPENMP=0  #使用OPENMP设置为1，否则为0
DEBUG=0  #使用DEBUG设置为1，否则为0


ARCH= -gencode arch=compute_80,code=sm_80 
#      -gencode arch=compute_20,code=[sm_20,sm_21] \ This one is deprecated?

# This is what I use, uncomment if you know your arch and want to specify
# ARCH= -gencode arch=compute_52,code=compute_52

VPATH=./src/:./examples
SLIB=libdarknet.so
ALIB=libdarknet.a
EXEC=darknet
OBJDIR=./obj/

CC=gcc
CPP=g++
NVCC=/usr/local/cuda-11.1/bin/nvcc #修改为自己的路径

4.2 开始编译

make

编译小结

原因：计算机显卡型号和算力需要在makefile里进行匹配
解决方法：本文为RTX3080，对应computer_80，makefile里没有写，添加-gencode arch=compute_80,code=sm_80并把之前的都删去。

# Tesla A100 (GA100), DGX-A100, RTX 3080
# ARCH= -gencode arch=compute_80,code=sm_80

# Tesla V100
# ARCH= -gencode arch=compute_70,code=sm_70

# GeForce RTX 2080 Ti, RTX 2080, RTX 2070, Quadro RTX 8000, Quadro RTX 6000, Quadro RTX 5000, Tesla T4, XNOR Tensor Cores
# ARCH= -gencode arch=compute_75,code=sm_75

# Jetson XAVIER
# ARCH= -gencode arch=compute_72,code=sm_72

# GTX 1080, GTX 1070, GTX 1060, GTX 1050, GTX 1030, Titan Xp, Tesla P40, Tesla P4
# ARCH= -gencode arch=compute_61,code=sm_61

# GP100/Tesla P100 - DGX-1
# ARCH= -gencode arch=compute_60,code=sm_60

# For Jetson TX1, Tegra X1, DRIVE CX, DRIVE PX - uncomment:
# ARCH= -gencode arch=compute_53,code=sm_53

# For Jetson Tx2 or Drive-PX2 uncomment:
# ARCH= -gencode arch=compute_62,code=sm_62

5.训练

./darknet detector train cfg/voc.data cfg/yolov3.cfg darknet53.conv.74 -gpus 0

训练小结

可能是GPU被占用或者需要修改yolov3.cfg文件中的subdivision（变大）、width和height（变小）的大小

6.测试

测试

./darknet detector valid cfg/voc.data cfg/yolov3.cfg darknet53.conv.74 -gpus 0

会在results文件夹下生成各个类别对应的.txt文件
计算mAP
（1）下载voc_eval.py：https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/datasets/voc_eval.py
（2）新建compute_AP.py（计算单类别AP）

from voc_eval import voc_eval
 
rec,prec,ap=voc_eval('/data/darknet/results/{}.txt', '/data/darknet/scripts/VOCdevkit/VOC2007/Annotations/{}.xml', '/data/darknet/scripts/VOCdevkit/VOC2007/ImageSets/Main/test.txt', 'smoke', '.')

注意修改voc_eval中对应的路径！应使用python2运行，若使用python3，需要把cPickle改为pickle!
（3）新建compute_mAP.py（计算多类别mAP）

from voc_eval import voc_eval
import os

current_path = os.getcwd()
results_path = current_path+"/results"
sub_files = os.listdir(results_path)

mAP = []
for i in range(len(sub_files)):
    class_name = sub_files[i].split(".txt")[0]
    rec, prec, ap = voc_eval('/data/darknet/results/{}.txt', '/data/darknet/scripts/VOCdevkit/VOC2007/Annotations/{}.xml', '/data/darknet/scripts/VOCdevkit/VOC2007/ImageSets/Main/test.txt', class_name, '.')
    print("{} :\t {} ".format(class_name, ap))
    mAP.append(ap)

mAP = tuple(mAP)

print("***************************")
print("mAP :\t {}".format( float( sum(mAP)/len(mAP)) ))