yolov5+torchlib1.6-cuda+vs2017+opencv4.4-cuda

最近在使用yolov5,网上找了些资料,发现torchlib用来部署yolov5相对更方便,效率更高,

github上也有调用torchlib的代码,但是后处理时间都花费过多,经过我的一番研究后,已经彻底解决了(测试结果为前处理(不包括图片从硬盘加载)8ms,推理15ms(2080ti),后处理7ms),也在github上的torchlib-yolov5--issue后提交了我的回复,

yolov5源码上export.py 导出的是cpu版本的,要改成gpu版本的,github上的链接:https://github.com/yasenh/libtorch-yolov5

下面是我原始的代码:

#include <opencv2/opencv.hpp>
#include<opencv2\dnn\dnn.hpp>
#include <torch/script.h>
#include<iostream>
#include<string>
#include<io.h>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <stdio.h>
#include<fstream>

using namespace std;
using namespace cv;
using namespace dnn;

struct files {//搜索文件的结构体
	String filepath;
	time_t time;
	string filename;
};


vector<files> fileSearch(string path)
{
	int result;
	long long hFile = 0;
	struct stat buf;
	vector<files> allfile;
	files  filed;
	struct _finddata_t fileInfo;
	vector<String> strpath;
	string pathName;
	if ((hFile = _findfirst(pathName.assign(path).append("\\*").c_str(), &fileInfo)) == -1)
		cout << "错误" << endl;
	do {
		if (strcmp(fileInfo.name, "..") == 0 || strcmp(fileInfo.name, ".") == 0) {
		}
		else {
			cout << path + "\\" + fileInfo.name << ":time:" << fileInfo.time_create << endl;
			if (strstr(fileInfo.name, "png") == NULL && strstr(fileInfo.name, "jpg") == NULL) {
				continue;
			}
			pathName = path + "\\" + fileInfo.name;
			result = stat((char*)pathName.data(), &buf);
			if (result != 0) {
				perror("显示文件状态信息出错");
			}
			else {
				filed.filepath = pathName;
				filed.time = buf.st_atime;
				filed.filename = fileInfo.name;
				allfile.push_back(filed);
			}
		}
	} while (_findnext(hFile, &fileInfo) == 0);
	_findclose(hFile);
	return allfile;
}


int main()
{
	// Loading  Module
	torch::jit::script::Module module = torch::jit::load(R"(E:\deeplearning\yolov5\runs\temp\best3.torchscript.pt)");//torchscript
	torch::DeviceType device_type;
	device_type = torch::kCUDA;
	torch::Device device0(device_type);
	module.to(device0);
	module.to(torch::kHalf);
	module.eval();

	Mat frame, img;
	vector<files> allf = fileSearch(R"(D:\图片)");//文件夹路径
	for (int n = 0; n < allf.size(); n++) {
		clock_t start = clock();
		frame = imread(allf[n].filepath);
		cout << clock() - start << "ms-read" << endl;
		if (frame.empty())
		{
			std::cout << "Read frame failed!" << std::endl;
			break;
		}

		// Preparing input tensor
		resize(frame, img, Size(640, 640));
		cvtColor(img, img, COLOR_BGR2RGB);
		img.convertTo(img, CV_32FC3, 1.0f / 255.0f);  // normalization 1/255
		auto tensor_img = torch::from_blob(img.data, { 1, img.rows, img.cols, img.channels() }).to(device0);
		tensor_img = tensor_img.permute({ 0, 3, 1, 2 }).contiguous();  // BHWC -> BCHW (Batch, Channel, Height, Width)
		tensor_img = tensor_img.to(torch::kHalf);
		std::vector<torch::jit::IValue> inputs;
		inputs.emplace_back(tensor_img);
		torch::jit::IValue output = module.forward(inputs);
		auto detections = output.toTuple()->elements()[0].toTensor();
		auto conf_mask = detections.select(2, 4).ge(0.5).unsqueeze(2);//过滤了score为0.5以下的
		detections = torch::masked_select(detections[0], conf_mask[0]).view({ -1, 85 });//类别+5,这里类别有80个
		detections = detections.to(torch::kFloat);
		detections = detections.to(torch::kCPU);//因为结果是在gpu上的,先拷到cpu,减少通信时间


		clock_t starrrr = clock();
		vector<Rect> boxes;
		vector<float> confidences;
		float* ptr = (float*)detections.data_ptr();//主要是通过这里的指针取值,缩短了后处理时间
		for (int i = 0; i < (int)detections.size(0); ++i, ptr += 85)//类别+5,这里80个类别
		{
			    float confidence = ptr[4];
				int centerX = (int)((ptr[0] / 640) * frame.cols);
				int centerY = (int)((ptr[1] / 640) * frame.rows);
				int width = (int)((ptr[2] / 640) * frame.cols);
				int height = (int)((ptr[3] / 640) * frame.rows);
				int left = (centerX - width / 2);
				int top = (centerY - height / 2);
				confidences.push_back(confidence);
				boxes.push_back(Rect(max(left, 0), max(top, 0), min(width, frame.cols - left), min(height, frame.rows - top)));

		}
		vector<int> indices;
		NMSBoxes(boxes, confidences, 0.5, 0.2, indices);
		for (size_t i = 0; i < indices.size(); ++i)
		{
			int idx = indices[i];
			Rect box = boxes[idx];
			rectangle(frame, box, Scalar(0, 0, 250), 2, 8, 0);

	}


	return 0;
}

 

评论 14
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值