ubisoft-laforge-ZeroEGGS-main/ZEGGS/data_pipeline.py

import json
import logging
import os
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
from omegaconf import DictConfig
from scipy import interpolate
from scipy.interpolate import griddata

from anim import bvh, quat
from audio.audio_files import read_wavfile, write_wavefile
from audio.spectrograms import extract_mel_spectrogram_for_tts

FILE_ROOT = os.path.dirname(os.path.realpath(__file__))
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

_logger = logging.getLogger(__name__)
_logger.propagate = False
warnings.simplefilter("ignore")


# ===============================================
#                   Audio
# ===============================================
def extract_energy(mel_spec):
    energy = np.linalg.norm(mel_spec, axis=0)
    return energy


def preprocess_audio(audio_data, anim_fs, anim_length, params, feature_type):
    if params.normalize_loudness:
        import pyloudnorm as pyln
        meter = pyln.Meter(params.sampling_rate)  # create BS.1770 meter
        loudness = meter.integrated_loudness(audio_data)
        # loudness normalize audio to -20 dB LUFS
        audio_data = pyln.normalize.loudness(audio_data, loudness, -20.0)

    resample_method = params.resample_method
    audio_feature = []
    # Extract MEL spectrogram
    mel_spec = extract_mel_spectrogram_for_tts(
        wav_signal=audio_data,
        fs=params.sampling_rate,
        n_fft=params.filter_length,
        step_size=params.hop_length,
        n_mels=params.n_mel_channels,
        mel_fmin=params.mel_fmin,
        mel_fmax=params.mel_fmax,
        min_amplitude=params.min_clipping,
        pre_emphasis=params.pre_emphasis,
        pre_emph_coeff=params.pre_emph_coeff,
        dynamic_range=None,
        real_amplitude=params.real_amplitude,
        centered=params.centered,
        normalize_mel_bins=params.normalize_mel_bins,
        normalize_range=params.normalize_range,
        logger=_logger,
    )[0].T
    mel_spec = 10 ** (mel_spec / 20)
    mel_spec = np.log(mel_spec)

    if "mel_spec" in feature_type:
        mel_spec_interp = interpolate.griddata(
            np.arange(len(mel_spec)),
            mel_spec,
            ((params.sampling_rate / params.hop_length) / anim_fs) * np.arange(anim_length),
            method=resample_method,
        ).astype(np.float32)
        audio_feature.append(mel_spec_interp)

    if "energy" in feature_type:
        energy = extract_energy(np.exp(mel_spec).T)
        f = interpolate.interp1d(np.arange(len(energy)), energy, kind=resample_method, fill_value="extrapolate")
        energy_interp = f(
            ((params.sampling_rate / params.hop_length) / anim_fs) * np.arange(anim_length)
        ).astype(np.float32)
        audio_feature.append(energy_interp[:, np.newaxis])

    audio_feature = np.concatenate(audio_feature, axis=1)

    return audio_feature


# ===============================================
#                   Animation
# ===============================================
def preprocess_animation(anim_data, conf=dict(), animation_path=None, info_df=None, i=0):
    nframes = len(anim_data["rotations"])
    njoints = len(anim_data["parents"])
    dt = anim_data["frametime"]

    lrot = quat.unroll(quat.from_euler(np.radians(anim_data["rotations"]), anim_data["order"]))

    lpos = anim_data["positions"]

    grot, gpos = quat.fk(lrot, lpos, anim_data["parents"])

    # Find root (Projected hips on the ground)
    root_pos = gpos[:, anim_data["names"].index("Spine2")] * np.array([1, 0, 1])
    # root_pos = signal.savgol_filter(root_pos, 31, 3, axis=0, mode="interp")

    # Root direction
    root_fwd = quat.mul_vec(grot[:, anim_data["names"].index("Hips")], np.array([[0, 0, 1]]))
    root_fwd[:, 1] = 0
    root_fwd = root_fwd / np.sqrt(np.sum(root_fwd * root_fwd, axis=-1))[..., np.newaxis]

    # root_fwd = signal.savgol_filter(root_fwd, 61, 3, axis=0, mode="interp")
    # root_fwd = root_fwd / np.sqrt(np.sum(root_fwd * root_fwd, axis=-1))[..., np.newaxis]

    # Root rotation
    root_rot = quat.normalize(
        quat.between(np.array([[0, 0, 1]]).repeat(len(root_fwd), axis=0), root_fwd)
    )

    # Find look at direction
    gaze_lookat = quat.mul_vec(grot[:, anim_data["names"].index("Head")], np.array([0, 0, 1]))
    gaze_lookat[:, 1] = 0
    gaze_lookat = gaze_lookat / np.sqrt(np.sum(np.square(gaze_lookat), axis=-1))[..., np.newaxis]

    # Find gaze position
    gaze_distance = 100  # Assume other actor is one meter away
    gaze_pos_all = root_pos + gaze_distance * gaze_lookat
    gaze_pos = np.median(gaze_pos_all, axis=0)
    gaze_pos = gaze_pos[np.newaxis].repeat(nframes, axis=0)

    # Visualize Gaze Pos
    if conf.get("visualize_gaze", False):
        import matplotlib.pyplot as plt

        plt.scatter(gaze_pos_all[:, 0], gaze_pos_all[:, 2], s=0.1, marker=".")
        plt.scatter(gaze_pos[0, 0], gaze_pos[0, 2])
        plt.scatter(root_pos[:, 0], root_pos[:, 2], s=0.1, marker=".")
        plt.quiver(root_pos[::60, 0], root_pos[::60, 2], root_fwd[::60, 0], root_fwd[::60, 2])
        plt.gca().set_aspect("equal")
        plt.show()

    # Compute local gaze dir
    gaze_dir = gaze_pos - root_pos
    # gaze_dir = gaze_dir / np.sqrt(np.sum(np.square(gaze_dir), axis=-1))[..., np.newaxis]
    gaze_dir = quat.mul_vec(quat.inv(root_rot), gaze_dir)

    # Make relative to root
    lrot[:, 0] = quat.mul(quat.inv(root_rot), lrot[:, 0])
    lpos[:, 0] = quat.mul_vec(quat.inv(root_rot), lpos[:, 0] - root_pos)

    # Local velocities
    lvel = np.zeros_like(lpos)
    lvel[1:] = (lpos[1:] - lpos[:-1]) / dt
    lvel[0] = lvel[1] - (lvel[3] - lvel[2])

    lvrt = np.zeros_like(lpos)
    lvrt[1:] = quat.to_helical(quat.abs(quat.mul(lrot[1:], quat.inv(lrot[:-1])))) / dt
    lvrt[0] = lvrt[1] - (lvrt[3] - lvrt[2])

    # Root velocities
    root_vrt = np.zeros_like(root_pos)
    root_vrt[1:] = quat.to_helical(quat.abs(quat.mul(root_rot[1:], quat.inv(root_rot[:-1])))) / dt
    root_vrt[0] = root_vrt[1] - (root_vrt[3] - root_vrt[2])
    root_vrt[1:] = quat.mul_vec(quat.inv(root_rot[:-1]), root_vrt[1:])
    root_vrt[0] = quat.mul_vec(quat.inv(root_rot[0]), root_vrt[0])

    root_vel = np.zeros_like(root_pos)
    root_vel[1:] = (root_pos[1:] - root_pos[:-1]) / dt
    root_vel[0] = root_vel[1] - (root_vel[3] - root_vel[2])
    root_vel[1:] = quat.mul_vec(quat.inv(root_rot[:-1]), root_vel[1:])
    root_vel[0] = quat.mul_vec(quat.inv(root_rot[0]), root_vel[0])

    # Compute character space
    crot, cpos, cvrt, cvel = quat.fk_vel(lrot, lpos, lvrt, lvel, anim_data["parents"])

    # Compute 2-axis transforms
    ltxy = np.zeros(dtype=np.float32, shape=[len(lrot), njoints, 2, 3])
    ltxy[..., 0, :] = quat.mul_vec(lrot, np.array([1.0, 0.0, 0.0]))
    ltxy[..., 1, :] = quat.mul_vec(lrot, np.array([0.0, 1.0, 0.0]))

    ctxy = np.zeros(dtype=np.float32, shape=[len(crot), njoints, 2, 3])
    ctxy[..., 0, :] = quat.mul_vec(crot, np.array([1.0, 0.0, 0.0]))
    ctxy[..., 1, :] = quat.mul_vec(crot, np.array([0.0, 1.0, 0.0]))

    if conf.get("save_normalized_animations", False):
        anim_data["positions"] = lpos
        anim_data["rotations"] = np.degrees(quat.to_euler(lrot, order=anim_data["order"]))

        normalized_animations_path = animation_path / "processed" / "normalized_animations"
        normalized_animations_path.mkdir(exist_ok=True)
        animation_norm_file = str(
            normalized_animations_path / info_df.iloc[i].anim_bvh).replace(
            ".bvh", "_norm.bvh"
        )

        bvh.save(animation_norm_file, anim_data)

        lpos_denorm = lpos.copy()
        lpos_denorm[:, 0] = quat.mul_vec(root_rot, lpos_denorm[:, 0]) + root_pos

        lrot_denorm = lrot.copy()
        lrot_denorm[:, 0] = quat.mul(root_rot, lrot_denorm[:, 0])

        anim_data["positions"] = lpos_denorm
        anim_data["rotations"] = np.degrees(quat.to_euler(lrot_denorm, order=anim_data["order"]))

        animation_denorm_file = str(
            animation_path / "processed" / "normalized_animations" / info_df.iloc[i].anim_bvh
        ).replace(".bvh", "_denorm.bvh")

        bvh.save(animation_denorm_file, anim_data)

    return (
        root_pos,
        root_rot,
        root_vel,
        root_vrt,
        lpos,
        lrot,
        ltxy,
        lvel,
        lvrt,
        cpos,
        crot,
        ctxy,
        cvel,
        cvrt,
        gaze_pos,
        gaze_dir,
    )


# ===============================================
#                   Pipeline
# ===============================================
def data_pipeline(conf):
    """Prepare Audio and Animation data for training

    Args:
        conf: config file

    Returns:
        processed_data, data_definition
    """
    from rich.progress import track
    from rich.console import Console
    from rich.table import Table
    console = Console(record=True)
    console.print("This may take a little bit of time ...")
    len_ratios = conf["len_ratios"]
    base_path = Path(conf["base_path"])
    processed_data_path = base_path / conf["processed_data_path"]
    processed_data_path.mkdir(exist_ok=True)
    info_filename = base_path / "info.csv"
    animation_path = base_path / "original"
    audio_path = base_path / "original"

    with open(str(processed_data_path / "data_pipeline_conf.json"), "w") as f:
        json.dump(conf, f, indent=4)
    conf = DictConfig(conf)

    info_df = pd.read_csv(info_filename)
    num_of_samples = len(info_df)
    audio_desired_fs = conf.audio_conf["sampling_rate"]

    X_audio_features = []

    Y_root_pos = []
    Y_root_rot = []
    Y_root_vrt = []
    Y_root_vel = []

    Y_lpos = []
    Y_lrot = []
    Y_ltxy = []
    Y_lvel = []
    Y_lvrt = []

    Y_gaze_pos = []
    Y_gaze_dir = []

    current_start_frame = 0
    ranges_train = []
    ranges_valid = []
    ranges_train_labels = []
    ranges_valid_labels = []

    # for i in track(range(num_of_samples), description="Processing...", complete_style="magenta"):
    for i in range(num_of_samples):
        animation_file = str(animation_path / info_df.iloc[i].anim_bvh)
        audio_file = audio_path / info_df.iloc[i].audio_filename

        # Load Animation #
        original_anim_data = bvh.load(animation_file)
        anim_fps = int(np.ceil(1 / original_anim_data["frametime"]))
        assert anim_fps == 60

        # Load Audio #
        audio_sr, original_audio_data = read_wavfile(
            audio_file,
            rescale=True,
            desired_fs=audio_desired_fs,
            desired_nb_channels=None,
            out_type="float32",
            logger=_logger,
        )

        # Silence Audio #
        speacker_timing_df = pd.read_csv(audio_file.with_suffix(".csv"))

        # Mark regions that don't need silencing
        mask = np.zeros_like(original_audio_data)
        for ind, row in speacker_timing_df.iterrows():

            if "R" in row["#"]:
                start_time = [int(num) for num in row["Start"].replace(".", ":").rsplit(":")]
                end_time = [int(num) for num in row["End"].replace(".", ":").rsplit(":")]

                start_time = (
                        start_time[0] * 60 * audio_desired_fs
                        + start_time[1] * audio_desired_fs
                        + int(start_time[2] * (audio_desired_fs / 1000))
                )

                end_time = (
                        end_time[0] * 60 * audio_desired_fs
                        + end_time[1] * audio_desired_fs
                        + int(end_time[2] * (audio_desired_fs / 1000))
                )

                mask[start_time:end_time] = 1.0

        # Silence unmarked regions
        original_audio_data = original_audio_data * mask

        # Sync & Trim #
        # Get mark-ups
        audio_start_time = info_df.iloc[i].audio_start_time
        audio_start_time = [int(num) for num in audio_start_time.rsplit(":")]
        anim_start_time = info_df.iloc[i].anim_start_time
        anim_start_time = [int(num) for num in anim_start_time.rsplit(":")]
        acting_start_time = info_df.iloc[i].acting_start_time
        acting_start_time = [int(num) for num in acting_start_time.rsplit(":")]
        acting_end_time = info_df.iloc[i].acting_end_time
        acting_end_time = [int(num) for num in acting_end_time.rsplit(":")]

        # Compute Timings (This is assuming that audio timing is given in 30fps)
        audio_start_time_in_thirds = (
                audio_start_time[0] * 216000
                + audio_start_time[1] * 3600
                + audio_start_time[2] * 60
                + audio_start_time[3] * 2
        )

        anim_start_time_in_thirds = (
                anim_start_time[0] * 216000
                + anim_start_time[1] * 3600
                + anim_start_time[2] * 60
                + anim_start_time[3] * 1
        )

        acting_start_time_in_thirds = (
                acting_start_time[0] * 216000
                + acting_start_time[1] * 3600
                + acting_start_time[2] * 60
                + acting_start_time[3] * 1
        )

        acting_end_time_in_thirds = (
                acting_end_time[0] * 216000
                + acting_end_time[1] * 3600
                + acting_end_time[2] * 60
                + acting_end_time[3] * 1
        )

        acting_start_in_audio_ref = int(
            np.round(
                (acting_start_time_in_thirds - audio_start_time_in_thirds) * (audio_sr / 60)
            )
        )

        acting_end_in_audio_ref = int(
            np.round((acting_end_time_in_thirds - audio_start_time_in_thirds) * (audio_sr / 60))
        )

        acting_start_in_anim_ref = int(
            np.round(
                (acting_start_time_in_thirds - anim_start_time_in_thirds) * (anim_fps / 60)
            )
        )

        acting_end_in_anim_ref = int(
            np.round((acting_end_time_in_thirds - anim_start_time_in_thirds) * (anim_fps / 60))
        )

        if (
                acting_start_in_audio_ref < 0
                or acting_start_in_anim_ref < 0
                or acting_end_in_audio_ref < 0
                or acting_end_in_anim_ref < 0
        ):
            raise ValueError("The timings are incorrect!")

        # Trim to equal length
        original_audio_data = original_audio_data[acting_start_in_audio_ref:acting_end_in_audio_ref]

        original_anim_data["rotations"] = original_anim_data["rotations"][
                                          acting_start_in_anim_ref:acting_end_in_anim_ref
                                          ]

        original_anim_data["positions"] = original_anim_data["positions"][
                                          acting_start_in_anim_ref:acting_end_in_anim_ref
                                          ]
        for len_ratio in len_ratios:
            anim_data = original_anim_data.copy()
            audio_data = original_audio_data.copy()
            if len_ratio != 1.0:
                n_anim_frames = len(original_anim_data["rotations"])
                nbones = anim_data["positions"].shape[1]
                original_times = np.linspace(0, n_anim_frames - 1, n_anim_frames)
                sample_times = np.linspace(0, n_anim_frames - 1, int(len_ratio * (n_anim_frames)))
                anim_data["positions"] = griddata(original_times, anim_data["positions"].reshape([n_anim_frames, -1]),
                                                  sample_times, method='cubic').reshape([len(sample_times), nbones, 3])

                rotations = quat.unroll(quat.from_euler(np.radians(anim_data['rotations']), order=anim_data['order']))
                rotations = griddata(original_times, rotations.reshape([n_anim_frames, -1]), sample_times,
                                     method='cubic').reshape([len(sample_times), nbones, 4])
                rotations = quat.normalize(rotations)
                anim_data["rotations"] = np.degrees(quat.to_euler(rotations, order=anim_data["order"]))

                n_audio_frames = len(audio_data)
                original_times = np.linspace(0, n_audio_frames - 1, n_audio_frames)
                sample_times = np.linspace(0, n_audio_frames - 1, int(len_ratio * (n_audio_frames)))
                audio_data = griddata(original_times, audio_data, sample_times, method='cubic')
                # assert len(audio_data) / audio_sr == len(anim_data["rotations"]) / anim_fps

            # Saving Trimmed Files
            folder = "valid" if info_df.iloc[i].validation else "train"
            trimmed_filename = info_df.iloc[i].anim_bvh.split(".")[0]
            trimmed_filename = trimmed_filename + "_x_" + str(len_ratio).replace(".", "_")

            if conf["save_trimmed_audio"]:
                target_path = processed_data_path / "trimmed" / folder
                target_path.mkdir(exist_ok=True, parents=True)
                write_wavefile(target_path / (trimmed_filename + ".wav"), audio_data, audio_sr)

            if conf["save_trimmed_animation"]:
                target_path = processed_data_path / "trimmed" / folder
                target_path.mkdir(exist_ok=True, parents=True)

                # Centering the character. Comment if you want the original global position and orientation
                output = anim_data.copy()
                lrot = quat.from_euler(np.radians(output["rotations"]), output["order"])
                offset_pos = output["positions"][0:1, 0:1].copy() * np.array([1, 0, 1])
                offset_rot = lrot[0:1, 0:1].copy() * np.array([1, 0, 1, 0])

                root_pos = quat.mul_vec(quat.inv(offset_rot), output["positions"][:, 0:1] - offset_pos)
                output["positions"][:, 0:1] = quat.mul_vec(quat.inv(offset_rot),
                                                           output["positions"][:, 0:1] - offset_pos)
                output["rotations"][:, 0:1] = np.degrees(
                    quat.to_euler(quat.mul(quat.inv(offset_rot), lrot[:, 0:1]), order=output["order"]))

                bvh.save(target_path / (trimmed_filename + ".bvh"), anim_data)

            # Extracting Audio Features #
            audio_features = preprocess_audio(
                audio_data,
                anim_fps,
                len(anim_data["rotations"]),
                conf.audio_conf,
                feature_type=conf.audio_feature_type,
            )

            # Check if the lengths are correct and no NaNs
            assert len(audio_features) == len(anim_data["rotations"])
            assert not np.any(np.isnan(audio_features))

            if conf["visualize_spectrogram"]:
                import matplotlib.pyplot as plt

                plt.imshow(audio_features.T, interpolation="nearest")
                plt.show()

            # Extracting Animation Features
            nframes = len(anim_data["rotations"])
            dt = anim_data["frametime"]
            (
                root_pos,
                root_rot,
                root_vel,
                root_vrt,
                lpos,
                lrot,
                ltxy,
                lvel,
                lvrt,
                cpos,
                crot,
                ctxy,
                cvel,
                cvrt,
                gaze_pos,
                gaze_dir,
            ) = preprocess_animation(anim_data, conf, animation_path, info_df, i)

            # Appending Data
            X_audio_features.append(audio_features)

            Y_root_pos.append(root_pos)
            Y_root_rot.append(root_rot)
            Y_root_vel.append(root_vel)
            Y_root_vrt.append(root_vrt)
            Y_lpos.append(lpos)
            Y_lrot.append(lrot)
            Y_ltxy.append(ltxy)
            Y_lvel.append(lvel)
            Y_lvrt.append(lvrt)

            Y_gaze_pos.append(gaze_pos)
            Y_gaze_dir.append(gaze_dir)

            # Append to Ranges
            current_end_frame = nframes + current_start_frame

            if info_df.iloc[i].validation:
                ranges_valid.append([current_start_frame, current_end_frame])
                ranges_valid_labels.append(info_df.iloc[i].style)
            else:
                ranges_train.append([current_start_frame, current_end_frame])
                ranges_train_labels.append(info_df.iloc[i].style)

            current_start_frame = current_end_frame

    # Processing Labels
    ranges_train = np.array(ranges_train, dtype=np.int32)
    ranges_valid = np.array(ranges_valid, dtype=np.int32)

    label_names = list(set(ranges_train_labels + ranges_valid_labels))

    ranges_train_labels = np.array(
        [label_names.index(label) for label in ranges_train_labels], dtype=np.int32
    )
    ranges_valid_labels = np.array(
        [label_names.index(label) for label in ranges_valid_labels], dtype=np.int32
    )

    # Concatenating Data
    X_audio_features = np.concatenate(X_audio_features, axis=0).astype(np.float32)

    Y_root_pos = np.concatenate(Y_root_pos, axis=0).astype(np.float32)
    Y_root_rot = np.concatenate(Y_root_rot, axis=0).astype(np.float32)
    Y_root_vel = np.concatenate(Y_root_vel, axis=0).astype(np.float32)
    Y_root_vrt = np.concatenate(Y_root_vrt, axis=0).astype(np.float32)

    Y_lpos = np.concatenate(Y_lpos, axis=0).astype(np.float32)
    Y_lrot = np.concatenate(Y_lrot, axis=0).astype(np.float32)
    Y_ltxy = np.concatenate(Y_ltxy, axis=0).astype(np.float32)
    Y_lvel = np.concatenate(Y_lvel, axis=0).astype(np.float32)
    Y_lvrt = np.concatenate(Y_lvrt, axis=0).astype(np.float32)

    Y_gaze_pos = np.concatenate(Y_gaze_pos, axis=0).astype(np.float32)
    Y_gaze_dir = np.concatenate(Y_gaze_dir, axis=0).astype(np.float32)

    # Compute Means & Stds
    # Filter out start and end frames
    ranges_mask = np.zeros(len(X_audio_features), dtype=bool)
    for s, e in ranges_train:
        ranges_mask[s + 2: e - 2] = True

    # Compute Means
    Y_root_vel_mean = Y_root_vel[ranges_mask].mean(axis=0)
    Y_root_vrt_mean = Y_root_vrt[ranges_mask].mean(axis=0)

    Y_lpos_mean = Y_lpos[ranges_mask].mean(axis=0)
    Y_ltxy_mean = Y_ltxy[ranges_mask].mean(axis=0)
    Y_lvel_mean = Y_lvel[ranges_mask].mean(axis=0)
    Y_lvrt_mean = Y_lvrt[ranges_mask].mean(axis=0)

    Y_gaze_dir_mean = Y_gaze_dir[ranges_mask].mean(axis=0)

    audio_input_mean = X_audio_features[ranges_mask].mean(axis=0)

    anim_input_mean = np.hstack(
        [
            Y_root_vel_mean.ravel(),
            Y_root_vrt_mean.ravel(),
            Y_lpos_mean.ravel(),
            Y_ltxy_mean.ravel(),
            Y_lvel_mean.ravel(),
            Y_lvrt_mean.ravel(),
            Y_gaze_dir_mean.ravel(),
        ]
    )

    # Compute Stds
    Y_root_vel_std = Y_root_vel[ranges_mask].std() + 1e-10
    Y_root_vrt_std = Y_root_vrt[ranges_mask].std() + 1e-10

    Y_lpos_std = Y_lpos[ranges_mask].std() + 1e-10
    Y_ltxy_std = Y_ltxy[ranges_mask].std() + 1e-10
    Y_lvel_std = Y_lvel[ranges_mask].std() + 1e-10
    Y_lvrt_std = Y_lvrt[ranges_mask].std() + 1e-10

    Y_gaze_dir_std = Y_gaze_dir[ranges_mask].std() + 1e-10

    audio_input_std = X_audio_features[ranges_mask].std() + 1e-10

    anim_input_std = np.hstack(
        [
            Y_root_vel_std.repeat(len(Y_root_vel_mean.ravel())),
            Y_root_vrt_std.repeat(len(Y_root_vrt_mean.ravel())),
            Y_lpos_std.repeat(len(Y_lpos_mean.ravel())),
            Y_ltxy_std.repeat(len(Y_ltxy_mean.ravel())),
            Y_lvel_std.repeat(len(Y_lvel_mean.ravel())),
            Y_lvrt_std.repeat(len(Y_lvrt_mean.ravel())),
            Y_gaze_dir_std.repeat(len(Y_gaze_dir_mean.ravel())),
        ]
    )

    # Compute Output Means
    anim_output_mean = np.hstack(
        [
            Y_root_vel_mean.ravel(),
            Y_root_vrt_mean.ravel(),
            Y_lpos_mean.ravel(),
            Y_ltxy_mean.ravel(),
            Y_lvel_mean.ravel(),
            Y_lvrt_mean.ravel(),
        ]
    )

    # Compute Output Stds
    Y_root_vel_out_std = Y_root_vel[ranges_mask].std(axis=0)
    Y_root_vrt_out_std = Y_root_vrt[ranges_mask].std(axis=0)

    Y_lpos_out_std = Y_lpos[ranges_mask].std(axis=0)
    Y_ltxy_out_std = Y_ltxy[ranges_mask].std(axis=0)
    Y_lvel_out_std = Y_lvel[ranges_mask].std(axis=0)
    Y_lvrt_out_std = Y_lvrt[ranges_mask].std(axis=0)

    anim_output_std = np.hstack(
        [
            Y_root_vel_out_std.ravel(),
            Y_root_vrt_out_std.ravel(),
            Y_lpos_out_std.ravel(),
            Y_ltxy_out_std.ravel(),
            Y_lvel_out_std.ravel(),
            Y_lvrt_out_std.ravel(),
        ]
    )

    processed_data = dict(
        X_audio_features=X_audio_features,
        Y_root_pos=Y_root_pos,
        Y_root_rot=Y_root_rot,
        Y_root_vel=Y_root_vel,
        Y_root_vrt=Y_root_vrt,
        Y_lpos=Y_lpos,
        Y_ltxy=Y_ltxy,
        Y_lvel=Y_lvel,
        Y_lvrt=Y_lvrt,
        Y_gaze_pos=Y_gaze_pos,
        ranges_train=ranges_train,
        ranges_valid=ranges_valid,
        ranges_train_labels=ranges_train_labels,
        ranges_valid_labels=ranges_valid_labels,
        audio_input_mean=audio_input_mean,
        audio_input_std=audio_input_std,
        anim_input_mean=anim_input_mean,
        anim_input_std=anim_input_std,
        anim_output_mean=anim_output_mean,
        anim_output_std=anim_output_std,
    )

    stats = dict(
        ranges_train=ranges_train,
        ranges_valid=ranges_valid,
        ranges_train_labels=ranges_train_labels,
        ranges_valid_labels=ranges_valid_labels,
        audio_input_mean=audio_input_mean,
        audio_input_std=audio_input_std,
        anim_input_mean=anim_input_mean,
        anim_input_std=anim_input_std,
        anim_output_mean=anim_output_mean,
        anim_output_std=anim_output_std,
    )

    data_definition = dict(
        dt=dt,
        label_names=label_names,
        parents=anim_data["parents"].tolist(),
        bone_names=anim_data["names"],
    )

    # Save Data
    if conf["save_final_data"]:
        np.savez(processed_data_path / "processed_data.npz", **processed_data)

        np.savez(processed_data_path / "stats.npz", **stats)

        with open(str(processed_data_path / "data_definition.json"), "w") as f:
            json.dump(data_definition, f, indent=4)

    # Data Stats:
    nlabels = len(label_names)
    df = pd.DataFrame()
    df["Dataset"] = ["Train", "Validation", "Total"]
    pd.set_option("display.max_rows", None, "display.max_columns", None)
    table = Table(title="Data Info", show_lines=True, row_styles=["magenta"])
    table.add_column("Dataset")
    data_len = 0
    for i in range(nlabels):
        ind_mask = ranges_train_labels == i
        ranges = ranges_train[ind_mask]
        num_train_frames = (
                np.sum(ranges[:, 1] - ranges[:, 0]) / 2
        )  # It is divided by two as we have mirrored versions too
        ind_mask = ranges_valid_labels == i
        ranges = ranges_valid[ind_mask]
        num_valid_frames = np.sum(ranges[:, 1] - ranges[:, 0]) / 2
        total = num_train_frames + num_valid_frames
        df[label_names[i]] = [
            f"{num_train_frames} frames - {num_train_frames / 60:.1f} secs",
            f"{num_valid_frames} frames - {num_valid_frames / 60:.1f} secs",
            f"{total} frames - {total / 60:.1f} secs",
        ]
        table.add_column(label_names[i])
        data_len += total

    for i in range(3):
        table.add_row(*list(df.iloc[i]))
    console.print(table)
    console.print(f"Total length of dataset is {data_len} frames - {data_len / 60:.1f} seconds")
    console_print_file = processed_data_path / "data_info.html"
    console.print(dict(conf))
    console.save_html(str(console_print_file))

    return processed_data, data_definition


if __name__ == "__main__":
    config_file = "../configs/data_pipeline_conf_v1.json"
    with open(config_file, "r") as f:
        conf = json.load(f)

    data_pipeline(conf)