libtorch推理
#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include <torch/script.h>
#include <torch/torch.h>
const std::vector<std::string> class_names = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" }; //类别名称
const int input_width = 640;
const int input_height = 640;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int num_classes = class_names.size();
const int output_numprob = 5 + num_classes;
const int output_numbox = 3 * (input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32);
torch::DeviceType device = at::kCUDA;
//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
//预处理
void pre_process(cv::Mat & image, std::vector<torch::jit::IValue> & inputs)
{
cv::Vec4d params;
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
cv::cvtColor(letterbox, letterbox, cv::COLOR_BGR2RGB);
letterbox.convertTo(letterbox, CV_32FC3, 1.0f / 255.0f);
torch::Tensor input = torch::from_blob(letterbox.data, { 1, letterbox.rows, letterbox.cols, letterbox.channels() }).to(device);
input = input.permute({ 0, 3, 1, 2 }).contiguous();
inputs.emplace_back(input);
}
//网络推理
void process(std::string model, std::vector<torch::jit::IValue> & inputs, torch::jit::IValue & outputs)
{
torch::jit::script::Module module = torch::jit::load(model);
module.to(device);
outputs = module.forward(inputs);
}
//NMS
void nms(std::vector<cv::Rect> & boxes, std::vector<float> & scores, float score_threshold, float nms_threshold, std::vector<int> & indices)
{
assert(boxes.size() == scores.size());
struct BoxScore
{
cv::Rect box;
float score;
int id;
};
std::vector<BoxScore> boxes_scores;
for (size_t i = 0; i < boxes.size(); i++)
{
BoxScore box_conf;
box_conf.box = boxes[i];
box_conf.score = scores[i];
box_conf.id = i;
if (scores[i] > score_threshold) boxes_scores.push_back(box_conf);
}
std::sort(boxes_scores.begin(), boxes_scores.end(), [](BoxScore a, BoxScore b) { return a.score > b.score; });
std::vector<float> area(boxes_scores.size());
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
area[i] = boxes_scores[i].box.width * boxes_scores[i].box.height;
}
std::vector<bool> isSuppressed(boxes_scores.size(), false);
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
if (isSuppressed[i]) continue;
for (size_t j = i + 1; j < boxes_scores.size(); ++j)
{
if (isSuppressed[j]) continue;
float x1 = (std::max)(boxes_scores[i].box.x, boxes_scores[j].box.x);
float y1 = (std::max)(boxes_scores[i].box.y, boxes_scores[j].box.y);
float x2 = (std::min)(boxes_scores[i].box.x + boxes_scores[i].box.width, boxes_scores[j].box.x + boxes_scores[j].box.width);
float y2 = (std::min)(boxes_scores[i].box.y + boxes_scores[i].box.height, boxes_scores[j].box.y + boxes_scores[j].box.height);
float w = (std::max)(0.0f, x2 - x1);
float h = (std::max)(0.0f, y2 - y1);
float inter = w * h;
float ovr = inter / (area[i] + area[j] - inter);
if (ovr >= nms_threshold) isSuppressed[j] = true;
}
}
for (int i = 0; i < boxes_scores.size(); ++i)
{
if (!isSuppressed[i]) indices.push_back(boxes_scores[i].id);
}
}
//box缩放到原图尺寸
void scale_boxes(cv::Rect & box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width - size.width * gain) / 2;
int pad_h = (input_height - size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}
//可视化函数
void draw_result(cv::Mat & image, std::string label, cv::Rect box)
{
cv::rectangle(image, box, cv::Scalar(255, 0, 0), 1);
int baseLine;
cv::Size label_size = cv::getTextSize(label, 1, 1, 1, &baseLine);
cv::Point tlc = cv::Point(box.x, box.y);
cv::Point brc = cv::Point(box.x, box.y + label_size.height + baseLine);
cv::putText(image, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 1);
}
//后处理
void post_process(cv::Mat & image, cv::Mat & result, torch::jit::IValue & outputs)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
torch::Tensor preds = outputs.toTuple()->elements()[0].toTensor().to(at::kCPU); //1*25200*85
float* data = new float[preds[0].numel()];
std::copy(preds[0].data_ptr<float>(), preds[0].data_ptr<float>() + preds[0].numel(), data);
for (int i = 0; i < output_numbox; ++i)
{
float* ptr = data + i * output_numprob;
float objness = ptr[4];
if (objness < confidence_threshold)
continue;
float* classes_scores = ptr + 5;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) - classes_scores;
float max_class_score = classes_scores[class_id];
float score = max_class_score * objness;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);
for (int i = 0; i < indices.size(); i++)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
std::string label = class_names[class_ids[idx]] + ":" + cv::format("%.2f", scores[idx]); //class_ids[idx]是class_id
draw_result(result, label, box);
}
//std::vector<torch::Tensor> output;
//for (size_t i = 0; i < preds.sizes()[0]; ++i)
//{
// torch::Tensor pred = preds.select(0, i);
// // Filter by scores
// torch::Tensor scores = pred.select(1, 4) * std::get<0>(torch::max(pred.slice(1, 5, pred.sizes()[1]), 1));
// pred = torch::index_select(pred, 0, torch::nonzero(scores > score_threshold).select(1, 0));
// if (pred.sizes()[0] == 0) continue;
// // (center_x, center_y, w, h) to (left, top, right, bottom)
// pred.select(1, 0) = pred.select(1, 0) - pred.select(1, 2) / 2;
// pred.select(1, 1) = pred.select(1, 1) - pred.select(1, 3) / 2;
// pred.select(1, 2) = pred.select(1, 0) + pred.select(1, 2);
// pred.select(1, 3) = pred.select(1, 1) + pred.select(1, 3);
// // Computing scores and classes
// std::tuple<torch::Tensor, torch::Tensor> max_tuple = torch::max(pred.slice(1, 5, pred.sizes()[1]), 1);
// pred.select(1, 4) = pred.select(1, 4) * std::get<0>(max_tuple);
// pred.select(1, 5) = std::get<1>(max_tuple);
// torch::Tensor dets = pred.slice(1, 0, 6);
// torch::Tensor keep = torch::empty({ dets.sizes()[0] });
// torch::Tensor areas = (dets.select(1, 3) - dets.select(1, 1)) * (dets.select(1, 2) - dets.select(1, 0));
// std::tuple<torch::Tensor, torch::Tensor> indexes_tuple = torch::sort(dets.select(1, 4), 0, 1);
// torch::Tensor indexes = std::get<1>(indexes_tuple);
// int count = 0;
// while (indexes.sizes()[0] > 0)
// {
// keep[count] = (indexes[0].item().toInt());
// count += 1;
// // Computing overlaps
// torch::Tensor lefts = torch::empty(indexes.sizes()[0] - 1);
// torch::Tensor tops = torch::empty(indexes.sizes()[0] - 1);
// torch::Tensor rights = torch::empty(indexes.sizes()[0] - 1);
// torch::Tensor bottoms = torch::empty(indexes.sizes()[0] - 1);
// torch::Tensor widths = torch::empty(indexes.sizes()[0] - 1);
// torch::Tensor heights = torch::empty(indexes.sizes()[0] - 1);
// for (size_t i = 0; i < indexes.sizes()[0] - 1; ++i)
// {
// lefts[i] = std::max(dets[indexes[0]][0].item().toFloat(), dets[indexes[i + 1]][0].item().toFloat());
// tops[i] = std::max(dets[indexes[0]][1].item().toFloat(), dets[indexes[i + 1]][1].item().toFloat());
// rights[i] = std::min(dets[indexes[0]][2].item().toFloat(), dets[indexes[i + 1]][2].item().toFloat());
// bottoms[i] = std::min(dets[indexes[0]][3].item().toFloat(), dets[indexes[i + 1]][3].item().toFloat());
// widths[i] = std::max(float(0), rights[i].item().toFloat() - lefts[i].item().toFloat());
// heights[i] = std::max(float(0), bottoms[i].item().toFloat() - tops[i].item().toFloat());
// }
// torch::Tensor overlaps = widths * heights;
// // FIlter by IOUs
// torch::Tensor ious = overlaps / (areas.select(0, indexes[0].item().toInt()) + torch::index_select(areas, 0, indexes.slice(0, 1, indexes.sizes()[0])) - overlaps);
// indexes = torch::index_select(indexes, 0, torch::nonzero(ious <= nms_threshold).select(1, 0) + 1);
// }
// keep = keep.toType(torch::kInt64);
// output.push_back(torch::index_select(dets, 0, keep.slice(0, 0, count)));
//}
//for (size_t i = 0; i < output[0].sizes()[0]; i++)
//{
// float left = output[0][i][0].item().toFloat();
// float top = output[0][i][1].item().toFloat();
// float right = output[0][i][2].item().toFloat();
// float down = output[0][i][3].item().toFloat();
// float score = output[0][i][4].item().toFloat();
// int class_id = output[0][i][5].item().toInt();
// cv::Rect box = cv::Rect(left, top, right - left, down - top);
// scale_boxes(box, image.size());
// std::string label = class_name[class_id] + ":" + cv::format("%.2f", score);
// draw_result(result, label, box);
//}
}
int main(int argc, char* argv[])
{
cv::Mat image = cv::imread("bus.jpg");
std::vector<torch::jit::IValue> inputs;
pre_process(image, inputs);
std::string model = "yolov5n.torchscript";
torch::jit::IValue outputs;
process(model, inputs, outputs);
cv::Mat result = image.clone();
post_process(image, result, outputs);
cv::imwrite("result.jpg", result);
return 0;
}
onnxruntime推理
import cv2
import numpy as np
import onnxruntime
class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush'] #coco80类别
input_shape = (640, 640)
score_threshold = 0.2
nms_threshold = 0.5
confidence_threshold = 0.2
def nms(boxes, scores, score_threshold, nms_threshold):
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[0]
keep.append(i)
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= nms_threshold)[0]
index = index[idx + 1]
return keep
def xywh2xyxy(x):
y = np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
def filter_box(outputs): #过滤掉无用的框
outputs = np.squeeze(outputs)
outputs = outputs[outputs[..., 4] > confidence_threshold]
classes_scores = outputs[..., 5:]
boxes = []
scores = []
class_ids = []
for i in range(len(classes_scores)):
class_id = np.argmax(classes_scores[i])
outputs[i][4] *= classes_scores[i][class_id]
outputs[i][5] = class_id
if outputs[i][4] > score_threshold:
boxes.append(outputs[i][:6])
scores.append(outputs[i][4])
class_ids.append(outputs[i][5])
boxes = np.array(boxes)
boxes = xywh2xyxy(boxes)
scores = np.array(scores)
indices = nms(boxes, scores, score_threshold, nms_threshold)
output = boxes[indices]
return output
def letterbox(im, new_shape=(416, 416), color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = (new_shape[1] - new_unpad[0])/2, (new_shape[0] - new_unpad[1])/2 # wh padding
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return im
def scale_boxes(boxes, shape):
# Rescale boxes (xyxy) from input_shape to shape
gain = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) # gain = old / new
pad = (input_shape[1] - shape[1] * gain) / 2, (input_shape[0] - shape[0] * gain) / 2 # wh padding
boxes[..., [0, 2]] -= pad[0] # x padding
boxes[..., [1, 3]] -= pad[1] # y padding
boxes[..., :4] /= gain
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
return boxes
def draw(image, box_data):
box_data = scale_boxes(box_data, image.shape)
boxes = box_data[...,:4].astype(np.int32)
scores = box_data[...,4]
classes = box_data[...,5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 1)
cv2.putText(image, '{0} {1:.2f}'.format(class_names[cl], score), (top, left), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1)
if __name__=="__main__":
image = cv2.imread('bus.jpg')
input = letterbox(image, input_shape)
input = input[:, :, ::-1].transpose(2, 0, 1).astype(dtype=np.float32) #BGR2RGB和HWC2CHW
input = input / 255.0
input = np.expand_dims(input, axis=0)
onnx_session = onnxruntime.InferenceSession('yolov5n.onnx', providers=['CPUExecutionProvider'])
input_name = []
for node in onnx_session.get_inputs():
input_name.append(node.name)
output_name = []
for node in onnx_session.get_outputs():
output_name.append(node.name)
inputs = {}
for name in input_name:
inputs[name] = input
outputs = onnx_session.run(None, inputs)
boxes = filter_box(outputs)
draw(image, boxes)
cv2.imwrite('result.jpg', image)
#include <iostream>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
const std::vector<std::string> class_names = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" }; //类别名称
const int input_width = 640;
const int input_height = 640;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int num_classes = class_names.size();
const int output_numprob = 5 + num_classes;
const int output_numbox = 3 * (input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32);
//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
//预处理
void pre_process(cv::Mat & image, std::vector<float> & inputs)
{
cv::Vec4d params;
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
//inputs.resize(letterbox.channels() * letterbox.cols * letterbox.rows);
//for (int c = 0; c < letterbox.channels(); c++)
//{
// for (int i = 0; i < letterbox.cols; i++)
// {
// for (int j = 0; j < letterbox.rows; j++)
// {
// float pix = letterbox.ptr<uchar>(i)[j * 3 + 2 - c];//转换通道，输入onnx模型的图片通道顺序是RGB，但是opencv存储默认是BGR
// inputs[c * letterbox.cols * letterbox.rows + i * letterbox.rows + size_t(j)] = pix / 255.0;//归一化
// }
// }
//}
cv::cvtColor(letterbox, letterbox, cv::COLOR_BGR2RGB);
letterbox.convertTo(letterbox, CV_32FC3, 1.0f / 255.0f);
std::vector<cv::Mat> split_images;
cv::split(letterbox, split_images);
for (size_t i = 0; i < letterbox.channels(); ++i)
{
std::vector<float> split_image_data = split_images[i].reshape(1, 1);
inputs.insert(inputs.end(), split_image_data.begin(), split_image_data.end());
}
}
//网络推理
void process(const wchar_t* model, std::vector<float> & inputs, std::vector<Ort::Value> & outputs)
{
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "yolov5n");
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);//设置线程数
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);//启用模型优化策略
//CUDA option set
OrtCUDAProviderOptions cuda_option;
cuda_option.device_id = 0;
cuda_option.arena_extend_strategy = 0;
cuda_option.cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive;
cuda_option.gpu_mem_limit = SIZE_MAX;
cuda_option.do_copy_in_default_stream = 1;
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
session_options.AppendExecutionProvider_CUDA(cuda_option);
Ort::Session session(env, model, session_options);
Ort::AllocatorWithDefaultOptions allocator;//通过使用Ort::AllocatorWithDefaultOptions类，可以方便地进行内存分配和管理，而无需手动实现内存分配和释放的细节。
std::vector<const char*> input_node_names;
for (size_t i = 0; i < session.GetInputCount(); i++)
{
input_node_names.push_back(session.GetInputName(i, allocator));
}
std::vector<const char*> output_node_names;
for (size_t i = 0; i < session.GetOutputCount(); i++)
{
output_node_names.push_back(session.GetOutputName(i, allocator));
}
// create input tensor object from data values
std::vector<int64_t> input_node_dims = { 1, 3, input_width, input_height };
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, inputs.data(), inputs.size(), input_node_dims.data(), input_node_dims.size());
std::vector<Ort::Value> ort_inputs;
ort_inputs.push_back(std::move(input_tensor));//右值引用，避免不必要的拷贝和内存分配操作
// score model & input tensor, get back output tensor
outputs = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), ort_inputs.data(), input_node_names.size(), output_node_names.data(), output_node_names.size());
}
//NMS
void nms(std::vector<cv::Rect>& boxes, std::vector<float>& scores, float score_threshold, float nms_threshold, std::vector<int>& indices)
{
assert(boxes.size() == scores.size());
struct BoxScore
{
cv::Rect box;
float score;
int id;
};
std::vector<BoxScore> boxes_scores;
for (size_t i = 0; i < boxes.size(); i++)
{
BoxScore box_conf;
box_conf.box = boxes[i];
box_conf.score = scores[i];
box_conf.id = i;
if (scores[i] > score_threshold) boxes_scores.push_back(box_conf);
}
std::sort(boxes_scores.begin(), boxes_scores.end(), [](BoxScore a, BoxScore b) { return a.score > b.score; });
std::vector<float> area(boxes_scores.size());
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
area[i] = boxes_scores[i].box.width * boxes_scores[i].box.height;
}
std::vector<bool> isSuppressed(boxes_scores.size(), false);
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
if (isSuppressed[i]) continue;
for (size_t j = i + 1; j < boxes_scores.size(); ++j)
{
if (isSuppressed[j]) continue;
float x1 = (std::max)(boxes_scores[i].box.x, boxes_scores[j].box.x);
float y1 = (std::max)(boxes_scores[i].box.y, boxes_scores[j].box.y);
float x2 = (std::min)(boxes_scores[i].box.x + boxes_scores[i].box.width, boxes_scores[j].box.x + boxes_scores[j].box.width);
float y2 = (std::min)(boxes_scores[i].box.y + boxes_scores[i].box.height, boxes_scores[j].box.y + boxes_scores[j].box.height);
float w = (std::max)(0.0f, x2 - x1);
float h = (std::max)(0.0f, y2 - y1);
float inter = w * h;
float ovr = inter / (area[i] + area[j] - inter);
if (ovr >= nms_threshold) isSuppressed[j] = true;
}
}
for (int i = 0; i < boxes_scores.size(); ++i)
{
if (!isSuppressed[i]) indices.push_back(boxes_scores[i].id);
}
}
//box缩放到原图尺寸
void scale_boxes(cv::Rect & box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width - size.width * gain) / 2;
int pad_h = (input_height - size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}
//可视化函数
void draw_result(cv::Mat & image, std::string label, cv::Rect box)
{
cv::rectangle(image, box, cv::Scalar(255, 0, 0), 1);
int baseLine;
cv::Size label_size = cv::getTextSize(label, 1, 1, 1, &baseLine);
cv::Point tlc = cv::Point(box.x, box.y);
cv::Point brc = cv::Point(box.x, box.y + label_size.height + baseLine);
cv::putText(image, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 1);
}
//后处理
void post_process(cv::Mat & image, cv::Mat & result, std::vector<Ort::Value> & outputs)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
for (int i = 0; i < output_numbox; ++i)
{
float* ptr = const_cast<float*> (outputs[0].GetTensorData<float>()) + i * output_numprob;
float objness = ptr[4];
if (objness < confidence_threshold)
continue;
float* classes_scores = ptr + 5;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) - classes_scores;
float max_class_score = classes_scores[class_id];
float score = max_class_score * objness;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);
for (int i = 0; i < indices.size(); i++)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
std::string label = class_names[class_ids[idx]] + ":" + cv::format("%.2f", scores[idx]); //class_ids[idx]是class_id
draw_result(result, label, box);
}
}
int main(int argc, char* argv[])
{
cv::Mat image = cv::imread("bus.jpg");
std::vector<float> inputs;
pre_process(image, inputs);
const wchar_t* model = L"yolov5n.onnx";
std::vector<Ort::Value> outputs;
process(model, inputs, outputs);
cv::Mat result = image.clone();
post_process(image, result, outputs);
cv::imwrite("result.jpg", result);
return 0;
}
opencv推理
import cv2
import numpy as np
class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush'] #coco80类别
input_shape = (640, 640)
score_threshold = 0.2
nms_threshold = 0.5
confidence_threshold = 0.2
def letterbox(im, new_shape=(416, 416), color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = (new_shape[1] - new_unpad[0])/2, (new_shape[0] - new_unpad[1])/2 # wh padding
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return im
def scale_boxes(boxes, shape):
# Rescale boxes (xyxy) from input_shape to shape
gain = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) # gain = old / new
pad = (input_shape[1] - shape[1] * gain) / 2, (input_shape[0] - shape[0] * gain) / 2 # wh padding
boxes[..., [0, 2]] -= pad[0] # x padding
boxes[..., [1, 3]] -= pad[1] # y padding
boxes[..., :4] /= gain
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
return boxes
def draw(image, box_data):
box_data = scale_boxes(box_data, image.shape)
boxes = box_data[...,:4].astype(np.int32)
scores = box_data[...,4]
classes = box_data[...,5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
print('class: {}, score: {}, coordinate: [{}, {}, {}, {}]'.format(class_names[cl], score, top, left, right, bottom))
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 1)
cv2.putText(image, '{0} {1:.2f}'.format(class_names[cl], score), (top, left), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1)
def filter_box(outputs):
boxes = []
scores = []
class_ids = []
for i in range(outputs.shape[1]):
data = outputs[0][i]
objness = data[4]
if objness > confidence_threshold:
score = data[5:] * objness
_, _, _, max_score_index = cv2.minMaxLoc(score)
max_id = max_score_index[1]
if score[max_id] > score_threshold:
x, y, w, h = data[0].item(), data[1].item(), data[2].item(), data[3].item()
boxes.append(np.array([x-w/2, y-h/2, x+w/2, y+h/2]))
scores.append(score[max_id]*objness)
class_ids.append(max_id)
indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold, nms_threshold)
output = []
for i in indices:
output.append(np.array([boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3], scores[i], class_ids[i]]))
output = np.array(output)
return output
if __name__=="__main__":
image = cv2.imread("bus.jpg")
input = letterbox(image, input_shape)
blob = cv2.dnn.blobFromImage(input, 1/255., size=input_shape, swapRB=True, crop=False)
net = cv2.dnn.readNet("yolov5n.onnx")
net.setInput(blob)
outputs = net.forward()
boxes = filter_box(outputs)
draw(image, boxes)
cv2.imwrite('result.jpg', image)
#include <iostream>
#include <opencv2/opencv.hpp>
const std::vector<std::string> class_names = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" }; //类别名称
const int input_width = 640;
const int input_height = 640;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int num_classes = class_names.size();
const int output_numprob = 5 + num_classes;
const int output_numbox = 3 * (input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32);
//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
//预处理
void pre_process(cv::Mat & image, cv::Mat & inputs)
{
cv::Vec4d params;
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
cv::dnn::blobFromImage(letterbox, inputs, 1. / 255., cv::Size(input_width, input_height), cv::Scalar(), true, false);
}
//网络推理
void process(std::string model, cv::Mat & inputs, std::vector<cv::Mat> & outputs)
{
cv::dnn::Net net = cv::dnn::readNet(model);
net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);//cv::dnn::DNN_BACKEND_OPENCV;
net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); //cv::dnn::DNN_TARGET_CPU;
net.setInput(inputs);
net.forward(outputs, net.getUnconnectedOutLayersNames());
}
//box缩放到原图尺寸
void scale_boxes(cv::Rect & box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width - size.width * gain) / 2;
int pad_h = (input_height - size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}
//可视化函数
void draw_result(cv::Mat & image, std::string label, cv::Rect box)
{
cv::rectangle(image, box, cv::Scalar(255, 0, 0), 1);
int baseLine;
cv::Size label_size = cv::getTextSize(label, 1, 1, 1, &baseLine);
cv::Point tlc = cv::Point(box.x, box.y);
cv::Point brc = cv::Point(box.x, box.y + label_size.height + baseLine);
cv::putText(image, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 1);
}
//后处理
void post_process(cv::Mat & image, cv::Mat & result, std::vector<cv::Mat> & outputs)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
for (int i = 0; i < output_numbox; ++i)
{
float* ptr = (float*)outputs[0].data + i * output_numprob;
float objness = ptr[4];
if (objness < confidence_threshold)
continue;
float* classes_scores = ptr + 5;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) - classes_scores;
float max_class_score = classes_scores[class_id];
float score = max_class_score * objness;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
cv::dnn::NMSBoxes(boxes, scores, score_threshold, nms_threshold, indices);
for (int i = 0; i < indices.size(); i++)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
std::string label = class_names[class_ids[idx]] + ":" + cv::format("%.2f", scores[idx]);
draw_result(result, label, box);
}
}
int main(int argc, char* argv[])
{
cv::Mat image = cv::imread("bus.jpg"), inputs;
pre_process(image, inputs);
std::string model = "yolov5n.onnx";
std::vector<cv::Mat> outputs;
process(model, inputs, outputs);
cv::Mat result = image.clone();
post_process(image, result, outputs);
cv::imwrite("result.jpg", result);
return 0;
}
openvino推理
import cv2
import numpy as np
from openvino.inference_engine import IECore
class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush"]
input_shape = (640, 640)
score_threshold = 0.2
nms_threshold = 0.5
confidence_threshold = 0.2
device = "CPU"
def nms(boxes, scores, score_threshold, nms_threshold):
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[0]
keep.append(i)
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= nms_threshold)[0]
index = index[idx + 1]
return keep
def xywh2xyxy(x):
y = np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
def filter_box(outputs): #过滤掉无用的框
outputs = np.squeeze(outputs)
outputs = outputs[outputs[..., 4] > confidence_threshold]
classes_scores = outputs[..., 5:]
boxes = []
scores = []
class_ids = []
for i in range(len(classes_scores)):
class_id = np.argmax(classes_scores[i])
outputs[i][4] *= classes_scores[i][class_id]
outputs[i][5] = class_id
if outputs[i][4] > score_threshold:
boxes.append(outputs[i][:6])
scores.append(outputs[i][4])
class_ids.append(outputs[i][5])
boxes = np.array(boxes)
boxes = xywh2xyxy(boxes)
scores = np.array(scores)
indices = nms(boxes, scores, score_threshold, nms_threshold)
output = boxes[indices]
return output
def letterbox(im, new_shape=(416, 416), color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = (new_shape[1] - new_unpad[0])/2, (new_shape[0] - new_unpad[1])/2 # wh padding
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return im
def scale_boxes(boxes, shape):
# Rescale boxes (xyxy) from input_shape to shape
gain = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) # gain = old / new
pad = (input_shape[1] - shape[1] * gain) / 2, (input_shape[0] - shape[0] * gain) / 2 # wh padding
boxes[..., [0, 2]] -= pad[0] # x padding
boxes[..., [1, 3]] -= pad[1] # y padding
boxes[..., :4] /= gain
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
return boxes
def draw(image, box_data):
box_data = scale_boxes(box_data, image.shape)
boxes = box_data[...,:4].astype(np.int32)
scores = box_data[...,4]
classes = box_data[...,5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 1)
cv2.putText(image, '{0} {1:.2f}'.format(class_names[cl], score), (top, left), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1)
if __name__=="__main__":
image = cv2.imread('bus.jpg')
input = letterbox(image, input_shape)
input = input[:, :, ::-1].transpose(2, 0, 1).astype(dtype=np.float32) #BGR2RGB和HWC2CHW
input = input / 255.0
input = np.expand_dims(input, axis=0)
model_path = "yolov5n-f32.onnx" #onnx推理支持fp32和fp16
model_xml = "yolov5n_int8.xml" #IR推理支持fp32、fp16和int8
model_bin = "yolov5n_int8.bin"
ie = IECore()
net = ie.read_network(model=model_xml) #net = ie.read_network(model=model_path)
exec_net = ie.load_network(network=net, device_name=device)
input_layer = next(iter(net.input_info))
infer_request_handle = exec_net.start_async(request_id=0, inputs={input_layer: input})
if infer_request_handle.wait(-1) == 0:
output_layer = infer_request_handle._outputs_list[0]
res = infer_request_handle.output_blobs[output_layer]
outputs = res.buffer
boxes = filter_box(outputs)
draw(image, boxes)
cv2.imwrite('result.jpg', image)
#include <iostream>
#include <opencv2/opencv.hpp>
#include <openvino/openvino.hpp>
const std::vector<std::string> class_names = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" }; //类别名称
const int input_width = 640;
const int input_height = 640;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int num_classes = class_names.size();
const int output_numprob = 5 + num_classes;
const int output_numbox = 3 * (input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32);
const std::string device = "CPU";
//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
//预处理
void pre_process(cv::Mat & image, cv::Mat & inputs)
{
cv::Vec4d params;
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
cv::dnn::blobFromImage(letterbox, inputs, 1. / 255., cv::Size(input_width, input_height), cv::Scalar(), true, false);
}
//网络推理
void process(std::string model, cv::Mat & inputs, ov::Tensor & outputs)
{
ov::Core core; //Initialize OpenVINO Runtime Core
auto compiled_model = core.compile_model(model, device); //Compile the Model
ov::InferRequest infer_request = compiled_model.create_infer_request(); //Create an Inference Request
auto input_port = compiled_model.input(); //Get input port for model with one input
ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), inputs.ptr(0)); //Create tensor from external memory
infer_request.set_input_tensor(input_tensor); // Set input tensor for model with one input
infer_request.infer(); //Start inference
outputs = infer_request.get_output_tensor(0); //Get the inference result
}
//NMS
void nms(std::vector<cv::Rect>& boxes, std::vector<float>& scores, float score_threshold, float nms_threshold, std::vector<int>& indices)
{
assert(boxes.size() == scores.size());
struct BoxScore
{
cv::Rect box;
float score;
int id;
};
std::vector<BoxScore> boxes_scores;
for (size_t i = 0; i < boxes.size(); i++)
{
BoxScore box_conf;
box_conf.box = boxes[i];
box_conf.score = scores[i];
box_conf.id = i;
if (scores[i] > score_threshold) boxes_scores.push_back(box_conf);
}
std::sort(boxes_scores.begin(), boxes_scores.end(), [](BoxScore a, BoxScore b) { return a.score > b.score; });
std::vector<float> area(boxes_scores.size());
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
area[i] = boxes_scores[i].box.width * boxes_scores[i].box.height;
}
std::vector<bool> isSuppressed(boxes_scores.size(), false);
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
if (isSuppressed[i]) continue;
for (size_t j = i + 1; j < boxes_scores.size(); ++j)
{
if (isSuppressed[j]) continue;
float x1 = (std::max)(boxes_scores[i].box.x, boxes_scores[j].box.x);
float y1 = (std::max)(boxes_scores[i].box.y, boxes_scores[j].box.y);
float x2 = (std::min)(boxes_scores[i].box.x + boxes_scores[i].box.width, boxes_scores[j].box.x + boxes_scores[j].box.width);
float y2 = (std::min)(boxes_scores[i].box.y + boxes_scores[i].box.height, boxes_scores[j].box.y + boxes_scores[j].box.height);
float w = (std::max)(0.0f, x2 - x1);
float h = (std::max)(0.0f, y2 - y1);
float inter = w * h;
float ovr = inter / (area[i] + area[j] - inter);
if (ovr >= nms_threshold) isSuppressed[j] = true;
}
}
for (int i = 0; i < boxes_scores.size(); ++i)
{
if (!isSuppressed[i]) indices.push_back(boxes_scores[i].id);
}
}
//box缩放到原图尺寸
void scale_boxes(cv::Rect & box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width - size.width * gain) / 2;
int pad_h = (input_height - size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}
//可视化函数
void draw_result(cv::Mat& image, std::string label, cv::Rect box)
{
cv::rectangle(image, box, cv::Scalar(255, 0, 0), 1);
int baseLine;
cv::Size label_size = cv::getTextSize(label, 1, 1, 1, &baseLine);
cv::Point tlc = cv::Point(box.x, box.y);
cv::Point brc = cv::Point(box.x, box.y + label_size.height + baseLine);
cv::putText(image, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 1);
}
//后处理
void post_process(cv::Mat & image, cv::Mat & result, ov::Tensor & outputs)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
for (int i = 0; i < output_numbox; ++i)
{
float* ptr = (float*)outputs.data() + i * output_numprob;
float objness = ptr[4];
if (objness < confidence_threshold)
continue;
float* classes_scores = ptr + 5;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) - classes_scores;
float max_class_score = classes_scores[class_id];
float score = max_class_score * objness;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);
for (int i = 0; i < indices.size(); i++)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
std::string label = class_names[class_ids[idx]] + ":" + cv::format("%.2f", scores[idx]); //class_ids[idx]是class_id
draw_result(result, label, box);
}
}
int main(int argc, char* argv[])
{
cv::Mat image = cv::imread("bus.jpg"), inputs;
pre_process(image, inputs);
std::string model = "yolov5n_fp32.onnx";
ov::Tensor outputs;
process(model, inputs, outputs);
cv::Mat result = image.clone();
post_process(image, result, outputs);
cv::imwrite("result.jpg", result);
return 0;
}
tensorrt推理
import cv2
import numpy as np
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush'] #coco80类别
input_shape = (640, 640)
score_threshold = 0.2
nms_threshold = 0.5
confidence_threshold = 0.2
def nms(boxes, scores, score_threshold, nms_threshold):
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[0]
keep.append(i)
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= nms_threshold)[0]
index = index[idx + 1]
return keep
def xywh2xyxy(x):
y = np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
def filter_box(outputs): #过滤掉无用的框
outputs = np.squeeze(outputs)
outputs = outputs[outputs[..., 4] > confidence_threshold]
classes_scores = outputs[..., 5:]
boxes = []
scores = []
class_ids = []
for i in range(len(classes_scores)):
class_id = np.argmax(classes_scores[i])
outputs[i][4] *= classes_scores[i][class_id]
outputs[i][5] = class_id
if outputs[i][4] > score_threshold:
boxes.append(outputs[i][:6])
scores.append(outputs[i][4])
class_ids.append(outputs[i][5])
boxes = np.array(boxes)
boxes = xywh2xyxy(boxes)
scores = np.array(scores)
indices = nms(boxes, scores, score_threshold, nms_threshold)
output = boxes[indices]
return output
def letterbox(im, new_shape=(416, 416), color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = (new_shape[1] - new_unpad[0])/2, (new_shape[0] - new_unpad[1])/2 # wh padding
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return im
def scale_boxes(input_shape, boxes, shape):
# Rescale boxes (xyxy) from input_shape to shape
gain = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) # gain = old / new
pad = (input_shape[1] - shape[1] * gain) / 2, (input_shape[0] - shape[0] * gain) / 2 # wh padding
boxes[..., [0, 2]] -= pad[0] # x padding
boxes[..., [1, 3]] -= pad[1] # y padding
boxes[..., :4] /= gain
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
return boxes
def draw(image, box_data):
box_data = scale_boxes(input_shape, box_data, image.shape)
boxes = box_data[...,:4].astype(np.int32)
scores = box_data[...,4]
classes = box_data[...,5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 1)
cv2.putText(image, '{0} {1:.2f}'.format(class_names[cl], score), (top, left), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1)
if __name__=="__main__":
logger = trt.Logger(trt.Logger.WARNING)
with open("yolov5n.engine", "rb") as f, trt.Runtime(logger) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
inputs_host = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
outputs_host = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)
inputs_device = cuda.mem_alloc(inputs_host.nbytes)
outputs_device = cuda.mem_alloc(outputs_host.nbytes)
stream = cuda.Stream()
image = cv2.imread('bus.jpg')
input = letterbox(image, input_shape)
input = input[:, :, ::-1].transpose(2, 0, 1).astype(dtype=np.float32) #BGR2RGB和HWC2CHW
input = input / 255.0
input = np.expand_dims(input, axis=0)
np.copyto(inputs_host, input.ravel())
with engine.create_execution_context() as context:
cuda.memcpy_htod_async(inputs_device, inputs_host, stream)
context.execute_async_v2(bindings=[int(inputs_device), int(outputs_device)], stream_handle=stream.handle)
cuda.memcpy_dtoh_async(outputs_host, outputs_device, stream)
stream.synchronize()
boxes = filter_box(outputs_host.reshape(context.get_binding_shape(1)))
draw(image, boxes)
cv2.imwrite('result.jpg', image)
#include <iostream>
#include <fstream>
#include <vector>
#include <opencv2/opencv.hpp>
#include <cuda_runtime.h>
#include <NvInfer.h>
#include <NvInferRuntime.h>
const std::vector<std::string> class_names = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" }; //类别名称
const int input_width = 640;
const int input_height = 640;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int input_numel = 1 * 3 * input_width * input_height;
const int num_classes = class_names.size();
const int output_numprob = 5 + num_classes;
const int output_numbox = 3 * (input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32);
const int output_numel = 1 * output_numprob * output_numbox;
inline const char* severity_string(nvinfer1::ILogger::Severity t)
{
switch (t)
{
case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
case nvinfer1::ILogger::Severity::kERROR: return "error";
case nvinfer1::ILogger::Severity::kWARNING: return "warning";
case nvinfer1::ILogger::Severity::kINFO: return "info";
case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
default: return "unknow";
}
}
class TRTLogger : public nvinfer1::ILogger
{
public:
virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override
{
if (severity <= Severity::kINFO)
{
if (severity == Severity::kWARNING)
printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
else if (severity <= Severity::kERROR)
printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
else
printf("%s: %s\n", severity_string(severity), msg);
}
}
} logger;
std::vector<unsigned char> load_file(const std::string & file)
{
std::ifstream in(file, std::ios::in | std::ios::binary);
if (!in.is_open())
return {};
in.seekg(0, std::ios::end);
size_t length = in.tellg();
std::vector<uint8_t> data;
if (length > 0)
{
in.seekg(0, std::ios::beg);
data.resize(length);
in.read((char*)& data[0], length);
}
in.close();
return data;
}
//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
//预处理
void pre_process(cv::Mat& image, float* input_data_host)
{
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
letterbox.convertTo(letterbox, CV_32FC3, 1.0f / 255.0f);
int image_area = letterbox.cols * letterbox.rows;
float* pimage = (float*)letterbox.data;
float* phost_b = input_data_host + image_area * 0;
float* phost_g = input_data_host + image_area * 1;
float* phost_r = input_data_host + image_area * 2;
for (int i = 0; i < image_area; ++i, pimage += 3)
{
*phost_r++ = pimage[0];
*phost_g++ = pimage[1];
*phost_b++ = pimage[2];
}
}
//网络推理
void process(std::string model, float* input_data_host, float* output_data_host)
{
TRTLogger logger;
auto engine_data = load_file(model);
auto runtime = nvinfer1::createInferRuntime(logger);
auto engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
auto execution_context = engine->createExecutionContext();
float* input_data_device = nullptr;
cudaMalloc(&input_data_device, sizeof(float) * input_numel);
cudaMemcpyAsync(input_data_device, input_data_host, sizeof(float) * input_numel, cudaMemcpyHostToDevice, stream);
float* output_data_device = nullptr;
cudaMalloc(&output_data_device, sizeof(float) * output_numel);
float* bindings[] = { input_data_device, output_data_device };
execution_context->enqueueV2((void**)bindings, stream, nullptr);
cudaMemcpyAsync(output_data_host, output_data_device, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);
cudaFree(input_data_device);
cudaFree(output_data_device);
}
//NMS
void nms(std::vector<cv::Rect>& boxes, std::vector<float>& scores, float score_threshold, float nms_threshold, std::vector<int>& indices)
{
assert(boxes.size() == scores.size());
struct BoxScore
{
cv::Rect box;
float score;
int id;
};
std::vector<BoxScore> boxes_scores;
for (size_t i = 0; i < boxes.size(); i++)
{
BoxScore box_conf;
box_conf.box = boxes[i];
box_conf.score = scores[i];
box_conf.id = i;
if (scores[i] > score_threshold) boxes_scores.push_back(box_conf);
}
std::sort(boxes_scores.begin(), boxes_scores.end(), [](BoxScore a, BoxScore b) { return a.score > b.score; });
std::vector<float> area(boxes_scores.size());
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
area[i] = boxes_scores[i].box.width * boxes_scores[i].box.height;
}
std::vector<bool> isSuppressed(boxes_scores.size(), false);
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
if (isSuppressed[i]) continue;
for (size_t j = i + 1; j < boxes_scores.size(); ++j)
{
if (isSuppressed[j]) continue;
float x1 = (std::max)(boxes_scores[i].box.x, boxes_scores[j].box.x);
float y1 = (std::max)(boxes_scores[i].box.y, boxes_scores[j].box.y);
float x2 = (std::min)(boxes_scores[i].box.x + boxes_scores[i].box.width, boxes_scores[j].box.x + boxes_scores[j].box.width);
float y2 = (std::min)(boxes_scores[i].box.y + boxes_scores[i].box.height, boxes_scores[j].box.y + boxes_scores[j].box.height);
float w = (std::max)(0.0f, x2 - x1);
float h = (std::max)(0.0f, y2 - y1);
float inter = w * h;
float ovr = inter / (area[i] + area[j] - inter);
if (ovr >= nms_threshold) isSuppressed[j] = true;
}
}
for (int i = 0; i < boxes_scores.size(); ++i)
{
if (!isSuppressed[i]) indices.push_back(boxes_scores[i].id);
}
}
//box缩放到原图尺寸
void scale_boxes(cv::Rect & box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width - size.width * gain) / 2;
int pad_h = (input_height - size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}
//可视化函数
void draw_result(cv::Mat& image, std::string label, cv::Rect box)
{
cv::rectangle(image, box, cv::Scalar(255, 0, 0), 1);
int baseLine;
cv::Size label_size = cv::getTextSize(label, 1, 1, 1, &baseLine);
cv::Point tlc = cv::Point(box.x, box.y);
cv::Point brc = cv::Point(box.x, box.y + label_size.height + baseLine);
cv::putText(image, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 1);
}
//后处理
void post_process(cv::Mat & image, cv::Mat & result, float* output_data_host)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
for (int i = 0; i < output_numbox; ++i)//25200
{
float* ptr = output_data_host + i * output_numprob;//85
float objness = ptr[4];
if (objness < confidence_threshold)
continue;
float* classes_scores = 5 + ptr;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) - classes_scores;
float max_class_score = classes_scores[class_id];
float score = max_class_score * objness;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);
for (int i = 0; i < indices.size(); i++)
{
int idx = indices[i];
cv::Rect box = boxes[idx];
std::string label = class_names[class_ids[idx]] + ":" + cv::format("%.2f", scores[idx]); //class_ids[idx]是class_id
draw_result(result, label, box);
}
}
int main(int argc, char* argv[])
{
float* inputs = nullptr;
float* outputs = nullptr;
cudaMallocHost(&inputs, sizeof(float) * input_numel);
cudaMallocHost(&outputs, sizeof(float) * output_numel);
cv::Mat image = cv::imread("bus.jpg");
pre_process(image, inputs);
std::string model = "yolov5n.engine";
process(model, inputs, outputs);
cv::Mat result = image.clone();
post_process(image, result, outputs);
cv::imwrite("result.jpg", result);
cudaFreeHost(inputs);
cudaFreeHost(outputs);
return 0;
}
各种格式的模型权重下载链接：yolov5n模型权重
更完整的工程见我的github：https://github.com/taifyang/yolov5-inference
扫一扫
1946
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
