使用Darknet框架训练目标检测模型

1.darknet环境准备

(1)安装darknet环境

git clone https://github.com/AlexeyAB/darknet.git

(2)修改Makefile文件

GPU=1
CUDNN=1
OPENCV=1

(3)编译darknet

cd darknet
Make

2.数据准备

以头部数据集为例

(1)数据集

JPEGImages

Annotations

(2) 制作yolo格式的标签（.txt）

在JPEGImages的同级目录下新建labels文件夹

将decode_xml.py中的img_path修改为JPEGImages的路径，这里我们提取Head类的标签，所以classes = ["Head"]

运行decode_xml.py文件，可以在labels文件夹下生成对应的txt文件

decode_xml.py

import os
import cv2
import xml.etree.ElementTree as ET

# classes to be trained
classes = ["Head"]

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(xml_path, img_h, img_w):
    in_file = open(xml_path)
    label_path = xml_path.replace('Annotations','labels').replace('.xml','.txt')
    out_file = open(label_path, 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()

    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes :
            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((img_w,img_h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

def xml2txt(img_path):
    img_names = os.listdir(img_path)
    for img_name in img_names:
        img_all_path = img_path + img_name

        #print(img_all_path)
        xml_path = img_all_path.replace('JPEGImages','Annotations').replace('jpg','xml')
        print(xml_path)
        
        img = cv2.imread(img_all_path)
        img_h = img.shape[0]
        img_w = img.shape[1]
        convert_annotation(xml_path, img_h,img_w)

if __name__ == "__main__":

    # Modify to your PEGImages path
    img_path = '/media/deepnorth/14b6945d-9936-41a8-aeac-505b96fc2be8/head-doc/JPEGImages/'
    xml2txt(img_path)

(3)生成train.txt文件 (JPEGImages文件中每个图片的路径)

在终端下执行以下命令生成train.txt

find /media/JPEGImages/ -name "*.jpg"  >  train.txt

3.训练文件准备

在darknet文件夹下新建一个head文件夹,包括以下文件

(1)新建head.data

classes= 1
train  = head/train.txt
valid  = head/test.txt
names =  head/head.names
backup = backup/

(2)计算数据集的anchor

./darknet detector calc_anchors head/head.data -num_of_clusters 6 -width 608 -height 608

(3)修改yolov3-tiny.cfg

batch=64
subdivisions=16
 
width=608
height=608
channels=3
 
[convolutional]
size=1
stride=1
pad=1
filters=18
activation=linear
 
[yolo]
mask = 3,4,5
anchors = 9,18,  13,27,  21,39,  35,60,  58,98,  99,174
classes=1
num=6
jitter=.3
ignore_thresh = .7
truth_thresh = 1
random=1

需要注意的点：

anchors 由上一个步骤(2)计算出

yolo层之前的convolutional的filters的计算公式：18 = 3*(5+class)

4.开始训练

./darknet detector train head/head.data head/yolov3-tiny.cfg -dont_show -map -gpus 0,1,2,3

5.测试

测试单张图像

./darknet detector test head/head.data head/yolov3-tiny.cfg  yolov3-tiny.weight img/head-test.jpg

测试视频

./darknet detector demo head/head.data  head/yolov3-tiny.cfg  yolov3-tiny.weights  <video file>

测试模型的map      mAP@IoU=0.5

./darknet detector map head/head.data head/yolov3-tiny.cfg backup/yolov3-tiny-30000.weights -iou_thresh 0.5

6.模型选择

方法一：

当命令行中添加了参数-map  将会生成下图map曲线，选择MAP最高的作为最终的模型

Method 2: Calculate the map through step5, and then select the model with the highest map as the final model.

方法二：根据第五步中的命令计算map, 选择map最好的模型