|
|
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()
|
