yolo-fastestv1

YoloDet.h

#ifndef NCNN_H_
#define NCNN_H_

#include "net.h"
#include "benchmark.h"

#include <opencv2/opencv.hpp>

//Model input image size
#define IW 320  //width
#define IH 320  //height

//cpu num threads
#define NUMTHREADS 8

typedef struct _TargetBox {
    int x1;        //left
    int y1;        //top
    int x2;        //right
    int y2;        //bottom
 
    int cate;      //category
    float score;   //Confidence level
}TargetBox;

class YoloDet {
public:

    //model init
    int init(const char* paramPath, const char* binPath);
    //body detect
    int detect(cv::Mat& srcImg, std::vector<TargetBox>& dstBoxes);

private:

};


#endif  //NCNN_H_

YoloDet.cpp

#include "YoloDet.h"

ncnn::Net DNet;

static ncnn::PoolAllocator g_blob_pool_allocator;
static ncnn::PoolAllocator g_workspace_pool_allocator;

int YoloDet::init(const char* paramPath, const char* binPath)
{
    printf("Ncnn mode init:\n %s \n %s\n", paramPath, binPath);
    
    g_blob_pool_allocator.clear();
    g_workspace_pool_allocator.clear();

    g_blob_pool_allocator.set_size_compare_ratio(0.0f);
    g_workspace_pool_allocator.set_size_compare_ratio(0.0f);

    DNet.load_param(paramPath);
    DNet.load_model(binPath);
    return 0;
}

int YoloDet::detect(cv::Mat& srcImg, std::vector<TargetBox>& dstBoxes)
{
    ncnn::Mat in = ncnn::Mat::from_pixels_resize(srcImg.data, ncnn::Mat::PIXEL_BGR2RGB,\
                                                 srcImg.cols, srcImg.rows, IW, IH);
    
    //Normalization of input image data
    const float mean_vals[3] = {0.f, 0.f, 0.f};
    const float norm_vals[3] = {1/255.f, 1/255.f, 1/255.f};
    in.substract_mean_normalize(mean_vals, norm_vals);
     
    //Forward
    ncnn::Mat out;
    ncnn::Extractor ex = DNet.create_extractor();
    ex.set_num_threads(NUMTHREADS);
    ex.input("data", in);
    ex.extract("output", out);
    
    //doresult
    dstBoxes.resize(out.h);
    for (int i = 0; i < out.h; i++)
    {
        const float* values = out.row(i);

        dstBoxes[i].cate = values[0];
        dstBoxes[i].score = values[1];

        dstBoxes[i].x1 = values[2] * srcImg.cols;
        dstBoxes[i].y1 = values[3] * srcImg.rows;
        dstBoxes[i].x2 = values[4] * srcImg.cols;
        dstBoxes[i].y2 = values[5] * srcImg.rows;
    }
    return 0;
}

demo.cpp

#include "YoloDet.h"

int drawBoxes(cv::Mat srcImg, std::vector<TargetBox> boxes)
{   
    printf("Detect box num: %d\n", boxes.size());
    for (int i = 0; i < boxes.size(); i++)
    {
       cv::rectangle (srcImg, 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 cate =std::to_string(boxes[i].cate);
        std::string Ttext = "Category:" + cate;
        cv::Point Tp = cv::Point(boxes[i].x1, boxes[i].y1-20);
        cv::putText(srcImg, Ttext, Tp, cv::FONT_HERSHEY_TRIPLEX, 0.5, 
                                   cv::Scalar(0, 255, 0), 1, CV_AA);

        std::string score =std::to_string(boxes[i].score);
        std::string Stext = "Score:" + score;
        cv::Point Sp = cv::Point(boxes[i].x1, boxes[i].y1-5);
        cv::putText(srcImg, Stext, Sp, cv::FONT_HERSHEY_TRIPLEX, 0.5, 
                                   cv::Scalar(0, 0, 255), 1, CV_AA);

    }
    return 0;
}


int testCam() {
    YoloDet api;
    //Init model
    api.init("model/yolo-fastest-1.1_body.param", 
              "model/yolo-fastest-1.1_body.bin");

    cv::Mat frame;    
    std::vector<TargetBox> output;

    cv::VideoCapture cap(0);

    while (true) {
        printf("=========================\n");
        cap >> frame;
        if (frame.empty()) break; //如果某帧为空则退出循环

        double start = ncnn::get_current_time();
        api.detect(frame, output);
        double end = ncnn::get_current_time();
        double time = end - start;
        printf("Detect Time:%7.2f \n",time);

        drawBoxes(frame, output);
        output.clear();
        
        cv::imshow("demo", frame);
        cv::waitKey(20);
    }

    cap.release();//释放资源
    return 0;
}

int main() {
    testCam();
    return 0;
}