使用onnxruntime加载YOLOv8生成的onnx文件进行目标检测

      在网上下载了60多幅包含西瓜和冬瓜的图像组成melon数据集,使用 LabelMe  工具进行标注,然后使用 labelme2yolov8 脚本将json文件转换成YOLOv8支持的.txt文件,并自动生成YOLOv8支持的目录结构,包括melon.yaml文件,其内容如下:

path: ../datasets/melon # dataset root dir
train: images/train # train images (relative to 'path')
val: images/val  # val images (relative to 'path')
test: # test images (optional)# Classes
names:0: watermelon1: wintermelon

      使用以下python脚本进行训练生成onnx文件

import argparse
import colorama
from ultralytics import YOLOdef parse_args():parser = argparse.ArgumentParser(description="YOLOv8 train")parser.add_argument("--yaml", required=True, type=str, help="yaml file")parser.add_argument("--epochs", required=True, type=int, help="number of training")parser.add_argument("--task", required=True, type=str, choices=["detect", "segment"], help="specify what kind of task")args = parser.parse_args()return argsdef train(task, yaml, epochs):if task == "detect":model = YOLO("yolov8n.pt") # load a pretrained modelelif task == "segment":model = YOLO("yolov8n-seg.pt") # load a pretrained modelelse:print(colorama.Fore.RED + "Error: unsupported task:", task)raiseresults = model.train(data=yaml, epochs=epochs, imgsz=640) # train the modelmetrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings rememberedmodel.export(format="onnx") #, dynamic=True) # export the model, cannot specify dynamic=True, opencv does not support# model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)model.export(format="torchscript") # libtorchif __name__ == "__main__":colorama.init()args = parse_args()train(args.task, args.yaml, args.epochs)print(colorama.Fore.GREEN + "====== execution completed ======")

      以下是使用onnxruntime接口加载onnx文件进行目标检测的实现代码:

namespace {constexpr bool cuda_enabled{ false };
constexpr int image_size[2]{ 640, 640 }; // {height,width}, input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 6, 8400)
constexpr float model_score_threshold{ 0.45 }; // confidence threshold
constexpr float model_nms_threshold{ 0.50 }; // iou threshold#ifdef _MSC_VER
constexpr char* onnx_file{ "../../../data/best.onnx" };
constexpr char* torchscript_file{ "../../../data/best.torchscript" };
constexpr char* images_dir{ "../../../data/images/predict" };
constexpr char* result_dir{ "../../../data/result" };
constexpr char* classes_file{ "../../../data/images/labels.txt" };
#else
constexpr char* onnx_file{ "data/best.onnx" };
constexpr char* torchscript_file{ "data/best.torchscript" };
constexpr char* images_dir{ "data/images/predict" };
constexpr char* result_dir{ "data/result" };
constexpr char* classes_file{ "data/images/labels.txt" };
#endifstd::vector<std::string> parse_classes_file(const char* name)
{std::vector<std::string> classes;std::ifstream file(name);if (!file.is_open()) {std::cerr << "Error: fail to open classes file: " << name << std::endl;return classes;}std::string line;while (std::getline(file, line)) {auto pos = line.find_first_of(" ");classes.emplace_back(line.substr(0, pos));}file.close();return classes;
}auto get_dir_images(const char* name)
{std::map<std::string, std::string> images; // image name, image path + image namefor (auto const& dir_entry : std::filesystem::directory_iterator(name)) {if (dir_entry.is_regular_file())images[dir_entry.path().filename().string()] = dir_entry.path().string();}return images;
}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);}//cv::imshow("Inference", frame);//cv::waitKey(-1);std::string path(result_dir);path += "/" + name;cv::imwrite(path, frame);
}std::wstring ctow(const char* str)
{constexpr size_t len{ 128 };wchar_t wch[len];swprintf(wch, len, L"%hs", str);return std::wstring(wch);
}float image_preprocess(const cv::Mat& src, cv::Mat& dst)
{cv::cvtColor(src, dst, cv::COLOR_BGR2RGB);float resize_scales{ 1. };if (src.cols >= src.rows) {resize_scales = src.cols * 1.f / image_size[1];cv::resize(dst, dst, cv::Size(image_size[1], static_cast<int>(src.rows / resize_scales)));} else {resize_scales = src.rows * 1.f / image_size[0];cv::resize(dst, dst, cv::Size(static_cast<int>(src.cols / resize_scales), image_size[0]));}cv::Mat tmp = cv::Mat::zeros(image_size[0], image_size[1], CV_8UC3);dst.copyTo(tmp(cv::Rect(0, 0, dst.cols, dst.rows)));dst = tmp;return resize_scales;
}template<typename T>
void image_to_blob(const cv::Mat& src, T* blob)
{for (auto c = 0; c < 3; ++c) {for (auto h = 0; h < src.rows; ++h) {for (auto w = 0; w < src.cols; ++w) {blob[c * src.rows * src.cols + h * src.cols + w] = (src.at<cv::Vec3b>(h, w)[c]) / 255.f;}}}
}void post_process(const float* data, int rows, int stride, float xfactor, float yfactor, const std::vector<std::string>& classes,cv::Mat& frame, const std::string& name)
{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 > model_score_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, model_score_threshold, model_nms_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);
}} // namespaceint test_yolov8_detect_onnxruntime()
{// reference: ultralytics/examples/YOLOv8-ONNXRuntime-CPPtry {Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo");Ort::SessionOptions session_option;if (cuda_enabled) {OrtCUDAProviderOptions cuda_option;cuda_option.device_id = 0;session_option.AppendExecutionProvider_CUDA(cuda_option);}session_option.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);session_option.SetIntraOpNumThreads(1);session_option.SetLogSeverityLevel(3);Ort::Session session(env, ctow(onnx_file).c_str(), session_option);Ort::AllocatorWithDefaultOptions allocator;std::vector<const char*> input_node_names, output_node_names;std::vector<std::string> input_node_names_, output_node_names_;for (auto i = 0; i < session.GetInputCount(); ++i) {Ort::AllocatedStringPtr input_node_name = session.GetInputNameAllocated(i, allocator);input_node_names_.emplace_back(input_node_name.get());}for (auto i = 0; i < session.GetOutputCount(); ++i) {Ort::AllocatedStringPtr output_node_name = session.GetOutputNameAllocated(i, allocator);output_node_names_.emplace_back(output_node_name.get());}for (auto i = 0; i < input_node_names_.size(); ++i)input_node_names.emplace_back(input_node_names_[i].c_str());for (auto i = 0; i < output_node_names_.size(); ++i)output_node_names.emplace_back(output_node_names_[i].c_str());Ort::RunOptions options(nullptr);std::unique_ptr<float[]> blob(new float[image_size[0] * image_size[1] * 3]);std::vector<int64_t> input_node_dims{ 1, 3, image_size[1], image_size[0] };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;}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;}auto tstart = std::chrono::high_resolution_clock::now();cv::Mat rgb;auto resize_scales = image_preprocess(frame, rgb);image_to_blob(rgb, blob.get());Ort::Value input_tensor = Ort::Value::CreateTensor<float>(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob.get(), 3 * image_size[1] * image_size[0], input_node_dims.data(), input_node_dims.size());auto output_tensors = session.Run(options, input_node_names.data(), &input_tensor, 1, output_node_names.data(), output_node_names.size());Ort::TypeInfo type_info = output_tensors.front().GetTypeInfo();auto tensor_info = type_info.GetTensorTypeAndShapeInfo();std::vector<int64_t> output_node_dims = tensor_info.GetShape();auto output = output_tensors.front().GetTensorMutableData<float>();int stride_num = output_node_dims[1];int signal_result_num = output_node_dims[2];cv::Mat raw_data = cv::Mat(stride_num, signal_result_num, CV_32F, output);raw_data = raw_data.t();float* data = (float*)raw_data.data;auto tend = std::chrono::high_resolution_clock::now();std::cout << "elapsed millisenconds: " << std::chrono::duration_cast<std::chrono::milliseconds>(tend - tstart).count() << " ms" << std::endl;post_process(data, signal_result_num, stride_num, resize_scales, resize_scales, classes, frame, key);}}catch (const std::exception& e) {std::cerr << "Error: " << e.what() << std::endl;return -1;}return 0;
}

      labels.txt文件内容如下:仅2类

watermelon 0
wintermelon 1

      说明

      1.这里使用的onnxruntime版本为1.18.0;

      2.windows下,onnxruntime库在debug和release为同一套库,在debug和release下均可执行;

      3.通过指定变量cuda_enabled判断走cpu还是gpu流程 ;

      4.windows下,onnxruntime中有些接口参数为wchar_t*,而linux下为char*,因此在windows下需要单独做转换,这里通过ctow函数实现从char*到wchar_t的转换;

      5.yolov8中提供的sample有问题,需要作调整。

      执行结果如下图所示:同样的预测图像集,与opencv dnn结果相似,它们具有相同的后处理流程;下面显示的耗时是在cpu下,gpu下仅20毫秒左右

      其中一幅图像的检测结果如下图所示:

      GitHub:https://github.com/fengbingchun/NN_Test

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/web/22400.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

干货!如何在Jmeter中实现对NCR响应的解析

最近做接口测试时发现了一个问题&#xff0c;部分请求的响应是通过NCR编码实现的&#xff0c;这样就导致了无法对这些请求进行断言&#xff0c;为了解决这个问题进行了如下调研&#xff0c;大家可以参考下面两篇文章&#xff1a; 使用Java apache commons包五分钟搞定NCR解析&…

CCIG 2024:大模型技术及其前沿应用论坛深度解析

一、CCIG论坛介绍 中国图象图形大会&#xff08;CCIG 2024&#xff09;是一场备受瞩目的学术盛会&#xff0c;近期在陕西省西安市曲江国际会议中心举行。这次会议以“图聚智生&#xff0c;象合慧成”为主题&#xff0c;由中国图象图形学学会主办&#xff0c;旨在汇聚图像图形领…

ABAP 长文本编辑器弹窗控件

前言 用户想在ALV上编辑长文本&#xff0c;但是ALV只有128个字符肯定是不够用的&#xff0c;所以需要用一个长文本编辑器来输入&#xff0c;本来想自己写的&#xff0c;发现有标准的函数&#xff0c;还挺好用的 代码 在用户双击ALV字段时&#xff0c;触发下述form&#xff0…

使用Rufus工具制作Ubuntu To Go——很详细

一、准备工作 准备工具&#xff1a; 1、下载Rufus(主角)软件 2、准备一个U盘或硬盘&#xff08;小白128G足够&#xff0c;装Ubuntu系统&#xff09; 3、下载Ubuntu系统镜像文件 1、下载软件Rufus 先来看一下官网介绍&#xff1a; Rufus 是一款格式化和创建 USB 启动盘的辅助工…

“GPT-4o深度解析:技术演进、能力评估与个人体验综述“

文章目录 每日一句正能量前言对比分析模型架构性能应用场景用户体验技术创新社区和生态系统总结 技术能力语言生成能力语言理解能力技术实现总结 个人感受关于GPT-4o的假设性观点&#xff1a;关于当前语言模型的一般性观点&#xff1a; 后记 每日一句正能量 又回到了原点&#…

“智能分析赋能等保:大数据技术在安全审计记录中的应用“

智能分析技术在信息安全领域&#xff0c;尤其是等保&#xff08;等级保护&#xff09;合规性方面&#xff0c;发挥了关键作用。特别是结合大数据技术&#xff0c;安全审计记录的处理和分析能力得到了显著增强。以下几点阐述了大数据技术是如何赋能等保安全审计的&#xff1a; …

【Python】 Python中的functools.wraps:装饰器的优雅包装

基本原理 在Python中&#xff0c;装饰器是一种非常强大的工具&#xff0c;它允许我们以一种非常灵活的方式修改或增强函数的行为。装饰器本质上是一个函数&#xff0c;它接收一个函数作为参数&#xff0c;并返回一个新的函数。然而&#xff0c;当我们使用装饰器时&#xff0c;…

【笔记】Sturctured Streaming笔记总结(Python版)

目录 相关资料 一、概述 1.1 基本概念 1.2 两种处理模型 &#xff08;1&#xff09;微批处理 &#xff08;2&#xff09;持续处理 1.3 Structured Streaming和Spark SQL、Spark Streaming关系 二、编写Structured Streaming程序的基本步骤 三、输入源 3.1 File源 &a…

Docker 基础使用 (1) 使用流程概览

文章目录 Docker 软件安装Docker 镜像仓库Docker 仓库指令Docker 镜像指令Docker 容器指令Docker 使用实例 —— 搭建 nginx 服务nginx 概念nginx 使用用 docker 启动 nginx Docker 基础使用&#xff08;0&#xff09;基础认识 Docker 基础使用 (1) 使用流程概览 Docker 基础使…

f-stack和DPDK

GPT-4 (OpenAI) f-stack和DPDK&#xff08;数据平面开发套件&#xff09;都是与高性能网络处理相关的技术。它们的目的是提高数据包的处理速度&#xff0c;优化网络I/O的性能。以下是对这两者的简要解释&#xff1a; 1. **DPDK (Data Plane Development Kit):** DPDK 是一个…

【软件安全国产化替代解决方案】亮相2024澳门万讯论坛

近日&#xff0c;2024万讯论坛在澳门成功举办。本次论坛由万讯电脑科技主办&#xff0c;旨在引进国内尖端科技厂商&#xff0c;提供全方位的信创解决方案&#xff0c;分享信创化过程中所面临的挑战及阶段性转换经验。开源网安作为拥有软件安全领域全链条产品的厂商&#xff0c;…

【网络安全的神秘世界】安装burpsuite

&#x1f31d;博客主页&#xff1a;泥菩萨 &#x1f496;专栏&#xff1a;Linux探索之旅 | 网络安全的神秘世界 | 专接本 java下载地址 burpsuite安装包 安装java 双击打开java安装包——>下一步 在java文件夹下创建jdk和jre文件夹&#xff08;不一定是C盘&#xff0c;根…

云原生下的数据协调艺术:etcd存储系统解析

目录 一、分布式存储简介 二、etcd介绍 三、etcd架构 四、etcd集成实践 一、分布式存储简介 随着云原生与容器化技术的兴起&#xff0c;分布式系统的复杂性大大增加。分布式系统面临一系列问题&#xff0c;比如部署复杂、响应时间慢、运维复杂等&#xff0c;其中最根本的问…

【Redis数据库】命令操作

文章目录 一、连接命令二、键命令 &#x1f308;你好呀&#xff01;我是 山顶风景独好 &#x1f495;欢迎来到我的博客&#xff0c;很高兴能够在这里和您见面&#xff01; &#x1f495;希望您在这里可以感受到一份轻松愉快的氛围&#xff01; &#x1f495;这里不仅可以获得有…

Arduino_ESP32_控制舵机运行【2024年】版

介绍 吧下面代码复制进去就可以使用了。 当然肯定不只是esp32基本上所有的Arduino都可以使用 然后接线的话信号线接到18号io口 然后其他的接电源正负极 图示 代码教程 #include <Arduino.h>// 定义舵机控制引脚 const int servoPin 18;// 定义舵机角度对应的脉宽范围 …

算法金 | 再见,PCA 主成分分析!

​大侠幸会&#xff0c;在下全网同名[算法金] 0 基础转 AI 上岸&#xff0c;多个算法赛 Top [日更万日&#xff0c;让更多人享受智能乐趣] 1. 概念&#xff1a;数据降维的数学方法 定义 主成分分析&#xff08;PCA&#xff09;是一种统计方法&#xff0c;通过正交变换将一组可…

RTOS笔记--任务状态与调度

任务状态 freertos中的任务分为四个状态&#xff1a;就绪状态&#xff08;ready&#xff09;、运行状态&#xff08;running&#xff09;、阻塞状态&#xff08;blocked&#xff09;、暂停状态&#xff08;suspended&#xff09; 完整的任务状态转换图&#xff1a; 在使用vTas…

24V转3V的高效解决方案-AH8610.sot23-6

24V转3V的高效解决方案-AH8610.sot23-6 ### 引言 随着电子设备向小型化、高性能和高效率发展&#xff0c;对电源管理芯片的要求也越来越高。AH8610.sot23-6是一款专为小型化电源转换设计的同步降压开关芯片&#xff0c;它能够将24V的输入电压转换为3V输出&#xff0c;适用于US…

风管静压箱的作用及选型

1.压力的种类 动压—由风速而产生的压力&#xff1b;空调厂家设计时均已经考虑&#xff0c;无需计算。静压—垂直作用于风管壁面的压力&#xff0c;用于克服风管阻力&#xff1b;所以&#xff0c;对于风管机组有零静压和带静压之分&#xff0c;零静压指静压为0Pa&#xff0c;不…

数学建模笔记

数学建模 定义角度 数学模型是针对参照某种事物系统的特征或数量依存关系&#xff0c;采用数学语言&#xff0c;概括地或近似地表述出的一种数学结构&#xff0c;这种数学结构是借助于数学符号刻画出来的某种系统的纯关系结构。从广义理解&#xff0c;数学模型包括数学中的各…