Update app.py

app.py

@@ -12,10 +12,8 @@ from utils import process_video, process_image, process_gif
 from database import get_db, ProcessedFile, PoseData, AnimationData
 
 def init_page():
-    """Initialize Streamlit page configuration with embedded styling."""
     st.set_page_config(layout="wide", page_title="Pose Detection & Animation Generator")
     
-    # Embedded CSS styles
     st.markdown("""
     <style>
     /* Base styling */
@@ -55,7 +53,6 @@ def init_page():
     </style>
     """, unsafe_allow_html=True)
 
-    # Detection settings
     st.sidebar.title("Settings")
     confidence_threshold = st.sidebar.slider(
         "Detection Confidence",
@@ -65,27 +62,27 @@ def init_page():
         step=0.1
     )
     
-    # Export format selection
     export_format = st.sidebar.selectbox(
         "Export Format",
         options=['uasset', 'fbx', 'bvh'],
         key='export_format'
     )
     
-    # Manual correction mode
+    # Чекбокс для ручной коррекции
     enable_corrections = st.sidebar.checkbox("Enable Manual Corrections")
     if enable_corrections:
         st.sidebar.info("Click on landmarks in the preview to adjust their positions")
     
+    # Сохранение состояния ручной коррекции в session_state
+    st.session_state.manual_correction = enable_corrections
+
     st.title("Pose Detection & Animation Generator")
     return confidence_threshold
 
 def init_components() -> Tuple[PoseDetector, SkeletonGenerator, AnimationExporter]:
-    """Initialize the main processing components."""
     return PoseDetector(), SkeletonGenerator(), AnimationExporter()
 
 def handle_upload(file_type: str, uploaded_file, components: Tuple, db_session) -> Optional[ProcessedFile]:
-    """Process uploaded file and store results in database."""
     processed_file = ProcessedFile(
         filename=uploaded_file.name,
         file_type='video' if uploaded_file.type == 'image/gif' else file_type,
@@ -165,11 +162,11 @@ def process_image_upload(uploaded_file, components, processed_file, db, col1, co
         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.copy()
-            
+        # Если включена ручная коррекция через чекбокс
         if st.session_state.get('manual_correction', False):
+            # Инициализируем состояние текущих координат, если ещё не задано
+            if 'current_landmarks' not in st.session_state:
+                st.session_state.current_landmarks = skeleton_data.copy()
             joints = st.session_state.current_landmarks
             selected_joint = st.selectbox("Select Joint to Adjust", list(joints.keys()))
             
@@ -189,8 +186,6 @@ def process_image_upload(uploaded_file, components, processed_file, db, col1, co
                 save_corrected_pose(db, processed_file.id, joints)
                 st.success("Corrections saved successfully!")
 
-    provide_download_button(animation_data_binary)
-
 def process_video_upload(uploaded_file, components, processed_file, db, is_gif, col1, col2):
     """Handle video/GIF file upload processing."""
     pose_detector, skeleton_generator, animation_exporter = components
@@ -228,20 +223,13 @@ def save_pose_data(db, file_id: int, skeleton_data: dict):
     db.commit()
 
 def save_animation_data(db, file_id: int, skeleton_data: dict):
-    animation_data = AnimationData(
-        file_id=file_id,
-        skeleton_data=skeleton_data
-    )
+    animation_data = AnimationData(file_id=file_id, skeleton_data=skeleton_data)
     db.add(animation_data)
     db.commit()
 
 def save_video_data(db, file_id: int, animation_frames: list):
     for frame_num, frame_data in enumerate(animation_frames):
-        pose_data = PoseData(
-            file_id=file_id,
-            frame_number=frame_num,
-            landmarks=frame_data
-        )
+        pose_data = PoseData(file_id=file_id, frame_number=frame_num, landmarks=frame_data)
         db.add(pose_data)
     db.commit()
 
@@ -253,6 +241,12 @@ def provide_download_button(animation_data_binary):
         mime="application/octet-stream"
     )
 
+# Stub для сохранения скорректированных данных (при необходимости можно расширить)
+def save_corrected_pose(db, file_id: int, joints: dict):
+    pose_data = PoseData(file_id=file_id, landmarks=joints, corrected_landmarks=joints, is_corrected=True)
+    db.add(pose_data)
+    db.commit()
+
 def show_instructions():
     with st.expander("Instructions"):
         st.markdown("""