You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

346 lines
12 KiB

2 weeks ago
import cv2
import numpy as np
import os
import sys
from ArithControl import Arith_EOController as Arith
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
from H30T_Reader import H30T_Reader
class VideoPlayer:
def __init__(self, video_path):
self.video_path = video_path
self.cap = cv2.VideoCapture(video_path)
self.h30t_reader = H30T_Reader()
self.h30t_reader.open(video_path)
if not self.cap.isOpened():
print(f"Error: Could not open video file {video_path}")
sys.exit(1)
# 获取视频属性
self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
self.fps = self.cap.get(cv2.CAP_PROP_FPS)
self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# 播放状态
self.current_frame = 0
self.is_playing = True
self.is_seeking = False
self.should_exit = False
# 鼠标绘制状态
self.is_drawing = False
self.start_point = None
self.end_point = None
self.tracking_boxes = [] # 存储所有绘制的矩形框
# 窗口设置
self.window_name = "Video Player - Press SPACE to pause/play, LEFT/RIGHT for frame control"
cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
cv2.resizeWindow(self.window_name, 800, 600)
# 创建highgui进度条
cv2.createTrackbar('Progress', self.window_name, 0, self.total_frames - 1, self.on_trackbar)
# 设置键盘回调
self.setup_keyboard_controls()
# 设置鼠标回调
cv2.setMouseCallback(self.window_name, self.mouse_callback)
self.ArithModule = Arith.EOController()
self.arithout = []
def on_trackbar(self, pos):
"""highgui进度条回调函数"""
if not self.is_seeking: # 避免循环更新
self.is_seeking = True
self.current_frame = pos
self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.current_frame)
self.is_seeking = False
def mouse_callback(self, event, x, y, flags, param):
"""鼠标回调函数 - 处理点击和拖动事件"""
# 检查鼠标是否在视频区域内(排除信息栏)
if y >= self.frame_height: # 鼠标在信息栏区域,忽略
return
# 边界检查
x = max(0, min(x, self.frame_width - 1))
y = max(0, min(y, self.frame_height - 1))
if event == cv2.EVENT_LBUTTONDOWN:
# 鼠标按下
self.is_drawing = True
self.start_point = (x, y)
self.end_point = (x, y)
print(f"鼠标按下: ({x}, {y})")
elif event == cv2.EVENT_MOUSEMOVE:
if self.is_drawing:
# 计算拖动距离
if self.start_point:
dx = x - self.start_point[0]
dy = y - self.start_point[1]
self.drag_distance = (dx * dx + dy * dy) ** 0.5
# 更新矩形终点
self.end_point = (x, y)
elif event == cv2.EVENT_LBUTTONUP:
if self.is_drawing:
self.end_point = (x, y)
# 根据框尺寸判断点还是框
w = self.end_point[0] - self.start_point[0]
h = self.end_point[1] - self.start_point[1]
if w > 5 and h > 5:
self.on_mouse_box([self.start_point[0], self.start_point[1], self.end_point[0] - self.start_point[0], self.end_point[1] - self.start_point[1]])
else:
self.on_mouse_click(self.start_point)
# 清除状态
self.is_drawing = False
self.start_point = None
self.end_point = None
def on_mouse_click(self, point):
"""鼠标点击事件处理"""
print(f"鼠标点击: ({point})")
def on_mouse_box(self, box):
"""鼠标框选事件处理"""
print(f"框选: ({box})")
# 调用自定义处理器
def setup_keyboard_controls(self):
"""设置键盘控制映射"""
# 默认键盘映射
self.key_handlers = {
27: self.handle_escape, # ESC
32: self.handle_space, # SPACE
81: self.handle_left_arrow, # LEFT ARROW
83: self.handle_right_arrow, # RIGHT ARROW
}
# 支持自定义键盘映射
self.custom_key_handlers = {}
# 支持自定义鼠标事件处理器
self.mouse_click_handler = None
self.mouse_drag_start_handler = None
self.mouse_drag_end_handler = None
def add_key_handler(self, key_code, handler_func):
"""添加自定义键盘处理器"""
self.custom_key_handlers[key_code] = handler_func
def remove_key_handler(self, key_code):
"""移除键盘处理器"""
if key_code in self.key_handlers:
del self.key_handlers[key_code]
if key_code in self.custom_key_handlers:
del self.custom_key_handlers[key_code]
def set_mouse_click_handler(self, handler_func):
"""设置鼠标点击事件处理器"""
self.mouse_click_handler = handler_func
def set_mouse_drag_start_handler(self, handler_func):
"""设置鼠标拖动开始事件处理器"""
self.mouse_drag_start_handler = handler_func
def set_mouse_drag_end_handler(self, handler_func):
"""设置鼠标拖动结束事件处理器"""
self.mouse_drag_end_handler = handler_func
def handle_keyboard_input(self, key):
"""处理键盘输入的回调函数"""
# 先检查自定义处理器
if key in self.custom_key_handlers:
self.custom_key_handlers[key]()
# 再检查默认处理器
elif key in self.key_handlers:
self.key_handlers[key]()
def handle_escape(self):
"""处理ESC键 - 退出程序"""
self.should_exit = True
print("Exiting...")
def handle_space(self):
"""处理空格键 - 播放/暂停切换"""
self.is_playing = not self.is_playing
status = "PLAYING" if self.is_playing else "PAUSED"
print(f"Play/Pause toggled: {status}")
def handle_left_arrow(self):
"""处理左方向键 - 上一帧"""
self.current_frame = max(self.current_frame - 1, 0)
print(f"Previous frame: {self.current_frame}")
def handle_right_arrow(self):
"""处理右方向键 - 下一帧"""
self.current_frame = min(self.current_frame + 1, self.total_frames - 1)
print(f"Next frame: {self.current_frame}")
def draw_info_overlay(self, frame):
"""在视频帧上绘制信息覆盖层"""
# 创建信息显示区域
info_height = 60
info_bar = np.zeros((info_height, self.frame_width, 3), dtype=np.uint8)
info_bar[:] = (0, 0, 0) # 黑色背景,半透明效果
# 计算时间信息
current_time = self.current_frame / self.fps if self.fps > 0 else 0
total_time = self.total_frames / self.fps if self.fps > 0 else 0
# 显示信息文本
time_text = f"Time: {current_time:.1f}s / {total_time:.1f}s"
frame_text = f"Frame: {self.current_frame} / {self.total_frames}"
status_text = "PLAYING" if self.is_playing else "PAUSED"
tracking_text = f"Tracking Boxes: {len(self.tracking_boxes)}"
# 添加文本到信息栏
cv2.putText(info_bar, time_text, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.putText(info_bar, frame_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.putText(info_bar, status_text, (self.frame_width - 150, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
(0, 255, 0) if self.is_playing else (0, 0, 255), 2)
cv2.putText(info_bar, tracking_text, (self.frame_width - 200, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
(255, 255, 0), 2) # 黄色显示跟踪框数量
# 将信息栏添加到帧的底部
combined_frame = np.vstack([frame, info_bar])
return combined_frame
def draw_tracking_boxes(self, frame):
"""在帧上绘制跟踪框类似selectROI的显示效果"""
# 绘制当前帧的跟踪框
for roi in self.tracking_boxes:
if roi['frame'] == self.current_frame:
x1, y1, x2, y2 = roi['bbox_corners']
# 绘制矩形框
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
# 绘制ID标签
cv2.putText(frame, f"ROI {roi['id']}", (x1, y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
# 绘制中心点
center = roi['center']
cv2.circle(frame, center, 3, (0, 255, 0), -1)
# 绘制尺寸信息
width, height = roi['bbox'][2], roi['bbox'][3]
cv2.putText(frame, f"{width}x{height}", (x1, y2+20),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
# 绘制正在绘制的矩形(实时预览)
if self.is_drawing and self.start_point and self.end_point:
x1, y1 = min(self.start_point[0], self.end_point[0]), min(self.start_point[1], self.end_point[1])
x2, y2 = max(self.start_point[0], self.end_point[0]), max(self.start_point[1], self.end_point[1])
# 绘制半透明矩形
overlay = frame.copy()
cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 0, 0), -1)
cv2.addWeighted(overlay, 0.3, frame, 0.7, 0, frame)
# 绘制边框
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
# 显示尺寸信息
width, height = x2 - x1, y2 - y1
cv2.putText(frame, f"ROI: {width} x {height}", (10, 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
return frame
def play(self):
while not self.should_exit:
if self.is_playing:
ret, frame = self.cap.read()
subtitle = self.h30t_reader.read()
if not ret:
# 视频结束,重新开始
self.current_frame = 0
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
continue
self.current_frame = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
# 获取当前帧
# self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.current_frame)
ret, frame = self.cap.read()
if not ret:
break
# 运行算法
self.arithout = self.ArithModule.run(frame)
# 绘制算法结果
frame = self.draw_arith_result(frame)
# 绘制信息覆盖层
display_frame = self.draw_info_overlay(frame)
# 显示帧
cv2.imshow(self.window_name, display_frame)
# 更新进度条位置(避免循环更新)
if not self.is_seeking:
cv2.setTrackbarPos('Progress', self.window_name, self.current_frame)
# 处理键盘输入
key = cv2.waitKey(int(1000/self.fps) if self.is_playing else 0) & 0xFF
self.handle_keyboard_input(key)
self.cleanup()
def draw_arith_result(self,frame):
for box in self.arithout:
x1, y1, x2, y2, cls, track_id = box
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f"ID: {track_id}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
return frame
def cleanup(self):
self.cap.release()
cv2.destroyAllWindows()
def main():
video_path = 'DJI_20250418150210_0006_S.MP4'
# 检查文件是否存在
if not os.path.exists(video_path):
print(f"Error: Video file {video_path} does not exist")
sys.exit(1)
# 创建并运行视频播放器
player = VideoPlayer(video_path)
player.play()
#
if __name__ == "__main__":
main()