YoloDet.h
#ifndef NCNN_H_
#define NCNN_H_
#include "net.h"
#include "benchmark.h"
#include
//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& 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& 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 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 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;
}