【YOLOv8 追踪检测器】

YOLOV8追踪检测类

detector.py

# 导入必要的库importnumpyasnpfromultralyticsimportYOLO# 导入 ultralytics 库中的 YOLO 类，这是 YOLOv8 的官方高级 API# 定义 YOLOv8 检测器类classYOLO_Detetor:def__init__(self):# 初始化时加载 YOLOv8s 预训练模型权重# ultralytics 会自动处理模型结构的加载和权重的下载/缓存self.model=YOLO("yolov8s.pt")defget_result(self,frame):""" 对输入图像进行目标检测，并返回格式化的检测结果 :param frame: 输入图像 (numpy array, BGR格式) :return: 包含检测框、置信度、类别的字典列表，若无检测结果则返回 None """out_results=[]# 用于存储最终格式化后的检测结果# 执行模型预测# conf=0.4: 置信度阈值，低于此值的框会被过滤# iou=0.45: NMS(非极大值抑制)的 IoU 阈值，用于去除重叠的冗余框# verbose=False: 关闭预测过程中的控制台打印信息# [0]: predict 返回的是一个列表，这里取第一个（也是唯一一个）图像的结果result=self.model.predict(frame,conf=0.4,iou=0.45,verbose=False)[0]# 从结果对象中提取边界框、置信度和类别索引，并转移到 CPU 转为 numpy 数组# xyxy: 返回 [x1, y1, x2, y2] 格式的左上角和右下角绝对坐标bboxes=result.boxes.xyxy.cpu().numpy()confs=result.boxes.conf.cpu().numpy()# 每个框的置信度分数cls_idxs=result.boxes.cls.cpu().numpy()# 每个框对应的类别索引 (如 0 代表 person)# 遍历每一个检测到的目标框fori,bboxinenumerate(bboxes):# 将浮点型坐标转换为整数，方便后续绘图或计算x1,y1,x2,y2=map(int,bbox)# 获取当前框的置信度，保留两位小数conf=round(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)imdata=cv2.imread("traffic.png")# 实例化 YOLOv8 检测器yolo_detetor=YOLO_Detetor()# 调用检测器获取结果result=yolo_detetor.get_result(imdata)# 打印检测结果到控制台print(result)

结果如图

iou_tracker.py

importnumpyasnpfromdetectorimportYOLO_Detecor# 导入之前封装的 YOLOv8 检测器importcv2classIOU_Tracker:def__init__(self):# 初始化 YOLOv8 检测器self.detector=YOLO_Detecor()# 全局跟踪 ID 计数器，每分配一个新目标就自增self.tracker_id=1# 存储上一帧的活跃轨迹，每个元素是一个字典，包含 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_x=np.maximum(bbox1[0],bbox2[0])l_y=np.maximum(bbox1[1],bbox2[1])r_x=np.minimum(bbox1[2],bbox2[2])r_y=np.minimum(bbox1[3],bbox2[3])# 计算交集的宽和高，如果无交集则为 0w=np.maximum(0,r_x-l_x)h=np.maximum(0,r_y-l_y)inter=w*h# 返回交并比，分母为并集面积returninter/(b1_area+b2_area-inter)defgenerate_tracker(self,frame_data):""" 核心跟踪逻辑：将当前帧的检测结果与历史轨迹进行匹配 """# 1. 获取当前帧的检测结果 (列表，若未检测到则为 None)new_det=self.detector.get_result(frame_data)# 【关键修复】如果当前帧没有检测到任何目标，直接清空历史轨迹并返回ifnew_detisNone:self.old_tracker=[]self.__vis(frame_data)returnkeep=[]# 存储当前帧匹配成功的轨迹# 2. 遍历历史轨迹，尝试在当前帧的检测结果中寻找匹配foroldinself.old_tracker:max_iou=-1best_det=None# 贪心匹配：找到与当前历史轨迹 IOU 最大的检测框fornewinnew_det:iou=self.__iou(old["bbox"],new["bbox"])ifiou>max_iou:max_iou=iou best_det=new# 如果找到了匹配且 IOU 大于阈值 (0.5)，则认为是同一个目标ifbest_detandmax_iou>0.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_id+=1# 新目标分配新的 ID# 4. 更新历史轨迹为当前帧的活跃轨迹self.old_tracker=keep# 5. 在当前帧上绘制跟踪结果self.__vis(frame_data)def__vis(self,frame_):""" 可视化函数：在图像上绘制边界框和跟踪 ID """fordetinself.old_tracker:x1,y1,x2,y2=det["bbox"]tr_id=det["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):""" 运行视频跟踪主循环 """cap=cv2.VideoCapture(video_path)whileTrue:ret,frame=cap.read()ifnotret:# 视频读取完毕breakself.generate_tracker(frame)cv2.imshow("img",frame)# 等待 10ms，按 'q' 键可提前退出 (优化建议)ifcv2.waitKey(10)&0xFF==ord('q'):breakcap.release()cv2.destroyAllWindows()if__name__=="__main__":video_path="a.mp4"iou_tracker=IOU_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.detector=YOLO_Detecor()# 存储上一帧的追踪结果（历史轨迹）self.old_tracker=[]# 追踪 ID 计数器，每出现一个新目标就自增self.tracker_id=1# 初始化各类别的计数统计字典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_x=np.maximum(bbox1[0],bbox2[0])l_y=np.maximum(bbox1[1],bbox2[1])r_x=np.minimum(bbox1[2],bbox2[2])r_y=np.minimum(bbox1[3],bbox2[3])# 计算交集的宽和高（如果为负数则取0，表示不相交）w=np.maximum(0,r_x-l_x)h=np.maximum(0,r_y-l_y)inter=w*h# 交集面积# 返回 IOU 值returninter/(b1_area+b2_area-inter)defgenerate_tracker(self,frame_data):""" 核心追踪逻辑：将当前帧的检测结果与历史轨迹进行匹配 """# 1. 获取当前帧的检测结果new_det=self.detector.get_result(frame_data)# 如果当前帧没有检测到目标，置为空列表new_det=new_detifnew_detisnotNoneelse[]# 用于存储当前帧匹配成功或新生成的轨迹keep=[]# --- 第一步：匹配历史轨迹 ---foroldinself.old_tracker:max_iou=-1# 记录当前历史目标与所有新检测框的最大 IOUbest_det=None# 记录最大 IOU 对应的新检测框# 遍历当前帧的所有新检测框fornewinnew_det:iou=self.__iou(old["bbox"],new["bbox"])# 计算 IOUifiou>max_iou:max_iou=iou best_det=new# 如果最大 IOU 大于阈值（0.5），说明是同一个目标ifbest_detandmax_iou>0.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_name=COCO_NAMES.get(det["cls_"],0)ifcls_name==0:# 如果是 person（索引为0），跳过计数continueself.cls_cout[f"{cls_name}"]+=1# 对应类别计数 +1self.tracker_id+=1# 全局追踪 ID 自增# 更新历史轨迹为当前帧的最终结果self.old_tracker=keep# 在当前帧画面上绘制追踪结果和统计信息self._vis(frame_data)def_vis(self,frame_):""" 可视化函数：在画面上绘制追踪框、ID 以及类别统计信息 """# 绘制每个追踪目标的框和 IDfordetinself.old_tracker:x1,y1,x2,y2=det["bbox"]tr_id=det["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,)# 在画面左上角显示各类别的累计计数y=22# 初始文本纵坐标forkey,valueinself.cls_cout.items():key_name=key values_=value cv2.putText(frame_,f"{key_name}:{values_}",(10,y),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1,)y+=22# 每显示一行，纵坐标下移defrun(self,video_path):""" 主运行函数：读取视频流并逐帧进行追踪 """cap=cv2.VideoCapture(video_path)whileTrue:ret,frame=cap.read()ifnotret:# 视频读取结束break# 对当前帧执行追踪逻辑self.generate_tracker(frame)# 显示画面，等待 10mscv2.imshow("img",frame)cv2.waitKey(10)# 释放视频捕获对象并关闭所有窗口cap.release()cv2.destroyAllWindows()# 程序入口if__name__=="__main__":video_path="b.mp4"# 指定视频文件路径iou_tracker=IOU_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",}model=YOLO(r"yolov8s.pt")cap=cv2.VideoCapture("a.mp4")# 2. 初始化全局计数器（ {类别ID: set([track_id_1, track_id_2, ...])}）tracked_ids={}whileTrue:ret,frame=cap.read()ifnotret:breakresults=model.track(frame,persist=True,# 跨帧保持追踪 ID 不重置verbose=False,tracker="bytetrack.yaml",# 选择追踪器配置# tracker="botsort.yaml",conf=0.5,# 检测置信度阈值iou=0.7,# NMS 的 IoU 阈值)# 提取当前帧的 Track ID 并更新全局计数ifresults[0].boxes.idisnotNone:# 获取当前帧所有目标的 track_id和 类别ID (转为 int 列表)current_ids=results[0].boxes.id.int().tolist()current_cls=results[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)# 绘制追踪框和轨迹annotated=results[0].plot()# 在左上角绘制分类别实时计数结果y_offset=40# 初始文本纵坐标forcls_id,ids_setintracked_ids.items():# 获取类别名称，如果不在 COCO_NAMES 中则显示未知cls_name=COCO_NAMES.get(cls_id,"unknown")count=len(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_offset+=22# 每显示一行，纵坐标下移 35 像素cv2.imshow("YOLOv8 Track",annotated)ifcv2.waitKey(10)&0xFF==ord("q"):# 按 'q' 键可提前退出breakcap.release()cv2.destroyAllWindows()

结果如图