使用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最好的模型
- 点赞
- 收藏
- 关注作者
评论(0)