【C++ yolov5 libtorch】GPU windows测试

yolov5LoadAndDetect.cpp

#include 
#include 
#include 

#include "include/detector.h"
#include "include/cxxopts.hpp"


using namespace std;

std::vector LoadNames(const std::string& path) {
	// load class names
	std::vector class_names;
	std::ifstream infile(path);
	if (infile.is_open()) {
		std::string line;
		while (getline(infile, line)) {
			class_names.emplace_back(line);
		}
		infile.close();
	}
	else {
		std::cerr << "Error loading the class names!\n";
	}

	return class_names;
}


void Demo(cv::Mat& img,
	const std::vector>& detections,
	const std::vector& class_names,
	bool label = true) {

	if (!detections.empty()) {
		for (const auto& detection : detections[0]) {
			const auto& box = detection.bbox;
			float score = detection.score;
			int class_idx = detection.class_idx;

			cv::rectangle(img, box, cv::Scalar(0, 0, 255), 2);

			if (label) {
				std::stringstream ss;
				ss << std::fixed << std::setprecision(2) << score;
				std::string s = class_names[class_idx] + " " + ss.str();

				auto font_face = cv::FONT_HERSHEY_DUPLEX;
				auto font_scale = 1.0;
				int thickness = 1;
				int baseline = 0;
				auto s_size = cv::getTextSize(s, font_face, font_scale, thickness, &baseline);
				cv::rectangle(img,
					cv::Point(box.tl().x, box.tl().y - s_size.height - 5),
					cv::Point(box.tl().x + s_size.width, box.tl().y),
					cv::Scalar(0, 0, 255), -1);
				cv::putText(img, s, cv::Point(box.tl().x, box.tl().y - 5),
					font_face, font_scale, cv::Scalar(255, 255, 255), thickness);
			}
		}
	}

	//cv::namedWindow("Result", cv::WINDOW_AUTOSIZE);
	cv::resize(img, img, cv::Size(717, 600));// cv::Size(956, 800)  cv::Size(717, 600)
	cv::imshow("Result", img);
	cv::waitKey(0);
}


int main(int argc, const char* argv[]) {
	cxxopts::Options parser(argv[0], "A LibTorch inference implementation of the yolov5");

	// TODO: add other args
	parser.allow_unrecognised_options().add_options()
		("weights", "model.torchscript.pt path", cxxopts::value()->default_value("weights/best.torchscript"))//torchscript格式 已验证
		("source", "source", cxxopts::value()->default_value("./test/1.bmp"))//bus.jpgv
		("conf-thres", "object confidence threshold", cxxopts::value()->default_value("0.4"))
		("iou-thres", "IOU threshold for NMS", cxxopts::value()->default_value("0.5"))
		("gpu", "Enable cuda device or cpu", cxxopts::value()->default_value("true"))
		("view-img", "display results", cxxopts::value()->default_value("true"))
		("h,help", "Print usage");

	auto opt = parser.parse(argc, argv);

	if (opt.count("help")) {
		std::cout << parser.help() << std::endl;
		exit(0);
	}

	// check if gpu flag is set
	bool is_gpu = opt["gpu"].as();

	// set device type - CPU/GPU
	torch::DeviceType device_type; //libtorch 版本应为GPU
	if (torch::cuda::is_available() && is_gpu) {// 链接器-命令行 /INCLUDE:?warp_size@cuda@at@@YAHXZ   /INCLUDE:?searchsorted_cuda@native@at@@YA?AVTensor@2@AEBV32@0_N1@Z
		device_type = torch::kCUDA;//重新编译 pytorch 使得编译时CUDA能够与运行时CUDA保持一致 https://blog.csdn.net/qq_36038453/article/details/120278523

	}
	else {
		device_type = torch::kCPU;
	}
	std::cout << "Device available :" << torch::cuda::is_available() << std::endl;
	// load class names from dataset for visualization
	std::vector class_names = LoadNames("weights/block.txt");
	if (class_names.empty()) {
		return -1;
	}

	// load network
	std::string weights = opt["weights"].as();
	auto detector = Detector(weights, device_type);





	// load input image
	std::string source = opt["source"].as();
	cv::Mat img = cv::imread(source);
	if (img.empty()) {
		std::cerr << "Error loading the image!\n";
		return -1;
	}

	// run once to warm up
	std::cout << "Run once on empty image" << std::endl;
	auto temp_img = cv::Mat::zeros(img.rows, img.cols, CV_32FC3);//
	detector.Run(temp_img, 1.0f, 1.0f);//temp_imag 是黑色

	// set up threshold
	float conf_thres = opt["conf-thres"].as();
	float iou_thres = opt["iou-thres"].as();

	// inference
	auto result = detector.Run(img, conf_thres, iou_thres);

	// visualize detections
	if (opt["view-img"].as()) {
		Demo(img, result, class_names);
	}

	cv::destroyAllWindows();
	return 0;
}

detector.cpp

#include "detector.h"


Detector::Detector(const std::string& model_path, const torch::DeviceType& device_type) : device_(device_type) {
	try {
		// Deserialize the ScriptModule from a file using torch::jit::load().
		LoadLibraryA("ATen_cuda.dll");
		LoadLibraryA("c10_cuda.dll");
		LoadLibraryA("torch_cuda.dll");
		LoadLibraryA("torchvision.dll");
		module_ = torch::jit::load(model_path);
	}
	catch (const c10::Error& e) {
		std::cerr << "Error loading the model!\n" << e.what();
		std::exit(EXIT_FAILURE);
	}

	half_ = (device_ != torch::kCPU);
	module_.to(device_);

	if (half_) {
		module_.to(torch::kHalf);
	}

	module_.eval();
}


std::vector>
Detector::Run(const cv::Mat& img, float conf_threshold, float iou_threshold) {
	torch::NoGradGuard no_grad;
	std::cout << "----------New Frame----------" << std::endl;

	// TODO: check_img_size()

	/*** Pre-process ***/

	auto start = std::chrono::high_resolution_clock::now();

	// keep the original image for visualization purpose
	cv::Mat img_input = img.clone();

	std::vector pad_info = LetterboxImage(img_input, img_input, cv::Size(416, 416));//cv::Size(640, 640)
	const float pad_w = pad_info[0];
	const float pad_h = pad_info[1];
	const float scale = pad_info[2];

	cv::cvtColor(img_input, img_input, cv::COLOR_BGR2RGB);  // BGR -> RGB     旧的cv::COLOR_BGR2RGB
	img_input.convertTo(img_input, CV_32FC3, 1.0f / 255.0f);  // normalization 1/255
	auto tensor_img = torch::from_blob(img_input.data, { 1, img_input.rows, img_input.cols, img_input.channels() }).to(device_);

	tensor_img = tensor_img.permute({ 0, 3, 1, 2 }).contiguous();  // BHWC -> BCHW (Batch, Channel, Height, Width)

	if (half_) {
		tensor_img = tensor_img.to(torch::kHalf);
	}

	std::vector inputs;
	inputs.emplace_back(tensor_img);

	auto end = std::chrono::high_resolution_clock::now();
	auto duration = std::chrono::duration_cast(end - start);
	// It should be known that it takes longer time at first time
	std::cout << "pre-process takes : " << duration.count() << " ms" << std::endl;

	/*** Inference ***/
	// TODO: add synchronize point
	start = std::chrono::high_resolution_clock::now();

	// inference
	torch::jit::IValue output = module_.forward(inputs);//
	//auto output = module_.forward(inputs);

	//auto detections = module_.forward(inputs).toTensor();
	end = std::chrono::high_resolution_clock::now();
	duration = std::chrono::duration_cast(end - start);
	// It should be known that it takes longer time at first time
	std::cout << "inference takes : " << duration.count() << " ms" << std::endl;

	/*** Post-process ***/

	start = std::chrono::high_resolution_clock::now();
	auto detections = output.toTuple()->elements()[0].toTensor();

	// result: n * 7
	// batch index(0), top-left x/y (1,2), bottom-right x/y (3,4), score(5), class id(6)
	auto result = PostProcessing(detections, pad_w, pad_h, scale, img.size(), conf_threshold, iou_threshold);

	end = std::chrono::high_resolution_clock::now();
	duration = std::chrono::duration_cast(end - start);
	// It should be known that it takes longer time at first time
	std::cout << "post-process takes : " << duration.count() << " ms" << std::endl;

	return result;
}


std::vector Detector::LetterboxImage(const cv::Mat& src, cv::Mat& dst, const cv::Size& out_size) {
	auto in_h = static_cast(src.rows);
	auto in_w = static_cast(src.cols);
	float out_h = out_size.height;
	float out_w = out_size.width;

	float scale = (std::min)(out_w / in_w, out_h / in_h);

	int mid_h = static_cast(in_h * scale);
	int mid_w = static_cast(in_w * scale);

	cv::resize(src, dst, cv::Size(mid_w, mid_h));

	int top = (static_cast(out_h) - mid_h) / 2;
	int down = (static_cast(out_h) - mid_h + 1) / 2;
	int left = (static_cast(out_w) - mid_w) / 2;
	int right = (static_cast(out_w) - mid_w + 1) / 2;

	cv::copyMakeBorder(dst, dst, top, down, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));

	std::vector pad_info{ static_cast(left), static_cast(top), scale };
	return pad_info;
}


std::vector> Detector::PostProcessing(const torch::Tensor& detections,
	float pad_w, float pad_h, float scale, const cv::Size& img_shape,
	float conf_thres, float iou_thres) {
	constexpr int item_attr_size = 5;
	int batch_size = detections.size(0);
	// number of classes, e.g. 80 for coco dataset
	auto num_classes = detections.size(2) - item_attr_size;

	// get candidates which object confidence > threshold
	auto conf_mask = detections.select(2, 4).ge(conf_thres).unsqueeze(2);

	std::vector> output;
	output.reserve(batch_size);

	// iterating all images in the batch
	for (int batch_i = 0; batch_i < batch_size; batch_i++) {
		// apply constrains to get filtered detections for current image
		auto det = torch::masked_select(detections[batch_i], conf_mask[batch_i]).view({ -1, num_classes + item_attr_size });

		// if none detections remain then skip and start to process next image
		if (0 == det.size(0)) {
			continue;
		}

		// compute overall score = obj_conf * cls_conf, similar to x[:, 5:] *= x[:, 4:5]
		det.slice(1, item_attr_size, item_attr_size + num_classes) *= det.select(1, 4).unsqueeze(1);

		// box (center x, center y, width, height) to (x1, y1, x2, y2)
		torch::Tensor box = xywh2xyxy(det.slice(1, 0, 4));

		// [best class only] get the max classes score at each result (e.g. elements 5-84)
		std::tuple max_classes = (torch::max)(det.slice(1, item_attr_size, item_attr_size + num_classes), 1);

		// class score
		auto max_conf_score = std::get<0>(max_classes);
		// index
		auto max_conf_index = std::get<1>(max_classes);

		max_conf_score = max_conf_score.to(torch::kFloat).unsqueeze(1);
		max_conf_index = max_conf_index.to(torch::kFloat).unsqueeze(1);

		// shape: n * 6, top-left x/y (0,1), bottom-right x/y (2,3), score(4), class index(5)
		det = torch::cat({ box.slice(1, 0, 4), max_conf_score, max_conf_index }, 1);

		// for batched NMS
		constexpr int max_wh = 4096;
		auto c = det.slice(1, item_attr_size, item_attr_size + 1) * max_wh;
		auto offset_box = det.slice(1, 0, 4) + c;

		std::vector offset_box_vec;
		std::vector score_vec;

		// copy data back to cpu
		auto offset_boxes_cpu = offset_box.cpu();
		auto det_cpu = det.cpu();
		const auto& det_cpu_array = det_cpu.accessor();

		// use accessor to access tensor elements efficiently
		Tensor2Detection(offset_boxes_cpu.accessor(), det_cpu_array, offset_box_vec, score_vec);

		// run NMS
		std::vector nms_indices;
		cv::dnn::NMSBoxes(offset_box_vec, score_vec, conf_thres, iou_thres, nms_indices);

		std::vector det_vec;
		for (int index : nms_indices) {
			Detection t;
			const auto& b = det_cpu_array[index];
			t.bbox =
				cv::Rect(cv::Point(b[Det::tl_x], b[Det::tl_y]),
					cv::Point(b[Det::br_x], b[Det::br_y]));
			t.score = det_cpu_array[index][Det::score];
			t.class_idx = det_cpu_array[index][Det::class_idx];
			det_vec.emplace_back(t);
		}

		ScaleCoordinates(det_vec, pad_w, pad_h, scale, img_shape);

		// save final detection for the current image
		output.emplace_back(det_vec);
	} // end of batch iterating

	return output;
}


void Detector::ScaleCoordinates(std::vector& data, float pad_w, float pad_h,
	float scale, const cv::Size& img_shape) {
	auto clip = [](float n, float lower, float upper) {
		return (std::max)(lower, (std::min)(n, upper));
	};

	std::vector detections;
	for (auto& i : data) {
		float x1 = (i.bbox.tl().x - pad_w) / scale;  // x padding
		float y1 = (i.bbox.tl().y - pad_h) / scale;  // y padding
		float x2 = (i.bbox.br().x - pad_w) / scale;  // x padding
		float y2 = (i.bbox.br().y - pad_h) / scale;  // y padding

		x1 = clip(x1, 0, (float)img_shape.width);
		y1 = clip(y1, 0, (float)img_shape.height);
		x2 = clip(x2, 0, (float)img_shape.width);
		y2 = clip(y2, 0, (float)img_shape.height);

		i.bbox = cv::Rect(cv::Point(x1, y1), cv::Point(x2, y2));
	}
}


torch::Tensor Detector::xywh2xyxy(const torch::Tensor& x) {
	auto y = torch::zeros_like(x);
	// convert bounding box format from (center x, center y, width, height) to (x1, y1, x2, y2)
	y.select(1, Det::tl_x) = x.select(1, 0) - x.select(1, 2).div(2);
	y.select(1, Det::tl_y) = x.select(1, 1) - x.select(1, 3).div(2);
	y.select(1, Det::br_x) = x.select(1, 0) + x.select(1, 2).div(2);
	y.select(1, Det::br_y) = x.select(1, 1) + x.select(1, 3).div(2);
	return y;
}


void Detector::Tensor2Detection(const at::TensorAccessor& offset_boxes,
	const at::TensorAccessor& det,
	std::vector& offset_box_vec,
	std::vector& score_vec) {

	for (int i = 0; i < offset_boxes.size(0); i++) {
		offset_box_vec.emplace_back(
			cv::Rect(cv::Point(offset_boxes[i][Det::tl_x], offset_boxes[i][Det::tl_y]),
				cv::Point(offset_boxes[i][Det::br_x], offset_boxes[i][Det::br_y]))
		);
		score_vec.emplace_back(det[i][Det::score]);
	}
}

detector.h

# pragma once

#include 

#include 
//#torch/torch
#include "torch/torch.h"
#include 
#include 

#include 
#include 
#include 
#include 
#include 
#include "utils.h"

class Detector {
public:
	/***
	 * @brief constructor
	 * @param model_path - path of the TorchScript weight file
	 * @param device_type - inference with CPU/GPU
	 */
	Detector(const std::string& model_path, const torch::DeviceType& device_type);

	/***
	 * @brief inference module
	 * @param img - input image
	 * @param conf_threshold - confidence threshold
	 * @param iou_threshold - IoU threshold for nms
	 * @return detection result - bounding box, score, class index
	 */
	std::vector>
		Run(const cv::Mat& img, float conf_threshold, float iou_threshold);

private:
	/***
	 * @brief Padded resize
	 * @param src - input image
	 * @param dst - output image
	 * @param out_size - desired output size
	 * @return padding information - pad width, pad height and zoom scale
	 */
	static std::vector LetterboxImage(const cv::Mat& src, cv::Mat& dst, const cv::Size& out_size = cv::Size(640, 640));

	/***
	 * @brief Performs Non-Maximum Suppression (NMS) on inference results
	 * @note For 640x640 image, 640 / 32(max stride) = 20, sum up boxes from each yolo layer with stride (8, 16, 32) and
	 *       3 scales at each layer, we can get total number of boxes - (20x20 + 40x40 + 80x80) x 3 = 25200
	 * @param detections - inference results from the network, example [1, 25200, 85], 85 = 4(xywh) + 1(obj conf) + 80(class score)
	 * @param conf_thres - object confidence(objectness) threshold
	 * @param iou_thres - IoU threshold for NMS algorithm
	 * @return detections with shape: nx7 (batch_index, x1, y1, x2, y2, score, classification)
	 */
	static std::vector> PostProcessing(const torch::Tensor& detections,
		float pad_w, float pad_h, float scale, const cv::Size& img_shape,
		float conf_thres = 0.4, float iou_thres = 0.6);

	/***
	 * @brief Rescale coordinates to original input image
	 * @param data - detection result after inference and nms
	 * @param pad_w - width padding
	 * @param pad_h - height padding
	 * @param scale - zoom scale
	 * @param img_shape - original input image shape
	 */
	static void ScaleCoordinates(std::vector& data, float pad_w, float pad_h,
		float scale, const cv::Size& img_shape);

	/***
	 * @brief box (center x, center y, width, height) to (x1, y1, x2, y2)
	 * @param x - input box with xywh format
	 * @return box with xyxy format
	 */
	static torch::Tensor xywh2xyxy(const torch::Tensor& x);

	/***
	 * @brief Convert data from Tensors to vectors
	 */
	static void Tensor2Detection(const at::TensorAccessor& offset_boxes,
		const at::TensorAccessor& det,
		std::vector& offset_box_vec,
		std::vector& score_vec);

	torch::jit::script::Module module_;
	torch::Device device_;
	bool half_;
};

utils.h

#pragma once

enum Det {
    tl_x = 0,
    tl_y = 1,
    br_x = 2,
    br_y = 3,
    score = 4,
    class_idx = 5
};

struct Detection {
    cv::Rect bbox;
    float score;
    int class_idx;
};

cxxopts.hpp

/*

Copyright (c) 2014, 2015, 2016, 2017 Jarryd Beck

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

*/

#ifndef CXXOPTS_HPP_INCLUDED
#define CXXOPTS_HPP_INCLUDED

#include 
#include 
#include 
#include 
#include 
#include 

属性：

include:

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1\include
D:\Software\?\cudaopencv\opencv4.5.0\build\install\include
C:\Program Files\PCL 1.10.0\3rdParty\Eigen\eigen3
D:\Software\?\cudaopencv\opencv4.5.0\build\install\include\opencv2

 连接器-命令行

/INCLUDE:?warp_size@cuda@at@@YAHXZ 

连接器-输入

asmjit.lib
torch_cuda_cpp.lib
c10.lib
c10_cuda.lib
caffe2_nvrtc.lib
clog.lib
cpuinfo.lib
dnnl.lib
kineto.lib
fbgemm.lib
libprotobuf.lib
libprotobuf-lite.lib
libprotoc.lib
pthreadpool.lib
torch.lib
torch_cpu.lib
torch_cuda.lib
torch_cuda_cu.lib
XNNPACK.lib

C/C++ 常规 附加包含目录：

D:\Software\vs2019+pcl+opencv\libtorch\include
D:\Software\vs2019+pcl+opencv\libtorch\include\torch\csrc\api\include

C/C++ 所有选项符合模式-是 (/permissive-)

不使用预编译头

库目录：D:\Software\vs2019+pcl+opencv\cudaopencv\opencv4.5.0\build\install\x64\vc16\lib

系统变量：

D:\Software\vs2019+pcl+opencv\libtorch\lib

D:\Software\vs2019+pcl+opencv\cudaopencv\opencv4.5.0\build\install\x64\vc16\bin

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1\bin

 D:\Software\vs2019+pcl+opencv\libtorch\bin