Update main.cpp

This commit is contained in:
hpc203
2023-07-30 16:37:12 +08:00
committed by GitHub
parent 5767b8bb19
commit f7c7606aee

294
main.cpp
View File

@@ -1,56 +1,56 @@
#define _CRT_SECURE_NO_WARNINGS #define _CRT_SECURE_NO_WARNINGS
#include <iostream> #include <iostream>
#include <fstream> #include <fstream>
#include <numeric> #include <numeric>
#include <opencv2/imgproc.hpp> #include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp> #include <opencv2/highgui.hpp>
//#include <cuda_provider_factory.h> ///nvidia-cuda<EFBFBD><EFBFBD><EFBFBD><EFBFBD> //#include <cuda_provider_factory.h> ///nvidia-cuda加速
#include <onnxruntime_cxx_api.h> #include <onnxruntime_cxx_api.h>
using namespace cv; using namespace cv;
using namespace std; using namespace std;
using namespace Ort; using namespace Ort;
typedef struct BoxInfo typedef struct BoxInfo
{ {
int xmin; int xmin;
int ymin; int ymin;
int xmax; int xmax;
int ymax; int ymax;
float score; float score;
string name; string name;
} BoxInfo; } BoxInfo;
class Detic class Detic
{ {
public: public:
Detic(string modelpath); Detic(string modelpath);
vector<BoxInfo> detect(Mat cv_image); vector<BoxInfo> detect(Mat cv_image);
private: private:
void preprocess(Mat srcimg); void preprocess(Mat srcimg);
vector<float> input_image_; vector<float> input_image_;
int inpWidth; int inpWidth;
int inpHeight; int inpHeight;
vector<string> class_names; vector<string> class_names;
const int max_size = 800; const int max_size = 800;
//<EFBFBD><EFBFBD><EFBFBD>ʼ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>õĿ<EFBFBD>ִ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> //存储初始化获得的可执行网络
Env env = Env(ORT_LOGGING_LEVEL_ERROR, "Head Pose Estimation"); Env env = Env(ORT_LOGGING_LEVEL_ERROR, "Detic");
Ort::Session *ort_session = nullptr; Ort::Session *ort_session = nullptr;
SessionOptions sessionOptions = SessionOptions(); SessionOptions sessionOptions = SessionOptions();
vector<char*> input_names; vector<char*> input_names;
vector<char*> output_names; vector<char*> output_names;
vector<vector<int64_t>> input_node_dims; // >=1 outputs vector<vector<int64_t>> input_node_dims; // >=1 outputs
vector<vector<int64_t>> output_node_dims; // >=1 outputs vector<vector<int64_t>> output_node_dims; // >=1 outputs
}; };
Detic::Detic(string model_path) Detic::Detic(string model_path)
{ {
//OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0); ///nvidia-cuda<EFBFBD><EFBFBD><EFBFBD><EFBFBD> //OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0); ///nvidia-cuda加速
sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC); sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
std::wstring widestr = std::wstring(model_path.begin(), model_path.end()); ///<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>windowsϵͳ<EFBFBD><EFBFBD><EFBFBD><EFBFBD>ôд std::wstring widestr = std::wstring(model_path.begin(), model_path.end()); ///如果在windows系统就这么写
ort_session = new Session(env, widestr.c_str(), sessionOptions); ///<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>windowsϵͳ<EFBFBD><EFBFBD><EFBFBD><EFBFBD>ôд ort_session = new Session(env, widestr.c_str(), sessionOptions); ///如果在windows系统就这么写
///ort_session = new Session(env, model_path.c_str(), sessionOptions); ///<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>linuxϵͳ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ôд ///ort_session = new Session(env, model_path.c_str(), sessionOptions); ///如果在linux系统就这么写
size_t numInputNodes = ort_session->GetInputCount(); size_t numInputNodes = ort_session->GetInputCount();
size_t numOutputNodes = ort_session->GetOutputCount(); size_t numOutputNodes = ort_session->GetOutputCount();
@@ -70,46 +70,46 @@ Detic::Detic(string model_path)
auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo(); auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
auto output_dims = output_tensor_info.GetShape(); auto output_dims = output_tensor_info.GetShape();
output_node_dims.push_back(output_dims); output_node_dims.push_back(output_dims);
} }
ifstream ifs("imagenet_21k_class_names.txt"); ifstream ifs("imagenet_21k_class_names.txt");
string line; string line;
while (getline(ifs, line)) while (getline(ifs, line))
{ {
this->class_names.push_back(line); ///<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ÿ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>RGBֵ this->class_names.push_back(line); ///你可以用随机数给每个类别分配RGB
} }
} }
void Detic::preprocess(Mat srcimg) void Detic::preprocess(Mat srcimg)
{ {
Mat dstimg; Mat dstimg;
cvtColor(srcimg, dstimg, COLOR_BGR2RGB); cvtColor(srcimg, dstimg, COLOR_BGR2RGB);
int im_h = srcimg.rows; int im_h = srcimg.rows;
int im_w = srcimg.cols; int im_w = srcimg.cols;
float oh, ow, scale; float oh, ow, scale;
if (im_h < im_w) if (im_h < im_w)
{ {
scale = (float)max_size / (float)im_h; scale = (float)max_size / (float)im_h;
oh = max_size; oh = max_size;
ow = scale * (float)im_w; ow = scale * (float)im_w;
} }
else else
{ {
scale = (float)max_size / (float)im_h; scale = (float)max_size / (float)im_h;
oh = scale * (float)im_h; oh = scale * (float)im_h;
ow = max_size; ow = max_size;
} }
float max_hw = std::max(oh, ow); float max_hw = std::max(oh, ow);
if (max_hw > max_size) if (max_hw > max_size)
{ {
scale = (float)max_size / max_hw; scale = (float)max_size / max_hw;
oh *= scale; oh *= scale;
ow *= scale; ow *= scale;
} }
resize(dstimg, dstimg, Size(int(ow + 0.5), int(oh + 0.5)), INTER_LINEAR); resize(dstimg, dstimg, Size(int(ow + 0.5), int(oh + 0.5)), INTER_LINEAR);
this->inpHeight = dstimg.rows; this->inpHeight = dstimg.rows;
this->inpWidth = dstimg.cols; this->inpWidth = dstimg.cols;
this->input_image_.resize(this->inpWidth * this->inpHeight * dstimg.channels()); this->input_image_.resize(this->inpWidth * this->inpHeight * dstimg.channels());
int k = 0; int k = 0;
for (int c = 0; c < 3; c++) for (int c = 0; c < 3; c++)
@@ -123,71 +123,71 @@ void Detic::preprocess(Mat srcimg)
k++; k++;
} }
} }
} }
} }
vector<BoxInfo> Detic::detect(Mat srcimg) vector<BoxInfo> Detic::detect(Mat srcimg)
{ {
int im_h = srcimg.rows; int im_h = srcimg.rows;
int im_w = srcimg.cols; int im_w = srcimg.cols;
this->preprocess(srcimg); this->preprocess(srcimg);
array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth }; array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth };
auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size()); Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
// <EFBFBD><EFBFBD>ʼ<EFBFBD><EFBFBD><EFBFBD><EFBFBD> // 开始推理
vector<Value> ort_outputs = ort_session->Run(RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size()); vector<Value> ort_outputs = ort_session->Run(RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size());
const float *pred_boxes = ort_outputs[0].GetTensorMutableData<float>(); const float *pred_boxes = ort_outputs[0].GetTensorMutableData<float>();
const float *scores = ort_outputs[1].GetTensorMutableData<float>(); const float *scores = ort_outputs[1].GetTensorMutableData<float>();
const int *pred_classes = ort_outputs[2].GetTensorMutableData<int>(); const int *pred_classes = ort_outputs[2].GetTensorMutableData<int>();
//const float *pred_masks = ort_outputs[3].GetTensorMutableData<float>(); //const float *pred_masks = ort_outputs[3].GetTensorMutableData<float>();
int num_box = ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape()[0]; int num_box = ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape()[0];
const float scale_x = float(im_w) / float(inpWidth); const float scale_x = float(im_w) / float(inpWidth);
const float scale_y = float(im_h) / float(inpHeight); const float scale_y = float(im_h) / float(inpHeight);
vector<BoxInfo> preds; vector<BoxInfo> preds;
for (int i = 0; i < num_box; i++) for (int i = 0; i < num_box; i++)
{ {
float xmin = pred_boxes[i * 4] * scale_x; float xmin = pred_boxes[i * 4] * scale_x;
float ymin = pred_boxes[i * 4 + 1] * scale_y; float ymin = pred_boxes[i * 4 + 1] * scale_y;
float xmax = pred_boxes[i * 4 + 2] * scale_x; float xmax = pred_boxes[i * 4 + 2] * scale_x;
float ymax = pred_boxes[i * 4 + 3] * scale_y; float ymax = pred_boxes[i * 4 + 3] * scale_y;
xmin = std::min(std::max(xmin, 0.f), float(im_w)); xmin = std::min(std::max(xmin, 0.f), float(im_w));
ymin = std::min(std::max(ymin, 0.f), float(im_h)); ymin = std::min(std::max(ymin, 0.f), float(im_h));
xmax = std::min(std::max(xmax, 0.f), float(im_w)); xmax = std::min(std::max(xmax, 0.f), float(im_w));
ymax = std::min(std::max(ymax, 0.f), float(im_h)); ymax = std::min(std::max(ymax, 0.f), float(im_h));
const float threshold = 0; const float threshold = 0;
const float width = xmax - xmin; const float width = xmax - xmin;
const float height = ymax - ymin; const float height = ymax - ymin;
if (width > threshold && height > threshold) if (width > threshold && height > threshold)
{ {
preds.push_back({ int(xmin), int(ymin), int(xmax), int(ymax), scores[i], class_names[pred_classes[i]] }); preds.push_back({ int(xmin), int(ymin), int(xmax), int(ymax), scores[i], class_names[pred_classes[i]] });
} }
} }
return preds; return preds;
} }
int main() int main()
{ {
Detic mynet("weights/Detic_C2_R50_640_4x_in21k.onnx"); Detic mynet("weights/Detic_C2_R50_640_4x_in21k.onnx");
string imgpath = "desk.jpg"; string imgpath = "desk.jpg";
Mat srcimg = imread(imgpath); Mat srcimg = imread(imgpath);
vector<BoxInfo> preds = mynet.detect(srcimg); vector<BoxInfo> preds = mynet.detect(srcimg);
for (size_t i = 0; i < preds.size(); ++i) for (size_t i = 0; i < preds.size(); ++i)
{ {
rectangle(srcimg, Point(preds[i].xmin, preds[i].ymin), Point(preds[i].xmax, preds[i].ymax), Scalar(0, 0, 255), 2); rectangle(srcimg, Point(preds[i].xmin, preds[i].ymin), Point(preds[i].xmax, preds[i].ymax), Scalar(0, 0, 255), 2);
string label = format("%.2f", preds[i].score); string label = format("%.2f", preds[i].score);
label = preds[i].name + " :" + label; label = preds[i].name + " :" + label;
putText(srcimg, label, Point(preds[i].xmin, preds[i].ymin - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 1); putText(srcimg, label, Point(preds[i].xmin, preds[i].ymin - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 1);
} }
//imwrite("result.jpg", srcimg); //imwrite("result.jpg", srcimg);
static const string kWinName = "Deep learning object detection in ONNXRuntime"; static const string kWinName = "Deep learning object detection in ONNXRuntime";
namedWindow(kWinName, WINDOW_NORMAL); namedWindow(kWinName, WINDOW_NORMAL);
imshow(kWinName, srcimg); imshow(kWinName, srcimg);
waitKey(0); waitKey(0);
destroyAllWindows(); destroyAllWindows();
} }