Yolo-Fastestv2训练及部署自己的数据集

欢天喜地小姐姐

已于 2024-05-22 09:38:23 修改

阅读量1.5k

点赞数 7

分类专栏：目标检测文章标签： python 计算机视觉人工智能

于 2024-05-16 16:17:30 首次发布

本文链接：https://blog.csdn.net/qq_41496108/article/details/138851982

版权

目标检测专栏收录该内容

5 篇文章 0 订阅

订阅专栏

Yolo-Fastestv2训练及部署自己的数据集

训练流程

一、配置环境

1.下载代码

git clone https://github.com/dog-qiuqiu/Yolo-FastestV2.git

2.创建虚拟环境，在虚拟环境下安装相关python依赖包

pip install -r requirements.txt -i https://pypi.mirrors.ustc.edu.cn/simple #采用国内源，速度较快

3.测试能否正常使用

python test.py --data data/coco.data --weights modelzoo/coco2017-0.241078ap-model.pth --img img/000139.jpg

二、数据集建立

1.数据集格式

数据集的格式是Yolo格式，每个图像对应一个.txt标签文件cx，cy，w，h。cx，cy是标签框的中心点的坐标，w，h是标签框的宽度和高度。示例如下：

0 0.410156 0.317708 0.807813 0.631250

2.图像与标签位置

数据集图像与对应的标签文件具有同样的名字，存储在同一目录中。结构如下：

 -train
   - 001.jpg
   - 001.txt
   - 002.jpg
   - 002.txt
 -val
   - 003.jpg
   - 003.txt
 -test
   - 004.jpg
   - 004.txt

3.生成数据集路径.txt文件

路径是图像的绝对路径，示例如下：

/home/dataset/train/001.jpg

python实现：

import os
images_path='/home/dataset/train/' 
txt_path='/home/dataset/train.txt'
with open(txt_path,'w') as f:
    files_name=os.listdir(images_path)
    f.truncate()
    for name in files_name:
        if name.endswith('.jpg'):
            f.write(images_path+name)
            f.write('\n')
f.close()

4.生成.names类别标签

./data/coco.names修改成自己数据集的类别。

5.最终构建的数据集目录结构

 -train
   - 001.jpg
   - 001.txt
   - 002.jpg
   - 002.txt
 -train.txt
 -val
   - 003.jpg
   - 003.txt
 -val.txt
 -test
   - 004.jpg
   - 004.txt
 -test.txt

三、训练

1.根据当前数据集生成锚点anchors

计算适应数据集的anchors，将输出值替换coco.data中的anchors

python genanchors.py --traintxt /home/dataset/train.txt

anchors6.txt文件将在当前目录中生成，anchors6.txt的示例内容如下：

12.64,19.39, 37.88,51.48, 55.71,138.31, 126.91,78.23, 131.57,214.55, 279.92,258.87  # anchor bias
0.636158       # iou

注意：anchor6.txt很重要

2.构建训练 .data配置文件

可以直接在./data/coco.data中修改（classes、anchors、train、val)

[name]
model_name=coco           # model name

[train-configure]
epochs=300          
steps=150,250             # Declining learning rate steps
batch_size=64             # batch size
subdivisions=1          
learning_rate=0.001      

[model-configure]
pre_weights=None          #预训练模型
classes=80                # 类别
width=352              
height=352                # The height of the model input image
anchor_num=3              # anchor num
anchors=12.64,19.39, 37.88,51.48, 55.71,138.31,
 126.91,78.23, 131.57,214.55, 279.92,258.87 #anchor bias

[data-configure]
train=/media/qiuqiu/D/coco/train2017.txt   # train dataset path .txt file
val=/media/qiuqiu/D/coco/val2017.txt       # val dataset path .txt file 
names=./data/coco.names                    # .names category label file

3.训练

python train.py --data data/coco.data

4.评估

计算map, F1

python evaluation.py --data data/coco.data --weights weights/.pth

5.预测

预测单张图片

python test.py --data data/coco.data --weights weights/.pth --img img/005.jpg

批量预测多张图片

python test.py --data data/coco.data --weights weights/.pth --img img/ --save save/

test_images.py:在test.py的基础上修改，代码如下：

import os
import cv2
import time
import argparse
import torch
import model.detector
import utils.utils
import os
#!/usr/bin/python3.7
if __name__ == '__main__':
    #指定训练配置文件
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='',
                        help='Specify training profile *.data')
    parser.add_argument('--weights', type=str, default='',
                        help='The path of the .pth model to be transformed')
    parser.add_argument('--img', type=str, default='',
                        help='The path of test image')
    parser.add_argument('--save',type=str,default='')
    opt = parser.parse_args()
    cfg = utils.utils.load_datafile(opt.data)
    assert os.path.exists(opt.weights), "请指定正确的模型路径"
    assert os.path.exists(opt.img), "请指定正确的测试图像路径"
    assert os.path.exists(opt.save),"请指定正确的测试图像路径"

    #模型加载
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.detector.Detector(cfg["classes"], cfg["anchor_num"], True).to(device)
    model.load_state_dict(torch.load(opt.weights, map_location=device))
    #sets the module in eval node
    model.eval()
    img_names=os.listdir(opt.img)
    for img_name in img_names:
         #数据预处理
         if img_name.endswith('.jpg'):
           ori_img = cv2.imread(opt.img+'/'+img_name)
         res_img = cv2.resize(ori_img, (cfg["width"], cfg["height"]), interpolation = cv2.INTER_LINEAR)
         img = res_img.reshape(1, cfg["height"], cfg["width"], 3)
         img = torch.from_numpy(img.transpose(0,3, 1, 2))
         img = img.to(device).float() / 255.0
         #模型推理
         preds = model(img)
         #特征图后处理
         output = utils.utils.handel_preds(preds, cfg, device)
         output_boxes = utils.utils.non_max_suppression(output, conf_thres = 0.3, iou_thres = 0.4)
         #加载label names
         LABEL_NAMES = []
         with open(cfg["names"], 'r') as f:
             for line in f.readlines():
                 LABEL_NAMES.append(line.strip())
         h, w, _ = ori_img.shape
         scale_h, scale_w = h / cfg["height"], w / cfg["width"]
    #绘制预测框
         for box in output_boxes[0]:
             box = box.tolist()
             obj_score = box[4]
             category = LABEL_NAMES[int(box[5])]
             x1, y1 = int(box[0] * scale_w), int(box[1] * scale_h)
             x2, y2 = int(box[2] * scale_w), int(box[3] * scale_h)
             cv2.rectangle(ori_img, (x1, y1), (x2, y2), (255, 255, 0), 2)
             cv2.putText(ori_img, '%.2f' % obj_score, (x1, y1 - 5), 0, 0.7, (0, 255, 0), 2)
             cv2.putText(ori_img, category, (x1, y1 - 25), 0, 0.7, (0, 255, 0), 2)
         cv2.imwrite(opt.save+'/'+img_name, ori_img)

NCNN部署流程

1.将.pth模型转化为onnx格式

python pytorch2onnx.py --data data/coco.data --weights weights/.pth --output ./ncnn/build/tools/onnx/yolo-fastestv2.onnx

2.onnx-sim

python -m onnxsim ./ncnn/build/tools/onnx/yolo-fastestv2.onnx ./ncnn/build/tools/onnx/yolo-fastestv2-opt.onnx

3.搭建NCNN

git clone https://github.com/Tencent/ncnn.git
cd ncnn
mkdir build
cd build
cmake ..
make
make install
cp -rf ./ncnn/build/install/* ./sample/ncnn

其中 cp参数含义：
-R/r：递归处理，将指定目录下的所有文件与子目录一并处理；
-f：强行复制文件或目录，不论目标文件或目录是否已存在；

4.将 onnx 转换成param 和 bin

ncnnoptimize 工具会自动将无用的 MemoryData 删除，并且自动将最终的 blob count 设置为合适的数量，顺便转为 fp16 存储减小模型体积

cd ./ncnn/build/tools/onnx
./onnx2ncnn yolo-fastestv2-opt.onnx yolo-fastestv2.param yolo-fastestv2.bin
cp yolo-fastestv2* ../
cd ../
./ncnnoptimize yolo-fastestv2.param yolo-fastestv2.bin yolo-fastestv2-opt.param yolo-fastestv2-opt.bin 1
cp yolo-fastestv2-opt* ../../../sample/ncnn/model

5.修改demo.cpp和yolo-fastestv2.cpp

./sample/ncnn/demo.cpp 修改
(1) static const char* class_names[ ] = { } 类别名
(2) api.loadModel中模型路径
(3) 测试图片路径 cv::Mat cvImg = cv::imread(“058.jpg”);
./sample/ncnn/src/yolo-fastestv2.cpp 修改
(1) numCategory = 80, 修改类别，否则报错
(2) std::vector bias{ } 修改为和anchors一样

6.运行

cd ./sample/ncnn
sh build.sh
./demo

7.批量运行，demo.cpp代码修改

部分代码如下：

   #include "yolo-fastestv2.h"
   #include <iostream>
   #include "opencv2/core.hpp"
   #include "opencv2/imgproc.hpp"
   #include <opencv2/core/utils/filesystem.hpp>
   #include <string>
   #include <typeinfo>
  //获取绝对文件路径的文件名
   std::string GetFileNameFromPath(const std::string & path){
        size_t found=path.find_last_of("/");
        if (found!=std::string::npos){
                return path.substr(found+1);
        }
        return path;

}
...
    api.loadModel(); 
    
    std::string img_dir="path/test/";
    std::string save_dir="path/save_test/";
    if (cv::utils::fs::exists(img_dir))
    {
        std::cout<<"该文件存在"<<std::endl;
     }
    //获取当前文件夹下指定格式的文件
    std::vector <cv::String> img_lists;
    cv::utils::fs::glob(img_dir,"*.jpg",img_lists);
    for(auto name:img_lists)
    {
    cv::Mat cvImg=cv::imread(name);
    
    std::vector<TargetBox> boxes;
    api.detection(cvImg, boxes);
    for (int i = 0; i < boxes.size(); i++) {         
        char text[256];
        sprintf(text, "%s %.1f%%", class_names[boxes[i].cate], boxes[i].score * 100);
        int baseLine = 0;
        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        int x = boxes[i].x1;
        int y = boxes[i].y1 - label_size.height - baseLine;
        if (y < 0)
            y = 0;
        if (x + label_size.width > cvImg.cols)
            x = cvImg.cols - label_size.width;
       cv::rectangle(cvImg, cv::Rect(cv::Point(x, y)，
                     cv::Size(label_size.width, label_size.height + baseLine)),
                     cv::Scalar(255, 255, 255), -1)
       cv::putText(cvImg, text, cv::Point(x, y + label_size.height),
                   cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
       cv::rectangle (cvImg, cv::Point(boxes[i].x1, boxes[i].y1),
                      cv::Point(boxes[i].x2, boxes[i].y2), cv::Scalar(255, 255, 0), 2, 2, 0);
    }
    
     std::string filename=GetFileNameFromPath(name);
     cv::imwrite(save_dir+filename,cvImg);
}
    return 0;
}