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 / conf["info_filename"]
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"):
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)
