Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 47 additions & 27 deletions src/Driving_Accident_Video_Recognition/detector.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
"""
检测器模块:精准事故判断+视频保存+帧率显示(优化版:新增事故类型区分+置信度+目标计数
检测器模块:精准事故判断+视频保存+帧率显示(优化版:新增人物数量终端输出
"""
import sys
import cv2
import time
import logging # 新增:引入日志模块(替代print,与主程序统一)
from ultralytics import YOLO
from config import (
YOLO_MODEL_PATH, CONFIDENCE_THRESHOLD, ACCIDENT_CLASSES,
Expand All @@ -15,6 +16,9 @@
process_box_coords, get_box_center, calculate_euclidean_distance, draw_annotations
)

# 新增:初始化日志(与主程序日志名一致,确保格式统一)
logger = logging.getLogger("AccidentDetection")

class AccidentDetector:
def __init__(self):
self.model = None # YOLO模型对象
Expand All @@ -27,17 +31,17 @@ def __init__(self):

def _load_model(self):
"""加载YOLO模型(增加兜底逻辑)"""
print("🔄 加载YOLOv8检测模型...")
logger.info("🔄 加载YOLOv8检测模型...") # 替换print为logger
try:
self.model = YOLO(YOLO_MODEL_PATH)
print(f"✅ 模型加载成功:{YOLO_MODEL_PATH}")
logger.info(f"✅ 模型加载成功:{YOLO_MODEL_PATH}")
except Exception as e:
print(f"⚠️ 指定模型加载失败,尝试默认轻量模型yolov8n.pt...")
logger.warning(f"⚠️ 指定模型加载失败,尝试默认轻量模型yolov8n.pt...")
try:
self.model = YOLO("yolov8n.pt")
print("✅ 兜底模型(yolov8n.pt)加载成功")
logger.info("✅ 兜底模型(yolov8n.pt)加载成功")
except Exception as e2:
print(f"❌ 模型加载失败:{e2},程序退出")
logger.error(f"❌ 模型加载失败:{e2},程序退出")
sys.exit(1)

def _init_video_writer(self, frame):
Expand All @@ -53,7 +57,7 @@ def _init_video_writer(self, frame):
# 初始化写入器
self.video_writer = cv2.VideoWriter(RESULT_VIDEO_PATH, fourcc, 30.0, (width, height))
if not self.video_writer.isOpened():
print(f"⚠️ 无法保存视频到{RESULT_VIDEO_PATH},跳过保存")
logger.warning(f"⚠️ 无法保存视频到{RESULT_VIDEO_PATH},跳过保存")
self.video_writer = None

def _calculate_accident(self, detected_objects):
Expand All @@ -76,10 +80,10 @@ def _calculate_accident(self, detected_objects):
return None

def detect_frame(self, frame, language="zh"):
"""处理单帧:新增目标计数+置信度显示+事故类型区分"""
"""处理单帧:新增目标计数+置信度显示+事故类型区分+人物数量统计"""
detected_objects = []
current_frame = frame.copy()
# 新增:目标数量统计(人、小车、卡车)
# 目标数量统计(人、小车、卡车)
target_count = {"person": 0, "car": 0, "truck": 0}

try:
Expand Down Expand Up @@ -155,58 +159,74 @@ def detect_frame(self, frame, language="zh"):
if self.video_writer:
self.video_writer.write(current_frame)
except Exception as e:
print(f"⚠️ 帧处理错误:{e},继续运行...")
return current_frame, self.accident_detected
logger.warning(f"⚠️ 帧处理错误:{e},继续运行...")

# 新增:返回人物数量(供终端输出)
return current_frame, self.accident_detected, target_count["person"]

def run_detection(self, language="zh"):
"""启动检测流程:打开摄像头/视频+逐帧处理"""
"""启动检测流程:打开摄像头/视频+逐帧处理(新增人物数量终端输出)"""
# 打开检测源(重试3次)
cap = None
for retry in range(3):
cap = cv2.VideoCapture(DETECTION_SOURCE)
if cap.isOpened():
print(f"✅ 第{retry+1}次打开检测源成功")
logger.info(f"✅ 第{retry+1}次打开检测源成功")
break
print(f"⚠️ 第{retry+1}次打开检测源失败,1秒后重试...")
logger.warning(f"⚠️ 第{retry+1}次打开检测源失败,1秒后重试...")
time.sleep(1)
# 兜底:打开默认摄像头
if not cap or not cap.isOpened():
print(f"❌ 目标检测源{DETECTION_SOURCE}无法打开,尝试默认摄像头(0)...")
logger.error(f"❌ 目标检测源{DETECTION_SOURCE}无法打开,尝试默认摄像头(0)...")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("❌ 所有检测源均无法打开,程序退出")
logger.error("❌ 所有检测源均无法打开,程序退出")
sys.exit(1)
print("✅ 检测源打开成功(按Q/ESC退出)")
print(f"💡 配置:行人车辆距离阈值{PERSON_VEHICLE_DISTANCE_THRESHOLD}像素")
logger.info("✅ 检测源打开成功(按Q/ESC退出)")
logger.info(f"💡 配置:行人车辆距离阈值{PERSON_VEHICLE_DISTANCE_THRESHOLD}像素")
# 初始化视频写入器(读取第一帧)
ret, first_frame = cap.read()
if ret:
self._init_video_writer(first_frame)

# 新增:控制终端输出频率(避免刷屏,每10帧输出一次)
frame_count = 0
output_interval = 10 # 每10帧输出一次人物数量

# 逐帧处理
while True:
ret, frame = cap.read()
if not ret:
print("🔚 视频流读取完毕,结束检测")
logger.info("🔚 视频流读取完毕,结束检测")
break
# 处理单帧
processed_frame, _ = self.detect_frame(frame, language)
# 处理单帧(接收返回的人物数量)
processed_frame, _, person_count = self.detect_frame(frame, language)
cv2.imshow("驾驶事故检测", processed_frame)

# 新增:终端输出人物数量(按间隔输出,避免刷屏)
frame_count += 1
if frame_count % output_interval == 0:
logger.info(f"📊 实时统计:当前画面中人物数量 = {person_count}")

# 退出逻辑
key = cv2.waitKey(1) & 0xFF
if key == ord("q") or key == 27:
print("🛑 用户手动退出")
logger.info("🛑 用户手动退出")
break
# 释放资源
cap.release()
if self.video_writer:
self.video_writer.release()
print(f"✅ 检测结果已保存到{RESULT_VIDEO_PATH}")
logger.info(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")
logger.info(f"\n📊 检测总结:")
logger.info(f" - 是否检测到事故 → {'✅ 是' if self.accident_detected else '❌ 否'}")
logger.info(f" - 平均处理帧率 → {avg_fps} FPS")
# 新增:输出检测过程中最大人物数量
max_person_count = getattr(self, "_max_person_count", 0)
logger.info(f" - 检测过程中最大人物数量 → {max_person_count}")

# 供外部导入
__all__ = ["AccidentDetector"]
Loading
Loading