Update app.py

app.py

@@ -4,6 +4,8 @@ import numpy as np
 import tempfile
 from typing import Optional, Tuple
 from datetime import datetime
+from PIL import Image
+from streamlit_drawable_canvas import st_canvas
 
 from pose_detector import PoseDetector
 from skeleton_generator import SkeletonGenerator
@@ -68,14 +70,13 @@ def init_page():
         key='export_format'
     )
     
-    # Инициализируем значение manual_correction только при первом запуске
     if "manual_correction" not in st.session_state:
         st.session_state.manual_correction = st.sidebar.checkbox("Enable Manual Corrections")
     else:
         st.session_state.manual_correction = st.sidebar.checkbox("Enable Manual Corrections", value=st.session_state.manual_correction)
         
     if st.session_state.manual_correction:
-        st.sidebar.info("Click on landmarks in the preview to adjust their positions")
+        st.sidebar.info("Click on the preview canvas to select a joint")
     
     st.title("Pose Detection & Animation Generator")
     return confidence_threshold
@@ -118,7 +119,7 @@ def main():
         with col1:
             st.subheader("Original")
         with col2:
-            st.subheader("Processed")
+            st.subheader("Processed (click to select joint)")
 
         try:
             if file_type == 'image' and not is_gif:
@@ -140,52 +141,94 @@ def main():
         db.close()
 
 def process_image_upload(uploaded_file, components, processed_file, db, col1, col2):
-    """Handle image file upload processing."""
+    """Handle image file upload processing with persistent state and interactive joint selection."""
     pose_detector, skeleton_generator, animation_exporter = components
     
-    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-    image = cv2.imdecode(file_bytes, 1)
+    # Сохраняем исходное изображение в session_state
+    if "uploaded_image" not in st.session_state:
+        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
+        image = cv2.imdecode(file_bytes, 1)
+        st.session_state.uploaded_image = image
+    else:
+        image = st.session_state.uploaded_image
 
     with col1:
         st.image(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), use_column_width=True)
+    
+    # Обработка изображения для получения позы и скелета
+    if "original_skeleton_data" not in st.session_state:
+        processed_image, skeleton_data = process_image(image, pose_detector, skeleton_generator)
+        st.session_state.original_skeleton_data = skeleton_data
+        st.session_state.processed_image = processed_image
+    else:
+        skeleton_data = st.session_state.original_skeleton_data
+        processed_image = st.session_state.processed_image
 
-    processed_image, skeleton_data = process_image(image, pose_detector, skeleton_generator)
-
-    if not skeleton_data:
-        raise ValueError("No pose detected in the image")
-
+    # Сохраняем данные поз в БД
     save_pose_data(db, processed_file.id, skeleton_data)
     animation_data_binary = animation_exporter.export_pose(skeleton_data)
     save_animation_data(db, processed_file.id, skeleton_data)
 
     with col2:
+        height, width = processed_image.shape[:2]
         processed_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
-        canvas_container = st.empty()
-        canvas_container.image(processed_rgb, use_column_width=True)
-
-        if st.session_state.get('manual_correction', False):
-            # Инициализируем текущие координаты, если ещё не заданы
-            if "current_landmarks" not in st.session_state or st.session_state.current_landmarks is None:
-                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()), key="selected_joint")
-            
-            # Используем уникальные ключи для слайдеров
-            x_pos = st.slider("X Position", 0.0, 1.0, float(joints[selected_joint]['position'][0]), 0.01, key=f"x_{selected_joint}")
-            y_pos = st.slider("Y Position", 0.0, 1.0, float(joints[selected_joint]['position'][1]), 0.01, key=f"y_{selected_joint}")
-            
-            apply_clicked = st.button("Apply Changes", key="apply_changes")
-            save_clicked = st.button("Save Corrections", key="save_corrections")
-            
-            if apply_clicked:
-                joints[selected_joint]['position'] = [x_pos, y_pos]
-                st.session_state.current_landmarks = joints  # обновляем значение в session_state
-                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)
+        pil_image = Image.fromarray(processed_rgb)
+        
+        # Отображаем изображение через st_canvas, чтобы можно было кликать по нему.
+        canvas_result = st_canvas(
+            fill_color="rgba(0, 0, 0, 0)",  # прозрачный фон для рисования
+            stroke_width=5,
+            stroke_color="#FF0000",
+            background_image=pil_image,
+            update_streamlit=True,
+            height=height,
+            width=width,
+            drawing_mode="point",  # режим для регистрации кликов
+            key="canvas"
+        )
+        
+        # Если пользователь сделал клик, в canvas_result.json_data появятся объекты.
+        if st.session_state.get('manual_correction', False) and canvas_result.json_data is not None:
+            objects = canvas_result.json_data.get("objects", [])
+            if objects:
+                # Берём последний добавленный объект как клик.
+                last_obj = objects[-1]
+                click_x = last_obj.get("left")
+                click_y = last_obj.get("top")
                 
-            if save_clicked:
+                # Находим ближайший joint (учитывая, что координаты суставов нормализованы)
+                min_dist = float("inf")
+                selected_joint = None
+                for joint_name, data in skeleton_data.items():
+                    joint_px = data['position'][0] * width
+                    joint_py = data['position'][1] * height
+                    dist = ((joint_px - click_x)**2 + (joint_py - click_y)**2)**0.5
+                    if dist < min_dist:
+                        min_dist = dist
+                        selected_joint = joint_name
+                threshold = 20  # порог в пикселях для выбора сустава
+                if min_dist < threshold:
+                    st.session_state.active_joint = selected_joint
+                    st.write(f"Selected joint: **{selected_joint}**")
+                else:
+                    st.write("Click closer to a joint to select it.")
+        
+        # Если активный joint выбран, показываем слайдеры для его корректировки.
+        if st.session_state.get("active_joint"):
+            active_joint = st.session_state.active_joint
+            joints = st.session_state.get("current_landmarks", skeleton_data.copy())
+            st.write(f"Active joint for adjustment: **{active_joint}**")
+            x_pos = st.slider("Adjust X", 0.0, 1.0, float(joints[active_joint]['position'][0]), 0.01, key=f"adj_x_{active_joint}")
+            y_pos = st.slider("Adjust Y", 0.0, 1.0, float(joints[active_joint]['position'][1]), 0.01, key=f"adj_y_{active_joint}")
+            if st.button("Apply Adjustment", key=f"apply_adj_{active_joint}"):
+                joints[active_joint]['position'] = [x_pos, y_pos]
+                st.session_state.current_landmarks = joints
+                corrected_image = pose_detector.draw_corrected_pose(image, joints)
+                st.session_state.processed_image = corrected_image
+                corrected_rgb = cv2.cvtColor(corrected_image, cv2.COLOR_BGR2RGB)
+                st.image(corrected_rgb, use_column_width=True)
+                st.write(f"Updated {active_joint}: {joints[active_joint]['position']}")
+            if st.button("Save Corrections", key=f"save_adj_{active_joint}"):
                 save_corrected_pose(db, processed_file.id, joints)
                 st.success("Corrections saved successfully!")
 
@@ -252,11 +295,13 @@ def save_corrected_pose(db, file_id: int, joints: dict):
 def show_instructions():
     with st.expander("Instructions"):
         st.markdown("""
-        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
-
+        1. Upload an image/video using the file uploader.
+        2. Wait for processing to complete.
+        3. In the **Processed** panel, click on the preview canvas near a joint to select it.
+        4. Adjust the selected joint using the sliders and click **Apply Adjustment**.
+        5. Click **Save Corrections** to store the changes.
+        6. Download animation data.
+        
         Supported formats:
         - Images: JPG, PNG
         - Videos: MP4, GIF