create_annotator.py

import json
import os
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Literal, Optional, Sequence, Union

import joblib
import numpy as np
import pandas as pd
from sklearn.pipeline import Pipeline
from tap import Tap
from tqdm import tqdm

import config
from components import (
    AddFrameRangeColumnsToAnnotations,
    CreateStyleWords,
    CreateTextures,
    FeatureAverager,
    MinMedianMaxScaler,
    NormalizeMinus1To1,
    ToDF,
)
from functional import expand_annotations_df
from loadbvh import Skeleton, loadbvh


@dataclass
class LMAAnnotator:
    joints_to_lma_pipeline: Pipeline

    @classmethod
    def create(
        cls,
        fps: int,
        downsample_step: int,
        window_style_word: int,
        step_style_word: int,
    ) -> "LMAAnnotator":
        create_textures_component = CreateTextures(
            step=downsample_step, joint_indices=config.JOINT_INDICES
        )

        style_words_fps = fps // downsample_step
        create_style_words_component = CreateStyleWords(
            window=window_style_word, step=step_style_word, fps=style_words_fps
        )

        feature_averager = FeatureAverager(
            groups=[
                config.LMA_BODY_INDICES,
                config.LMA_EFFORT_WEIGHT_STRONG_INDICES,
                config.LMA_EFFORT_TIME_SUDDEN_INDICES,
                config.LMA_EFFORT_FLOW_BOUND_INDICES,
                config.LMA_SHAPE_INDICES,
                config.LMA_SPACE_INDICES,
            ]
        )
        to_df = ToDF(
            [
                "BODY",
                "EFFORT_WEIGHT_STRONG",
                "EFFORT_TIME_SUDDEN",
                "EFFORT_FLOW_BOUND",
                "SHAPE",
                "SPACE",
            ]
        )

        joints_to_lma = Pipeline(
            [
                ("create_textures", create_textures_component),
                ("create_style_words", create_style_words_component),
                ("scaler", MinMedianMaxScaler()),
                ("averager", feature_averager),
                ("to_dataframe", to_df),
                ("normalize_to_0_1", NormalizeMinus1To1()),
                (
                    "add_frame_range_columns",
                    AddFrameRangeColumnsToAnnotations(
                        step=step_style_word, window=window_style_word
                    ),
                ),
            ]
        )
        return cls(joints_to_lma)

    def annotate(
        self, joints: np.ndarray, return_expanded: bool = False
    ) -> pd.DataFrame:
        """joints should be"""
        if not return_expanded:
            return self.joints_to_lma(joints)
        else:
            textures = self.joints_to_textures(joints)
            annotations = self.texture_to_lma(textures)
            downsampled_frames = textures.shape[-1]
            expanded_annotations = expand_annotations_df(
                annotations, pad_to_length=downsampled_frames
            )
            return expanded_annotations

    def annotate_bvh(
        self, bvh_file: Union[os.PathLike, str], return_expanded: bool = False
    ) -> pd.DataFrame:
        skeleton, _, _, fps = loadbvh(bvh_file)
        return self.annotate(skeleton.d_xyz, return_expanded)

    def fit_to_joints(self, x: np.ndarray) -> pd.DataFrame:
        return self.joints_to_lma_pipeline.fit_transform(x)

    def fit_to_textures(self, textures: np.ndarray) -> pd.DataFrame:
        return self.textures_to_lma_pipeline.fit_transform(textures)

    def fit_to_style_words(self, x: np.ndarray) -> pd.DataFrame:
        return self.style_words_to_lma_pipeline.fit_transform(x)

    @property
    def scaler(self) -> MinMedianMaxScaler:
        return self.joints_to_lma_pipeline.steps[2][1]

    def joints_to_lma(self, x: np.ndarray) -> pd.DataFrame:
        return self.joints_to_lma_pipeline.transform(x)

    @property
    def joints_to_textures_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[:1]

    def joints_to_textures(self, x: np.ndarray) -> np.ndarray:
        return self.joints_to_textures_pipeline.transform(x)

    @property
    def textures_to_style_words_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[1:2]

    def textures_to_style_words(self, x: np.ndarray) -> np.ndarray:
        return self.textures_to_style_words_pipeline.transform(x)

    @property
    def joints_to_style_words_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[:2]

    def joints_to_style_words(self, x: np.ndarray) -> np.ndarray:
        return self.joints_to_style_words_pipeline.transform(x)

    @property
    def style_words_to_lma_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[2:]

    def style_words_to_lma(self, x: np.ndarray) -> pd.DataFrame:
        return self.style_words_to_lma_pipeline.transform(x)

    @property
    def textures_to_lma_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[1:]

    def texture_to_lma(self, x: np.ndarray) -> pd.DataFrame:
        return self.textures_to_lma_pipeline.transform(x)

    def save(self, save_path: Union[os.PathLike, str]) -> None:
        joblib.dump(self.joints_to_lma_pipeline, save_path)

    @classmethod
    def load(cls, save_path: Union[os.PathLike, str]) -> "LMAAnnotator":
        save_path = Path(save_path)
        joints_to_lma = joblib.load(save_path)
        return cls(joints_to_lma)

    @property
    def style_words_to_scaled_style_words_pipeline(self) -> Pipeline:
        return self.joints_to_lma_pipeline[2:3]

    def scale_style_words(self, style_words: np.ndarray) -> np.ndarray:
        return self.style_words_to_scaled_style_words_pipeline.transform(style_words)

    def _repr_mimebundle_(self, **kwargs) -> dict:
        return self.joints_to_lma_pipeline._repr_mimebundle_(**kwargs)

    def __str__(self):
        return str(self.joints_to_lma_pipeline)


# fmt: off
class ArgumentParser(Tap):
    bvh_dir: Path = Path("BVHs")  # the directory with the bvh files to fit the model to.
    window_style_word: int = 16  # window size for style words.
    step_style_word: int = 4  # step size for style words.
    method: Literal["StyleWords", "MotionWords"] = "StyleWords"  # Select which method to use for creating words.
    create_annotations: bool = False  # If set, extracts annotations for each motion
    downsample_step: int = 24  # pick a frame every `downsample_step` frames from the original sequence.
    fps: int = 120  # expected fps for original sequences.
# fmt: on


def main(
    args: Optional[Sequence[str]] = None,
    known_only: bool = False,
):
    parser = ArgumentParser(underscores_to_dashes=True)
    args = parser.parse_args(args, known_only)

    annotator = LMAAnnotator.create(
        args.fps, args.downsample_step, args.window_style_word, args.step_style_word
    )
    print("Annotator Pipeline:", annotator, sep="\n")

    bvh_dir = args.bvh_dir

    timestamp = timestamp_str()
    save_dir = Path("save") / f"{timestamp}-{bvh_dir.stem}"
    save_dir.mkdir(exist_ok=True, parents=True)
    print(f"Results will be saved in: {save_dir}")

    args_json = save_dir / "args.json"
    with open(args_json, "w") as fp:
        args_dict = make_serializable(args.as_dict())
        json.dump(args_dict, fp, indent=2)

    textures_save_dir = save_dir / "textures"
    textures_save_dir.mkdir(exist_ok=True, parents=True)

    style_words_save_dir = save_dir / "style-words"
    style_words_save_dir.mkdir(exist_ok=True, parents=True)

    joints_save_dir = save_dir / "joints"
    joints_save_dir.mkdir(exist_ok=True, parents=True)

    skeletons: List[Skeleton] = []
    all_joints: List[np.ndarray] = []
    all_style_words: List[np.ndarray] = []

    results = {}
    print(f"Loading Motion data from '{bvh_dir}'.")
    bvh_files = sorted(bvh_dir.glob("*.bvh"))
    for i, bvh_file in enumerate(tqdm(bvh_files)):
        skeleton, frame_times, duration, fps = loadbvh(bvh_file)
        skeletons.append(skeleton)
        assert fps == args.fps
        print(f"Loaded BVH file: {bvh_file}")
        print(f"fps: {fps}")
        print(f"Duration: {duration} seconds")
        print(f"FPS: {frame_times.size // duration if duration > 0 else 'N/A'}")
        print("*" * 30)

        joints = skeleton.d_xyz
        textures = annotator.joints_to_textures(joints)
        style_words = annotator.textures_to_style_words(textures)
        all_joints.append(joints)
        all_style_words.append(style_words)

        joints_save_path = joints_save_dir / f"{bvh_file.stem}.npy"
        np.save(joints_save_path, joints)

        textures_save_path = textures_save_dir / f"{bvh_file.stem}.npy"
        np.save(textures_save_path, textures)

        style_words_save_path = style_words_save_dir / f"{bvh_file.stem}.npy"
        np.save(style_words_save_path, style_words)

        results[bvh_file.stem] = {
            "bvh_file": bvh_file,
            "joint_file": joints_save_path,
            "style_words_file": style_words_save_path,
            "fps": fps,
            "duration": duration,
            "total_frames": skeleton.d_xyz.shape[-1],
            "downsampled_frames": textures.shape[-1],
        }
    print(f"Loaded {len(skeletons)} motions")

    style_words_stacked = np.vstack(all_style_words)
    print("All style words stacked shape", style_words_stacked.shape)

    print("Fitting model on all style words...", sep=" ")
    annotator.fit_to_style_words(style_words_stacked)
    print("Done")

    annotator_save_path = save_dir / "lma-annotator.joblib"
    annotator.save(annotator_save_path)
    print("Annotator model saved as", annotator_save_path)

    if args.create_annotations:
        annotations_save_dir = save_dir / "annotations"
        annotations_save_dir.mkdir(exist_ok=True, parents=True)
        print(
            f"Creating annotation for each motions. Annotations will be saved in '{annotations_save_dir}'..."
        )
        for bvh_file, joints in tqdm(zip(bvh_files, all_joints), total=len(bvh_files)):
            annotations = annotator.annotate(joints, return_expanded=True)
            annotations_file = annotations_save_dir / f"{bvh_file.stem}.csv"
            annotations.to_csv(annotations_file, index=False)
            results[bvh_file.stem]["annotations_file"] = annotations_file

    results_csv_file = save_dir / "results.csv"
    results_df = pd.DataFrame.from_dict(results).T
    results_df.to_csv(results_csv_file, index=True)

    return os.EX_OK


def timestamp_str(date: Optional[datetime] = None) -> str:
    date = date or datetime.now()
    return date.strftime("%y%m%d-%H%M")


def make_serializable(d: Dict[str, Union[int, bool, str, Path]]):
    d = {k: str(v) if isinstance(v, Path) else v for k, v in d.items()}
    return d


if __name__ == "__main__":
    sys.exit(main())