Add pad_or_trim function to handle segment before encoding (#705)

This commit is contained in:
trungkienbkhn
2024-02-29 23:08:28 +07:00
committed by GitHub
parent 06d32bf0c1
commit 16141e65d9
2 changed files with 17 additions and 1 deletions

View File

@@ -102,3 +102,18 @@ def _resample_frames(frames, resampler):
# Add None to flush the resampler.
for frame in itertools.chain(frames, [None]):
yield from resampler.resample(frame)
def pad_or_trim(array, length: int, *, axis: int = -1):
"""
Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
"""
if array.shape[axis] > length:
array = array.take(indices=range(length), axis=axis)
if array.shape[axis] < length:
pad_widths = [(0, 0)] * array.ndim
pad_widths[axis] = (0, length - array.shape[axis])
array = np.pad(array, pad_widths)
return array