基于darknet框架训练yolov4模型,获取推理结果的bbox信息,导出yolo标注文件,实现图像的自动标注,方便图像标注。
1. Darknet 安装
按照官网https://pjreddie.com/darknet/ 的指引安装即可。
报错与解决:
(1)使用GPU时,修改makefile GPU=1 cudnn=1,make报错darknet make include/darknet.h:16:23: 致命错误:cudnn.h
找到cudnn.h复制到/usr/include下,我的cudnn.h在anaconda/include下;
(2)然后make,再报一个错误:/usr/bin/ld: 找不到 -lcudnncollect2
找到libcudnn.so文件,复制到/usr/local/cuda-10.1/lib64,我的libcudnn.so文件在anaconda/lib下
2. 训练模型--数据准备
参照github的指南准备数据,我用lalbelImg标注,保存图像及标注,分成train和val,即可训练
3. 训练模型--修改cfg配置文件
4. 获取模型推理的bbox信息,导出yolo标注文件
import darknet
import darknet_images
import cv2
import time
import os
config_file='cfg/yolov4-tiny-circle.cfg'
data_file='data/circle.data'
# weights='backup/11-23/yolov4-tiny_best.weights'
weights='backup/12-02/yolov4-tiny_best.weights'
img_path='/camera_12022010/camera_log'
result_path = '/camera_12022010/result'
label_path='/camera_11_19/-labels/'
save_label = False #是否保存成yolo TXT标注,不保存的话,只做推理
#load model
network, class_names, colors = darknet.load_network(config_file, data_file, weights, batch_size=1)
FileList = os.listdir(img_path)
for filename in FileList:
img = os.path.join(img_path, filename)
image = cv2.imread(img)
start = time.time()
image, detections = darknet_images.image_detection(img, network, class_names, colors, 0.1)
cv2.imwrite(result_path + f're_{filename}', image)
#识别结果写入label文件
if save_label:
if len(detections)>0:
bbox = detections[0][2]
x1, y1, w, h = bbox[0],bbox[1],bbox[2],bbox[3]
x2 = x1 + w
y2 = y1 + h
label_x = x1 / image.shape[1]
label_y = y1 / image.shape[0]
label_w = w / image.shape[1]
label_h = h / image.shape[0]
#这里只有一类,所以class标号只有‘0’,可根据class.txt对应类别写入标号
msg = '0' + ' ' + str(label_x) + ' '+ str(label_y) + ' '+ str(label_w) + ' '+ str(label_h)
txt_path = label_path + filename[:-4] + '.txt'
file = open(txt_path, 'w')
file.write(msg)
file.close()
print(f'inference time:{time.time()-start}')