diff --git a/src/Driving_Accident_Video_Recognition/detector.py b/src/Driving_Accident_Video_Recognition/detector.py index d281ccaf87..f840b486de 100644 --- a/src/Driving_Accident_Video_Recognition/detector.py +++ b/src/Driving_Accident_Video_Recognition/detector.py @@ -1,5 +1,5 @@ """ -检测器模块：精准事故判断+视频保存+帧率显示（无报错版） +检测器模块：精准事故判断+视频保存+帧率显示（优化版：新增事故类型区分+置信度+目标计数） """ import sys import cv2 @@ -15,7 +15,6 @@ process_box_coords, get_box_center, calculate_euclidean_distance, draw_annotations ) - class AccidentDetector: def __init__(self): self.model = None # YOLO模型对象 @@ -24,7 +23,6 @@ def __init__(self): # 帧率计算（滑动平均，避免波动） self.fps_history = [] self.prev_time = time.time() - self._load_model() # 初始化时加载模型 def _load_model(self): @@ -59,35 +57,38 @@ def _init_video_writer(self, frame): self.video_writer = None def _calculate_accident(self, detected_objects): - """精准判断事故：多车/人车接触""" + """精准判断事故类型：返回None/多车事故/人车接触事故""" persons = [obj for obj in detected_objects if obj[0] == "person"] vehicles = [obj for obj in detected_objects if obj[0] in ["car", "truck"]] - - # 条件1：车辆数量≥配置阈值 + + # 条件1：多车事故（车辆数量≥配置阈值） if len(vehicles) >= MIN_VEHICLE_COUNT: - return True - # 条件2：行人和车辆距离≤阈值 + return "multi_vehicle" + # 条件2：人车接触事故（行人和车辆距离≤阈值） if PERSON_VEHICLE_CONTACT and len(persons) >= 1 and len(vehicles) >= 1: p_centers = [get_box_center(*obj[1:]) for obj in persons] v_centers = [get_box_center(*obj[1:]) for obj in vehicles] for p in p_centers: for v in v_centers: if calculate_euclidean_distance(p, v) <= PERSON_VEHICLE_DISTANCE_THRESHOLD: - return True - return False + return "person_vehicle" + # 无事故 + return None def detect_frame(self, frame, language="zh"): - """处理单帧：检测+标注+帧率计算""" + """处理单帧：新增目标计数+置信度显示+事故类型区分""" detected_objects = [] current_frame = frame.copy() - + # 新增：目标数量统计（人、小车、卡车） + target_count = {"person": 0, "car": 0, "truck": 0} + try: # 缩放帧（适配YOLO输入） frame_resized = cv2.resize(current_frame, (RESIZE_WIDTH, RESIZE_HEIGHT)) # 模型推理（关闭冗余日志） results = self.model(frame_resized, conf=CONFIDENCE_THRESHOLD, verbose=False) - - # 解析检测结果 + + # 解析检测结果（新增置信度提取） for r in results: if not hasattr(r, "boxes") or r.boxes is None: continue @@ -97,36 +98,64 @@ def detect_frame(self, frame, language="zh"): cls_idx = int(box.cls[0]) if cls_idx in ACCIDENT_CLASSES: cls_name = self.model.names[cls_idx] + # 新增：获取检测置信度（保留2位小数） + conf = round(float(box.conf[0]), 2) # 坐标缩放回原始帧 scale_x = current_frame.shape[1] / RESIZE_WIDTH scale_y = current_frame.shape[0] / RESIZE_HEIGHT x1, y1, x2, y2 = process_box_coords(box, scale_x, scale_y) - detected_objects.append((cls_name, x1, y1, x2, y2)) - - # 判断事故 - self.accident_detected = self._calculate_accident(detected_objects) - # 绘制标注 - current_frame = draw_annotations(current_frame, detected_objects, self.accident_detected, language) - - # 计算滑动平均帧率 + detected_objects.append((cls_name, conf, x1, y1, x2, y2)) # 新增conf参数 + # 统计目标数量 + target_count[cls_name] += 1 + + # 判定事故类型（替代原布尔值判断） + accident_type = self._calculate_accident(detected_objects) + self.accident_detected = accident_type is not None + + # 绘制标注（适配新增的置信度和事故类型） + font = cv2.FONT_HERSHEY_SIMPLEX + # 1. 绘制目标框+标签（含置信度） + for obj in detected_objects: + cls_name, conf, x1, y1, x2, y2 = obj + # 类别名称映射（保留原逻辑） + class_map = { + "person": "Ren（人）" if language == "zh" else "Person", + "car": "Xiao Che（小车）" if language == "zh" else "Car", + "truck": "Ka Che（卡车）" if language == "zh" else "Truck" + } + display_name = f"{class_map.get(cls_name, cls_name)}({conf})" # 新增置信度显示 + # 绘制绿色框（原逻辑不变） + cv2.rectangle(current_frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + # 绘制标签（避免超出画面） + label_y = y1 - 10 if y1 > 20 else y1 + 20 + cv2.putText(current_frame, display_name, (x1, label_y), font, 0.8, (0, 255, 0), 2) + + # 2. 绘制事故提示（按类型区分颜色） + if accident_type == "multi_vehicle": + accident_text = "Duo Che Shi Gu!（多车事故！）" if language == "zh" else "Multi-Vehicle Accident!" + cv2.putText(current_frame, accident_text, (50, 50), font, 1.2, (0, 255, 255), 3) # 黄色 + elif accident_type == "person_vehicle": + accident_text = "Ren Che Jie Chu!（人车接触！）" if language == "zh" else "Person-Vehicle Contact!" + cv2.putText(current_frame, accident_text, (50, 50), font, 1.2, (0, 0, 255), 3) # 红色 + + # 3. 绘制目标数量统计（新增） + count_text = f"Ren: {target_count['person']} | Xiao Che: {target_count['car']} | Ka Che: {target_count['truck']}" if language == "zh" else f"Person: {target_count['person']} | Car: {target_count['car']} | Truck: {target_count['truck']}" + cv2.putText(current_frame, count_text, (50, 150), font, 0.8, (255, 255, 0), 2) # 青色 + + # 4. 绘制帧率（调整位置避免重叠） current_time = time.time() self.fps_history.append(1 / (current_time - self.prev_time)) self.prev_time = current_time - # 只保留最近10帧的帧率（避免波动） if len(self.fps_history) > 10: self.fps_history.pop(0) avg_fps = int(sum(self.fps_history) / len(self.fps_history)) if self.fps_history else 0 - # 绘制帧率 - cv2.putText(current_frame, f"FPS: {avg_fps}", (50, 100), - cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2) - - # 保存视频帧 + cv2.putText(current_frame, f"FPS: {avg_fps}", (50, 100), font, 1, (255, 0, 0), 2) + + # 保存视频帧（原逻辑不变） if self.video_writer: self.video_writer.write(current_frame) - except Exception as e: print(f"⚠️ 帧处理错误：{e}，继续运行...") - return current_frame, self.accident_detected def run_detection(self, language="zh"): @@ -140,7 +169,6 @@ def run_detection(self, language="zh"): break print(f"⚠️ 第{retry+1}次打开检测源失败，1秒后重试...") time.sleep(1) - # 兜底：打开默认摄像头 if not cap or not cap.isOpened(): print(f"❌ 目标检测源{DETECTION_SOURCE}无法打开，尝试默认摄像头（0）...") @@ -148,45 +176,37 @@ def run_detection(self, language="zh"): if not cap.isOpened(): print("❌ 所有检测源均无法打开，程序退出") sys.exit(1) - print("✅ 检测源打开成功（按Q/ESC退出）") print(f"💡 配置：行人车辆距离阈值{PERSON_VEHICLE_DISTANCE_THRESHOLD}像素") - # 初始化视频写入器（读取第一帧） ret, first_frame = cap.read() if ret: self._init_video_writer(first_frame) - # 逐帧处理 while True: ret, frame = cap.read() if not ret: print("🔚 视频流读取完毕，结束检测") break - # 处理单帧 processed_frame, _ = self.detect_frame(frame, language) cv2.imshow("驾驶事故检测", processed_frame) - # 退出逻辑 key = cv2.waitKey(1) & 0xFF if key == ord("q") or key == 27: print("🛑 用户手动退出") break - # 释放资源 cap.release() if self.video_writer: self.video_writer.release() print(f"✅ 检测结果已保存到{RESULT_VIDEO_PATH}") cv2.destroyAllWindows() - # 检测总结 avg_fps = int(sum(self.fps_history) / len(self.fps_history)) if self.fps_history else 0 print(f"\n📊 检测总结：") print(f" - 是否检测到事故 → {'✅ 是' if self.accident_detected else '❌ 否'}") print(f" - 平均处理帧率 → {avg_fps} FPS") - # 供外部导入 __all__ = ["AccidentDetector"]