posedet/main_app.py

100 lines
4.0 KiB
Python
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import streamlit as st
import os
import cv2
import tempfile
import torch
# Import the main app class and config flags
from motion_app import MotionComparisonApp
from config import REALSENSE_AVAILABLE, PYGAME_AVAILABLE
def main():
"""Main function to run the Streamlit app."""
st.set_page_config(page_title="Motion Comparison", page_icon="🏃", layout="wide")
st.title("🏃 Motion Comparison & Pose Analysis System")
st.markdown("---")
# Initialize the app object in session state
if 'app' not in st.session_state:
st.session_state.app = MotionComparisonApp()
app = st.session_state.app
# --- Sidebar UI ---
with st.sidebar:
st.header("🎛️ Control Panel")
# Display settings
resolution_options = {"High (1280x800)": "high", "Medium (960x720)": "medium", "Standard (640x480)": "low"}
selected_res = st.selectbox("Display Resolution", list(resolution_options.keys()), index=1)
app.display_settings['resolution_mode'] = resolution_options[selected_res]
st.markdown("---")
# Video Source Selection
video_source = st.radio("Video Source", ["Preset Video", "Upload Video"])
video_path = None
if video_source == "Preset Video":
preset_path = "preset_videos/liuzi.mp4"
if os.path.exists(preset_path):
st.success("'六字诀' video found.")
video_path = preset_path
else:
st.error("❌ Preset video not found. Please place 'liuzi.mp4' in 'preset_videos' folder.")
else:
uploaded_file = st.file_uploader("Upload a video", type=['mp4', 'avi', 'mov', 'mkv'])
if uploaded_file:
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
tmp_file.write(uploaded_file.read())
video_path = tmp_file.name
st.markdown("---")
# System Initialization
st.subheader("⚙️ System Initialization")
if st.button("🚀 Initialize System", use_container_width=True):
with st.spinner("Initializing detectors and cameras..."):
app.initialize_detector()
app.initialize_camera()
# System Status
st.subheader(" System Status")
st.info(f"Computation: {'GPU (CUDA)' if torch.cuda.is_available() else 'CPU'}")
st.info(f"Camera: {'RealSense' if REALSENSE_AVAILABLE else 'USB Webcam'}")
st.info(f"Audio: {'Enabled' if PYGAME_AVAILABLE else 'Disabled'}")
# --- Main Page UI ---
if video_path:
# Display video info and control buttons
# This part is identical to your original `main` function's logic
# It sets up the "Preview Camera" and "Start Comparison" buttons
# And calls app.start_comparison(video_path) when clicked.
# Example of how you would structure the main page:
if st.button("🚀 Start Comparison", use_container_width=True):
if not app.body_detector:
st.error("⚠️ Please initialize the system from the sidebar first!")
else:
# The start_comparison method now contains the main display loop
app.start_comparison(video_path)
else:
st.info("👈 Please select or upload a standard video from the sidebar to begin.")
with st.expander("📖 Usage Guide", expanded=True):
st.markdown("""
1. **Select Video**: Choose a preset or upload your own video in the sidebar.
2. **Initialize**: Click 'Initialize System' to prepare the camera and AI model.
3. **Start**: Click 'Start Comparison' to begin the analysis.
""")
if __name__ == "__main__":
# Set environment variables for performance
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
try:
import torch
torch.set_num_threads(1)
except ImportError:
pass
main()