Darknet 训练yolov4模型,做图像自动标注(代码)

基于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配置文件

参考:GitHub - AlexeyAB/darknet: YOLOv4 / Scaled-YOLOv4 / YOLO - Neural Networks for Object Detection (Windows and Linux version of Darknet )

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}')

  • 0
    点赞
  • 9
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值