Update app.py

app.py

@@ -12,24 +12,51 @@ from utils import process_video, process_image, process_gif
 from database import get_db, ProcessedFile, PoseData, AnimationData
 
 def init_page():
-    """Initialize Streamlit page configuration and styling."""
+    """Initialize Streamlit page configuration with embedded styling."""
     st.set_page_config(layout="wide", page_title="Pose Detection & Animation Generator")
-    with open('static/style.css') as f:
-        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
     
-    # Theme selection
-    theme = st.sidebar.selectbox("Theme", ["Light", "Dark"], key="theme")
-    if theme == "Dark":
-        st.markdown("""
-        <style>
-        .stApp { background-color: #1E1E1E; color: #FFFFFF; }
-        </style>
-        """, unsafe_allow_html=True)
+    # Embedded CSS styles
+    st.markdown("""
+    <style>
+    /* Base styling */
+    .stApp {
+        max-width: 100% !important;
+        padding: 2rem;
+        transition: background-color 0.3s;
+    }
     
-    # Detection settings
-    st.sidebar.title("Settings")
+    /* Dark theme */
+    [data-theme="dark"] .stApp {
+        background-color: #1E1E1E;
+        color: #FFFFFF;
+    }
+    
+    /* Widget styling */
+    .stButton>button {
+        background-color: #4CAF50;
+        color: white;
+        border-radius: 4px;
+        padding: 0.5rem 1rem;
+    }
+    
+    .stDownloadButton>button {
+        background-color: #008CBA;
+    }
+    
+    /* Progress bars */
+    .stProgress > div > div > div {
+        background-color: #4CAF50;
+    }
     
+    /* Columns spacing */
+    .stColumn {
+        padding: 0 1rem;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
     # Detection settings
+    st.sidebar.title("Settings")
     confidence_threshold = st.sidebar.slider(
         "Detection Confidence",
         min_value=0.0,
@@ -50,11 +77,6 @@ def init_page():
     if enable_corrections:
         st.sidebar.info("Click on landmarks in the preview to adjust their positions")
     
-    # Custom skeleton mapping
-    show_mapping = st.sidebar.expander("Skeleton Mapping")
-    with show_mapping:
-        st.text_area("Custom Mapping (JSON)", value="{}", key="custom_mapping")
-    
     st.title("Pose Detection & Animation Generator")
     return confidence_threshold
 
@@ -64,8 +86,6 @@ def init_components() -> Tuple[PoseDetector, SkeletonGenerator, AnimationExporte
 
 def handle_upload(file_type: str, uploaded_file, components: Tuple, db_session) -> Optional[ProcessedFile]:
     """Process uploaded file and store results in database."""
-    pose_detector, skeleton_generator, animation_exporter = components
-    
     processed_file = ProcessedFile(
         filename=uploaded_file.name,
         file_type='video' if uploaded_file.type == 'image/gif' else file_type,
@@ -74,50 +94,20 @@ def handle_upload(file_type: str, uploaded_file, components: Tuple, db_session)
     db_session.add(processed_file)
     db_session.commit()
     db_session.refresh(processed_file)
-    
     return processed_file
 
 def main():
     init_page()
     components = init_components()
     
-    try:
-        uploaded_file = st.file_uploader(
-            "Choose an image or video file (max 50MB)", 
-            type=['jpg', 'jpeg', 'png', 'mp4', 'avi', 'gif']
-        )
-        if uploaded_file is not None:
-            st.cache_data.clear()  # Clear cache to prevent stale data
-    except Exception as e:
-        st.error("Network error occurred. Please try uploading again.")
-        return
-
-    if uploaded_file is None:
+    uploaded_file = st.file_uploader(
+        "Choose an image or video file (max 50MB)", 
+        type=['jpg', 'jpeg', 'png', 'mp4', 'avi', 'gif']
+    )
+    
+    if not uploaded_file:
         st.warning("Please upload a file to begin.")
         return
-        
-    if uploaded_file.type == 'video/mp4':
-        try:
-            st.info("Processing video... This may take a moment.")
-            file_size = len(uploaded_file.getvalue()) / (1024 * 1024)  # Size in MB
-            if file_size > 50:
-                st.error("Video file size must be under 50MB. Please upload a smaller file.")
-                return
-                
-            # Validate video file
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
-                tfile.write(uploaded_file.getvalue())
-                cap = cv2.VideoCapture(tfile.name)
-                if not cap.isOpened():
-                    st.error("Invalid video file. Please try a different file.")
-                    return
-                cap.release()
-        except Exception as e:
-            st.error(f"Error processing video: {str(e)}")
-            return
-
-    if uploaded_file is None:
-        return
 
     db = next(get_db())
     try:
@@ -135,7 +125,7 @@ def main():
         try:
             if file_type == 'image' and not is_gif:
                 process_image_upload(uploaded_file, components, processed_file, db, col1, col2)
-            elif file_type == 'video' or is_gif:
+            else:
                 process_video_upload(uploaded_file, components, processed_file, db, is_gif, col1, col2)
         except Exception as e:
             st.error(f"Processing error: {str(e)}")
@@ -171,39 +161,33 @@ def process_image_upload(uploaded_file, components, processed_file, db, col1, co
     save_animation_data(db, processed_file.id, skeleton_data)
 
     with col2:
-        # Create a canvas for manual corrections
-        canvas_container = st.empty()
         processed_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
-        
-        # Add manual correction controls
+        canvas_container = st.empty()
+        canvas_container.image(processed_rgb, use_column_width=True)
+
         if st.button("Enable Manual Correction"):
             st.session_state.manual_correction = True
-            st.session_state.current_landmarks = skeleton_data
+            st.session_state.current_landmarks = skeleton_data.copy()
             
         if st.session_state.get('manual_correction', False):
-            # Display current joint positions
             joints = st.session_state.current_landmarks
-            
             selected_joint = st.selectbox("Select Joint to Adjust", list(joints.keys()))
             
-            col1, col2 = st.columns(2)
-            with col1:
+            col_x, col_y = st.columns(2)
+            with col_x:
                 x_pos = st.slider("X Position", 0.0, 1.0, float(joints[selected_joint]['position'][0]), 0.01)
-            with col2:
+            with col_y:
                 y_pos = st.slider("Y Position", 0.0, 1.0, float(joints[selected_joint]['position'][1]), 0.01)
             
             if st.button("Apply Changes"):
-                joints[selected_joint]['position'][0] = x_pos
-                joints[selected_joint]['position'][1] = y_pos
-                st.session_state.current_landmarks = joints
+                joints[selected_joint]['position'] = [x_pos, y_pos]
                 processed_image = pose_detector.draw_corrected_pose(image, joints)
                 processed_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+                canvas_container.image(processed_rgb, use_column_width=True)
                 
             if st.button("Save Corrections"):
-                save_corrected_pose(db, processed_file.id, st.session_state.current_landmarks)
+                save_corrected_pose(db, processed_file.id, joints)
                 st.success("Corrections saved successfully!")
-                
-        canvas_container.image(processed_rgb, use_column_width=True)
 
     provide_download_button(animation_data_binary)
 
@@ -211,15 +195,15 @@ def process_video_upload(uploaded_file, components, processed_file, db, is_gif,
     """Handle video/GIF file upload processing."""
     pose_detector, skeleton_generator, animation_exporter = components
     progress_bar = st.progress(0)
-    status_text = st.empty()
     
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.gif' if is_gif else '.mp4') as tfile:
+    with tempfile.NamedTemporaryFile() as tfile:
         tfile.write(uploaded_file.read())
-        video_path = tfile.name
-
+        tfile.seek(0)
+        
         with col1:
-            st.video(video_path)
+            st.video(tfile.read())
 
+        video_path = tfile.name
         if is_gif:
             processed_video_path, animation_frames = process_gif(video_path, pose_detector, skeleton_generator)
         else:
@@ -233,20 +217,17 @@ def process_video_upload(uploaded_file, components, processed_file, db, is_gif,
 
         with col2:
             if processed_video_path:
-                st.video(processed_video_path)
+                with open(processed_video_path, "rb") as f:
+                    st.video(f.read())
 
         provide_download_button(animation_data_binary)
-        
-        cleanup_temp_files(video_path, processed_video_path)
 
 def save_pose_data(db, file_id: int, skeleton_data: dict):
-    """Save pose data to database."""
     pose_data = PoseData(file_id=file_id, landmarks=skeleton_data)
     db.add(pose_data)
     db.commit()
 
 def save_animation_data(db, file_id: int, skeleton_data: dict):
-    """Save animation data to database."""
     animation_data = AnimationData(
         file_id=file_id,
         skeleton_data=skeleton_data
@@ -255,7 +236,6 @@ def save_animation_data(db, file_id: int, skeleton_data: dict):
     db.commit()
 
 def save_video_data(db, file_id: int, animation_frames: list):
-    """Save video frame data to database."""
     for frame_num, frame_data in enumerate(animation_frames):
         pose_data = PoseData(
             file_id=file_id,
@@ -266,7 +246,6 @@ def save_video_data(db, file_id: int, animation_frames: list):
     db.commit()
 
 def provide_download_button(animation_data_binary):
-    """Provide download button for animation data."""
     st.download_button(
         label="Download Animation Data",
         data=animation_data_binary,
@@ -274,30 +253,19 @@ def provide_download_button(animation_data_binary):
         mime="application/octet-stream"
     )
 
-def cleanup_temp_files(*file_paths):
-    """Clean up temporary files."""
-    for file_path in file_paths:
-        if file_path:
-            try:
-                import os
-                os.unlink(file_path)
-            except Exception:
-                pass
-
 def show_instructions():
-    """Show usage instructions."""
     with st.expander("Instructions"):
         st.markdown("""
-        1. Upload an image or video file using the file uploader above
-        2. Wait for the pose detection and skeleton generation to complete
-        3. Preview the results in the right column
-        4. Download the animation data for use in Unreal Engine
-
-        Supported file formats:
-        - Images: JPG, JPEG, PNG
-        - Videos: MP4, AVI, GIF
+        1. Upload an image/video using the file uploader
+        2. Wait for processing to complete
+        3. Preview results in the right column
+        4. Download animation data
+
+        Supported formats:
+        - Images: JPG, PNG
+        - Videos: MP4, GIF
         """)
 
 if __name__ == "__main__":
     main()
-    show_instructions()
+    show_instructions()