两百行C++代码实现yolov5车辆计数部署(通俗易懂版)

本文是文章传统图像处理方法实现车辆计数的后续。这里用OpenCV实现了基于yolov5检测器的单向车辆计数功能,方法是撞线计数。该代码只能演示视频demo效果,一些功能未完善,离实际工程应用还有距离。
实现流程:
(1)训练yolov5模型,这里就没有自己训练了,直接使用官方的开源模型yolov5s.pt;
(2)运行yolov5工程下面的export.py,将pt模型转成onnx模型;
(3)编写yolov5部署的C++工程,包括前处理、推理和后处理部分;
(4)读取视频第一帧,用yolov5检测第一帧图像的车辆目标,计算这些检测框的中心点,
(5)读取视频的后续帧,用yolov5检测每帧图像上的车辆目标,计算新目标和上一帧图像中检测框中心点的距离矩阵;
(6)通过距离矩阵确定新旧目标检测框之间的对应关系;
(7)计算对应新旧目标检测框中心点之间的连线,判断和事先设置的虚拟撞线是否相交,若相交则计数加1;
(8)重复(5)-(7)。
由于程序在CPU端运行,为了提速,实际实现的时候采取的是隔帧判断而不是使用相邻帧,v1的代码实现如下:

#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>


// 常量
const float INPUT_WIDTH = 640.0;
const float INPUT_HEIGHT = 640.0;
const float SCORE_THRESHOLD = 0.5;
const float NMS_THRESHOLD = 0.45;
const float CONFIDENCE_THRESHOLD = 0.45;

const std::vector<std::string> class_name = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush" };


// 画框函数
void draw_label(cv::Mat& input_image, std::string label, int left, int top)
{
	int baseLine;
	cv::Size label_size = cv::getTextSize(label, 0.7, 0.7, 1, &baseLine);
	top = std::max(top, label_size.height);
	cv::Point tlc = cv::Point(left, top);
	cv::Point brc = cv::Point(left , top + label_size.height + baseLine);
	cv::putText(input_image, label, cv::Point(left, top + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 255), 1);
}


// 预处理
std::vector<cv::Mat> preprocess(cv::Mat& input_image, cv::dnn::Net& net)
{
	cv::Mat blob;
	cv::dnn::blobFromImage(input_image, blob, 1. / 255., cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::Scalar(), true, false);

	net.setInput(blob);

	std::vector<cv::Mat> preprcess_image;
	net.forward(preprcess_image, net.getUnconnectedOutLayersNames());

	return preprcess_image;
}


// 后处理
std::vector<cv::Rect> postprocess(std::vector<cv::Mat>& preprcess_image, cv::Mat& output_image)
{
	std::vector<int> class_ids;
	std::vector<float> confidences;
	std::vector<cv::Rect> boxes;
	std::vector<cv::Rect> boxes_nms;

	float x_factor = output_image.cols / INPUT_WIDTH;
	float y_factor = output_image.rows / INPUT_HEIGHT;

	float* data = (float*)preprcess_image[0].data;

	const int dimensions = 85;
	const int rows = 25200;
	for (int i = 0; i < rows; ++i)
	{
		float confidence = data[4];
		if (confidence >= CONFIDENCE_THRESHOLD)
		{
			float* classes_scores = data + 5;
			cv::Mat scores(1, class_name.size(), CV_32FC1, classes_scores);
			cv::Point class_id;
			double max_class_score;
			cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
			if (max_class_score > SCORE_THRESHOLD)
			{
				confidences.push_back(confidence);
				class_ids.push_back(class_id.x);

				float cx = data[0];
				float cy = data[1];
				float w = data[2];
				float h = data[3];
				int left = int((cx - 0.5 * w) * x_factor);
				int top = int((cy - 0.5 * h) * y_factor);
				int width = int(w * x_factor);
				int height = int(h * y_factor);
				boxes.push_back(cv::Rect(left, top, width, height));
			}
		}
		data += 85;
	}

	std::vector<int> indices;
	cv::dnn::NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, indices);
	for (size_t i = 0; i < indices.size(); i++)
	{
		int idx = indices[i];
		cv::Rect box = boxes[idx];
		boxes_nms.push_back(box);

		int left = box.x;
		int top = box.y;
		int width = box.width;
		int height = box.height;
		cv::rectangle(output_image, cv::Point(left, top), cv::Point(left + width, top + height), cv::Scalar(255, 0, 0), 1);

		std::string label = cv::format("%.2f", confidences[idx]);
		label = class_name[class_ids[idx]] + ":" + label;
		draw_label(output_image, label, left, top);
	}
	return boxes_nms;
}


std::vector<cv::Point> get_centers(std::vector<cv::Rect> detections)
{
	std::vector<cv::Point> detections_centers(detections.size());
	for (size_t i = 0; i < detections.size(); i++)
	{
		detections_centers[i] = cv::Point(detections[i].x + detections[i].width / 2, detections[i].y + detections[i].height / 2);
	}

	return detections_centers;
}


float get_distance(cv::Point p1, cv::Point p2)
{
	return sqrt(pow(p1.x - p2.x, 2) + pow(p1.y - p2.y, 2));
}


bool is_cross(cv::Point p1, cv::Point p2)
{
	if (p1.x == p2.x) return false;

	int y = 500;  //line1: y = 500
	float k = (p1.y - p2.y) / (p1.x - p2.x);  //
	float b = p1.y - k * p1.x; //line2: y = kx + b
	float x = (y - b) / k;
	return (x > std::min(p1.x, p2.x) && x < std::max(p1.x, p2.x));
}


int main(int argc, char** argv)
{
	cv::VideoCapture capture("test.mp4");
	cv::Mat frame;
	cv::dnn::Net net = cv::dnn::readNet("yolov5s-f32.onnx");

	int frame_num = 0;
	int count = 0;
	std::vector<cv::Point> detections_centers_old;
	std::vector<cv::Point> detections_centers_new;

	while(cv::waitKey(1) < 0)
	{
	    capture >> frame;
		if (frame.empty())
			break;

		std::cout << "******************************************************************* frame_num: " << frame_num << std::endl;

		cv::Mat image = frame.clone();
		std::vector<cv::Mat> preprcess_image = preprocess(image, net);

		std::vector<cv::Rect> detections = postprocess(preprcess_image, image);

		if (frame_num == 0)
		{
			detections_centers_old = get_centers(detections);

			std::cout << "detections_center:" << std::endl;
			for (size_t i = 0; i < detections_centers_old.size(); i++)
			{
				std::cout << detections_centers_old[i] << std::endl;
			}
		}
		else if (frame_num % 2 == 0)
		{
			detections_centers_new = get_centers(detections);

			std::cout << "detections_center:" << std::endl;
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				std::cout << detections_centers_new[i] << std::endl;
			}

			std::vector<std::vector<float>> distance_matrix(detections_centers_new.size(), std::vector<float>(detections_centers_old.size()));
			std::cout << "distance_matrix:" << std::endl;
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				for (size_t j = 0; j < detections_centers_old.size(); j++)
				{
					distance_matrix[i][j] = get_distance(detections_centesr_new[i], detections_centers_old[j]); //
					std::cout << distance_matrix[i][j] << " ";
				}
				std::cout << std::endl;
			}

			std::cout << "min_index:" << std::endl;
			std::vector<float> min_indices(detections_centers_new.size());
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				std::vector<float> distance_vector = distance_matrix[i];
				int min_index = std::min_element(distance_vector.begin(), distance_vector.end()) - distance_vector.begin();
				min_indices[i] = min_index;
				std::cout << min_index << " ";
			}
			std::cout << std::endl;

			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				cv::Point p1 = detections_centers_new[i];
				cv::Point p2 = detections_centers_old[min_indices[i]];
				std::cout << p1 << " " << p2 << std::endl;

				if (is_cross(p1, p2))
				{
					std::cout << "is_cross" << p1 << " " << p2 << std::endl;
					count++;
				}
			}
			detections_centers_old = detections_centers_new;
		}

		frame_num++;

		cv::putText(image, "car num: " + std::to_string(count), cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 255), 1);
		cv::line(image, cv::Point(0, 500), cv::Point(1280, 500) , cv::Scalar(0, 0, 255));
		cv::imshow("output", image);
		cv::imwrite(std::to_string(frame_num) + ".jpg", image);
	}

	capture.release();
	return 0;
}

在调试中,发现v1的实现存在如下问题:出现新目标的时候,计算新旧检测框的对应关系出现匹配错误,导致计数偏多。因此在v2中设置匹配的距离阈值,并简化了判断检测框中心点连线和撞线是否相交的方法。
v2的代码实现如下:

#include <iostream>
#include <opencv2/opencv.hpp>


//#define DEBUG


// 常量
const float INPUT_WIDTH = 640.0;
const float INPUT_HEIGHT = 640.0;
const float SCORE_THRESHOLD = 0.5;
const float NMS_THRESHOLD = 0.25;
const float CONFIDENCE_THRESHOLD = 0.5;

const std::vector<std::string> class_name = {
	"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
	"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
	"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
	"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
	"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
	"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
	"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
	"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
	"hair drier", "toothbrush" };

const int IMAGE_WIDTH = 1280;
const int IMAGE_HEIGHT = 720;
const int LINE_HEIGHT = IMAGE_HEIGHT / 2;


//画出检测框和标签
void draw_label(cv::Mat& input_image, std::string label, int left, int top)
{
	int baseLine;
	cv::Size label_size = cv::getTextSize(label, 0.7, 0.7, 1, &baseLine);
	top = std::max(top, label_size.height);
	cv::Point tlc = cv::Point(left, top);
	cv::Point brc = cv::Point(left , top + label_size.height + baseLine);
	cv::putText(input_image, label, cv::Point(left, top + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 255), 1);
}


//预处理
std::vector<cv::Mat> preprocess(cv::Mat& input_image, cv::dnn::Net& net)
{
	cv::Mat blob;
	cv::dnn::blobFromImage(input_image, blob, 1. / 255., cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::Scalar(), true, false);

	net.setInput(blob);

	std::vector<cv::Mat> preprcess_image;
	net.forward(preprcess_image, net.getUnconnectedOutLayersNames());

	return preprcess_image;
}


//后处理
std::vector<cv::Rect> postprocess(std::vector<cv::Mat>& preprcess_image, cv::Mat& output_image)
{
	std::vector<int> class_ids;
	std::vector<float> confidences;
	std::vector<cv::Rect> boxes;
	std::vector<cv::Rect> boxes_nms;

	float x_factor = output_image.cols / INPUT_WIDTH;
	float y_factor = output_image.rows / INPUT_HEIGHT;

	float* data = (float*)preprcess_image[0].data;

	const int dimensions = 85;
	const int rows = 25200;
	for (int i = 0; i < rows; ++i)
	{
		float confidence = data[4];
		if (confidence >= CONFIDENCE_THRESHOLD)
		{
			float* classes_scores = data + 5;
			cv::Mat scores(1, class_name.size(), CV_32FC1, classes_scores);
			cv::Point class_id;
			double max_class_score;
			cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
			if (max_class_score > SCORE_THRESHOLD)
			{
				confidences.push_back(confidence);
				class_ids.push_back(class_id.x);

				float cx = data[0];
				float cy = data[1];
				float w = data[2];
				float h = data[3];
				int left = int((cx - 0.5 * w) * x_factor);
				int top = int((cy - 0.5 * h) * y_factor);
				int width = int(w * x_factor);
				int height = int(h * y_factor);
				boxes.push_back(cv::Rect(left, top, width, height));
			}
		}
		data += 85;
	}

	std::vector<int> indices;
	cv::dnn::NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, indices);
	for (size_t i = 0; i < indices.size(); i++)
	{
		int idx = indices[i];
		cv::Rect box = boxes[idx];
		boxes_nms.push_back(box);

		int left = box.x;
		int top = box.y;
		int width = box.width;
		int height = box.height;
		cv::rectangle(output_image, cv::Point(left, top), cv::Point(left + width, top + height), cv::Scalar(255, 0, 0), 1);

		std::string label = cv::format("%.2f", confidences[idx]);
		//label = class_name[class_ids[idx]] + ":" + label;
		label = "car";
		draw_label(output_image, label, left, top);
	}

	return boxes_nms;
}


//计算检测框的中心
std::vector<cv::Point> get_centers(std::vector<cv::Rect> detections)
{
	std::vector<cv::Point> detections_centers(detections.size());
	for (size_t i = 0; i < detections.size(); i++)
	{
		detections_centers[i] = cv::Point(detections[i].x + detections[i].width / 2, detections[i].y + detections[i].height / 2);
	}

	return detections_centers;
}


//计算两点间距离
float get_distance(cv::Point p1, cv::Point p2)
{
	return sqrt(pow(p1.x - p2.x, 2) + pow(p1.y - p2.y, 2));
}


//判断连接相邻两帧对应检测框中心的线段是否与红线相交
bool is_cross(cv::Point p1, cv::Point p2)
{
	return (p1.y <= LINE_HEIGHT && p2.y > LINE_HEIGHT) || (p1.y > LINE_HEIGHT && p2.y <= LINE_HEIGHT);
}


int main(int argc, char** argv)
{
	cv::VideoCapture capture("test.mp4");
	cv::Mat frame;
	cv::dnn::Net net = cv::dnn::readNet("yolov5s-f32.onnx");

	int frame_num = 0;
	int count = 0;
	std::vector<cv::Point> detections_centers_old;
	std::vector<cv::Point> detections_centers_new;

	while(cv::waitKey(1) < 0)
	{
	    capture >> frame;
		if (frame.empty())
			break;

		std::cout << "******************************************************************* frame_num: " << frame_num << std::endl;

		cv::Mat image = frame.clone();
		std::vector<cv::Mat> preprcess_image = preprocess(image, net);

		std::vector<cv::Rect> detections = postprocess(preprcess_image, image);

		if (frame_num == 0)
		{
			detections_centers_old = get_centers(detections);

#ifdef DEBUG
			std::cout << "detections_center:" << std::endl;
			for (size_t i = 0; i < detections_centers_old.size(); i++)
			{
				std::cout << detections_centers_old[i] << std::endl;
			}
#endif // DEBUG
		}
		else if (frame_num % 2 == 0)
		{
			detections_centers_new = get_centers(detections);

#ifdef DEBUG
			std::cout << "detections_center:" << std::endl;
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				std::cout << detections_centers_new[i] << std::endl;
			}
#endif // DEBUG

			std::vector<std::vector<float>> distance_matrix(detections_centers_new.size(), std::vector<float>(detections_centers_old.size())); //距离矩阵
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				for (size_t j = 0; j < detections_centers_old.size(); j++)
				{
					distance_matrix[i][j] = get_distance(detections_centers_new[i], detections_centers_old[j]); 
				}
			}

#ifdef DEBUG
			std::cout << "min_index:" << std::endl;
#endif // DEBUG

			std::vector<float> min_indices(detections_centers_new.size());
			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				std::vector<float> distance_vector = distance_matrix[i];
				float min_val = *std::min_element(distance_vector.begin(), distance_vector.end());
				int min_index = -1;
				if (min_val < LINE_HEIGHT / 5)
					 min_index = std::min_element(distance_vector.begin(), distance_vector.end()) - distance_vector.begin();
				
				min_indices[i] = min_index;
#ifdef DEBUG
				std::cout << min_index << " ";
#endif // DEBUG
			}
			std::cout << std::endl;

			for (size_t i = 0; i < detections_centers_new.size(); i++)
			{
				if (min_indices[i] < 0)
					continue;

				cv::Point p1 = detections_centers_new[i];
				cv::Point p2 = detections_centers_old[min_indices[i]];

#ifdef DEBUG
				std::cout << p1 << " " << p2 << std::endl;
#endif // DEBUG

				if (is_cross(p1, p2))
				{
#ifdef DEBUG
					std::cout << "is_cross" << p1 << " " << p2 << std::endl;
#endif // DEBUG
					count++;
				}
			}

			detections_centers_old = detections_centers_new;
		}

		cv::putText(image, "car num: " + std::to_string(count), cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 0, 255), 1);
		cv::line(image, cv::Point(0, LINE_HEIGHT), cv::Point(IMAGE_WIDTH, LINE_HEIGHT), cv::Scalar(0, 0, 255));
		cv::imshow("output", image);

#ifdef DEBUG
		if (frame_num % 2 == 0)
			cv::imwrite(std::to_string(frame_num) + ".jpg", image);
#endif // DEBUG

		frame_num++;
	}

	capture.release();
	return 0;
}

检测效果如下图,效果还是可以的,比传统方法有大幅提升。完整视频中有一次计数异常(总共52辆向图像上方行驶的车辆,检出53辆),是因为检测器不准导致车辆检测框位置漂移,可以后续优化。注:由于官方提供的coco80类的开源权重文件用于车辆检测效果不是很好,LZ把检测出的类别直接固定为car,实际应自己重新训练一个车辆检测的模型。
在这里插入图片描述

更详细注释的代码、测试视频和转好的权重文件放在下载链接:点击跳转

后续文章:
目标跟踪算法实现车辆计数
五十行python代码实现yolov8车辆计数

评论 13
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

给算法爸爸上香

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值