OpenVINO加载YOLOv8模型进行Classify/Detect/Segment/OBB C++实现
·
之前在 https://blog.csdn.net/fengbingchun/article/details/159862438 中对OpenVINO进行了介绍，这里通过OpenVINO加载已训练好的YOLOv8模型实现Classify/Detect/Segment/OBB，两类(watermelon, wintermelon，OBB除外)数据集。
Classify主要测试代码如下：
typedef struct OVContext {
ov::Core core{};
ov::CompiledModel compiled_model{};
ov::InferRequest infer_request{};
ov::Shape input_shape{};
} OVContext;
OVContext openvino_init()
{
OVContext ctx{};
std::cout << "available devices: ";
for (const auto& dev : ctx.core.get_available_devices())
std::cout << dev << " ";
std::cout << std::endl;
std::string device_name{ "CPU" };
if (cuda_enabled)
device_name = "GPU";
auto model = ctx.core.read_model(openvino_file);
ctx.compiled_model = ctx.core.compile_model(model, device_name);
ctx.infer_request = ctx.compiled_model.create_infer_request();
ctx.input_shape = model->input().get_shape(); // [1,3,H,W]
return ctx;
}
std::vector<ov::Tensor> openvino_infer(const cv::Mat& input, OVContext& ctx)
{
ov::Tensor input_tensor(ov::element::f32, ctx.input_shape, input.data);
ctx.infer_request.set_input_tensor(input_tensor);
ctx.infer_request.infer();
std::vector<ov::Tensor> outputs; // classify: [1, num_classes]; detect: [1, 4 + num_classes, num_boxes]; segment: output0: [1, 4 + num_classes + mask_dim, num_boxes] output1: [1, mask_dim, mask_h, mask_w]
for (size_t i = 0; i < ctx.compiled_model.outputs().size(); ++i) {
outputs.push_back(ctx.infer_request.get_output_tensor(i));
}
return outputs;
}
int test_yolov8_classify_openvino()
{
auto classes = parse_classes_file(classes_file);
if (classes.size() == 0) {
std::cerr << "Error: fail to parse classes file: " << classes_file << std::endl;
return -1;
}
OVContext ctx = openvino_init();
const auto net_w{ ctx.input_shape[3] }, net_h{ ctx.input_shape[2] };
for (const auto& [key, val] : get_dir_images(images_dir)) {
cv::Mat frame = cv::imread(val, cv::IMREAD_COLOR);
if (frame.empty()) {
std::cerr << "Warning: unable to load image: " << val << std::endl;
continue;
}
cv::resize(frame, frame, cv::Size(net_w, net_h));
cv::Mat blob;
cv::dnn::blobFromImage(frame, blob, 1.0 / 255.0, cv::Size(net_w, net_h), cv::Scalar(), true, false);
auto outputs = openvino_infer(blob, ctx);
const float* data = outputs[0].data<const float>();
int num_classes = outputs[0].get_shape()[1];
auto class_id = std::max_element(data, data + num_classes) - data;
std::cout << "image name: " << key << ", class id: " << class_id << ", class name: " << classes[class_id] << std::endl;
}
return 0;
}执行结果如下图所示：
Detect主要测试代码如下：
void draw_boxes(const std::vector<std::string>& classes, const std::vector<int>& ids, const std::vector<float>& confidences,
const std::vector<cv::Rect>& boxes, const std::string& name, cv::Mat& frame)
{
if (ids.size() != confidences.size() || ids.size() != boxes.size() || confidences.size() != boxes.size()) {
std::cerr << "Error: their lengths are inconsistent: " << ids.size() << ", " << confidences.size() << ", " << boxes.size() << std::endl;
return;
}
std::cout << "image name: " << name << ", number of detections: " << ids.size() << std::endl;
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<int> dis(100, 255);
for (auto i = 0; i < ids.size(); ++i) {
auto color = cv::Scalar(dis(gen), dis(gen), dis(gen));
cv::rectangle(frame, boxes[i], color, 2);
std::string class_string = classes[ids[i]] + ' ' + std::to_string(confidences[i]).substr(0, 4);
cv::Size text_size = cv::getTextSize(class_string, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
cv::Rect text_box(boxes[i].x, boxes[i].y - 40, text_size.width + 10, text_size.height + 20);
cv::rectangle(frame, text_box, color, cv::FILLED);
cv::putText(frame, class_string, cv::Point(boxes[i].x + 5, boxes[i].y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
}
std::string path(result_dir);
path += "/" + name;
cv::imwrite(path, frame);
}
void post_process(const float* data, int rows, int stride, float xfactor, float yfactor, const std::vector<std::string>& classes,
const std::string& name, cv::Mat& frame)
{
std::vector<int> class_ids;
std::vector<float> confidences;
std::vector<cv::Rect> boxes;
for (auto i = 0; i < rows; ++i) {
const float* classes_scores = data + 4;
cv::Mat scores(1, classes.size(), CV_32FC1, (float*)classes_scores);
cv::Point class_id;
double max_class_score;
cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
if (max_class_score > confidence_threshold) {
confidences.push_back(max_class_score);
class_ids.push_back(class_id.x);
float x = data[0];
float y = data[1];
float w = data[2];
float h = data[3];
int left = int((x - 0.5 * w) * xfactor);
int top = int((y - 0.5 * h) * yfactor);
int width = int(w * xfactor);
int height = int(h * yfactor);
boxes.push_back(cv::Rect(left, top, width, height));
}
data += stride;
}
std::vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, iou_threshold, nms_result);
std::vector<int> ids;
std::vector<float> confs;
std::vector<cv::Rect> rects;
for (size_t i = 0; i < nms_result.size(); ++i) {
ids.emplace_back(class_ids[nms_result[i]]);
confs.emplace_back(confidences[nms_result[i]]);
rects.emplace_back(boxes[nms_result[i]]);
}
draw_boxes(classes, ids, confs, rects, name, frame);
}
int test_yolov8_detect_openvino()
{
namespace fs = std::filesystem;
if (!fs::exists(result_dir)) {
fs::create_directories(result_dir);
}
auto classes = parse_classes_file(classes_file);
if (classes.size() == 0) {
std::cerr << "Error: fail to parse classes file: " << classes_file << std::endl;
return -1;
}
OVContext ctx = openvino_init();
const int net_w{ static_cast<int>(ctx.input_shape[3]) }, net_h{ static_cast<int>(ctx.input_shape[2]) };
for (const auto& [key, val] : get_dir_images(images_dir)) {
cv::Mat frame = cv::imread(val, cv::IMREAD_COLOR);
if (frame.empty()) {
std::cerr << "Warning: unable to load image: " << val << std::endl;
continue;
}
cv::Mat blob{};
cv::Mat bgr = modify_image_size(frame);
cv::dnn::blobFromImage(bgr, blob, 1.0 / 255.0, cv::Size(net_w, net_h), cv::Scalar(), true, false);
auto outputs = openvino_infer(blob, ctx);
auto output_shape = outputs[0].get_shape();
const cv::Mat src(cv::Size(static_cast<int>(output_shape[2]), static_cast<int>(output_shape[1])), CV_32FC1, (float*)outputs[0].data<const float>());
cv::Mat transposed{};
cv::transpose(src, transposed);
float scalex = bgr.cols * 1.f / net_w;
float scaley = bgr.rows * 1.f / net_h;
post_process((float*)transposed.data, transposed.rows, transposed.cols, scalex, scaley, classes, key, frame);
}
return 0;
}执行结果如下图所示：
Segment主要测试代码如下：
void draw_boxes_mask(const std::vector<std::string>& classes, const std::vector<int>& ids, const std::vector<float>& confidences,
const std::vector<cv::Rect>& boxes, const std::vector<cv::Mat>& masks, const std::string& name, cv::Mat& frame)
{
std::cout << "image name: " << name << ", number of detections: " << ids.size() << std::endl;
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<int> dis(100, 255);
cv::Mat mk = frame.clone();
std::vector<cv::Scalar> colors;
for (auto i = 0; i < classes.size(); ++i)
colors.emplace_back(cv::Scalar(dis(gen), dis(gen), dis(gen)));
for (auto i = 0; i < ids.size(); ++i) {
cv::rectangle(frame, boxes[i], colors[ids[i]], 2);
std::string class_string = classes[ids[i]] + ' ' + std::to_string(confidences[i]).substr(0, 4);
cv::Size text_size = cv::getTextSize(class_string, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
cv::Rect text_box(boxes[i].x, boxes[i].y - 40, text_size.width + 10, text_size.height + 20);
cv::rectangle(frame, text_box, colors[ids[i]], cv::FILLED);
cv::putText(frame, class_string, cv::Point(boxes[i].x + 5, boxes[i].y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
mk(boxes[i]).setTo(colors[ids[i]], masks[i]);
}
cv::addWeighted(frame, 0.5, mk, 0.5, 0, frame);
std::string path(result_dir);
cv::imwrite(path + "/" + name, frame);
}
void post_process_mask(const cv::Mat& output0, const cv::Mat& output1, const std::vector<int>& output1_sizes, const std::vector<std::string>& classes, const std::string& name, cv::Mat& frame)
{
std::vector<int> class_ids;
std::vector<float> confidences;
std::vector<cv::Rect> boxes;
std::vector<std::vector<float>> masks;
float scalex = frame.cols * 1.f / input_size[1]; // note: image_preprocess function
float scaley = frame.rows * 1.f / input_size[0];
auto scale = (scalex > scaley) ? scalex : scaley;
const float* data = (float*)output0.data;
for (auto i = 0; i < output0.rows; ++i) {
cv::Mat scores(1, classes.size(), CV_32FC1, (float*)data + 4);
cv::Point class_id;
double max_class_score;
cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
if (max_class_score > confidence_threshold) {
confidences.emplace_back(max_class_score);
class_ids.emplace_back(class_id.x);
masks.emplace_back(std::vector<float>(data + 4 + classes.size(), data + output0.cols)); // 32
float x = data[0];
float y = data[1];
float w = data[2];
float h = data[3];
int left = std::max(0, std::min(int((x - 0.5 * w) * scale), frame.cols));
int top = std::max(0, std::min(int((y - 0.5 * h) * scale), frame.rows));
int width = std::max(0, std::min(int(w * scale), frame.cols - left));
int height = std::max(0, std::min(int(h * scale), frame.rows - top));
boxes.emplace_back(cv::Rect(left, top, width, height));
}
data += output0.cols;
}
std::vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, iou_threshold, nms_result);
cv::Mat proto = output1.reshape(0, { output1_sizes[1], output1_sizes[2] * output1_sizes[3] });
std::vector<int> ids;
std::vector<float> confs;
std::vector<cv::Rect> rects;
std::vector<cv::Mat> mks;
for (size_t i = 0; i < nms_result.size(); ++i) {
auto index = nms_result[i];
ids.emplace_back(class_ids[index]);
confs.emplace_back(confidences[index]);
boxes[index] = boxes[index] & cv::Rect(0, 0, frame.cols, frame.rows);
cv::Mat mk;
get_masks(cv::Mat(masks[index]).t(), proto, output1_sizes, frame, boxes[index], mk);
mks.emplace_back(mk);
rects.emplace_back(boxes[index]);
}
draw_boxes_mask(classes, ids, confs, rects, mks, name, frame);
}
int test_yolov8_segment_openvino()
{
namespace fs = std::filesystem;
if (!fs::exists(result_dir)) {
fs::create_directories(result_dir);
}
auto classes = parse_classes_file(classes_file);
if (classes.size() == 0) {
std::cerr << "Error: fail to parse classes file: " << classes_file << std::endl;
return -1;
}
OVContext ctx = openvino_init();
const int net_w{ static_cast<int>(ctx.input_shape[3]) }, net_h{ static_cast<int>(ctx.input_shape[2]) };
for (const auto& [key, val] : get_dir_images(images_dir)) {
cv::Mat frame = cv::imread(val, cv::IMREAD_COLOR);
if (frame.empty()) {
std::cerr << "Warning: unable to load image: " << val << std::endl;
continue;
}
cv::Mat blob{};
cv::Mat bgr = modify_image_size(frame);
cv::dnn::blobFromImage(bgr, blob, 1.0 / 255.0, cv::Size(net_w, net_h), cv::Scalar(), true, false);
auto outputs = openvino_infer(blob, ctx);
auto output1_shape = outputs[0].get_shape();
auto output2_shape = outputs[1].get_shape();
const cv::Mat src(cv::Size(static_cast<int>(output1_shape[2]), static_cast<int>(output1_shape[1])), CV_32FC1, (float*)outputs[0].data<const float>());
cv::Mat transposed{};
cv::transpose(src, transposed);
std::vector<int> sizes;
for (int i = 0; i < 4; ++i)
sizes.emplace_back(output2_shape[i]);
cv::Mat mask = cv::Mat(sizes, CV_32F, (float*)outputs[1].data<const float>());
post_process_mask(transposed, mask, sizes, classes, key, frame);
}
return 0;
}执行结果如下图所示：
OBB主要测试代码如下：
void draw_rotated_rect(const std::vector<std::string>classes, const std::vector<int>& ids, const std::vector<float>& confs, const std::vector<cv::RotatedRect>& boxes, const std::string& name, cv::Mat& frame)
{
std::cout << "image name: " << name << ", number of detections: " << ids.size() << std::endl;
//std::random_device rd;
std::mt19937 gen(66); // gen(rd())
std::uniform_int_distribution<int> dis(100, 255);
std::vector<cv::Scalar> colors;
for (auto i = 0; i < classes.size(); ++i)
colors.emplace_back(cv::Scalar(dis(gen), dis(gen), dis(gen)));
for (auto i = 0; i < ids.size(); ++i) {
auto rotated_rect = boxes[i];
cv::Point2f pts[4];
rotated_rect.points(pts);
for (auto j = 0; j < 4; ++j)
cv::line(frame, pts[j], pts[(j + 1) % 4], colors[ids[i]], 2);
auto center = boxes[i].center;
std::string text = classes[ids[i]] + "," + std::to_string(confs[i]).substr(0, 4);
cv::Size text_size = cv::getTextSize(text, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
cv::Point text_org(static_cast<int>(center.x - text_size.width / 2), static_cast<int>(center.y + text_size.height / 2));
cv::putText(frame, text, text_org, cv::FONT_HERSHEY_DUPLEX, 1, colors[ids[i]], 2);
}
std::string path(result_dir);
cv::imwrite(path + "/" + name, frame);
}
void post_process2(const float* data, int rows, int stride, float xfactor, float yfactor, const std::vector<std::string>classes, const std::string& image_name, cv::Mat& frame)
{
const double PI{ std::acos(-1) }; // 3.1415926...
std::vector<int> class_ids{};
std::vector<float> confidences{};
std::vector<cv::RotatedRect> boxes{};
for (auto i = 0; i < rows; ++i) {
const float* classes_scores = data + 4;
cv::Mat scores(1, classes.size(), CV_32FC1, (float*)classes_scores);
cv::Point class_id{};
double max_class_score{};
cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
if (max_class_score > confidence_threshold) {
confidences.push_back(max_class_score);
class_ids.push_back(class_id.x);
float x = data[0] * xfactor;
float y = data[1] * yfactor;
float w = data[2] * xfactor;
float h = data[3] * yfactor;
float angle = data[stride - 1];
if (angle >= 0.5 * PI && angle <= 0.75 * PI)
angle = angle - PI;
cv::RotatedRect box = cv::RotatedRect(cv::Point2f(x, y), cv::Size2f(w, h), angle * 180 / PI);
boxes.push_back(box);
}
data += stride;
}
std::vector<int> nms_result{};
cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, iou_threshold, nms_result);
std::vector<int> ids;
std::vector<float> confs;
std::vector<cv::RotatedRect> result{};
for (size_t i = 0; i < nms_result.size(); ++i) {
ids.emplace_back(class_ids[nms_result[i]]);
confs.emplace_back(confidences[nms_result[i]]);
result.emplace_back(boxes[nms_result[i]]);
}
draw_rotated_rect(classes, ids, confs, result, image_name, frame);
}
int test_yolov8_obb_openvino()
{
namespace fs = std::filesystem;
if (!fs::exists(result_dir)) {
fs::create_directories(result_dir);
}
OVContext ctx = openvino_init();
const int net_w{ static_cast<int>(ctx.input_shape[3]) }, net_h{ static_cast<int>(ctx.input_shape[2]) };
for (const auto& [key, val] : get_dir_images(images_dir)) {
cv::Mat frame = cv::imread(val, cv::IMREAD_COLOR);
if (frame.empty()) {
std::cerr << "Warning: unable to load image: " << val << std::endl;
continue;
}
cv::Mat blob{};
cv::Mat bgr = modify_image_size(frame);
cv::dnn::blobFromImage(bgr, blob, 1.0 / 255.0, cv::Size(net_w, net_h), cv::Scalar(), true, false);
auto outputs = openvino_infer(blob, ctx);
auto output_shape = outputs[0].get_shape();
const cv::Mat src(cv::Size(static_cast<int>(output_shape[2]), static_cast<int>(output_shape[1])), CV_32FC1, (float*)outputs[0].data<const float>());
cv::Mat transposed{};
cv::transpose(src, transposed);
float scalex = bgr.cols * 1.f / net_w;
float scaley = bgr.rows * 1.f / net_h;
post_process2((float*)transposed.data, transposed.rows, transposed.cols, scalex, scaley, obb_class_names, key, frame);
}
return 0;
}执行结果如下图所示：
注：
1.Classify,Detect,Segment,OBB共用模型加载初始化函数openvino_init以及推理函数openvino_infer。
2.OpenVINO的前处理和后处理操作与使用OpenCV DNN完全一致。
3.加载YOLOv8还是YOLO11以上测试代码通用。
Agent 垂直技术社区，欢迎活跃、内容共建。
更多推荐
4
0- 0
已为社区贡献4条内容
所有评论(0)