Method Predict

lesson7-Detection/src/Detector.cpp:274–317

}

void Detector::Predict(cv::Mat image, bool show, float conf_thresh, float nms_thresh) {
	int origin_width = image.cols;
	int origin_height = image.rows;
	cv::resize(image, image, { width,height });
	auto img_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte);
	img_tensor = img_tensor.permute({ 2, 0, 1 }).unsqueeze(0).to(torch::kFloat) / 255.0;

	float anchor[12] = { 10,14,  23,27,  37,58,  81,82,  135,169,  344,319 };
	auto anchors_ = torch::from_blob(anchor, { 6,2 }, torch::TensorOptions(torch::kFloat32));
	int image_size[2] = { width,height };
	img_tensor = img_tensor.to(device);

	auto outputs = detector->forward(img_tensor);
	std::vector<torch::Tensor> output_list = {};
	auto tensor_input = outputs[1];
	auto output_decoded = DecodeBox(tensor_input, anchors_.narrow(0, 0, 3), name_list.size(), image_size);
	output_list.push_back(output_decoded);

	tensor_input = outputs[0];
	output_decoded = DecodeBox(tensor_input, anchors_.narrow(0, 3, 3), name_list.size(), image_size);
	output_list.push_back(output_decoded);

	//std::cout << tensor_input << anchors_.narrow(0, 3, 3);

	auto output = torch::cat(output_list, 1);
	auto detection = non_maximum_suppression(output, name_list.size(), conf_thresh, nms_thresh);

	float w_scale = float(origin_width) / width;
	float h_scale = float(origin_height) / height;
	for (int i = 0; i < detection.size(); i++) {
		for (int j = 0; j < detection[i].size(0) / 7; j++)
		{
			detection[i].select(0, 7 * j + 0) *= w_scale;
			detection[i].select(0, 7 * j + 1) *= h_scale;
			detection[i].select(0, 7 * j + 2) *= w_scale;
			detection[i].select(0, 7 * j + 3) *= h_scale;
		}
	}

	cv::resize(image, image, { origin_width,origin_height });
	if(show)
		show_bbox(image, detection[0], name_list);
	return;
}