Ubuntu下使用yolo训练自己的数据集

Ubuntu下使用Yolov3训练自己的数据集

主要内容:
使用yolov3卷积神经网络训练自己的数据集,从而实现道路环境检测。主要检测的目标有三类,分别是行人、路坑以及道路上的石头。

目录
1.采集数据集
2.搭建训练环境
3.下载darknet源码
4.下载预训练模型
5.处理数据集
6.修改网络配置文件
7 . 训练模型
8. 训练注意的问题

1.采集数据集
采集含有目标物的图片,制作成数据集。数据集的数量尽可能的多,包含的场景尽量多,这样训练出的模型效果会更好。 使用labelme软件进行数据集的标注,可以标注多个类别。 最终的数据集格式是一张图片对应一个xml,数据格式如下图所示:所有图片存在一个文件夹,所有xml存在一个文件夹,且图片与xml一一对应。注意事项,图片的名称和标注的类别不能含有中文字符。

2.搭建训练环境
原本打算在虚拟机上训练,由于显卡驱动安装失败,放弃。选择装双系统:win10+Ubuntu16.04。本机GPU为gtx1650,官网找到对应显卡驱动并安装。配置cuda10.0和cudnn,opencv4.0。在配置cuda时注意,由于以及安装了显卡驱动,故有一项关于驱动的安装应选择NO。

3.下载darknet源码
终端输入命令

git clone https://github.com/pjreddie/darknet

进入到darknet目录,修改makefile文件。需要使能GPU,cudnn和opencv。安装opencv4.0版本,make时报错,显示不支持c++11。这里将不使能opencv。
Ubuntu下使用yolo训练自己的数据集_第1张图片
修改cuda路径,即nvcc文件路径。
Ubuntu下使用yolo训练自己的数据集_第2张图片
开始make编译,编译成功如下
Ubuntu下使用yolo训练自己的数据集_第3张图片4.下载预训练模型

wget https://pjreddie.com/media/files/darknet53.conv.74
./darknet detect cfg/yolov3.cfg darknet53.conv.74 data/dog.jpg

显示如下错误提示
在这里插入图片描述终端输入,解决问题,我也不知道啥意思

export LD_LIBRARY_PATH="/usr/local/cuda-10.0/lib64"

继续以上操作,显示如下错误提示,显存不足
Ubuntu下使用yolo训练自己的数据集_第4张图片

解决办法:
需要修改所使用的模型cfg文件中的subdivision的参数。
由subdivisions=8改成subdivisions=64。
subdivision:这个参数很有意思的,它会让你的每一个batch不是一下子都丢到网络里。而是分成subdivision对应数字的份数,一份一份的跑完后,在一起打包算作完成一次iteration。这样会降低对显存的占用情况。如果设置这个参数为1的话就是一次性把所有batch的图片都丢到网络里,如果为2的话就是一次丢一半。

cd cfg
gedit yolov3.cfg
继续之前操作,可以看到YOLO的运行结果。到这里,YOLOV3已经走通了,是时候开始加入自己的数据了。

5.处理数据集
5.1按照voc2007数据集格式制作文件夹,如下图所示。Annotatios为xml文件夹,JPEGImages为图片文件夹。
Ubuntu下使用yolo训练自己的数据集_第5张图片ImageSets里Main文件夹,用到4个文件:

  • train.txt 是用来训练的图片文件的文件名列表 ,约trainval的50%

  • val.txt是用来验证的图片文件的文件名列表 ,约trainval的50%

  • trianval.txt是用来训练和验证的图片文件的文件名列表 ,约整个数据集的50%

  • test.txt 是用来测试的图片文件的文件名列表 ,约整个数据集的50%

使用Matlab做VOC里的txt文件

clear;clc;  
%%%%%%%%% 第一类检测目标pedestrian 制作TXT文件 %%%%%%%%%%%%%%%
file1 = dir('G:\My_Project\train_yolov3_bymyself\VOC2007\Annotations\pedestrian');  % pedestrian xml文件地址xml文件地址
path1 = 'G:\My_Project\train_yolov3_bymyself\VOC2007\ImageSets\Main\pedestrian';     % txt文件存放的地址
len1 = length(file1)-2;  

num1_trainval=sort(randperm(len1, floor(1*len1/2)));      %trainval集占所有数据的1/2,可以根据需要设置  
num1_train=sort(num1_trainval(randperm(length(num1_trainval), floor(1*length(num1_trainval)/2))));    %train集占trainval集的1/2,可以根据需要设置  
num1_val=setdiff(num1_trainval,num1_train);  %trainval集剩下的作为val集  
num1_test=setdiff(1:len1,num1_trainval);%剩下的作为test集  

fid=fopen(strcat(path1, '_trainval.txt'),'a+');  
for i=1:length(num1_trainval)  
    s = sprintf('%s',file1(num1_trainval(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path1, '_train.txt'),'a+');  
for i=1:length(num1_train)  
    s = sprintf('%s',file1(num1_train(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path1, '_val.txt'),'a+');  
for i=1:length(num1_val)  
    s = sprintf('%s',file1(num1_val(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path1, '_test.txt'),'a+');  
for i=1:length(num1_test)  
    s = sprintf('%s',file1(num1_test(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

%%%%%%%%% 第二类检测目标chuckhole 制作TXT文件 %%%%%%%%%%%%%%%
file2 = dir('G:\My_Project\train_yolov3_bymyself\VOC2007\Annotations\chuckhole');  % chuckhole xml文件地址xml文件地址
path2 = 'G:\My_Project\train_yolov3_bymyself\VOC2007\ImageSets\Main\chuckhole';     % txt文件存放的地址
len2 = length(file2)-2;  

num2_trainval=sort(randperm(len2, floor(1*len2/2)));      %trainval集占所有数据的1/2,可以根据需要设置  
num2_train=sort(num2_trainval(randperm(length(num2_trainval), floor(1*length(num2_trainval)/2))));    %train集占trainval集的1/2,可以根据需要设置  
num2_val=setdiff(num2_trainval,num2_train);  %trainval集剩下的作为val集  
num2_test=setdiff(1:len2,num2_trainval);%剩下的作为test集  

fid=fopen(strcat(path2, '_trainval.txt'),'a+');  
for i=1:length(num2_trainval)  
    s = sprintf('%s',file2(num2_trainval(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path2, '_train.txt'),'a+');  
for i=1:length(num2_train)  
    s = sprintf('%s',file2(num2_train(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path2, '_val.txt'),'a+');  
for i=1:length(num2_val)  
    s = sprintf('%s',file2(num2_val(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path2, '_test.txt'),'a+');  
for i=1:length(num2_test)  
    s = sprintf('%s',file2(num2_test(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

%%%%%%%%% 第三类检测目标stone 制作TXT文件 %%%%%%%%%%%%%%%
file3 = dir('G:\My_Project\train_yolov3_bymyself\VOC2007\Annotations\stone');  % stone xml文件地址xml文件地址
path3 = 'G:\My_Project\train_yolov3_bymyself\VOC2007\ImageSets\Main\stone';     % txt文件存放的地址
len3 = length(file3)-2;  

num3_trainval=sort(randperm(len3, floor(1*len3/2)));      %trainval集占所有数据的1/2,可以根据需要设置  
num3_train=sort(num3_trainval(randperm(length(num3_trainval), floor(1*length(num3_trainval)/2))));    %train集占trainval集的1/2,可以根据需要设置  
num3_val=setdiff(num3_trainval,num3_train);      %trainval集剩下的作为val集  
num3_test=setdiff(1:len3,num3_trainval);         %剩下的作为test集  

fid=fopen(strcat(path3, '_trainval.txt'),'a+');  
for i=1:length(num3_trainval)  
    s = sprintf('%s',file3(num3_trainval(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path3, '_train.txt'),'a+');  
for i=1:length(num3_train)  
    s = sprintf('%s',file3(num3_train(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path3, '_val.txt'),'a+');  
for i=1:length(num3_val)  
    s = sprintf('%s',file3(num3_val(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid);  

fid=fopen(strcat(path3, '_test.txt'),'a+');  
for i=1:length(num3_test)  
    s = sprintf('%s',file3(num3_test(i)+2).name);  
    fprintf(fid,[s(1:length(s)-4) '\n']);  
end  
fclose(fid); 

5.2加入自己的数据集
在代码的darknet目录下新建VOCdevkit文件夹,然后把刚才制作的VOC2007文件夹拷贝到该文件夹下。有的读者可能了解过YOLOV3的label,YOLOV3的label标注的一行五个数分别代表类别(从 0 开始编号), BoundingBox 中心 X 坐标,中心 Y 坐标,宽,高。这些坐标都是 0~1 的相对坐标。和我们刚才标注的label不同,因此我们需要下面的py文件帮我们转换label。

wget https://pjreddie.com/media/files/voc_label.py

这里需要修改两个地方,sets和classes,classes根据自己需要修改。
Ubuntu下使用yolo训练自己的数据集_第6张图片
由于报错出现width 和hight有为0,我就在siza[0] 和size[1]后加0.00000001。不知道对后面训练有什么影响,暂时先这样处理吧。运行没有报错。

6.修改配置文件
6.1 修改cfg/voc.data。
位于darknet/cfg文件夹下,分别将train、val的路径改为scripts文件夹下边的2007_train.txt、2007_val.txt的路径。back 修改为自己的路径,输出的权重信息将存储这个文件内。
Ubuntu下使用yolo训练自己的数据集_第7张图片
6.2修改参数文件cfg/yolov3-voc.cfg,
Class是类别数
Filters = 3*(class + 5)
该文件共有六处需要修改
位于605行 611行,689行 695行,773行 779行

参数文件开头的地方可以选训练的batchsize,要注意! Testing注释掉 打开Training
Ubuntu下使用yolo训练自己的数据集_第8张图片
6.3修改voc.names文件
在这里插入图片描述
7.训练模型

./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg darknet53.conv.74

发生一个错误,并解决
Ubuntu下使用yolo训练自己的数据集_第9张图片又发生错误,cuda显存不足
Ubuntu下使用yolo训练自己的数据集_第10张图片
需要修改所使用的模型cfg文件中的subdivision的参数。

解决方法:由subdivisions=8改成subdivisions=16,64等,继续运行

哈哈,又报错了
Ubuntu下使用yolo训练自己的数据集_第11张图片
继续改呀,这是因为voc.data文件中的路径不是绝对路径,找不到文件。只要打开voc.data(在cfg文件夹下)把两个路径改为如下形式(根据自己实际情况)即可:
train = /home/你的用户名/darknet/scripts/2018_train.txt
valid = /home/你的用户名/darknet/scripts/2018_val.txt

然而还报错,还是一大堆,我,,,我再坚持下。。
Ubuntu下使用yolo训练自己的数据集_第12张图片发现问题所在了,被自己蠢哭了。我的数据集图片里有一部分是自己收集的,我做的是jpg格式。还有一部分是直接引用别人的,当时没注意观察是png格式的。而前面生成的文件名都是需要jpg格式的,故这里有一部分文件打不开。

8 训练注意问题
8.1 训练参数说明
终于可以正常训练了,上张训练过程的截图。训练默认的是前10000轮每1000轮保存一次模型,10000轮后每10000轮保存一次模型。可以修改examples/detector.c文件的138行。修改完重新编译一下,在darknet目录下执行make。
在这里插入图片描述
Region xx:cfg文件中yolo-layer的索引;
Avg IOU: 当前迭代中,预测的box与标注的box的平均交并比,越大越好,期望数值为1;
Class:标注物体的分类准确率,越大越好,期望数值为1;
obj:越大越好,期望数值为1;
No obj:越小越好;
.5R:以IOU=0.5为阈值时候的recall; recall = 检出的正样本/实际的正样本
0.75R:以IOU=0.75为阈值时候的recall;
count:正样本数目

9216: 指示当前训练的迭代次数

0.071332: 是总体的Loss(损失)

0.42692 avg: 是平均Loss,这个数值应该越低越好,一般来说,一旦这个数值低于0.060730 avg就可以终止训练了。也可当看到平均损失0.xxxxxx avg在多次迭代时不再减少时,停止训练。

0.000100 rate: 代表当前的学习率,是在.cfg文件中定义的。

11.317444 seconds: 表示当前批次训练花费的总时间。

589824 images: 这一行最后的这个数值表示到目前为止,参与训练的图片的总量。

在Region 82 Avg IOU、Region 94 Avg IOU、Region 106 Avg IOU中出现很多nan

在训练过程中,nan的屏幕占比30%是正常的,如果太大,全是nan,则就是训练出了问题
解决方法一:在显存允许的情况下,可以适当增加batch(darknet-master/yolov3-voc.cfg中的batch)的大小(要视自己数据集 的大小情况来增加batch,不能盲目的改大),这样能够一定程度减少nan的出现。
解决方法二:增加数据集的规模。若是对于10类以内的图片,500张以内的训练集未必是太少了,因此可以增加数据集,实在不 行的话就进行数据增强,把数据集扩展到原来的几倍到几十倍不等。

8.2 训练的暂定与继续
不保存训练日志的训练执行命令

./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg scripts/darknet53.conv.74

保存训练日志的训练执行命令

./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg scripts/darknet53.conv.74 | tee train_yolov3-voc.log

ubuntu系统下,在训练终端处,使用ctrl+c终止训练。
在终端输入一下命令,接着上一步停止处的训练状态继续训练

./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg backup/yolov3-voc.backup

临时暂停: ctrl+z 暂停训练,fg继续训练。

8.3 训练日志的可视化
 (1)如果在训练的时候已经保存了训练日志,可以使用以下python脚本对训练日志进行可视化,得loss变化曲线和Avg IOU曲线(Avg IOU是交并比,标记的框和预测的框重复的部分除以他们的和,这个值越接近1越好)
visualization_train_yolov3-voc_log.py

# -*- coding: utf-8 -*-
# @Func    :yolov3 训练日志可视化,把该脚本和日志文件放在同一目录下运行。
 
import pandas as pd
import matplotlib.pyplot as plt
import os
 
# ==================可能需要修改的地方=====================================#
g_log_path = "train_yolov3-voc.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()
 

(2)将上述python脚本文件和训练日志放在同一目录下,打开此目录下的终端,运行上述.py文件可以得到loss变化曲线和Avg IOU变化曲线。同时,在当前目录下生成了train_log_iou.txt和train_log_loss.txt文件。

9 转化c++代码
9.1 图像转换和缩放
yolo采用自定义的image格式进行图像读取和处理,而一般我们工程中使用较多的是OpenCV或者指向图像数据的指针,因此此处先对图像转换和缩放操作进行修改,代码如下:
imprcoess.h

#ifndef IMPROCESS_H
#define IMPROCESS_H
 
#include
 
void imgConvert(const cv::Mat& img, float* dst);
 
void imgResize(float* src, float* dst,int srcWidth,int srcHeight,int dstWidth,int dstHeight);
 
void resizeInner(float *src, float* dst,int srcWidth,int srcHeight,int dstWidth,int dstHeight);
 
#endif // IMPROCESS_H

imprcoess.cpp

#include

void imgConvert(const cv::Mat& img, float* dst){
    uchar *data = img.data;
    int h = img.rows;
    int w = img.cols;
    int c = img.channels();

    for(int k= 0; k < c; ++k){
        for(int i = 0; i < h; ++i){
            for(int j = 0; j < w; ++j){
                dst[k*w*h+i*w+j] = data[(i*w + j)*c + k]/255.;
            }
        }
    }
}

void imgResize(float *src, float* dst,int srcWidth,int srcHeight,int dstWidth,int dstHeight){
    int new_w = srcWidth;
    int new_h = srcHeight;
    if (((float)dstWidth/srcWidth) < ((float)dstHeight/srcHeight)) {
        new_w = dstWidth;
        new_h = (srcHeight * dstWidth)/srcWidth;
    } else {
        new_h = dstHeight;
        new_w = (srcWidth * dstHeight)/srcHeight;
    }

    float* ImgReInner;
    size_t sizeInner=new_w*new_h*3*sizeof(float);
    ImgReInner=(float*)malloc(sizeInner);
    resizeInner(src,ImgReInner,srcWidth,srcHeight,new_w,new_h);

    for(int i=0;i<dstWidth*dstHeight*3;i++){
        dst[i]=0.5;
    }

    for(int k = 0; k < 3; ++k){
        for(int y = 0; y < new_h; ++y){
            for(int x = 0; x < new_w; ++x){
                float val = ImgReInner[k*new_w*new_h+y*new_w+x];
                dst[k*dstHeight*dstWidth + ((dstHeight-new_h)/2+y)*dstWidth + (dstWidth-new_w)/2+x]=val;
            }
        }
    }
    free(ImgReInner);
}

void resizeInner(float *src, float* dst,int srcWidth,int srcHeight,int dstWidth,int dstHeight){
    float* part;
    size_t sizePa=dstWidth*srcHeight*3*sizeof(float);
    part=(float*)malloc(sizePa);

    float w_scale = (float)(srcWidth - 1) / (dstWidth - 1);
    float h_scale = (float)(srcHeight - 1) / (dstHeight - 1);

    for(int k = 0; k < 3; ++k){
        for(int r = 0; r < srcHeight; ++r){
            for(int c = 0; c < dstWidth; ++c){
                float val = 0;
                if(c == dstWidth-1 || srcWidth == 1){
                    val=src[k*srcWidth*srcHeight+r*srcWidth+srcWidth-1];
                } else {
                    float sx = c*w_scale;
                    int ix = (int) sx;
                    float dx = sx - ix;
                    val=(1 - dx) * src[k*srcWidth*srcHeight+r*srcWidth+ix] + dx * src[k*srcWidth*srcHeight+r*srcWidth+ix+1];
                }
                part[k*srcHeight*dstWidth + r*dstWidth + c]=val;
            }
        }
    }

    for(int k = 0; k < 3; ++k){
        for(int r = 0; r < dstHeight; ++r){
            float sy = r*h_scale;
            int iy = (int) sy;
            float dy = sy - iy;
            for(int c = 0; c < dstWidth; ++c){
                float val = (1-dy) * part[k*dstWidth*srcHeight+iy*dstWidth+c];
                dst[k*dstWidth*dstHeight + r*dstWidth + c]=val;
            }
            if(r == dstHeight-1 || srcHeight == 1)
                continue;
            for(int c = 0; c < dstWidth; ++c){
                float val = dy * part[k*dstWidth*srcHeight+(iy+1)*dstWidth+c];
                dst[k*dstWidth*dstHeight + r*dstWidth + c]+=val;
            }
        }
    }
    free(part);
}


其中,imgConvert函数将OpenCV的图像由RGBRGBRGB…转化为yolo的RRRGGGBBB…格式(由代码可知,yolo输入图像的像素取值范围为0~1)。imgResize函数将图像缩放到cfg指定的网络输入的大小。代码修改自yolo的源码,将其image格式改为我们需要的指针形式。

9.2 调用darknet
为了让代码跑通,我们首先用OpenCV读取视频,然后将OpenCV的图像转为指针指向的数据格式(如果想直接采用OpenCV可自行修改)。模型是darknet官网上的模型,当然你可以换成自己训练的模型。代码如下:
darknet.h 文件在include目录下已包含,此处不再列出。

main.cpp

在这里插入代码片
```#include<iostream>
#include
#include
#include

using namespace std;
using namespace cv;

float colors[6][3] = { {1,0,1}, {0,0,1},{0,1,1},{0,1,0},{1,1,0},{1,0,0} };

float get_color(int c, int x, int max)
{
    float ratio = ((float)x/max)*5;
    int i = floor(ratio);
    int j = ceil(ratio);
    ratio -= i;
    float r = (1-ratio) * colors[i][c] + ratio*colors[j][c];
    return r;
}


int main()
{
    string cfgfile = "/home/oliver/darknet-master/cfg/yolov3.cfg";//读取模型文件,请自行修改相应路径
    string weightfile = "/home/oliver/darknet-master/yolov3.weights";
    float thresh=0.5;//参数设置
    float nms=0.35;
    int classes=80;

    network *net=load_network((char*)cfgfile.c_str(),(char*)weightfile.c_str(),0);//加载网络模型
    set_batch_network(net, 1);
    VideoCapture capture(0);//读取视频,请自行修改相应路径
    capture.set(CV_CAP_PROP_FRAME_WIDTH,1920);
    capture.set(CV_CAP_PROP_FRAME_HEIGHT,1080);
    Mat frame;
    Mat rgbImg;

    vector<string> classNamesVec;
    ifstream classNamesFile("/home/oliver/darknet-master/data/coco.names");//标签文件coco有80类

    if (classNamesFile.is_open()){
        string className = "";
        while (getline(classNamesFile, className))
            classNamesVec.push_back(className);
    }

    bool stop=false;
    while(!stop)
    {
        cout<<frame.size<<endl;
        if (!capture.read(frame))
        {
            printf("fail to read.\n");
            return 0;
        }
        cvtColor(frame, rgbImg, cv::COLOR_BGR2RGB);

        float* srcImg;
        size_t srcSize=rgbImg.rows*rgbImg.cols*3*sizeof(float);
        srcImg=(float*)malloc(srcSize);

        imgConvert(rgbImg,srcImg);//将图像转为yolo形式

        float* resizeImg;
        size_t resizeSize=net->w*net->h*3*sizeof(float);
        resizeImg=(float*)malloc(resizeSize);
        imgResize(srcImg,resizeImg,frame.cols,frame.rows,net->w,net->h);//缩放图像

        network_predict(net,resizeImg);//网络推理
        int nboxes=0;
        detection *dets=get_network_boxes(net,rgbImg.cols,rgbImg.rows,thresh,0.5,0,1,&nboxes);

        if(nms){
            do_nms_sort(dets,nboxes,classes,nms);
        }

        vector<cv::Rect>boxes;
        boxes.clear();
        vector<int>classNames;

        for (int i = 0; i < nboxes; i++){
            bool flag=0;
            int className;
            for(int j=0;j<classes;j++){
                if(dets[i].prob[j]>thresh){
                    if(!flag){
                        flag=1;
                        className=j;
                    }
                }
            }
            if(flag)
            {
                int left = (dets[i].bbox.x - dets[i].bbox.w / 2.)*frame.cols;
                int right = (dets[i].bbox.x + dets[i].bbox.w / 2.)*frame.cols;
                int top = (dets[i].bbox.y - dets[i].bbox.h / 2.)*frame.rows;
                int bot = (dets[i].bbox.y + dets[i].bbox.h / 2.)*frame.rows;

                if (left < 0)
                    left = 0;
                if (right > frame.cols - 1)
                    right = frame.cols - 1;
                if (top < 0)
                    top = 0;
                if (bot > frame.rows - 1)
                    bot = frame.rows - 1;

                Rect box(left, top, fabs(left - right), fabs(top - bot));
                boxes.push_back(box);
                classNames.push_back(className);
            }
        }
        free_detections(dets, nboxes);

        for(int i=0;i<boxes.size();i++)
        {
            int offset = classNames[i]*123457 % 80;
            float red = 255*get_color(2,offset,80);
            float green = 255*get_color(1,offset,80);
            float blue = 255*get_color(0,offset,80);

            rectangle(frame,boxes[i],Scalar(blue,green,red),2);

            String label = String(classNamesVec[classNames[i]]);
            int baseLine = 0;
            Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
            putText(frame, label, Point(boxes[i].x, boxes[i].y + labelSize.height),
                    FONT_HERSHEY_SIMPLEX, 1, Scalar(red, blue, green),2);
        }
//        Mat resize_img;
//        resize(frame,resize_img,cv::Size(f_width,f_height),(0,0),(0,0),cv::INTER_LINEAR);
//        cout<
        namedWindow("video",0);
        imshow("video",frame);


        int c=waitKey(30);
              if((char)c==27)
                  break;
              else if(c>=0)
                  waitKey(0);

        free(srcImg);
        free(resizeImg);
    }
    free_network(net);
    capture.release();
    return 1;
}

待续。。。。。

你可能感兴趣的:(Ubuntu下使用yolo训练自己的数据集)