app.py

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))