import streamlit as st import cv2 import time import os import numpy as np import torch import threading import queue from rtmlib import Body, draw_skeleton from audio_player import AudioPlayer from pose_analyzer import PoseSimilarityAnalyzer from config import REALSENSE_AVAILABLE def draw_skeleton_with_similarity(img, keypoints, scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=2): """ 自定义骨骼绘制函数，根据关节相似度设置颜色 相似度 > 90: 绿色 相似度 80-90: 黄色 相似度 < 80: 红色 支持多人检测，将绘制所有检测到的人体骨骼 """ if keypoints is None or len(keypoints) == 0: return img # OpenPose连接关系 skeleton = [ [1, 0], [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [0, 14], [0, 15], [14, 16], [15, 17] ] # 关节与关节角度的映射 joint_to_angle_mapping = { (2, 3): 'left_shoulder', (3, 4): 'left_elbow', (5, 6): 'right_shoulder', (6, 7): 'right_elbow', (8, 9): 'left_hip', (9, 10): 'left_knee', (11, 12): 'right_hip', (12, 13): 'right_knee' } # 骨骼的默认颜色 default_color = (0, 255, 255) # 黄色 # 确保keypoints和scores的形状正确，处理多人情况 if len(keypoints.shape) == 2: keypoints = keypoints[None, :, :] scores = scores[None, :, :] # 遍历所有人 num_instances = keypoints.shape[0] for person_idx in range(num_instances): person_kpts = keypoints[person_idx] person_scores = scores[person_idx] # 绘制骨骼 for limb_id, limb in enumerate(skeleton): joint_a, joint_b = limb if joint_a >= len(person_scores) or joint_b >= len(person_scores): continue if person_scores[joint_a] < kpt_thr or person_scores[joint_b] < kpt_thr: continue x_a, y_a = person_kpts[joint_a] x_b, y_b = person_kpts[joint_b] # 确定线条颜色 color = default_color if joint_similarities is not None: # 检查这个连接是否有对应的关节角度 if (joint_a, joint_b) in joint_to_angle_mapping: angle_name = joint_to_angle_mapping[(joint_a, joint_b)] if angle_name in joint_similarities: similarity = joint_similarities[angle_name] if similarity > 90: color = (0, 255, 0) # 绿色 elif similarity > 80: color = (0, 255, 255) # 黄色 else: color = (0, 0, 255) # 红色 cv2.line(img, (int(x_a), int(y_a)), (int(x_b), int(y_b)), color, thickness=line_width) # 绘制关键点 for kpt_id, (x, y) in enumerate(person_kpts): if person_scores[kpt_id] < kpt_thr: continue cv2.circle(img, (int(x), int(y)), 3, (255, 0, 255), -1) return img if REALSENSE_AVAILABLE: import pyrealsense2 as rs class MotionComparisonApp: """Main application class for motion comparison.""" def __init__(self): self.body_detector = None self.is_running = False self.standard_video_path = None self.webcam_cap = None self.standard_cap = None self.similarity_analyzer = PoseSimilarityAnalyzer() self.frame_counter = 0 self.audio_player = AudioPlayer() self.display_settings = {'resolution_mode': 'high', 'target_width': 960, 'target_height': 720} self.realsense_pipeline = None self.is_realsense_active = False self.last_error_time = 0 self.error_count = 0 # Async processing components self.pose_data_queue = queue.Queue(maxsize=50) self.similarity_thread = None self.similarity_stop_flag = threading.Event() if 'comparison_state' not in st.session_state: st.session_state.comparison_state = {'is_running': False, 'should_stop': False, 'should_restart': False} def get_display_resolution(self): modes = {'high': (1024, 576), 'medium': (960, 720), 'low': (640, 480)} mode = self.display_settings.get('resolution_mode', 'medium') return modes.get(mode, (960, 720)) def initialize_detector(self): if self.body_detector is None: try: device = 'cuda' if torch.cuda.is_available() else 'cpu' self.body_detector = Body(mode='lightweight', to_openpose=True, backend='onnxruntime', device=device) st.success(f"Keypoint detector initialized on device: {device}") return True except Exception as e: st.error(f"Detector initialization failed: {e}") return False return True def initialize_camera(self): if REALSENSE_AVAILABLE: try: self.realsense_pipeline = rs.pipeline() config = rs.config() width, height = self.get_display_resolution() config.enable_stream(rs.stream.color, width, height, rs.format.bgr8, 30) profile = self.realsense_pipeline.start(config) device = profile.get_device().get_info(rs.camera_info.name) st.success(f"✅ RealSense摄像头初始化成功: {device} ({width}x{height})") self.is_realsense_active = True return True except Exception as e: st.warning(f"RealSense初始化失败: {e}. 切换到USB摄像头.") return self._initialize_webcam() else: return self._initialize_webcam() def _initialize_webcam(self): try: self.webcam_cap = cv2.VideoCapture(0) if self.webcam_cap.isOpened(): self.webcam_cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')) width, height = self.get_display_resolution() self.webcam_cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) self.webcam_cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) self.webcam_cap.set(cv2.CAP_PROP_FPS, 30) actual_w = int(self.webcam_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) actual_h = int(self.webcam_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) st.success(f"✅ USB摄像头初始化成功 ({actual_w}x{actual_h})") return True else: st.error("❌ 无法打开USB摄像头") return False except Exception as e: st.error(f"❌ USB摄像头初始化失败: {e}") return False def read_camera_frame(self): if self.is_realsense_active and self.realsense_pipeline: try: frames = self.realsense_pipeline.wait_for_frames(timeout_ms=1000) color_frame = frames.get_color_frame() if not color_frame: return False, None return True, np.asanyarray(color_frame.get_data()) except Exception: return False, None elif self.webcam_cap and self.webcam_cap.isOpened(): return self.webcam_cap.read() return False, None def get_camera_preview_frame(self): ret, frame = self.read_camera_frame() if not ret or frame is None: return None frame = cv2.flip(frame, 1) if self.body_detector: try: keypoints, scores = self.body_detector(frame) frame = draw_skeleton_with_similarity(frame.copy(), keypoints, scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=1) except Exception: pass return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) def cleanup(self): """Cleans up all resources.""" # Stop similarity calculation thread if self.similarity_thread and self.similarity_thread.is_alive(): self.similarity_stop_flag.set() self.similarity_thread.join(timeout=2) if self.standard_cap: self.standard_cap.release() if self.webcam_cap: self.webcam_cap.release() if self.is_realsense_active and self.realsense_pipeline: self.realsense_pipeline.stop() self.audio_player.cleanup() self.is_running = False st.session_state.comparison_state['is_running'] = False def show_final_statistics(self): """Displays final statistics after the comparison ends.""" st.markdown("---") # 添加分隔线 history = self.similarity_analyzer.similarity_history if not history: st.warning("没有足够的比较数据来生成分析结果。") if st.button("返回主菜单", use_container_width=True): st.rerun() return final_avg = sum(history) / len(history) # 我们将评级文本和颜色分开定义，逻辑更清晰 if final_avg >= 80: level = "非常棒! 👏" md_color = "green" # 使用 'green' 颜色 elif final_avg >= 60: level = "整体不错! 👍" md_color = "blue" # 使用 'blue' 颜色 else: level = "需要改进! 💪" md_color = "orange" # 使用 'orange' 颜色替代 'warning' st.success("🎉 比较完成!") st.markdown(f"**整体表现**: :{md_color}[{level}]") col1, col2, col3 = st.columns(3) col1.metric("平均相似度", f"{final_avg:.1f}%") col2.metric("最高相似度", f"{max(history):.1f}%") col3.metric("最低相似度", f"{min(history):.1f}%") if final_avg < 60: with st.expander("💡 改善建议"): st.markdown("- 确保您的全身在摄像头画面中清晰可见\n" "- 尽量匹配标准视频的节奏和动作幅度\n" "- 确保光线充足且稳定") # 添加返回按钮 st.markdown("---") if st.button("返回主菜单", use_container_width=True): st.session_state.comparison_state['is_running'] = False # 确保状态重置 st.rerun() def _similarity_calculation_worker(self, start_time, video_fps): """Background thread for similarity calculation and plot updates.""" while not self.similarity_stop_flag.is_set(): try: # Get pose data from queue with timeout pose_data = self.pose_data_queue.get(timeout=1.0) if pose_data is None: # Poison pill to stop break elapsed_time, standard_angles, webcam_angles = pose_data if standard_angles and webcam_angles: current_similarity = self.similarity_analyzer.calculate_similarity(standard_angles, webcam_angles) self.similarity_analyzer.add_similarity_score(current_similarity, elapsed_time) self.pose_data_queue.task_done() except queue.Empty: continue except Exception as e: # Log error but continue processing continue def start_comparison(self, video_path): """The main loop for comparing motion, synchronized with real-time.""" self.is_running = True st.session_state.comparison_state.update({'is_running': True, 'should_stop': False, 'should_restart': False}) self.standard_video_path = video_path self.frame_counter = 0 self.similarity_analyzer.reset() audio_loaded = self.audio_player.load_audio(video_path) if audio_loaded: st.success("✅ 音频加载成功") else: st.info("ℹ️ 将不会播放音频") self.standard_cap = cv2.VideoCapture(video_path) if not self.standard_cap.isOpened(): st.error("无法打开标准视频文件") return if not self.is_realsense_active and (not self.webcam_cap or not self.webcam_cap.isOpened()): if not self.initialize_camera(): return # --- 获取视频信息和设置变量 --- total_frames = int(self.standard_cap.get(cv2.CAP_PROP_FRAME_COUNT)) video_fps = self.standard_cap.get(cv2.CAP_PROP_FPS) if video_fps == 0: video_fps = 30 # 防止除零错误 video_duration = total_frames / video_fps target_width, target_height = self.get_display_resolution() # --- UI 占位符 --- st.markdown("### 📺 视频比较") video_col1, video_col2 = st.columns(2, gap="small") standard_placeholder = video_col1.empty() webcam_placeholder = video_col2.empty() with video_col1: st.markdown("#### 🎯 标准动作视频") with video_col2: st.markdown(f"#### 📹 {'RealSense' if self.is_realsense_active else 'USB'}实时影像") st.markdown("---") control_col1, control_col2 = st.columns(2) stop_button = control_col1.button("⏹️ 停止", use_container_width=True, key="stop_comparison") restart_button = control_col2.button("🔄 重新开始", use_container_width=True, key="restart_comparison") st.markdown("---") st.markdown("### 📊 动作相似度分析") sim_col1, sim_col2, sim_col3 = st.columns([1, 1, 2]) similarity_score_placeholder = sim_col1.empty() avg_score_placeholder = sim_col2.empty() similarity_plot_placeholder = sim_col3.empty() progress_bar = st.progress(0) status_text = st.empty() # --- 按钮逻辑 --- if stop_button: st.session_state.comparison_state['should_stop'] = True if restart_button: st.session_state.comparison_state['should_restart'] = True # --- 主循环 (基于时间同步，专注于实时显示) --- start_time = time.time() # --- 启动异步相似度计算线程 --- self.similarity_stop_flag.clear() self.similarity_thread = threading.Thread( target=self._similarity_calculation_worker, args=(start_time, video_fps), daemon=True ) self.similarity_thread.start() current_similarity = 0 last_plot_update = 0 if audio_loaded: self.audio_player.play() while True: # --- 检查退出条件 --- elapsed_time = time.time() - start_time if elapsed_time >= video_duration or st.session_state.comparison_state['should_stop']: # 停止音频播放 if audio_loaded: self.audio_player.stop() break # --- 检查重新开始 --- if st.session_state.comparison_state['should_restart']: self.similarity_analyzer.reset() self.standard_cap.set(cv2.CAP_PROP_POS_FRAMES, 0) # 重置视频到开头 start_time = time.time() # 重置计时器 if audio_loaded: self.audio_player.restart() st.session_state.comparison_state['should_restart'] = False continue # --- 基于真实时间获取标准视频帧 --- target_frame_idx = int(elapsed_time * video_fps) self.standard_cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_idx) ret_standard, standard_frame = self.standard_cap.read() if not ret_standard: continue # 如果读取失败，跳过这一轮 # --- 读取当前摄像头帧 --- ret_webcam, webcam_frame = self.read_camera_frame() if not ret_webcam or webcam_frame is None: status_text.warning("摄像头画面读取失败，请检查连接...") time.sleep(0.1) # 短暂等待 continue self.frame_counter += 1 webcam_frame = cv2.flip(webcam_frame, 1) standard_frame = cv2.resize(standard_frame, (target_width, target_height)) webcam_frame = cv2.resize(webcam_frame, (target_width, target_height)) # --- 关键点检测 (实时处理) --- try: standard_keypoints, standard_scores = self.body_detector(standard_frame) webcam_keypoints, webcam_scores = self.body_detector(webcam_frame) # --- 异步相似度计算 (发送给后台线程) --- if standard_keypoints is not None and webcam_keypoints is not None: standard_angles = self.similarity_analyzer.extract_joint_angles(standard_keypoints, standard_scores) webcam_angles = self.similarity_analyzer.extract_joint_angles(webcam_keypoints, webcam_scores) # 非阻塞地将数据放入队列 try: self.pose_data_queue.put_nowait((elapsed_time, standard_angles, webcam_angles)) except queue.Full: # 如果队列满了，丢弃最旧的数据 try: self.pose_data_queue.get_nowait() self.pose_data_queue.put_nowait((elapsed_time, standard_angles, webcam_angles)) except queue.Empty: pass except Exception: standard_keypoints, webcam_keypoints = None, None # --- 绘制与显示 (可以保持原来的逻辑，只修改UI更新部分) --- try: # 计算关节相似度 joint_similarities = None if standard_keypoints is not None and webcam_keypoints is not None: standard_angles = self.similarity_analyzer.extract_joint_angles(standard_keypoints, standard_scores) webcam_angles = self.similarity_analyzer.extract_joint_angles(webcam_keypoints, webcam_scores) if standard_angles and webcam_angles: joint_similarities = self.similarity_analyzer.calculate_joint_similarities(standard_angles, webcam_angles) # 绘制骨骼 (线条更窄，根据相似度设置颜色) standard_with_keypoints = draw_skeleton_with_similarity(standard_frame.copy(), standard_keypoints, standard_scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=1) webcam_with_keypoints = draw_skeleton_with_similarity(webcam_frame.copy(), webcam_keypoints, webcam_scores, joint_similarities=joint_similarities, openpose_skeleton=True, kpt_thr=0.43, line_width=1) # 更新视频画面 (主线程专注于实时显示) standard_placeholder.image(cv2.cvtColor(standard_with_keypoints, cv2.COLOR_BGR2RGB), use_container_width=True) webcam_placeholder.image(cv2.cvtColor(webcam_with_keypoints, cv2.COLOR_BGR2RGB), use_container_width=True) # 更新UI (在主线程中，降低更新频率) if self.frame_counter % 30 == 0: # 每10帧更新一次UI # 更新状态信息 if self.similarity_analyzer.similarity_history: latest_similarity = self.similarity_analyzer.similarity_history[-1] avg_sim = np.mean(self.similarity_analyzer.similarity_history) similarity_score_placeholder.metric("当前相似度", f"{latest_similarity:.1f}%") avg_score_placeholder.metric("平均相似度", f"{avg_sim:.1f}%") # 更新图表 (每50帧更新一次) if self.frame_counter % 30 == 0: plot = self.similarity_analyzer.get_similarity_plot() if plot: similarity_plot_placeholder.plotly_chart(plot, use_container_width=True) # 更新进度条和状态文本 progress = min(elapsed_time / video_duration, 1.0) progress_bar.progress(progress) processing_fps = self.frame_counter / elapsed_time if elapsed_time > 0 else 0 status_text.text(f"时间: {elapsed_time:.1f}s / {video_duration:.1f}s | " f"处理帧率: {processing_fps:.1f} FPS | " f"标准视频帧: {target_frame_idx}/{total_frames}") except Exception as e: # 如果UI更新失败，避免程序崩溃 pass # 停止相似度计算线程 self.similarity_stop_flag.set() if self.similarity_thread and self.similarity_thread.is_alive(): # 发送毒丸信号停止线程 try: self.pose_data_queue.put_nowait(None) except queue.Full: pass self.similarity_thread.join(timeout=2) # 确保音频已停止 if audio_loaded: self.audio_player.stop() self.cleanup() self.show_final_statistics()