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)