libtorch实现yolo

cv::Mat origin_image, resized_image;
int input_image_size = 416;
origin_image = cv::imread("004.bmp");
resized_image = origin_image.clone();
cv::resize(resized_image, resized_image, cv::Size(input_image_size, input_image_size));
cv::Mat img_float;
resized_image.convertTo(img_float, CV_32F, 1.0 / 255);

torch::Tensor img_tensor = torch::from_blob(img_float.data, { 1, input_image_size,input_image_size, 3 }).permute({ 0, 3, 1, 2 });
img_tensor = (img_tensor - 0.5) / 0.5;

InputImage.push_back(img_tensor);
auto pred = model->forward(InputImage).toTensor();
//pred格式为2535X6，13x13x3+26x26x3=2535
auto flag = torch::sigmoid(torch::slice(pred, 1, 4, 5)) > 0.5; //含有物体的概率小于0.5的框舍去
auto num_rect = torch::sum(flag).item<int>();
torch::Tensor m_RectTensor = torch::empty({ num_rect, 5 });
auto pp = torch::sum(flag);
int m_rectcount = 0;
for (int k = 0; k < pred.sizes()[0]; k++)
{
	if (flag[k].item<double>() == 1)
	{
		double bx, by, bw, bh;
		getboundrect(pred, k, bx, by, bw, bh);
		float data[5] = { bx - bw / 2, by - bh / 2, bx + bw / 2, by + bh / 2, pred[k][4].item<double>() };
		torch::Tensor rect = torch::from_blob(data, { 5 });
		m_RectTensor[m_rectcount] = rect;
		m_rectcount++;
	}
}
auto sort_index = torch::argsort(m_RectTensor.transpose(1, 0)[4], 0, true); //将矩形按含有物体的概率大小从高到低排列
m_RectTensor = torch::index_select(m_RectTensor, 0, sort_index);
auto bbox = torch::slice(m_RectTensor, 1, 0, 4);
std::vector<torch::Tensor> m_keep;
non_max_suppression(bbox, 0.4, m_keep); //进行非最大值抑制，好像没有怎么起作用
for (int k = 0;k < m_keep.size();k++)
{
	double x1, y1, x2, y2;
	x1 = m_keep[k][0].item<double>();
	y1 = m_keep[k][1].item<double>();
	x2 = m_keep[k][2].item<double>();
	y2 = m_keep[k][3].item<double>();
	cv::rectangle(origin_image, cv::Point(x1, y1), cv::Point(x2, y2), (0, 0, 255), 1, 8);
}
imwrite("dd.jpg", origin_image);
system("pause");