#include "你的头文件.h"
#include <iostream>
#include<opencv2//opencv.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include<math.h>
#include <opencv2/dnn.hpp>
#include <fstream>
#define USE_CUDA false
using namespace std;
using namespace cv;
using namespace dnn;
bool Yolo::readModel(Net& net, string& netPath, bool isCuda = false) {
try {
net = readNet(netPath);
}
catch (const std::exception&) {
return false;
}
//cuda
if (isCuda) {
net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
}
//cpu
else {
net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
}
return true;
}
bool Yolo::Detect(Mat& SrcImg, Net& net, vector<Output>& output) {
Mat blob;
int col = SrcImg.cols;
int row = SrcImg.rows;
int maxLen = MAX(col, row);
Mat netInputImg = SrcImg.clone();
if (maxLen > 1.2 * col || maxLen > 1.2 * row) {
Mat resizeImg = Mat::zeros(maxLen, maxLen, CV_8UC3);
SrcImg.copyTo(resizeImg(Rect(0, 0, col, row)));
netInputImg = resizeImg;
}
vector<Ptr<Layer> > layer;
vector<string> layer_names;
layer_names = net.getLayerNames();
blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(0, 0, 0), true, false);
net.setInput(blob);
std::vector<cv::Mat> netOutputImg;
net.forward(netOutputImg, net.getUnconnectedOutLayersNames());
std::vector<int> classIds;
std::vector<float> confidences;
std::vector<cv::Rect> boxes;
float ratio_h = (float)netInputImg.rows / netHeight;
float ratio_w = (float)netInputImg.cols / netWidth;
int net_width = className.size() + 5;
float* pdata = (float*)netOutputImg[0].data;
for (int stride = 0; stride < strideSize; stride++) { //stride
int grid_x = (int)(netWidth / netStride[stride]);
int grid_y = (int)(netHeight / netStride[stride]);
for (int anchor = 0; anchor < 3; anchor++) { //anchors
for (int i = 0; i < grid_y; i++) {
for (int j = 0; j < grid_x; j++) {
float box_score = pdata[4];
if (box_score >= boxThreshold) {
cv::Mat scores(1, className.size(), CV_32FC1, pdata + 5);
Point classIdPoint;
double max_class_socre;
minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
max_class_socre = max_class_socre;
if (max_class_socre >= classThreshold) {
float x = pdata[0];//(sigmoid_x(pdata[0]) * 2.f - 0.5f + j) * m_Mark_Stride[stride];//x
float y = pdata[1];//y
float w = pdata[2];//powf(sigmoid_x(pdata[2]) * 2.f, 2.f) * anchor_w*ratio_c;//w
float h = pdata[3];//h
int left = round((x - 0.5 * w) * ratio_w);
int top = round((y - 0.5 * h) * ratio_h);
int right = round((x + 0.5 * w) * ratio_w);
int bottom = round((y + 0.5 * h) * ratio_h);
int width = round(ratio_w);
int height = round(ratio_h);
classIds.push_back(classIdPoint.x);
confidences.push_back(max_class_socre * box_score);
boxes.push_back(Rect(left, top, round(w * ratio_w), round(h * ratio_h)));
}
}
pdata += net_width;
}
}
}
}
vector<int> nms_result;
NMSBoxes(boxes, confidences, nmsScoreThreshold, nmsThreshold, nms_result);
for (int i = 0; i < nms_result.size(); i++) {
int idx = nms_result[i];
Output result;
result.id = classIds[idx];
result.confidence = confidences[idx];
result.box = boxes[idx];
output.push_back(result);
}
if (output.size())
return true;
else
return false;
}
void segNet(Mat& img) {
String modelFile = "你的分割模型.onnx";
if (img.empty()) {
cout << "未找到图片！" << endl;
}
// 记录原始图片大小
int org_width = img.cols;
int org_height = img.rows;
// 计算缩放比例
double scale = 1.0;
if (org_width >= org_height) {
scale = (double)512 / org_width;
}
else {
scale = (double)512 / org_height;
}
// 缩放图片，并保持长宽比
Mat resized_img;
resize(img, resized_img, Size(), scale, scale);
// 计算灰度填充的大小
int pad_x = (512 - resized_img.cols) / 2;
int pad_y = (512 - resized_img.rows) / 2;
// 加上灰色填充
Mat padded_img = Mat::zeros(512, 512, CV_8UC1);
copyMakeBorder(resized_img, padded_img, pad_y, pad_y, pad_x, pad_x, BORDER_CONSTANT, Scalar(128));
padded_img.convertTo(padded_img, CV_32FC3, 1.f / 255.f);
cv::cvtColor(padded_img, padded_img, cv::COLOR_BGR2RGB);
Mat inputBolb = blobFromImage(padded_img);
dnn::Net Unet = cv::dnn::readNetFromONNX(modelFile);
Unet.setInput(inputBolb);
Mat output = Unet.forward();
int N = inputBolb.size[0];
int C = inputBolb.size[2];
int H = inputBolb.size[2];
int W = inputBolb.size[3];
Mat predMat = Mat::zeros(512, 512, CV_32F);
for (int h = 0; h < H; h++) {
for (int w = 0; w < W; w++) {
float bg = output.ptr<float>(0, 0, h)[w];
float c1 = output.ptr<float>(0, 1, h)[w];
float c2 = output.ptr<float>(0, 2, h)[w];
if (bg >= c1 && bg >= c2) {//背景
predMat.at<float>(h, w) = 255;
}
else if (c1 >= c2) {//类别1
predMat.at<float>(h, w) = 180;
}
else {//类别二
predMat.at<float>(h, w) = 0;
}
}
}
// 去除灰色填充并还原原始大小
Rect roi(pad_x, pad_y, org_width * scale, org_height * scale);
Mat roi_img = predMat(roi);
resize(roi_img, roi_img, Size(org_width, org_height));
//保存语义分割图
imwrite("seg.png", roi_img);
}
void Yolo::drawPred(Mat& img, vector<Output> result) {
for (int i = 0; i < result.size(); i++) {
int x1, x2, y1, y2;
x1 = result[i].box.x;
x2 = result[i].box.x + result[i].box.width;
y1 = result[i].box.y;
y2 = result[i].box.y + result[i].box.height;
std::cout << x1 << "," << y1 << "," << x2 << "," << y2 << endl;
int center_x = result[i].box.x + 0.5 * result[i].box.width;
int center_y = result[i].box.y + 0.5 * result[i].box.height;
int color_num = i;
rectangle(img, result[i].box, (0, 255, 0), 2, 8);
string label = className[result[i].id] + ":" + to_string(result[i].confidence);
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
y1 = max(y1, labelSize.height);
putText(img, label, Point(x1, y1), FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2);
// 截取图像中指定位置的矩形区域
cv::Rect roi(x1, y1, result[i].box.width, result[i].box.height);
cv::Mat Detect = img(roi);
cv::imwrite("目标检测图.png", Detect);
segNet(Detect);
}
}
int main()
{
string img_path = "你的图片.png";
string model_path = "你的yolo文件路径best.onnx";
Yolo test;
Net net;
if (test.readModel(net, model_path, USE_CUDA)) {
cout << "read net ok!" << endl;
}
else {
cout << "read onnx model failed!";
return -1;
}
vector<Output> result;
Mat rgb_image = imread(img_path);
if (test.Detect(rgb_image, net, result)) {
test.drawPred(rgb_image, result);
}
else {
cout << "No target object!" << endl;
}
return 0;
}
