ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

【YOLOv8 追踪检测器】

【YOLOv8 追踪检测器】 YOLOV8追踪检测类detector.py# 导入必要的库importnumpyasnpfromultralyticsimportYOLO# 导入 ultralytics 库中的 YOLO 类这是 YOLOv8 的官方高级 API# 定义 YOLOv8 检测器类classYOLO_Detetor:def__init__(self):# 初始化时加载 YOLOv8s 预训练模型权重# ultralytics 会自动处理模型结构的加载和权重的下载/缓存self.modelYOLO(yolov8s.pt)defget_result(self,frame): 对输入图像进行目标检测并返回格式化的检测结果 :param frame: 输入图像 (numpy array, BGR格式) :return: 包含检测框、置信度、类别的字典列表若无检测结果则返回 None out_results[]# 用于存储最终格式化后的检测结果# 执行模型预测# conf0.4: 置信度阈值低于此值的框会被过滤# iou0.45: NMS(非极大值抑制)的 IoU 阈值用于去除重叠的冗余框# verboseFalse: 关闭预测过程中的控制台打印信息# [0]: predict 返回的是一个列表这里取第一个也是唯一一个图像的结果resultself.model.predict(frame,conf0.4,iou0.45,verboseFalse)[0]# 从结果对象中提取边界框、置信度和类别索引并转移到 CPU 转为 numpy 数组# xyxy: 返回 [x1, y1, x2, y2] 格式的左上角和右下角绝对坐标bboxesresult.boxes.xyxy.cpu().numpy()confsresult.boxes.conf.cpu().numpy()# 每个框的置信度分数cls_idxsresult.boxes.cls.cpu().numpy()# 每个框对应的类别索引 (如 0 代表 person)# 遍历每一个检测到的目标框fori,bboxinenumerate(bboxes):# 将浮点型坐标转换为整数方便后续绘图或计算x1,y1,x2,y2map(int,bbox)# 获取当前框的置信度保留两位小数confround(float(confs[i]),2)# 获取当前框的类别索引转为整数cls_int(cls_idxs[i])# 将当前目标的检测结果打包成字典追加到结果列表中out_results.append({bbox:[x1,y1,x2,y2],# 边界框坐标 [左上x, 左上y, 右下x, 右下y]conf:conf,# 置信度cls_:cls_# 类别索引})# 如果检测到了目标返回结果列表否则返回 Noneiflen(out_results)0:returnout_resultselse:returnNone# 导入 OpenCV 用于图像读取importcv2# 主程序入口if__name____main__:# 使用 OpenCV 读取待检测的图片 (traffic.png)imdatacv2.imread(traffic.png)# 实例化 YOLOv8 检测器yolo_detetorYOLO_Detetor()# 调用检测器获取结果resultyolo_detetor.get_result(imdata)# 打印检测结果到控制台print(result)结果如图iou_tracker.pyimportnumpyasnpfromdetectorimportYOLO_Detecor# 导入之前封装的 YOLOv8 检测器importcv2classIOU_Tracker:def__init__(self):# 初始化 YOLOv8 检测器self.detectorYOLO_Detecor()# 全局跟踪 ID 计数器每分配一个新目标就自增self.tracker_id1# 存储上一帧的活跃轨迹每个元素是一个字典包含 tracker_id, bbox, conf, cls_self.old_tracker[]def__iou(self,bbox1,bbox2): 计算两个边界框的交并比 (Intersection over Union) :param bbox1: [x1, y1, x2, y2] 格式 :param bbox2: [x1, y1, x2, y2] 格式 :return: IOU 值 (0~1) # 计算两个框的面积b1_area(bbox1[2]-bbox1[0])*(bbox1[3]-bbox1[1])b2_area(bbox2[2]-bbox2[0])*(bbox2[3]-bbox2[1])# 计算交集区域的左上角和右下角坐标l_xnp.maximum(bbox1[0],bbox2[0])l_ynp.maximum(bbox1[1],bbox2[1])r_xnp.minimum(bbox1[2],bbox2[2])r_ynp.minimum(bbox1[3],bbox2[3])# 计算交集的宽和高如果无交集则为 0wnp.maximum(0,r_x-l_x)hnp.maximum(0,r_y-l_y)interw*h# 返回交并比分母为并集面积returninter/(b1_areab2_area-inter)defgenerate_tracker(self,frame_data): 核心跟踪逻辑将当前帧的检测结果与历史轨迹进行匹配 # 1. 获取当前帧的检测结果 (列表若未检测到则为 None)new_detself.detector.get_result(frame_data)# 【关键修复】如果当前帧没有检测到任何目标直接清空历史轨迹并返回ifnew_detisNone:self.old_tracker[]self.__vis(frame_data)returnkeep[]# 存储当前帧匹配成功的轨迹# 2. 遍历历史轨迹尝试在当前帧的检测结果中寻找匹配foroldinself.old_tracker:max_iou-1best_detNone# 贪心匹配找到与当前历史轨迹 IOU 最大的检测框fornewinnew_det:iouself.__iou(old[bbox],new[bbox])ifioumax_iou:max_iouiou best_detnew# 如果找到了匹配且 IOU 大于阈值 (0.5)则认为是同一个目标ifbest_detandmax_iou0.5:# 保留原有的 tracker_id更新为当前帧的坐标和置信度old[bbox]best_det[bbox]old[conf]best_det[conf]keep.append(old)# 从当前检测结果列表中移除已匹配的框防止重复匹配new_det.remove(best_det)# 3. 处理剩余的未匹配检测框将它们作为新目标初始化fordetinnew_det:save_det{tracker_id:self.tracker_id,bbox:det[bbox],conf:det[conf],cls_:det[cls_],}keep.append(save_det)self.tracker_id1# 新目标分配新的 ID# 4. 更新历史轨迹为当前帧的活跃轨迹self.old_trackerkeep# 5. 在当前帧上绘制跟踪结果self.__vis(frame_data)def__vis(self,frame_): 可视化函数在图像上绘制边界框和跟踪 ID fordetinself.old_tracker:x1,y1,x2,y2det[bbox]tr_iddet[tracker_id]# 绘制绿色边界框cv2.rectangle(frame_,(x1,y1),(x2,y2),(0,255,0),1)# 在框上方绘制跟踪 IDcv2.putText(frame_,f{tr_id},(x1,y1-10),cv2.FONT_HERSHEY_COMPLEX,1,(0,255,255),1,)defrun(self,video_path): 运行视频跟踪主循环 capcv2.VideoCapture(video_path)whileTrue:ret,framecap.read()ifnotret:# 视频读取完毕breakself.generate_tracker(frame)cv2.imshow(img,frame)# 等待 10ms按 q 键可提前退出 (优化建议)ifcv2.waitKey(10)0xFFord(q):breakcap.release()cv2.destroyAllWindows()if__name____main__:video_patha.mp4iou_trackerIOU_Tracker()iou_tracker.run(video_path)效果图添加实时统计功能importnumpyasnpfromdetectorimportYOLO_Detecor# 导入自定义的 YOLO 检测器importcv2# 定义需要追踪的 COCO 数据集类别名称映射COCO_NAMES{0:person,1:bicycle,2:car,3:motorcycle,5:bus,7:truck,}# 实时统计与追踪类classIOU_Tracker:def__init__(self):# 初始化 YOLO 检测器self.detectorYOLO_Detecor()# 存储上一帧的追踪结果历史轨迹self.old_tracker[]# 追踪 ID 计数器每出现一个新目标就自增self.tracker_id1# 初始化各类别的计数统计字典self.cls_cout{person:0,bicycle:0,car:0,motorcycle:0,bus:0,truck:0,}def__iou(self,bbox1,bbox2): 计算两个边界框bbox的 IOU交并比 参数格式均为 [x1, y1, x2, y2]左上角和右下角坐标 返回值交集面积 / 并集面积 # 计算两个框的面积b1_area(bbox1[2]-bbox1[0])*(bbox1[3]-bbox1[1])b2_area(bbox2[2]-bbox2[0])*(bbox2[3]-bbox2[1])# 计算交集区域的左上角和右下角坐标l_xnp.maximum(bbox1[0],bbox2[0])l_ynp.maximum(bbox1[1],bbox2[1])r_xnp.minimum(bbox1[2],bbox2[2])r_ynp.minimum(bbox1[3],bbox2[3])# 计算交集的宽和高如果为负数则取0表示不相交wnp.maximum(0,r_x-l_x)hnp.maximum(0,r_y-l_y)interw*h# 交集面积# 返回 IOU 值returninter/(b1_areab2_area-inter)defgenerate_tracker(self,frame_data): 核心追踪逻辑将当前帧的检测结果与历史轨迹进行匹配 # 1. 获取当前帧的检测结果new_detself.detector.get_result(frame_data)# 如果当前帧没有检测到目标置为空列表new_detnew_detifnew_detisnotNoneelse[]# 用于存储当前帧匹配成功或新生成的轨迹keep[]# --- 第一步匹配历史轨迹 ---foroldinself.old_tracker:max_iou-1# 记录当前历史目标与所有新检测框的最大 IOUbest_detNone# 记录最大 IOU 对应的新检测框# 遍历当前帧的所有新检测框fornewinnew_det:iouself.__iou(old[bbox],new[bbox])# 计算 IOUifioumax_iou:max_iouiou best_detnew# 如果最大 IOU 大于阈值0.5说明是同一个目标ifbest_detandmax_iou0.5:# 保留原有的 track_id仅更新其坐标为最新检测框的坐标old[bbox]best_det[bbox]keep.append(old)# 从新检测列表中移除已匹配的目标避免重复分配new_det.remove(best_det)# --- 第二步处理剩余的新检测框即新出现的目标 ---fordetinnew_det:# 为新目标分配新的 tracker_idsave_det{tracker_id:self.tracker_id,bbox:det[bbox],conf:det[conf],cls_:det[cls_],}keep.append(save_det)# 统计新出现的目标类别数量cls_nameCOCO_NAMES.get(det[cls_],0)ifcls_name0:# 如果是 person索引为0跳过计数continueself.cls_cout[f{cls_name}]1# 对应类别计数 1self.tracker_id1# 全局追踪 ID 自增# 更新历史轨迹为当前帧的最终结果self.old_trackerkeep# 在当前帧画面上绘制追踪结果和统计信息self._vis(frame_data)def_vis(self,frame_): 可视化函数在画面上绘制追踪框、ID 以及类别统计信息 # 绘制每个追踪目标的框和 IDfordetinself.old_tracker:x1,y1,x2,y2det[bbox]tr_iddet[tracker_id]# 画红色矩形框cv2.rectangle(frame_,(x1,y1),(x2,y2),(0,0,255),1)# 在框上方显示追踪 IDcv2.putText(frame_,f{tr_id},(x1,y1-10),cv2.FONT_HERSHEY_COMPLEX,1,(0,255,255),1,)# 在画面左上角显示各类别的累计计数y22# 初始文本纵坐标forkey,valueinself.cls_cout.items():key_namekey values_value cv2.putText(frame_,f{key_name}:{values_},(10,y),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1,)y22# 每显示一行纵坐标下移defrun(self,video_path): 主运行函数读取视频流并逐帧进行追踪 capcv2.VideoCapture(video_path)whileTrue:ret,framecap.read()ifnotret:# 视频读取结束break# 对当前帧执行追踪逻辑self.generate_tracker(frame)# 显示画面等待 10mscv2.imshow(img,frame)cv2.waitKey(10)# 释放视频捕获对象并关闭所有窗口cap.release()cv2.destroyAllWindows()# 程序入口if__name____main__:video_pathb.mp4# 指定视频文件路径iou_trackerIOU_Tracker()# 创建追踪器实例iou_tracker.run(video_path)# 开始运行结果如图yolov8追踪推理并计数if__name____main__: 追踪不需要训练 直接调用的模型 importcv2fromultralyticsimportYOLO COCO_NAMES{0:person,1:bicycle,2:car,3:motorcycle,5:bus,7:truck,}modelYOLO(ryolov8s.pt)capcv2.VideoCapture(a.mp4)# 2. 初始化全局计数器 {类别ID: set([track_id_1, track_id_2, ...])}tracked_ids{}whileTrue:ret,framecap.read()ifnotret:breakresultsmodel.track(frame,persistTrue,# 跨帧保持追踪 ID 不重置verboseFalse,trackerbytetrack.yaml,# 选择追踪器配置# trackerbotsort.yaml,conf0.5,# 检测置信度阈值iou0.7,# NMS 的 IoU 阈值)# 提取当前帧的 Track ID 并更新全局计数ifresults[0].boxes.idisnotNone:# 获取当前帧所有目标的 track_id和 类别ID (转为 int 列表)current_idsresults[0].boxes.id.int().tolist()current_clsresults[0].boxes.cls.int().tolist()# 遍历当前帧检测到的每一个目标forcls_id,track_idinzip(current_cls,current_ids):# 如果该类别第一次出现先初始化一个空集合ifcls_idnotintracked_ids:tracked_ids[cls_id]set()# 将当前的 track_id 加入对应类别的集合中自动去重tracked_ids[cls_id].add(track_id)# 绘制追踪框和轨迹annotatedresults[0].plot()# 在左上角绘制分类别实时计数结果y_offset40# 初始文本纵坐标forcls_id,ids_setintracked_ids.items():# 获取类别名称如果不在 COCO_NAMES 中则显示未知cls_nameCOCO_NAMES.get(cls_id,unknown)countlen(ids_set)# 集合的长度即为该类别的唯一目标数cv2.putText(annotated,f{cls_name}:{count},(10,y_offset),# 左上角坐标 (x, y)cv2.FONT_HERSHEY_SIMPLEX,1.0,# 字体大小(0,255,0),# 绿色字体1,# 字体粗细))y_offset22# 每显示一行纵坐标下移 35 像素cv2.imshow(YOLOv8 Track,annotated)ifcv2.waitKey(10)0xFFord(q):# 按 q 键可提前退出breakcap.release()cv2.destroyAllWindows()结果如图
返回列表