diff --git a/utils/data.py b/utils/data.py index a38092e9..6c47d5eb 100644 --- a/utils/data.py +++ b/utils/data.py @@ -1,7 +1,7 @@ import numpy as np -def pad_data(x, length): +def _pad_data(x, length): _pad = 0 assert x.ndim == 1 return np.pad(x, (0, length - x.shape[0]), @@ -11,7 +11,33 @@ def pad_data(x, length): def prepare_data(inputs): max_len = max((len(x) for x in inputs)) - return np.stack([pad_data(x, max_len) for x in inputs]) + return np.stack([_pad_data(x, max_len) for x in inputs]) + + +def _pad_tensor(x, length): + _pad = 0 + assert x.ndim == 2 + x = np.pad(x, [[0, 0], [0, length - x.shape[1]]], mode='constant', constant_values=_pad) + return x + +def prepare_tensor(inputs, out_steps): + max_len = max((x.shape[1] for x in inputs)) + 1 # zero-frame + remainder = max_len % out_steps + pad_len = max_len + (out_steps - remainder) if remainder > 0 else max_len + return np.stack([_pad_tensor(x, pad_len) for x in inputs]) + + +def _pad_stop_target(x, length): + _pad = 1. + assert x.ndim == 1 + return np.pad(x, (0, length - x.shape[0]), mode='constant', constant_values=_pad) + + +def prepare_stop_target(inputs, out_steps): + max_len = max((x.shape[0] for x in inputs)) + 1 # zero-frame + remainder = max_len % out_steps + pad_len = max_len + (out_steps - remainder) if remainder > 0 else max_len + return np.stack([_pad_stop_target(x, pad_len) for x in inputs]) def pad_per_step(inputs, pad_len):