diff --git a/src/yolo11n_vehicle_counter/scripts/keyboard_handler.py b/src/yolo11n_vehicle_counter/scripts/keyboard_handler.py new file mode 100644 index 000000000..eb0e75477 --- /dev/null +++ b/src/yolo11n_vehicle_counter/scripts/keyboard_handler.py @@ -0,0 +1,54 @@ +import cv2 as cv +import os + +def handle_keyboard_events(key, frame, frame_count, cap, out, window_name): + """处理键盘事件 + + Args: + key: 按下的键 + frame: 当前帧 + frame_count: 当前帧计数 + cap: 视频捕获对象 + out: 视频输出对象 + window_name: 窗口名称 + + Returns: + bool: 是否继续运行 + """ + if key == ord('p'): # 按'p'键暂停 + # 进入暂停状态 + paused = True + while paused: + pause_key = cv.waitKey(0) & 0xff + if pause_key == ord('p'): # 再次按'p'键继续 + paused = False + elif pause_key == ord(' '): # 按空格键逐帧播放 + # 读取下一帧 + ret, frame = cap.read() + if not ret: + paused = False + break + frame_count += 1 + # 显示当前帧 + cv.imshow(window_name, frame) + # 写入帧到输出视频 + out.write(frame) + elif pause_key == ord('q'): # 按'q'键退出程序 + # 释放资源 + cap.release() + out.release() + cv.destroyAllWindows() + return False + elif key == ord('s'): # 按's'键保存当前帧截图 + # 创建截图目录（如果不存在） + screenshot_dir = "screenshots" + if not os.path.exists(screenshot_dir): + os.makedirs(screenshot_dir) + # 保存截图 + screenshot_path = os.path.join(screenshot_dir, f"frame_{frame_count}.jpg") + cv.imwrite(screenshot_path, frame) + print(f"✅ 截图已保存: {screenshot_path}") + elif key == ord('q'): # 按'q'键退出程序 + return False + + return True diff --git a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter.py b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter.py index 4a0d5fe19..b5cf6e26c 100644 --- a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter.py +++ b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter.py @@ -5,6 +5,7 @@ import os import json import time +from keyboard_handler import handle_keyboard_events # ==================== 配置路径 ==================== # 模型文件路径 @@ -657,7 +658,9 @@ def draw_tracks_and_count(frame, detections, total_counts, region, vehicle_state # 显示当前帧 cv.imshow("Camera", frame) - if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停 + # 键盘事件处理 + key = cv.waitKey(1) & 0xff + if not handle_keyboard_events(key, frame, frame_count, cap, out, "Camera"): break # 计算整体精度 diff --git a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_carla.py b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_carla.py index d45c911f0..48f7c0e55 100644 --- a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_carla.py +++ b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_carla.py @@ -2,6 +2,7 @@ from ultralytics import YOLO import numpy as np import supervision as sv +from keyboard_handler import handle_keyboard_events # ==================== 配置路径 ==================== # 模型文件路径 @@ -190,9 +191,10 @@ def draw_tracks_and_count(frame, detections, total_counts, limits): # 显示当前帧 cv.imshow("YOLO11n Vehicle Counter - CARLA", frame) - if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停 - print("用户暂停，按任意键继续...") - cv.waitKey(0) + # 键盘事件处理 + key = cv.waitKey(1) & 0xff + if not handle_keyboard_events(key, frame, frame_count, cap, out, "YOLO11n Vehicle Counter - CARLA"): + break # 释放资源 cap.release() diff --git a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_improved.py b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_improved.py index de6bbe3da..cc1023942 100644 --- a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_improved.py +++ b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_improved.py @@ -2,6 +2,7 @@ from ultralytics import YOLO import numpy as np import supervision as sv +from keyboard_handler import handle_keyboard_events # ==================== 配置路径 ==================== # 模型文件路径 @@ -166,7 +167,9 @@ def draw_tracks_and_count(frame, detections, total_counts, limits): # 显示当前帧 cv.imshow("Camera", frame) - if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停 + # 键盘事件处理 + key = cv.waitKey(1) & 0xff + if not handle_keyboard_events(key, frame, len(total_counts), cap, out, "Camera"): break # 释放资源 diff --git a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_region.py b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_region.py index 610ae2af1..40ba977d3 100644 --- a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_region.py +++ b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_region.py @@ -14,6 +14,7 @@ from ultralytics import YOLO import numpy as np import supervision as sv +from keyboard_handler import handle_keyboard_events # ==================== 配置路径 ==================== # 模型文件路径 @@ -197,9 +198,10 @@ def annotate_frame(frame, detections): # 显示当前帧 cv.imshow("YOLO11n Region Vehicle Counter", frame) - if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停 - print("用户暂停，按任意键继续...") - cv.waitKey(0) + # 键盘事件处理 + key = cv.waitKey(1) & 0xff + if not handle_keyboard_events(key, frame, frame_count, cap, out, "YOLO11n Region Vehicle Counter"): + break # 释放资源 cap.release() diff --git a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_updown.py b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_updown.py index a635b9262..95883a4ad 100644 --- a/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_updown.py +++ b/src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_updown.py @@ -2,6 +2,7 @@ from ultralytics import YOLO import numpy as np import supervision as sv +from keyboard_handler import handle_keyboard_events # ==================== 配置路径 ==================== # 模型文件路径 @@ -223,7 +224,9 @@ def draw_tracks_and_count(frame, detections, total_counts, limits): # 显示当前帧 cv.imshow("Camera", frame) - if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停 + # 键盘事件处理 + key = cv.waitKey(1) & 0xff + if not handle_keyboard_events(key, frame, len(total_counts), cap, out, "Camera"): break # 释放资源