yolov5+bytrack可以实现行人跟踪,每个人都会有一个ID,这就能很好的入侵人数做计算,或者能保存每个入侵者的图像
思路是:
准备一个与检测图像同样大小的attribute_bord,指定的禁入区域值为255,安全区值为0
每一帧都用行人目标的中心点坐标去attribute_bord上索要属性值,检查是否已进入禁入区
如进入禁入区,则将其ID加入入侵列表中。
效果:
关键代码:
*!
@Description : https://github.com/shaoshengsong/
@Author : shaoshengsong
@Date : 2022-09-23 02:52:22
*/
#include <fstream>
#include <sstream>
//#include <iostream>
//#include <cstring>
#include <opencv2/opencv.hpp>
#include "YOLOv5Detector.h"
//#include "FeatureTensor.h"
#include "BYTETracker.h" //bytetrack
#include "tracker.h"//deepsort
//Deep SORT parameter
// https://cloud.tencent.com/developer/article/2099504
using namespace std;
using namespace cv;
std::vector<cv::Point> pts;
void on_Mouse(int event, int x, int y, int flags, void* param){
if (event == EVENT_LBUTTONDOWN)// 左键按下
{
Point pt = Point(x, y);
pts.push_back(pt);
}
}
void invade_detect(Mat& img, vector<Point>& pts, vector<STrack>& bytrack_result){
int invaded = 255; //安全区标识板颜色值为0
vector<int> invade_IDs; // 入侵人员id容器
Mat attribute_board = Mat::zeros(img.rows, img.cols, CV_8UC1); // 标识板
if(pts.size() == 2){
line(img,pts[0], pts[1], Scalar(0,0,180), 2, 8);
}
else if (pts.size() > 2) {
fillPoly(attribute_board, pts, Scalar(invaded)); // 设置入侵区域标识板颜色值为255
polylines(img, pts, true, Scalar(0,0,180), 2, 8); // 画出入侵区域
}
for (size_t i=0; i<bytrack_result.size(); i++){
STrack temp_person = bytrack_result[i];
// perosn bbox center point
int left = temp_person.tlwh[0];
int top = temp_person.tlwh[1];
int width = temp_person.tlwh[2];
int height = temp_person.tlwh[3];
Point center = Point(left + width/2, top + height/2);
// whether invade
int whether_invade = attribute_board.at<uchar>(center.y, center.x);
if (whether_invade == invaded){ // 闯入划定区域
if(!std::count(invade_IDs.begin(), invade_IDs.end(), temp_person.track_id )) { // 新入侵者
invade_IDs.push_back(temp_person.track_id ); // 加入入侵者列表
rectangle(img, cv::Rect(temp_person.tlwh[0], temp_person.tlwh[1],
temp_person.tlwh[2], temp_person.tlwh[3]), Scalar(0,0,255),2, 8); // 画bbox,红框
string label = "Person ID:" + to_string(temp_person.track_id);
putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,255), 2, 8);
}
}
else{ // 为闯入划定区域
rectangle(img, cv::Rect(temp_person.tlwh[0], temp_person.tlwh[1],
temp_person.tlwh[2], temp_person.tlwh[3]), Scalar(255,255,255),2, 8); // 画bbox,白框
string label = "Person ID:" + to_string(temp_person.track_id);
putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255,255,255), 2, 8);
}
}
}
std::vector<STrack> test_bytetrack(cv::Mat& frame, std::vector<detect_result>& results,
BYTETracker& tracker, std::vector<std::string> & classes)
{
std::vector<detect_result> objects;
for (detect_result dr : results)
{
if(dr.classId == 0 ) // person
{
objects.push_back(dr);
}
}
// bytetrack主函数
std::vector<STrack> output_stracks = tracker.update(objects);
// // 对于track结果
// for (unsigned long i = 0; i < output_stracks.size(); i++)
// {
// std::vector<float> tlwh = output_stracks[i].tlwh;
// // 目标像素不能小于20个像素、宽高比需 < 1.6
// // 其实这里加一个判定不是很合理,因为 bytetrack主函数中ID已经增加了,只不过没有显示该目标而已
bool vertical = tlwh[2] / tlwh[3] > 1.6;
if (tlwh[2] * tlwh[3] > 20 && !vertical)
// if(1)
// {
// // 框出目标,并标注ID
// cv::Scalar s = tracker.get_color(output_stracks[i].track_id);
// cv::putText(frame, cv::format("%s--%d",classes[output_stracks[i].class_index].c_str(),output_stracks[i].track_id),
// cv::Point(tlwh[0], tlwh[1] - 5),
// 0, 0.6, cv::Scalar(0, 0, 255), 2, cv::LINE_AA);
// cv::rectangle(frame, cv::Rect(tlwh[0], tlwh[1], tlwh[2], tlwh[3]), s, 2);
// }
// }
return output_stracks;
}
int main(int argc, char *argv[])
{
// 加载类别名称
std::vector<std::string> classes;
std::string file="/home/jason/work/my-deploy/01-bytetrack-deepsort/coco_80_labels_list.txt";
std::ifstream ifs(file);
if (!ifs.is_open())
CV_Error(cv::Error::StsError, "File " + file + " not found");
std::string line;
while (std::getline(ifs, line))
{
classes.push_back(line);
}
// 检测器
std::cout<<"classes:"<<classes.size();
std::shared_ptr<YOLOv5Detector> detector(new YOLOv5Detector());
detector->init(k_detect_model_path);
//bytetrack设置
int fps=20;
BYTETracker bytetracker(fps, 30); // 后面的30是30帧没有发现,
// 读取视频
std::cout<<"begin read video"<<std::endl;
cv::VideoCapture capture("/home/jason/work/my-deploy/01-bytetrack-deepsort/1.mp4");
if (!capture.isOpened()) {
printf("could not read this video file...\n");
return -1;
}
std::cout<<"end read video"<<std::endl;
// yolo检测结果
std::vector<detect_result> results;
// 输出另存
cv::VideoWriter video("/home/jason/work/my-deploy/01-bytetrack-deepsort/out.avi",cv::VideoWriter::fourcc('M','J','P','G'),10, cv::Size(1920,1080));
//
namedWindow("bytrack && invade");
setMouseCallback("bytrack && invade", on_Mouse, 0);
int num_frames = 0;
cv::Mat frame;
while (true)
{
if (!capture.read(frame)) // if not success, break loop
{
std::cout<<"\n Cannot read the video file. please check your video.\n";
break;
}
num_frames ++;
//Second/Millisecond/Microsecond 秒s/毫秒ms/微秒us
auto start = std::chrono::system_clock::now();
// 获得检测结果
detector->detect(frame, results);
// 计算检测耗时
auto end = std::chrono::system_clock::now();
auto detect_time =std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();//ms
// std::cout<<classes.size()<<":"<<results.size()<<":"<<num_frames<<std::endl;
printf("视频尺寸:%d宽 * %d高\n", frame.cols, frame.rows );
printf("帧数:%d 检测器耗时:%dms ",num_frames, (int)detect_time);
// 进行跟踪
std::vector<STrack> temp_tracks;
auto start2 = std::chrono::system_clock::now();
temp_tracks = test_bytetrack(frame, results,bytetracker,classes);
auto end2 = std::chrono::system_clock::now();
// 计算跟踪器耗时
auto detect_time2 =std::chrono::duration_cast<std::chrono::milliseconds>(end2 - start2).count();//ms
printf("跟踪器耗时:%dms \n", (int)detect_time2);
// 入侵检测
invade_detect(frame, pts, temp_tracks);
cv::imshow("bytrack && invade", frame);
// 保存结果至指定路径
video.write(frame);
if(cv::waitKey(1) == 27) // Wait for 'esc' key press to exit
{
break;
}
results.clear();
}
capture.release();
video.release();
cv::destroyAllWindows();
}