Updates to the app such that downloading and visualization work more intuitively

app.py

@@ -7,7 +7,6 @@ from pose_format.pose_visualizer import PoseVisualizer
 from pathlib import Path
 from pyzstd import decompress
 from PIL import Image
-import cv2
 import mediapipe as mp
 
 mp_holistic = mp.solutions.holistic
@@ -18,6 +17,10 @@ FACEMESH_CONTOURS_POINTS = [
     )
 ]
 
+# Initialize session state
+# if "filtered_pose" not in st.session_state:
+#     st.session_state.filtered_pose = None
+
 
 def pose_normalization_info(pose_header):
     if pose_header.components[0].name == "POSE_LANDMARKS":
@@ -76,12 +79,12 @@ def get_pose_frames(pose: Pose, transparency: bool = False):
     return frames, images
 
 
-def get_pose_gif(pose: Pose, step: int = 1, fps: int = None):
+def get_pose_gif(pose: Pose, step: int = 1, start_frame:int=None, end_frame:int=None, fps: int = None):
     if fps is not None:
         pose.body.fps = fps
     v = PoseVisualizer(pose)
     frames = [frame_data for frame_data in v.draw()]
-    frames = frames[::step]
+    frames = frames[start_frame:end_frame:step]
     return v.save_gif(None, frames=frames)
 
 
@@ -110,28 +113,64 @@ if uploaded_file is not None:
         "How to select components?", options=["manual", "signclip"]
     )
 
+    component_names = [c.name for c in pose.header.components]
+    chosen_component_names = []
+    points_dict = {}
+    hide_legs = False
+
     if component_selection == "manual":
-        chosen_component_names = st.multiselect(
-            "Select components to visualize", options=[c.name for c in pose.header.components]
+        
+
+        chosen_component_names = st.pills(
+            "Select components to visualize", options=component_names, default=component_names,selection_mode="multi"
         )
-        if chosen_component_names:
-            pose = pose.get_components(chosen_component_names)
+        
+        for component in pose.header.components:
+            if component.name in chosen_component_names:
+                with st.expander(f"Points for {component.name}"):
+                    selected_points = st.multiselect(
+                        f"Select points for component {component.name}:",
+                        options=component.points,
+                        default=component.points,
+                    )
+                    if selected_points != component.points:  # Only add entry if not all points are selected
+                        points_dict[component.name] = selected_points
+        
+
 
     elif component_selection == "signclip":
         st.write("Selected landmarks used for SignCLIP.")
-        pose = pose.get_components(
-            ["POSE_LANDMARKS", "FACE_LANDMARKS", "LEFT_HAND_LANDMARKS", "RIGHT_HAND_LANDMARKS"]
-        )
-        pose = pose_hide_legs(pose)
+        chosen_component_names = ["POSE_LANDMARKS", "FACE_LANDMARKS", "LEFT_HAND_LANDMARKS", "RIGHT_HAND_LANDMARKS"]
+        points_dict = {"FACE_LANDMARKS": FACEMESH_CONTOURS_POINTS}
+        
+
 
     # Filter button logic
-    if st.button("Filter Components/Points"):
-        st.write("### Filtered .pose file")
+        # Filter section
+    st.write("### Filter .pose File")
+    filtered = st.button("Filter")
+    if filtered:
+        pose = pose.get_components(chosen_component_names, points=points_dict if points_dict else None)
+        if hide_legs:
+            pose = pose_hide_legs(pose)
+            
+        st.session_state.filtered_pose = pose
+
+    filtered_pose = st.session_state.get('filtered_pose', pose)
+    if filtered_pose:
+        filtered_pose = st.session_state.get('filtered_pose', pose)
+        st.write(f"#### Filtered .pose file")
+        st.write(f"Pose data shape: {filtered_pose.body.data.shape}")
         with st.expander("Show header"):
-            st.write(pose.header)
+            st.write(filtered_pose.header)
         with st.expander("Show body"):
-            st.write(pose.body)
-        
+            st.write(filtered_pose.body)
+        # with st.expander("Show data:"):
+        #     for frame in filtered_pose.body.data:
+        #         st.write(f"Frame:{frame}")
+        #         for person in frame:
+        #             st.write(person)
+
         pose_file_out = Path(uploaded_file.name).with_suffix(".pose")
         with pose_file_out.open("wb") as f:
             pose.write(f)
@@ -139,22 +178,19 @@ if uploaded_file is not None:
         with pose_file_out.open("rb") as f:
             st.download_button("Download Filtered Pose", f, file_name=pose_file_out.name)
 
+    
+    st.write("### Visualization")
+    step = st.select_slider("Step value to select every nth image", list(range(1, len(frames))), value=1)
+    fps = st.slider("FPS for visualization", min_value=1.0, max_value=filtered_pose.body.fps, value=filtered_pose.body.fps)
+    start_frame, end_frame = st.slider(
+        "Select Frame Range",
+        0,
+        len(frames),
+        (0, len(frames)),  # Default range
+    )
     # Visualization button logic
     if st.button("Visualize"):
-        st.write("### Visualization")
-        step = st.select_slider("Step value to select every nth image", list(range(1, len(frames))), value=1)
-        fps = st.slider("FPS for visualization", min_value=1.0, max_value=pose.body.fps, value=pose.body.fps)
-        st.image(get_pose_gif(pose=pose, step=step, fps=fps))
-
-
-            
-
-
-            # st.write(pose.body.data.shape)
-
-            # st.write(visualize_pose(pose=pose)) # bunch of ndarrays
-            # st.write([Image.fromarray(v.cv2.cvtColor(frame, cv_code)) for frame in frames])
-
-            # for i, image in enumerate(images[::n]):
-            #     print(f"i={i}")
-            #     st.image(image=image, width=width)
+        # Load filtered pose if it exists; otherwise, use the unfiltered pose
+        
+        
+        st.image(get_pose_gif(pose=filtered_pose, step=step, start_frame=start_frame, end_frame=end_frame, fps=fps))