首页 > 编程札记 > 编程

项目2：使用Yolov5和deepsort实现车辆和行人目标跟踪，实时计算目标运动速度和加速度（有检测超速功能）

阅读：评论：0

项目演示视频

项目获取地址及演示视频：.html

项目简介

本项目使用Yolov5+DeepSort实现车辆、行人跟踪，并实时统计各类别目标数量，以及测量目标运动速度、加速度，对于超速的车辆进行标记保存。

项目支持对高分辨率的视频进行检测，可以使用滑动窗口检测，具体的做法就是按照指定的滑动步长以及窗口大小，对每一帧的图片进行切割，例如切割成512*512的大小的切片输入到模型中进行推理，然后对所有切片的推理结果进行合并，合并时需要再进行一次非极大值抑制，以去掉不同切片检测到的重叠框。
本项目的预训练模型使用的是YOLOv5官方提供的yolov5s预训练权重，用户可以自行更换自己的模型权重文件。
本项目可以指定需要检测的类别，并实时统计每一帧中各类别的目标数量。
本项目可以实时统计各个目标的移动速度、加速度。
对于超速的车辆，可以将其进行标记保存，便于交通部门的管理。

主函数

if __name__ == '__main__':#Adding necessary input argumentsparser = argparse.ArgumentParser(description='test')parser.add_argument("--model_path", default="./weights/yolov5s.pt", type=str,help ='预训练模型的路径')parser.add_argument('--input_path',default='./mytest.mp4', type=str,help ='输入视频文件路径')parser.add_argument('--output_dir',default = './mytest', type=str,help='输出检测结果保存路径')parser.add_argument("--is_split",default=False, action="store_true",help="是否对视频的每一帧图片进行切割检测(自动合并)")parser.add_argument("--subsize",default=512, type=int, help="切割每一帧时指定的切片大小")parser.add_argument("--gap", default=100, type=int, help="滑动窗口的重叠部分的像素长度，值越大，滑动窗口步长越小")parser.add_argument("--num_process",default=8,type=int,help="使用的进程个数")parser.add_argument("--names", default=['bus', 'car', 'truck', "person"],type=list,help="需要检测的目标")parser.add_argument("--conf_thresh", default=0.2, type=float, help="合并切片时需要再次进行NMS去除重复框")parser.add_argument("--iou_thresh", default=0.4, type=float, help="合并切片时需要再次进行NMS去除重复框")parser.add_argument("--speed_thresh", default=10, type=int, help="设定车辆速度上限阈值，如果超过该阈值就会被记录下来, 单位是千米/小时，-1则表示关闭速度检测")parser.add_argument("--pro_speed_thresh", default=-1, type=int, help="设定车辆加速度上限阈值，如果超过该阈值就会被记录下来, 单位是米/平方秒，-1则表示关闭加速度检测")args = parser.parse_args()main(args)

主要函数

def update_tracker(args, target_detector, image, fps):new_faces = []allbboxes = []cls_idlist = []if args.is_split:# 首先将当前帧存入指定的临时文件夹中args.splitDir = os.path.join(args.output_dir,"splitDir")if not ists(args.splitDir):os.makedirs(args.splitDir)tmpdir = os.path.join(args.splitDir,"tmp")tmpdir2 = os.path.join(args.splitDir,"tmp_split")if not ists(tmpdir):os.makedirs(tmpdir)if not ists(tmpdir2):os.makedirs(tmpdir2)cv2.imwrite(os.path.join(tmpdir,"tmp.png"),image)split = splitbase(tmpdir,tmpdir2,gap=args.gap,subsize=args.subsize,num_process=args.num_process)split.splitdata(1) # 1表示不放缩原图进行裁剪for filename in os.listdir(tmpdir2):filepath = os.path.join(tmpdir2,filename) # tmp__1__0___0yshfit = int(filename.split("___")[1].split(".")[0])xshfit = int(filename.split("__")[2])img = cv2.imread(filepath)_, bboxes = target_detector.detect(img) # 检测器推理图片for x1, y1, x2, y2, cls_id, conf in bboxes:cls_idlist.append(cls_id)x1 += xshfity1 += yshfitx2 += xshfity2 += yshfitallbboxes.append([x1,y1,x2,y2,conf.cpu()])else:_, bboxes = target_detector.detect(image) # 检测器推理图片for x1, y1, x2, y2, cls_id, conf in bboxes:cls_idlist.append(cls_id)allbboxes.append([x1,y1,x2,y2,conf.cpu()])allbboxes = np.array(allbboxes)keep = list(range(allbboxes.shape[0])) if not args.is_split else py_cpu_nms(allbboxes,thresh=args.iou_thresh)bboxes = allbboxes[keep]            clss = []for idx in keep:clss.append(cls_idlist[idx])bbox_xywh = []confs = []for x1, y1, x2, y2, conf in bboxes:obj = [int((x1+x2)/2), int((y1+y2)/2),x2-x1, y2-y1]bbox_xywh.append(obj)confs.append(conf)# clss.append(cls_id)xywhs = torch.Tensor(bbox_xywh)confss = torch.Tensor(confs)outputs = deepsort.update(xywhs, confss, clss, image)bboxes2draw = []face_bboxes = []current_ids = []for value in list(outputs):x1, y1, x2, y2, cls_, track_id = valuebboxes2draw.append((x1, y1, x2, y2, cls_, track_id))current_ids.append(track_id)if cls_ == 'face':if not track_id in target_detector.faceTracker:target_detector.faceTracker[track_id] = 0face = image[y1:y2, x1:x2]new_faces.append((face, track_id))face_bboxes.append((x1, y1, x2, y2))# 计算每个目标的速度和加速度大小speed_list,speed_pro_list,speed_pro_change_list = get_speed_for_obj(bboxes2draw, fps)ids2delete = []for history_id in target_detector.faceTracker:if not history_id in current_ids:target_detector.faceTracker[history_id] -= 1if target_detector.faceTracker[history_id] < -5:ids2delete.append(history_id)for ids in ids2delete:target_detector.faceTracker.pop(ids)print('-[INFO] Delete track id:', ids)image = plot_bboxes(args, image, speed_list, speed_pro_list, speed_pro_change_list, bboxes2draw)return image, new_faces, face_bboxes

完整项目的获取方式请查看：.html