使用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: watermelon
  1: wintermelon

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

import argparse
import colorama
from ultralytics import YOLO

def 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 args

def train(task, yaml, epochs):
	if task == "detect":
		model = YOLO("yolov8n.pt") # load a pretrained model
	elif task == "segment":
		model = YOLO("yolov8n-seg.pt") # load a pretrained model
	else:
		print(colorama.Fore.RED + "Error: unsupported task:", task)
		raise

	results = model.train(data=yaml, epochs=epochs, imgsz=640) # train the model

	metrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings remembered

	model.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") # libtorch

if __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" };
#endif

std::vector parse_classes_file(const char* name)
{
	std::vector 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 images; // image name, image path + image name

	for (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& classes, const std::vector& ids, const std::vector& confidences,
	const std::vector& 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 dis(100, 255);

	for (auto i = 0; i = src.rows) {
		resize_scales = src.cols * 1.f / image_size[1];
		cv::resize(dst, dst, cv::Size(image_size[1], static_cast(src.rows / resize_scales)));
	} else {
		resize_scales = src.rows * 1.f / image_size[0];
		cv::resize(dst, dst, cv::Size(static_cast(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
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(h, w)[c]) / 255.f;
			}
		}
	}
}

void post_process(const float* data, int rows, int stride, float xfactor, float yfactor, const std::vector& classes,
	cv::Mat& frame, const std::string& name)
{
	std::vector class_ids;
	std::vector confidences;
	std::vector boxes;

	for (auto i = 0; i  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 nms_result;
	cv::dnn::NMSBoxes(boxes, confidences, model_score_threshold, model_nms_threshold, nms_result);

	std::vector ids;
	std::vector confs;
	std::vector 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);
}

} // namespace

int test_yolov8_detect_onnxruntime()
{
	// reference: ultralytics/examples/YOLOv8-ONNXRuntime-CPP
	try {
		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 input_node_names, output_node_names;
		std::vector 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 blob(new float[image_size[0] * image_size[1] * 3]);
		std::vector 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(
				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 output_node_dims = tensor_info.GetShape();
			auto output = output_tensors.front().GetTensorMutableData();
			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(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毫秒左右

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

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

使用onnxruntime加载YOLOv8生成的onnx文件进行目标检测插图(1)

GitHub：http://github.com/fengbingchun/NN_Test

一	二	三	四	五	六	日
						1
2	3	4	5	6	7	8
9	10	11	12	13	14	15
16	17	18	19	20	21	22
23	24	25	26	27	28	29
30

admin 钻石

相关推荐