app.py

import streamlit as st
from streamlit.runtime.uploaded_file_manager import UploadedFile
import pandas as pd
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 cv2
import mediapipe as mp
import torch

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


def preprocess_pose(pose):
    pose = pose.get_components(
        [
            "POSE_LANDMARKS",
            "FACE_LANDMARKS",
            "LEFT_HAND_LANDMARKS",
            "RIGHT_HAND_LANDMARKS",
        ],
        {"FACE_LANDMARKS": FACEMESH_CONTOURS_POINTS},
    )

    pose = pose.normalize(pose_normalization_info(pose.header))
    pose = pose_hide_legs(pose)

    # from sign_vq.data.normalize import pre_process_mediapipe, normalize_mean_std
    # from pose_anonymization.appearance import remove_appearance

    # pose = remove_appearance(pose)
    # pose = pre_process_mediapipe(pose)
    # pose = normalize_mean_std(pose)

    feat = np.nan_to_num(pose.body.data)
    feat = feat.reshape(feat.shape[0], -1)

    pose_frames = torch.from_numpy(np.expand_dims(feat, axis=0)).float()

    return pose_frames


# @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, 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]
    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(f"pose shape: {pose.body.data.shape}")

    header = pose.header
    st.write("### File Info")
    with st.expander(f"Show full Pose-format header from {uploaded_file.name}"):

        st.write(header)
    with st.expander(f"Show body information from {uploaded_file.name}"):
        st.write(pose.body)
    # st.write(pose.body.data.shape)
    # st.write(pose.body.fps)

    st.write(f"### Selection")

    components = pose.header.components

    component_names = [component.name for component in components]
    chosen_component_names = component_names

    component_selection = st.radio(
        "How to select components?", options=["manual", "signclip"]
    )
    if component_selection == "manual":
        st.write(f"### Component selection: ")
        chosen_component_names = st.pills(
            "Components to visualize",
            options=component_names,
            selection_mode="multi",
            default=component_names,
        )

        # st.write(chosen_component_names)

        st.write("### Point selection:")
        point_names = []
        new_chosen_components = []
        points_dict = {}
        for component in pose.header.components:
            with st.expander(f"points for {component.name}"):

                if component.name in chosen_component_names:

                    st.write(f"#### {component.name}")
                    selected_points = st.multiselect(
                        f"points for component {component.name}:",
                        options=component.points,
                        default=component.points,
                    )
                    if selected_points == component.points:
                        st.write(
                            f"All selected, no need to add a points dict entry for {component.name}"
                        )
                    else:
                        st.write(f"Adding dictionary for {component.name}")
                        points_dict[component.name] = selected_points

        # selected_points = st.multiselect("points to visualize", options=point_names, default=point_names)
        if chosen_component_names:

            if not points_dict:
                points_dict = None
            # else:
            #     st.write(points_dict)
            # st.write(chosen_component_names)

            pose = pose.get_components(chosen_component_names, points=points_dict)
            # st.write(pose.header)

    elif component_selection == "signclip":
        st.write("Selected landmarks used for SignCLIP. (Face countours only)")
        pose = pose.get_components(
            [
                "POSE_LANDMARKS",
                "FACE_LANDMARKS",
                "LEFT_HAND_LANDMARKS",
                "RIGHT_HAND_LANDMARKS",
            ],
            {"FACE_LANDMARKS": FACEMESH_CONTOURS_POINTS},
        )

        # pose = pose.normalize(pose_normalization_info(pose.header)) Visualization goes blank
        pose = pose_hide_legs(pose)
        with st.expander("Show facemesh contour points:"):
            st.write(f"{FACEMESH_CONTOURS_POINTS}")
        with st.expander(f"Show header:"):
            st.write(pose.header)
        # st.write(f"signclip selected, new header:")
        # st.write(pose.body.data.shape)
        # st.write(pose.header)
    else:
        pass

    st.write(f"### Visualization")
    width = st.select_slider(
        "select width of images",
        list(range(1, pose.header.dimensions.width + 1)),
        value=pose.header.dimensions.width / 2,
    )
    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,
    )
    visualize_clicked = st.button(f"Visualize!")

    if visualize_clicked:

        st.write(f"Generating gif...")

        # st.write(pose.body.data.shape)

        st.image(get_pose_gif(pose=pose, step=step, fps=fps))

        with st.expander("See header"):
            st.write(f"### header after filtering:")
            st.write(pose.header)

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