preprocessing.py 文件源码

def stack_and_pad(values: List[Union[np.ndarray, int, float]], pad=0) -> np.ndarray:
    """Pads a list of numpy arrays so that they have equal dimensions, then stacks them."""
    if isinstance(values[0], int) or isinstance(values[0], float):
        return np.array(values)

    dims = len(values[0].shape)
    max_shape = [max(sizes) for sizes in zip(*[v.shape for v in values])]

    padded_values = []

    for value in values:
        pad_width = [(0, max_shape[i] - value.shape[i])
                     for i in range(dims)]
        padded_value = np.lib.pad(value, pad_width, mode='constant',
                                  constant_values=pad)
        padded_values.append(padded_value)

    return np.stack(padded_values)