windows平台下darknet训练自己的数据集

写完前面关于python接口的文章发现忘了写训练自己的数据集的文章了,这里补上,用的还是AlexeyAB版本的darknet。

第一步:首先就是用labelimg做标注,这一步是最费时间的,注意这里一张图片是可以同时标注多个物体的,只是要做好对照关系,我之前网上找了一些教程都只标注了一个物体,这里说明一下标注多个物体也是可以的。

第二步:构建训练时可以读取的文件夹目录格式,如下所示:

windows平台下darknet训练自己的数据集_第1张图片

将xml文件放入Annotations文件夹,将图片放入JPEGImages文件夹,利用下面的代码1生成的mytrain_letter_train.txt文件就放在当前目录,代码1同时生成的表示坐标的txt文件放入labels文件夹,利用下面代码2将数据集分割为train、val、test、trainval四个文件。

# -*- coding: utf-8 -*-
"""
Created on Wed Nov 25 16:46:12 2020

@author: SS
"""

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
#源代码sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
sets=[('mytrain_letter', 'train')]  # 改成自己建立的myData
 
classes = ["B_A","B_B","B_C","B_D","B_E","B_F","B_G","blue_A","blue_B","blue_C","blue_D","blue_E","blue_F","blue_G"] # 改成自己的类别
 
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('E:/darknet/darknet-master/mytrain_letter/Annotations/%s.xml'%(image_id))  # 源代码VOCdevkit/VOC%s/Annotations/%s.xml
    out_file = open('E:/darknet/darknet-master/mytrain_letter/labels/%s.txt'%(image_id), 'w')  # 源代码VOCdevkit/VOC%s/labels/%s.txt
    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('E:/darknet/darknet-master/mytrain_letter/labels/'):  # 改成自己建立的myData
        os.makedirs('E:/darknet/darknet-master/mytrain_letter/labels/')
    image_ids = open('E:/darknet/darknet-master/mytrain_letter/ImageSets/Main/%s.txt'%(image_set)).read().strip().split()
    list_file = open('E:/darknet/darknet-master/mytrain_letter/%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('E:/darknet/darknet-master/mytrain_letter/JPEGImages/%s.jpg\n'%(image_id))
        convert_annotation(year, image_id)
    list_file.close()
import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'E:\\darknet\\darknet-master\\mytrain_letter\\Annotations'
txtsavepath = 'E:\\darknet\\darknet-master\\mytrain_letter\\ImageSets\\Main'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open('E:/darknet/darknet-master/mytrain_letter/ImageSets/Main/trainval.txt', 'w')
ftest = open('E:/darknet/darknet-master/mytrain_letter/ImageSets/Main/test.txt', 'w')
ftrain = open('E:/darknet/darknet-master/mytrain_letter/ImageSets/Main/train.txt', 'w')
fval = open('E:/darknet/darknet-master/mytrain_letter/ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

如果自己没有数据集,也可以拍摄一个视频的帧提取成图片进行自己的数据集的制作,可以用下面的代码3进行视频帧提取:

# -*- coding: utf-8 -*-
"""
Created on Fri Nov 20 14:02:07 2020

@author: SS
"""

# coding=utf-8

# 全局变量
VIDEO_PATH = 'C:\\Users\\SS\\Desktop\\sgcai_file\\test5.mp4' # 视频地址
EXTRACT_FOLDER = 'C:\\Users\\SS\Desktop\\sgcai_file\\result_video5\\pic' # 存放帧图片的位置
EXTRACT_FREQUENCY = 2 # 帧提取频率


def extract_frames(video_path, dst_folder, index):
    # 主操作
    import cv2
    video = cv2.VideoCapture()
    if not video.open(video_path):
        print("can not open the video")
        exit(1)
    count = 1
    while True:
        _, frame = video.read()
        if frame is None:
            break
        if count % EXTRACT_FREQUENCY == 0:
            save_path = "{}/{:>03d}.jpg".format(dst_folder, index)
            cv2.imwrite(save_path, frame)
            index += 1
        count += 1
    video.release()
    # 打印出所提取帧的总数
    print("Totally save {:d} pics".format(index-1))


def main():
    # 递归删除之前存放帧图片的文件夹,并新建一个
    import shutil
    try:
        shutil.rmtree(EXTRACT_FOLDER)
    except OSError:
        pass
    import os
    os.mkdir(EXTRACT_FOLDER)
    # 抽取帧图片,并保存到指定路径
    extract_frames(VIDEO_PATH, EXTRACT_FOLDER, 1)


if __name__ == '__main__':
    main()

第三步:修改参数:

①cfg文件:修改网络的参数,主要是训练的迭代次数、学习率和要检测的目标类别数量相关的参数

下面分为test和train两种模式,test的时候记得注释掉train的部分,train的时候注释掉test的部分

windows平台下darknet训练自己的数据集_第2张图片

filters=3*(5+len(classes))=3*(5+14)=57

下面总共有三处,每处都要修改。

windows平台下darknet训练自己的数据集_第3张图片

②data文件:检测类别数量、数据集的路径、类别对应名字、训练好的模型放置的路径

windows平台下darknet训练自己的数据集_第4张图片

③names文件:类别对应的名字

windows平台下darknet训练自己的数据集_第5张图片

第四步:下载预训练权重之后训练模型

下载:darknet53.conv.74

训练:darknet.exe detector train obj.data cfg/yolov3-obj.cfg darknet53.conv.74

第五步:python接口调用模型处理图片

# -*- coding: utf-8 -*-
"""
Created on Fri Jan  1 18:41:37 2021

@author: SS
"""

#coding:utf-8
import numpy as np
import cv2
import os
#%%
weightsPath='E:\\darknet\\darknet-master\\mytrain_letter\\yolov3-obj_final.weights'# 模型权重文件
configPath="E:\\darknet\\darknet-master\\build\\darknet\\x64\\cfg\\yolov3-obj.cfg"# 模型配置文件
labelsPath = "E:\\darknet\\darknet-master\\build\\darknet\\x64\\obj.names"# 模型类别标签文件
#%%
#初始化一些参数
LABELS = open(labelsPath).read().strip().split("\n")
print(LABELS)
#%%
boxes = []
confidences = []
classIDs = []

#加载 网络配置与训练的权重文件 构建网络
net = cv2.dnn.readNetFromDarknet(configPath,weightsPath)  
#读入待检测的图像
image = cv2.imread('E:\\darknet\\darknet-master\\mytrain_letter\\JPEGImages\\001.jpg')
#得到图像的高和宽
(H,W) = image.shape[0:2]
print(H,W)
#%%

# 得到 YOLO需要的输出层
ln = net.getLayerNames()
#%%
print(ln)
#%%
out = net.getUnconnectedOutLayers()#得到未连接层得序号  [[200] /n [267]  /n [400] ]
x = []
for i in out:   # 1=[200]
    x.append(ln[i[0]-1])    # i[0]-1    取out中的数字  [200][0]=200  ln(199)= 'yolo_82'
ln=x
print(ln)
# ln  =  ['yolo_82', 'yolo_94', 'yolo_106']  得到 YOLO需要的输出层
#%%


#从输入图像构造一个blob,然后通过加载的模型,给我们提供边界框和相关概率
#blobFromImage(image, scalefactor=None, size=None, mean=None, swapRB=None, crop=None, ddepth=None)
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416),swapRB=True, crop=False)#构造了一个blob图像,对原图像进行了图像的归一化,缩放了尺寸 ,对应训练模型

net.setInput(blob) #将blob设为输入??? 具体作用还不是很清楚
layerOutputs = net.forward(ln)  #ln此时为输出层名称  ,向前传播  得到检测结果

for output in layerOutputs:  #对三个输出层 循环
    for detection in output:  #对每个输出层中的每个检测框循环
        scores=detection[5:]  #detection=[x,y,h,w,c,class1,class2] scores取第6位至最后
        classID = np.argmax(scores)#np.argmax反馈最大值的索引
        confidence = scores[classID]
        if confidence >0.5:#过滤掉那些置信度较小的检测结果
            box = detection[0:4] * np.array([W, H, W, H])
            #print(box)
            (centerX, centerY, width, height)= box.astype("int")
            # 边框的左上角
            x = int(centerX - (width / 2))
            y = int(centerY - (height / 2))
            # 更新检测出来的框
            boxes.append([x, y, int(width), int(height)])
            confidences.append(float(confidence))
            classIDs.append(classID)


idxs=cv2.dnn.NMSBoxes(boxes, confidences, 0.2,0.3)
box_seq = idxs.flatten()#[ 2  9  7 10  6  5  4]
if len(idxs)>0:
    for seq in box_seq:
        (x, y) = (boxes[seq][0], boxes[seq][1])  # 框左上角
        (w, h) = (boxes[seq][2], boxes[seq][3])  # 框宽高
        if classIDs[seq]==0: #根据类别设定框的颜色
            color = [0,0,255]
        else:
            color = [0,255,0]
        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)  # 画框
        text = "{}: {:.4f}".format(LABELS[classIDs[seq]], confidences[seq])
        cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.3, color, 1)  # 写字
cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
cv2.imshow("Image", image)
cv2.waitKey(0)

 

你可能感兴趣的:(计算机视觉)