import streamlit as st import cv2 import time import os import numpy as np import torch from rtmlib import Body, draw_skeleton from audio_player import AudioPlayer from pose_analyzer import PoseSimilarityAnalyzer from config import REALSENSE_AVAILABLE 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 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': (1280, 800), '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(): 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(frame.copy(), keypoints, scores, openpose_skeleton=True, kpt_thr=0.43) except Exception: pass return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) def cleanup(self): """Cleans up all resources.""" 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.""" history = self.similarity_analyzer.similarity_history if not history: return final_avg = sum(history) / len(history) level, color = ("非常棒! 👏", "success") if final_avg >= 80 else \ ("整体不错! 👍", "info") if final_avg >= 60 else \ ("需要改进! 💪", "warning") st.success("🎉 比较完成!") st.markdown(f"**整体表现**: :{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" "- 确保光线充足且稳定") def start_comparison(self, video_path): """The main loop for comparing motion.""" # Setup and initialization... (abbreviated for clarity, logic is the same as original) 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 # Get video info and setup variables total_frames = int(self.standard_cap.get(cv2.CAP_PROP_FRAME_COUNT)) video_fps = self.standard_cap.get(cv2.CAP_PROP_FPS) frame_delay = 1.0 / video_fps if video_fps > 0 else 1.0 / 30.0 target_width, target_height = self.get_display_resolution() # UI Placeholders 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: camera_type = "RealSense摄像头" if self.is_realsense_active else "USB摄像头" st.markdown(f"#### 📹 {camera_type}实时影像") # 创建控制按钮区域（紧凑布局） st.markdown("---") control_container = st.container() with control_container: control_col1, control_col2, control_col3, control_col4 = st.columns([1, 1, 1, 1]) with control_col1: stop_button = st.button("⏹️ 停止", use_container_width=True, key="stop_comparison") with control_col2: restart_button = st.button("🔄 重新开始", use_container_width=True, key="restart_comparison") with control_col3: if audio_loaded: if st.button("🔊 音频状态", use_container_width=True, key="audio_status"): st.info(f"音频: {'播放中' if self.audio_player.is_playing else '已停止'}") with control_col4: # 分辨率切换按钮 if st.button("📐 切换分辨率", use_container_width=True, key="resolution_toggle"): modes = ['high', 'medium', 'low'] current_idx = modes.index(self.display_settings.get('resolution_mode', 'high')) next_idx = (current_idx + 1) % len(modes) self.display_settings['resolution_mode'] = modes[next_idx] st.info(f"分辨率模式: {modes[next_idx]}") # Similarity UI 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() # 处理按钮点击 if stop_button: st.session_state.comparison_state['should_stop'] = True if restart_button: st.session_state.comparison_state['should_restart'] = True # 状态显示区域 status_container = st.container() with status_container: progress_bar = st.progress(0) status_text = st.empty() # 主循环 video_frame_idx = 0 start_time = time.time() current_similarity = 0 last_plot_update = 0 # Start Audio if audio_loaded: self.audio_player.play() while (st.session_state.comparison_state['is_running'] and not st.session_state.comparison_state['should_stop']): loop_start = time.time() self.frame_counter += 1 # 检查重新开始 if st.session_state.comparison_state['should_restart']: self.standard_cap.set(cv2.CAP_PROP_POS_FRAMES, 0) self.similarity_analyzer.reset() start_time = time.time() video_frame_idx = 0 if audio_loaded: self.audio_player.restart() st.session_state.comparison_state['should_restart'] = False continue # 读取标准视频的当前帧 ret_standard, standard_frame = self.standard_cap.read() if not ret_standard: # 视频结束，重新开始 self.standard_cap.set(cv2.CAP_PROP_POS_FRAMES, 0) video_frame_idx = 0 self.similarity_analyzer.reset() start_time = time.time() if audio_loaded: self.audio_player.restart() continue # 读取摄像头当前帧 ret_webcam, webcam_frame = self.read_camera_frame() if not ret_webcam or webcam_frame is None: self.error_count += 1 if self.error_count > 100: # 连续错误过多时停止 st.error("摄像头连接出现问题，停止比较") break continue self.error_count = 0 # 重置错误计数 # 翻转摄像头画面（镜像效果） 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) except Exception as e: # 关键点检测失败时继续显示原始图像 standard_keypoints, standard_scores = None, None webcam_keypoints, webcam_scores = None, None # 计算相似度（每5帧计算一次以提高性能） if (self.frame_counter % 5 == 0 and standard_keypoints is not None and webcam_keypoints is not None): try: # 提取当前标准视频帧和摄像头帧的关节角度 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: current_similarity = self.similarity_analyzer.calculate_similarity( standard_angles, webcam_angles ) # 添加到历史记录 timestamp = time.time() - start_time self.similarity_analyzer.add_similarity_score(current_similarity, timestamp) except Exception as e: pass # 忽略相似度计算错误 # 绘制关键点 try: if standard_keypoints is not None and standard_scores is not None: standard_with_keypoints = draw_skeleton( standard_frame.copy(), standard_keypoints, standard_scores, openpose_skeleton=True, kpt_thr=0.43 ) else: standard_with_keypoints = standard_frame.copy() if webcam_keypoints is not None and webcam_scores is not None: webcam_with_keypoints = draw_skeleton( webcam_frame.copy(), webcam_keypoints, webcam_scores, openpose_skeleton=True, kpt_thr=0.43 ) else: webcam_with_keypoints = webcam_frame.copy() except Exception as e: # 如果绘制失败，使用原始帧 standard_with_keypoints = standard_frame.copy() webcam_with_keypoints = webcam_frame.copy() # 转换颜色空间 (BGR to RGB) standard_rgb = cv2.cvtColor(standard_with_keypoints, cv2.COLOR_BGR2RGB) webcam_rgb = cv2.cvtColor(webcam_with_keypoints, cv2.COLOR_BGR2RGB) # 添加帧信息到图像上 current_time = time.time() - start_time frame_info = f"时间: {current_time:.1f}s | 帧: {video_frame_idx}/{total_frames}" audio_info = f" | 音频: {'🔊' if self.audio_player.is_playing else '🔇'}" if audio_loaded else "" resolution_info = f" | {target_width}x{target_height}" # 显示大尺寸画面 with video_col1: standard_placeholder.image( standard_rgb, caption=f"标准动作 - {frame_info}{audio_info}{resolution_info}", use_container_width=True ) with video_col2: webcam_placeholder.image( webcam_rgb, caption=f"您的动作 - 实时画面{resolution_info}", use_container_width=True ) # 显示相似度信息（紧凑显示） if len(self.similarity_analyzer.similarity_history) > 0: try: avg_similarity = sum(self.similarity_analyzer.similarity_history) / len(self.similarity_analyzer.similarity_history) # 使用不同颜色显示相似度 if current_similarity >= 80: similarity_color = "🟢" level = "优秀" elif current_similarity >= 60: similarity_color = "🟡" level = "良好" else: similarity_color = "🔴" level = "需要改进" with sim_col1: similarity_score_placeholder.metric( "当前相似度", f"{similarity_color} {current_similarity:.1f}%", delta=f"{level}" ) with sim_col2: avg_score_placeholder.metric( "平均相似度", f"{avg_similarity:.1f}%" ) # 更新相似度图表（每20帧更新一次以提高性能） if (len(self.similarity_analyzer.similarity_history) >= 2 and self.frame_counter - last_plot_update >= 20): try: similarity_plot = self.similarity_analyzer.get_similarity_plot() if similarity_plot: with sim_col3: similarity_plot_placeholder.plotly_chart( similarity_plot, use_container_width=True, key=f"similarity_plot_{int(time.time() * 1000)}" # 使用时间戳避免重复ID ) last_plot_update = self.frame_counter except Exception as e: pass # 忽略图表更新错误 except Exception as e: pass # 忽略显示错误 # 更新进度和状态（紧凑显示） try: progress = min(video_frame_idx / total_frames, 1.0) if total_frames > 0 else 0 progress_bar.progress(progress) elapsed_time = time.time() - start_time fps_actual = self.frame_counter / elapsed_time if elapsed_time > 0 else 0 status_text.text( f"进度: {video_frame_idx}/{total_frames} | " f"实际FPS: {fps_actual:.1f} | " f"分辨率: {target_width}x{target_height} | " f"模式: {self.display_settings['resolution_mode']}" ) except Exception as e: pass # 忽略状态更新错误 video_frame_idx += 1 # 精确的帧率控制 loop_elapsed = time.time() - loop_start sleep_time = max(0, frame_delay - loop_elapsed) if sleep_time > 0: time.sleep(sleep_time) # 强制更新UI（每30帧一次） if self.frame_counter % 30 == 0: st.empty() # 触发UI更新 self.cleanup() self.show_final_statistics()