import cv2
import numpy as np
from ultralytics import YOLO
from filterpy.kalman import KalmanFilter
import time
from collections import deque
import pyttsx3
import threading
from PIL import Image, ImageDraw, ImageFont
# --- 更新后的配置参数 ---
CONFIG = {
'CLASS_NAMES': {2: "轿车", 5: "公交车", 7: "卡车"},
'MOVING_COLOR': (0, 255, 0), # 正常行驶车辆颜色 (绿色)
'ABNORMAL_COLOR': (0, 0, 255), # 异常状态车辆颜色 (红色)
'CONFIDENCE_THRESHOLD': 0.3, # YOLO检测置信度阈值
'COUNT_LINE_POSITION': 0.6, # 计数线位置(屏幕高度的比例)
'SPEED_CALCULATION_WINDOW_SECONDS': 1.0, # 速度计算时间窗口
'ABNORMAL_SPEED_THRESHOLD_KMH': 10.0, # 异常状态速度阈值 (20 km/h)
'ABNORMAL_DURATION_THRESHOLD': 3.0, # 异常状态持续时间阈值 (3秒)
'FOCAL_LENGTH_PX': 640, # 焦距(像素)
'REFERENCE_HEIGHTS': { # 各类车辆的实际参考高度(单位:米)
2: 1.5, # 轿车
5: 3.0, # 公交车
7: 3.5 # 卡车
},
'FONT_PATH': 'simhei.ttf', # 中文字体文件路径
# 移除了COUNTING_DIRECTION配置项
}
# 中文显示工具类(保持不变)
class ChineseDisplay:
def __init__(self, font_path='simhei.ttf', font_size=20):
try:
self.font = ImageFont.truetype(font_path, font_size)
self.large_font = ImageFont.truetype(font_path, 24)
self.small_font = ImageFont.truetype(font_path, 18)
print(f"成功加载字体: {font_path}")
except IOError:
print(f"警告: 无法加载字体 {font_path}, 将使用默认字体")
self.font = ImageFont.load_default()
self.large_font = ImageFont.load_default()
self.small_font = ImageFont.load_default()
def draw_chinese_text(self, frame, text, position, color=(255, 255, 255), font_size='normal'):
img_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
draw = ImageDraw.Draw(img_pil)
if font_size == 'large':
font = self.large_font
elif font_size == 'small':
font = self.small_font
else:
font = self.font
draw.text(position, text, font=font, fill=color)
return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
# 语音播报管理器(保持不变)
class VoiceAnnouncer:
def __init__(self):
self.engine = pyttsx3.init()
self.engine.setProperty('rate', 150)
self.engine.setProperty('volume', 0.8)
self.announced_abnormal = set()
self.lock = threading.Lock()
def announce_abnormal(self, track_id, vehicle_type, duration):
with self.lock:
if track_id not in self.announced_abnormal:
message = f"检测到车辆 {track_id} 号,{vehicle_type},状态异常"
print(f"【语音播报】: {message}")
def speak():
try:
self.engine.say(message)
self.engine.runAndWait()
except Exception as e:
print(f"语音播报错误: {e}")
thread = threading.Thread(target=speak)
thread.daemon = True
thread.start()
self.announced_abnormal.add(track_id)
def reset_announcements(self):
with self.lock:
self.announced_abnormal.clear()
# 卡尔曼滤波器初始化(保持不变)
def init_kalman(x, y, fps):
kf = KalmanFilter(dim_x=4, dim_z=2)
dt = 1.0 / fps
kf.F = np.array([[1, 0, dt, 0], [0, 1, 0, dt], [0, 0, 1, 0], [0, 0, 0, 1]])
kf.H = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])
kf.P *= 10.0
kf.R = np.array([[1, 0], [0, 1]]) * 5
kf.Q = np.eye(4) * 0.1
kf.x = np.array([x, y, 0, 0])
return kf
# 获取跟踪结果(保持不变)
def get_tracks(image, model):
boxes = []
results = model.track(image, persist=True, tracker="./bytetrack.yaml",
verbose=False, conf=CONFIG['CONFIDENCE_THRESHOLD'],
iou=0.5, classes=[2, 5, 7])
if results[0].boxes.id is None:
return boxes
result_boxes = results[0].boxes
for box in result_boxes:
if box.is_track:
bbox = box.xyxy.cpu().numpy()[0]
cls_id = int(box.cls.cpu().numpy()[0])
track_id = int(box.id.cpu().numpy()[0])
center_x = (bbox[0] + bbox[2]) / 2
center_y = (bbox[1] + bbox[3]) / 2
boxes.append({
'track_id': track_id,
'cls_id': cls_id,
'bbox': bbox,
'center': (center_x, center_y),
'bottom_center': (center_x, bbox[3]),
'height': bbox[3] - bbox[1]
})
return boxes
# 绘制速度和状态信息(保持不变)
def draw_vehicle_info(frame, chinese_display, bbox, speed_kmh, is_abnormal, abnormal_duration, cls_id, track_id):
x1, y1, x2, y2 = bbox.astype(int)
class_name = CONFIG['CLASS_NAMES'].get(cls_id, "车辆")
speed_text = f"{speed_kmh:.1f} km/h"
if is_abnormal:
info_text = f"ID:{track_id} {class_name} {speed_text}"
status_text = f"异常状态: {abnormal_duration:.1f}秒"
frame = chinese_display.draw_chinese_text(
frame,
status_text,
(x1, y1 - 40),
color=CONFIG['ABNORMAL_COLOR'],
font_size='small'
)
else:
info_text = f"ID:{track_id} {class_name} {speed_text}"
frame = chinese_display.draw_chinese_text(
frame,
info_text,
(x1, y1 - 15),
color=(255, 255, 255),
font_size='small'
)
return frame
# 计算实际距离(保持不变)
def calculate_distance(pixel_height, cls_id):
if cls_id not in CONFIG['REFERENCE_HEIGHTS']:
return None
reference_height = CONFIG['REFERENCE_HEIGHTS'][cls_id]
focal_length = CONFIG['FOCAL_LENGTH_PX']
distance = (reference_height * focal_length) / pixel_height
return distance
# 主函数
def main(source):
yolo = YOLO("yolo11n.pt")
cap = cv2.VideoCapture(source)
if not cap.isOpened():
print("无法打开视频源,请检查路径是否正确")
return
chinese_display = ChineseDisplay(CONFIG['FONT_PATH'])
announcer = VoiceAnnouncer()
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print(f"视频信息: {frame_width}x{frame_height}, {fps:.1f} FPS")
count_line_y = int(frame_height * CONFIG['COUNT_LINE_POSITION'])
track_states = {}
traffic_stats = {
'total_count': 0,
'car_count': 0,
'bus_count': 0,
'truck_count': 0,
'current_normal': 0,
'current_abnormal': 0
}
counted_vehicles = set()
speed_window_size = int(CONFIG['SPEED_CALCULATION_WINDOW_SECONDS'] * fps)
# 记录未计数的车辆(用于调试)
uncounted_vehicles = set()
while True:
ret, frame = cap.read()
if not ret:
break
current_time = time.time()
tracks = get_tracks(frame, yolo)
traffic_stats['current_normal'] = 0
traffic_stats['current_abnormal'] = 0
# 绘制计数线
cv2.line(frame, (0, count_line_y), (frame_width, count_line_y), (255, 255, 0), 2)
frame = chinese_display.draw_chinese_text(
frame,
f"计数线 ({count_line_y}px)",
(10, count_line_y - 15),
color=(255, 255, 0),
font_size='small'
)
for track in tracks:
track_id = track['track_id']
cls_id = track['cls_id']
bbox = track['bbox']
center = track['center']
bottom_center = track['bottom_center']
pixel_height = track['height']
x1, y1, x2, y2 = bbox.astype(int)
cx, cy = center
bx, by = bottom_center
# 初始化车辆状态跟踪
if track_id not in track_states:
track_states[track_id] = {
"kf": init_kalman(cx, cy, fps),
"position_history": deque(maxlen=speed_window_size),
"height_history": deque(maxlen=speed_window_size),
"abnormal_start_time": None,
"is_abnormal": False,
"counted": False,
"last_by": by, # 记录上一帧的底部y坐标
"crossed_line": False # 标记是否已经过线
}
# 显示新车辆ID
frame = chinese_display.draw_chinese_text(
frame,
f"新车辆: ID{track_id}",
(x1, y1 - 70),
color=(0, 255, 255),
font_size='small'
)
# 卡尔曼滤波更新
kf = track_states[track_id]["kf"]
kf.predict()
kf.update(np.array([cx, cy]))
filtered_center = (kf.x[0], kf.x[1])
# 记录位置和高度历史
pos_history = track_states[track_id]["position_history"]
height_history = track_states[track_id]["height_history"]
pos_history.append((current_time, filtered_center[0], filtered_center[1], by))
height_history.append((current_time, pixel_height))
# 计算速度
speed_kmh = 0.0
if len(pos_history) >= 2 and len(height_history) >= 2:
last_time, last_x, last_y, last_by = pos_history[-1]
prev_time, prev_x, prev_y, prev_by = pos_history[0]
time_diff = last_time - prev_time
if time_diff > 0.1:
pixel_dist = np.sqrt((last_x - prev_x) ** 2 + (last_y - prev_y) ** 2)
if cls_id in CONFIG['REFERENCE_HEIGHTS'] and pixel_height > 0:
last_distance = calculate_distance(height_history[-1][1], cls_id)
prev_distance = calculate_distance(height_history[0][1], cls_id)
if last_distance is not None and prev_distance is not None:
distance_diff = abs(last_distance - prev_distance)
speed_mps = distance_diff / time_diff
speed_kmh = speed_mps * 3.6
# 检测异常状态
is_abnormal = speed_kmh < CONFIG['ABNORMAL_SPEED_THRESHOLD_KMH']
abnormal_duration = 0.0
abnormal_start_time = track_states[track_id]["abnormal_start_time"]
color = (255, 255, 255)
if is_abnormal:
if abnormal_start_time is None:
abnormal_start_time = current_time
else:
abnormal_duration = current_time - abnormal_start_time
if abnormal_duration > CONFIG['ABNORMAL_DURATION_THRESHOLD']:
track_states[track_id]["is_abnormal"] = True
traffic_stats['current_abnormal'] += 1
color = CONFIG['ABNORMAL_COLOR']
class_name = CONFIG['CLASS_NAMES'].get(cls_id, "车辆")
announcer.announce_abnormal(track_id, class_name, abnormal_duration)
else:
track_states[track_id]["is_abnormal"] = False
traffic_stats['current_normal'] += 1
color = (0, 255, 255)
else:
abnormal_start_time = None
track_states[track_id]["is_abnormal"] = False
traffic_stats['current_normal'] += 1
color = CONFIG['MOVING_COLOR']
track_states[track_id]["abnormal_start_time"] = abnormal_start_time
# === 修改后的计数线逻辑:不再区分方向 ===
current_by = by
last_by = track_states[track_id]["last_by"]
# 判断车辆是否通过计数线(无论方向)
crossed = False
# 检测从上方穿过计数线(向下移动)
if last_by <= count_line_y and current_by > count_line_y:
crossed = True
# 检测从下方穿过计数线(向上移动)
if last_by >= count_line_y and current_by < count_line_y:
crossed = True
# 更新上一帧的底部y坐标
track_states[track_id]["last_by"] = current_by
# 如果检测到穿过计数线且未计数
if crossed and not track_states[track_id]["counted"]:
track_states[track_id]["counted"] = True
track_states[track_id]["crossed_line"] = True
counted_vehicles.add(track_id)
traffic_stats['total_count'] += 1
if cls_id == 2:
traffic_stats['car_count'] += 1
elif cls_id == 5:
traffic_stats['bus_count'] += 1
elif cls_id == 7:
traffic_stats['truck_count'] += 1
# 显示计数信息
frame = chinese_display.draw_chinese_text(
frame,
f"计数: ID{track_id}",
(x1, y1 - 85),
color=(0, 255, 0),
font_size='small'
)
# 显示计数线判断信息(用于调试)
count_status = "已计数" if track_states[track_id]["counted"] else "未计数"
frame = chinese_display.draw_chinese_text(
frame,
f"计数状态: {count_status}",
(x1, y1 - 100),
color=(200, 200, 255) if not track_states[track_id]["counted"] else (0, 255, 0),
font_size='small'
)
# 绘制车辆边界框和信息
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
frame = draw_vehicle_info(
frame,
chinese_display,
bbox,
speed_kmh,
track_states[track_id]["is_abnormal"],
abnormal_duration,
cls_id,
track_id
)
# 显示距离信息
if cls_id in CONFIG['REFERENCE_HEIGHTS']:
distance = calculate_distance(pixel_height, cls_id)
if distance is not None:
frame = chinese_display.draw_chinese_text(
frame,
f"距离: {distance:.1f}米",
(x1, y1 - 55),
color=(200, 200, 255),
font_size='small'
)
# 显示统计信息
stats_texts = [
f"总车辆: {traffic_stats['total_count']} (轿车: {traffic_stats['car_count']}, 公交: {traffic_stats['bus_count']}, 卡车: {traffic_stats['truck_count']})",
f"正常: {traffic_stats['current_normal']} | 异常: {traffic_stats['current_abnormal']}",
f"帧率: {fps:.1f} FPS"
]
for i, text in enumerate(stats_texts):
frame = chinese_display.draw_chinese_text(
frame,
text,
(10, 30 + i * 25),
color=(255, 255, 255),
font_size='small'
)
# 显示状态说明
frame = chinese_display.draw_chinese_text(
frame,
"绿色: 正常行驶 | 黄色: 低速警告 | 红色: 异常状态",
(10, frame_height - 20),
color=(255, 255, 255),
font_size='small'
)
# 显示计数线位置(不再显示方向)
count_line_text = f"计数线位置: y={count_line_y}px"
frame = chinese_display.draw_chinese_text(
frame,
count_line_text,
(frame_width - 200, frame_height - 20),
color=(255, 255, 0),
font_size='small'
)
cv2.imshow("高速公路车辆状态监测", frame)
# 键盘控制
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
break
elif key == ord(" "):
cv2.waitKey(0)
elif key == ord("r"):
traffic_stats = {k: 0 for k in traffic_stats}
counted_vehicles.clear()
track_states.clear()
announcer.reset_announcements()
print("所有计数器和状态已重置")
elif key == ord("c"):
# 调整计数线位置
CONFIG['COUNT_LINE_POSITION'] += 0.05
if CONFIG['COUNT_LINE_POSITION'] > 0.95:
CONFIG['COUNT_LINE_POSITION'] = 0.1
count_line_y = int(frame_height * CONFIG['COUNT_LINE_POSITION'])
print(f"计数线位置调整为: {CONFIG['COUNT_LINE_POSITION']} (y={count_line_y}px)")
# 更新实时帧率
fps = cap.get(cv2.CAP_PROP_FPS)
# 释放资源
cap.release()
cv2.destroyAllWindows()
print("\n=== 最终交通统计 ===")
print(f"总车辆数: {traffic_stats['total_count']}")
print(f"轿车: {traffic_stats['car_count']}")
print(f"公交车: {traffic_stats['bus_count']}")
print(f"卡车: {traffic_stats['truck_count']}")
print(f"检测到异常状态车辆: {traffic_stats['current_abnormal']}")
if __name__ == "__main__":
main("video/04-12_09_K221_1953.mp4")shh视频中车辆车速在100km/h左右,但识别结果只有60左右,修改
最新发布