代码地址:RefineDet
相关链接:
- 《RefineDet:Single-Shot Refinement Neural Network for Object Detection》论文笔记
- RefineDet:(1)训练脚本解析
- RefineDet:(2)检测部分网络解析
1. 概述
Python版本的测试代码在相应的文件夹下已经写好了,这里根据对应的Python文件写了一个测试Demo。
2. 测试Demo
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <vector>
#include <map>
#include <memory>
#include <string>
#include <opencv2/opencv.hpp>
#include <caffe/blob.hpp>
#include <caffe/common.hpp>
#include <caffe/layer.hpp>
#include <caffe/net.hpp>
int main(int argc, char* argv[])
{
//GPU设置
int device_id = 0;
caffe::Caffe::set_mode(caffe::Caffe::GPU);
caffe::Caffe::SetDevice(device_id);
std::cout << "Use GPU: " << device_id << std::endl;
//网络均值
std::vector<float> pixel_mean={104.0, 117.0, 123.0};
//模型文件路径
std::string deploy_file = "./models/deploy.prototxt";
std::string model_file = "./models/refine_det.caffemodel";
//模型所需的宽高尺寸
int det_img_width = 512;
int det_img_height = 704;
//测试图像
std::string test_file = "./test.jpg";
//检测的概率阈值
float det_obj_thresh = 0.3;
//网络加载
std::shared_ptr<caffe::Net<float> > det_net;
det_net.reset(new caffe::Net<float>(deploy_file, caffe::TEST));
det_net->CopyTrainedLayersFrom(model_file);
//加载测试图片
cv::Mat src_img = cv::imread(test_file, cv::IMREAD_ANYCOLOR|cv::IMREAD_ANYDEPTH);
if(!img.data)
{
std::cout << "no image data: " << test_file<< std::endl;
continue;
}
//输入数据准备
cv::Mat det_img = src_img.clone();
det_img.convertTo(det_img, CV_32FC3);
cv::resize(det_img, det_img, cv::Size(det_img_width, det_img_height), CV_INTER_LINEAR);
caffe::Blob<float>* input_blob = det_net->input_blobs()[0]; //get input blob
input_blob->Reshape(1, det_img.channels(), det_img.rows, det_img.cols);
float *blob_data = input_blob->mutable_cpu_data();
const int cols = det_img.cols;
const int rows = det_img.rows;
for (int i = 0; i < rows; ++i)
{
const cv::Vec3f* data = det_img.ptr<cv::Vec3f>(i); //此处用指针读取Mat数据,并存到blob中
for (int j = 0; j < cols; ++j)
{
blob_data[(0 * rows + i)*cols + j] = data[j][0] - pixel_mean[0];
blob_data[(1 * rows + i)*cols + j] = data[j][1] - pixel_mean[1];
blob_data[(2 * rows + i)*cols + j] = data[j][2] - pixel_mean[2];
}
}
// 网络前向传播,并获取检测结果
float loss;
det_net->Forward(&loss);
boost::shared_ptr<caffe::Blob<float> > det_output = this->det_net->blob_by_name("detection_out");
// 抽取网络的检测结果
std::cout << "model detect boxes num:" << det_output->height() << std::endl;
for(int i=0; i<det_output->height(); ++i)
{
int label = det_output->data_at(0, 0, i, 1);
if (label < 1) continue; //background type
float conf = det_output->data_at(0, 0, i, 2);
int x_min = static_cast<int>(det_output->data_at(0, 0, i, 3)*float(src_img.cols));
int y_min = static_cast<int>(det_output->data_at(0, 0, i, 4)*float(src_img.rows));
int x_max = static_cast<int>(det_output->data_at(0, 0, i, 5)*float(src_img.cols));
int y_max = static_cast<int>(det_output->data_at(0, 0, i, 6)*float(src_img.rows));
if (conf < this->det_obj_thresh) continue;
//格式化存储检测框
......
}
return 0;
}