yolo模型和pytorch yolo模型使用

目录

• step1 制作数据集
• step2 训练模型
• step3 测试
• step4 可视化训练日志

Darknet深度学习框架是由Joseph Redmon提出的一个用C和CUDA编写的开源神经网络框架，具体的环境搭建可以参考之前写的一篇文章：
基本环境搭建成功后，就可以使用自己制作的数据集训练自己的yolo模型了。

文中出现的使用的已标注好的数据集来自：

step1 制作数据集

1、
（1）按照
中制作数据集的1、（1）（2）两步制作好自己的数据集，
然后使用以下代码将xml格式的标注文件转换为txt格式：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
classes = ["people","front","side","back"]#这里是你的所有分类的名称
myRoot = r'E:\Tensorflow_Codes'#这里是你项目的根目录
xmlRoot = myRoot +r'\Annotations'
txtRoot = myRoot + r'\labels'
imageRoot = myRoot + r'\JPEGImages'
 
def getFile_name(file_dir):
    L=[]
    for root, dirs, files in os.walk(file_dir):
        print(files)
        for file in files:
            if os.path.splitext(file)[1] == '.jpg':
                L.append(os.path.splitext(file)[0]) #L.append(os.path.join(root, file))
    return L
 
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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(image_id):
    in_file = open(xmlRoot + '\%s.xml' % (image_id))
 
    out_file = open(txtRoot + '\%s.txt' % (image_id), 'w')  # 生成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'):
        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((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
 
#image_ids_train = open('D:/darknet-master/scripts/VOCdevkit/voc/list.txt').read().strip().split('\',\'')  # list格式只有000000 000001
image_ids_train = getFile_name(imageRoot)
# image_ids_val = open('/home/*****/darknet/scripts/VOCdevkit/voc/list').read().strip().split()
 
 
list_file_train = open(myRoot +r'\ImageSets\Main\train.txt', 'w')
#list_file_val = open('boat_val.txt', 'w')
 
for image_id in image_ids_train:
    list_file_train.write(imageRoot + '\%s.jpg\n' % (image_id))
    convert_annotation(image_id)
list_file_train.close()  # 只生成训练集，自己根据自己情况决定
 
# for image_id in image_ids_val:
 
#    list_file_val.write('/home/*****/darknet/boat_detect/images/%s.jpg\n'%(image_id))
#    convert_annotation(image_id)
# list_file_val.close()

（2）转换成功后，在darknet-master\build\darknet\x64\data\文件夹下，创建obj文件夹，用于存放原始训练图片数据和txt格式的标注数据。

2、

在darknet-master\build\darknet\x64\data\文件夹下新建train.txt，用于存放原始训练图片的路径（绝对路径和相对路径都可以），步骤1中代码的最后部分也包含了自动写入train.txt的代码（执行前需要先在myRoot\ImageSets\Main\下新建一个train.txt文件）：

3、

在darknet-master\build\darknet\x64文件夹下新建yolo-obj.cfg文件（可以直接复制yolov3.cfg，然后重命名为yolo-obj.cfg）：

搜索yolo，一共出现三次，对每个yolo块都作如下修改：

filter=(5+classes)*3 根据自己的类别情况进行计算

classes 根据自己的分类数设置

random 默认为1，如果显存小于4G建议改为0

4、

在darknet-master\build\darknet\x64\data\文件夹下新建obj.names文件，写入你自己数据集的类别：

5、

在darknet-master\build\darknet\x64\data\文件夹下新建obj.data文件，写入类别数、训练数据集路径文件、类别标签文件、训练模型保存路径：

6、修改darknet-master目录下的Makefile文件：

step2 训练模型

1、
Win+R打开运行终端，输入cmd，切换到darknet-master\build\darknet\x64目录下，输入以下命令开始训练：
tee是在Linux系统下用于输出训练日志文件的命令（要先在darknet-master\build\darknet\x64下先建立一个文件夹visualization），但是该命令在windows系统下不起作用，因此需要额外加入该命令的配置：
链接：https://pan.baidu.com/s/1R3bZCEqXGHNDNggs9m93Rw 提取码：ef8t
下载压缩包后，解压，将UnxUtils\usr\local\wbi文件夹下的tee.exe文件复制到C:\Windows\System32即可。

darknet.exe detector train data/obj.data yolo-obj.cfg 2>&1 | tee visualization/train_yolov3.log

2、
开始训练后，注意观察loss曲线的变化，最后生成的权重文件位于darknet-master\build\darknet\x64\backup目录下。

step3 测试

由于选取的样本较少（30张图片），迭代次数1000次，最后的loss在5左右，预测结果很一般。

step4 可视化训练日志

第三步训练过程中输出了训练日志文件train_yolov3.log，根据此文件可以进一步画出loss曲线、IOU曲线等用于模型分析。

实现代码如下：

import pandas as pd
import matplotlib.pyplot as plt
import os

g_log_path = "train_yolov3.log"  # 此处修改为你的训练日志文件名

def extract_log(log_file, new_log_file, key_word):
    '''
    :param log_file:日志文件
    :param new_log_file:挑选出可用信息的日志文件
    :param key_word:根据关键词提取日志信息
    :return:
    '''
    with open(log_file, "r") as f:
        with open(new_log_file, "w") as train_log:
            for line in f:
                # 去除多gpu的同步log
                if "Syncing" in line:
                    continue
                # 去除nan log
                if "nan" in line:
                    continue
                if key_word in line:
                    train_log.write(line)
    f.close()
    train_log.close()


def drawAvgLoss(loss_log_path):
    '''
    :param loss_log_path: 提取到的loss日志信息文件
    :return: 画loss曲线图
    '''
    line_cnt = 0
    for count, line in enumerate(open(loss_log_path, "rU")):
        line_cnt += 1
    result = pd.read_csv(loss_log_path, skiprows=[iter_num for iter_num in range(line_cnt) if ((iter_num < 500))],
                         error_bad_lines=False,
                         names=["loss", "avg", "rate", "seconds", "images"])
    result["avg"] = result["avg"].str.split(" ").str.get(1)
    result["avg"] = pd.to_numeric(result["avg"])

    fig = plt.figure(1, figsize=(6, 4))
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(result["avg"].values, label="Avg Loss", color="#ff7043")
    ax.legend(loc="best")
    ax.set_title("Avg Loss Curve")
    ax.set_xlabel("Batches")
    ax.set_ylabel("Avg Loss")


def drawIOU(iou_log_path):
    '''
    :param iou_log_path: 提取到的iou日志信息文件
    :return: 画iou曲线图
    '''
    line_cnt = 0
    for count, line in enumerate(open(iou_log_path, "rU")):
        line_cnt += 1
    result = pd.read_csv(iou_log_path, skiprows=[x for x in range(line_cnt) if (x % 39 != 0 | (x < 5000))],
                         error_bad_lines=False,
                         names=["Region Avg IOU", "Class", "Obj", "No Obj", "Avg Recall", "count"])
    result["Region Avg IOU"] = result["Region Avg IOU"].str.split(": ").str.get(1)

    result["Region Avg IOU"] = pd.to_numeric(result["Region Avg IOU"])

    result_iou = result["Region Avg IOU"].values
    # 平滑iou曲线
    for i in range(len(result_iou) - 1):
        iou = result_iou[i]
        iou_next = result_iou[i + 1]
        if abs(iou - iou_next) > 0.2:
            result_iou[i] = (iou + iou_next) / 2

    fig = plt.figure(2, figsize=(6, 4))
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(result_iou, label="Region Avg IOU", color="#ff7043")
    ax.legend(loc="best")
    ax.set_title("Avg IOU Curve")
    ax.set_xlabel("Batches")
    ax.set_ylabel("Avg IOU")


if __name__ == "__main__":
    loss_log_path = "train_log_loss.txt"
    iou_log_path = "train_log_iou.txt"
    if os.path.exists(g_log_path) is False:
        exit(-1)
    if os.path.exists(loss_log_path) is False:
        extract_log(g_log_path, loss_log_path, "images")
    if os.path.exists(iou_log_path) is False:
        extract_log(g_log_path, iou_log_path, "IOU")
    drawAvgLoss(loss_log_path)
    drawIOU(iou_log_path)
    plt.show()

由自己训练模型时的训练日志绘制出的loss曲线图和IOU曲线图：

刚刚接触模型训练，还有很多细节需要不断学习调整