import streamlit as st from streamlit.runtime.uploaded_file_manager import UploadedFile import numpy as np from pose_format import Pose from pose_format.pose_visualizer import PoseVisualizer from pathlib import Path from pyzstd import decompress from PIL import Image import mediapipe as mp mp_holistic = mp.solutions.holistic FACEMESH_CONTOURS_POINTS = [ str(p) for p in sorted( set([p for p_tup in list(mp_holistic.FACEMESH_CONTOURS) for p in p_tup]) ) ] def pose_normalization_info(pose_header): if pose_header.components[0].name == "POSE_LANDMARKS": return pose_header.normalization_info( p1=("POSE_LANDMARKS", "RIGHT_SHOULDER"), p2=("POSE_LANDMARKS", "LEFT_SHOULDER"), ) if pose_header.components[0].name == "BODY_135": return pose_header.normalization_info( p1=("BODY_135", "RShoulder"), p2=("BODY_135", "LShoulder") ) if pose_header.components[0].name == "pose_keypoints_2d": return pose_header.normalization_info( p1=("pose_keypoints_2d", "RShoulder"), p2=("pose_keypoints_2d", "LShoulder") ) def pose_hide_legs(pose): if pose.header.components[0].name == "POSE_LANDMARKS": point_names = ["KNEE", "ANKLE", "HEEL", "FOOT_INDEX"] # pylint: disable=protected-access points = [ pose.header._get_point_index("POSE_LANDMARKS", side + "_" + n) for n in point_names for side in ["LEFT", "RIGHT"] ] pose.body.confidence[:, :, points] = 0 pose.body.data[:, :, points, :] = 0 return pose else: raise ValueError("Unknown pose header schema for hiding legs") # @st.cache_data(hash_funcs={UploadedFile: lambda p: str(p.name)}) def load_pose(uploaded_file: UploadedFile) -> Pose: # with input_path.open("rb") as f_in: if uploaded_file.name.endswith(".zst"): return Pose.read(decompress(uploaded_file.read())) else: return Pose.read(uploaded_file.read()) @st.cache_data(hash_funcs={Pose: lambda p: np.array(p.body.data)}) def get_pose_frames(pose: Pose, transparency: bool = False): v = PoseVisualizer(pose) frames = [frame_data for frame_data in v.draw()] if transparency: cv_code = v.cv2.COLOR_BGR2RGBA else: cv_code = v.cv2.COLOR_BGR2RGB images = [Image.fromarray(v.cv2.cvtColor(frame, cv_code)) for frame in frames] return frames, images 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[start_frame:end_frame:step] return v.save_gif(None, frames=frames) st.write("# Pose-format explorer") st.write( "`pose-format` is a toolkit/library for 'handling, manipulation, and visualization of poses'. See [The documentation](https://pose-format.readthedocs.io/en/latest/)" ) st.write( "I made this app to help me visualize and understand the format, including different 'components' and 'points', and what they are named." ) uploaded_file = st.file_uploader("Upload a .pose file", type=[".pose", ".pose.zst"]) if uploaded_file is not None: with st.spinner(f"Loading {uploaded_file.name}"): pose = load_pose(uploaded_file) frames, images = get_pose_frames(pose=pose) st.success("Done loading!") st.write("### File Info") with st.expander(f"Show full Pose-format header from {uploaded_file.name}"): st.write(pose.header) st.write(f"### Selection") component_selection = st.radio( "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.pills( "Select components to visualize", options=component_names, default=component_names,selection_mode="multi" ) 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.") chosen_component_names = ["POSE_LANDMARKS", "FACE_LANDMARKS", "LEFT_HAND_LANDMARKS", "RIGHT_HAND_LANDMARKS"] points_dict = {"FACE_LANDMARKS": FACEMESH_CONTOURS_POINTS} # Filter button logic # Filter section st.write("### Filter .pose File") filtered = st.button("Apply 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(filtered_pose.header) with st.expander("Show 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) 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"): # 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))