YOLOv3-tiny 数据准备、模型训练以及测试过程记录

参考网站

https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects

源码准备和编译

git clone https://github.com/AlexeyAB/darknet.git
cd darknet-master

GPU运行，修改Makefile中的：

GPU=1
CUDNN=1
OPENCV=1
LIBSO=1  #可用于生成so文件

编译环境：CUDA10.2 cudnn 8.0.5 opencv 4.1.0 cmake 3.20.6
编译darknet:

sudo make 

网上下载配置文件、预训练模型文件后，可用以下命令进行测试：

./darknet detect cfg/yolov3-tiny.cfg yolov3-tiny.weights data/dog.jpg

准备自己的数据集

按照如下树状结构创建文件夹：

---VOCdevkit
	---VOC2007
		---Annotations
		---ImageSets
			---Main
		---JPEGImages

获取数据集视频，提帧

按间隔帧来提取视频中的帧（从第1帧取到第900帧，隔3帧取一次）

ffmpeg -i VID_20220527_093607.mp4 -vf "select=between(n\,1\,900)*not(mod(n\,3))" -vsync 0 image_%05d.jpg

重新修改下图像文件名

def rename3():
    pathlist = os.listdir(r'D:/PoseEstimationDatasets/can_image/img1/')
    path_old = r'D:/PoseEstimationDatasets/can_image/img1/'
    path = r'C:/Users/Chris/Desktop/VOCdevkit/VOC2007/JPEGImages/'
    #pathlist.sort(key=lambda x: int(x.split('.')[0]))
    i=0
    for filename in pathlist:
        img = path_old + filename
        res = cv2.imread(img)
        newname = str(i) + '.jpg'
        cv2.imwrite(path+newname, res)
        print(filename + ' change to ' + newname)
        i = i + 1

将标记好后的图像及xml文件放入images和labels文件夹下，然后运行以下代码，完成数据重命名。voc.py

from xml.etree.ElementTree import ElementTree
from os import walk, path
import cv2
import os


def read_xml(in_path):
    tree = ElementTree()
    tree.parse(in_path)
    return tree


def write_xml(tree, out_path):
    tree.write(out_path, encoding="utf-8", xml_declaration=True)

def get_path_prex(rootdir):
    data_path = []
    prefixs = []
    for root, dirs, files in walk(rootdir, topdown=True):
        for name in files:
            pre, ending = path.splitext(name)
            if ending != ".xml":
                continue
            else:
                data_path.append(path.join(root, name))
                prefixs.append(pre)

    return data_path, prefixs


if __name__ == "__main__":

    # build files which will be used in VOC2007
    if not os.path.exists("Annotations"):
        os.mkdir("Annotations")
    if not os.path.exists("JPEGImages"):
        os.mkdir("JPEGImages")

    xml_paths, prefixs = get_path_prex("labels")

    for i in range(len(xml_paths)):
        # rename and save the corresponding xml
        tree = read_xml(xml_paths[i])
        # save output xml, 000001.xml
        write_xml(tree, "Annotations/{}.xml".format("%06d" % (i + 1)))

        # rename and save the corresponding image
        img_pre = prefixs[i] + ".jpg"
        root = os.getcwd() + '/images/'
        img_path = path.join(root, img_pre)
        img = cv2.imread(img_path)
        # save output jpg, 000001.jpg
        cv2.imwrite('JPEGImages/{}.jpg'.format("%06d" % (i + 1)), img)

将命名好的图像和文件放到/VOCdevkit/VOC2007下，运行以下文件将数据分成trainval, train, test, val四类，并生成用于YOLO训练的txt文件：get_voc_2007_main.py

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]   #需按自己的需求进行修改

classes = ["box", "pen", "laptop"]    #需按自己的需求进行修改


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(year, image_id):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, 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 year, image_set in sets:
    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

os.system("cat 2007_train.txt 2007_val.txt > train.txt")
os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt > train.all.txt")

修改cfg/voc.data文件，确定类别数量和正确路径。
修改cfg/yolov3-tiny.cfg

change [filters=255] to filters=(classes + 5)x3 in the 3 [convolutional] before each
[yolo] layer, keep in mind that it only has to be the last [convolutional] before each 
of the [yolo] layers.

修改data/voc.names

训练模型

# yolo partial提取已经训练好的网络中的部分权重
./darknet partial ./cfg/yolov3-tiny.cfg ./yolov3-tiny.weights ./yolov3-tiny.conv.15 15
# 执行训练
sudo ./darknet detector train cfg/voc.data cfg/yolov3-tiny.cfg yolov3-tiny.conv.15

模型测试

单张图片测试：-thresh 0.5可用于置信度的调整

./darknet detector test cfg/voc.data cfg/yolov3-tiny.cfg backup/yolov3-tiny_10000.weights VOCdevkit/VOC2007/JPEGImages/000022.jpg

多张图片测试: （相关链接：https://github.com/pjreddie/darknet/issues/723）

./darknet detector test ./cfg/coco.data ./cfg/yolov3.cfg ./yolov3.weights -dont_show < data/train.txt > result.txt

data/train.txt 应包含待检测图片的路径
result.txt 将存入检测结果

主目录下进行评估：.

./darknet detector map cfg/voc.data cfg/yolov4-tiny.cfg backup/yolov4-tiny_60000.weights

对训练好的权重进行评价。权重在backup下，可以自己查看评价哪一个权重。