| |
| |
|
|
| from typing import List |
| from typing import Optional |
| from typing import Tuple |
| from typing import Union |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def chunk_audio(audio: np.ndarray, max_duration_seconds: int = 30, sample_rate: int = 16000) -> List[np.ndarray]: |
| """split long audio into chunks |
| |
| Args: |
| audio: |
| max_duration_seconds: |
| sample_rate: |
| |
| Returns: |
| chunks |
| """ |
| max_len = int(max_duration_seconds * sample_rate) |
|
|
| if len(audio) <= max_len: |
| return [audio] |
|
|
| chunks = [] |
| for i in range(0, len(audio), max_len): |
| chunk = audio[i : i + max_len] |
| chunks.append(chunk) |
|
|
| return chunks |
|
|
|
|
| def process_audio_batch( |
| audios: Union[np.ndarray, List[np.ndarray], List[List[np.ndarray]]], |
| feature_extractor, |
| sampling_rate: int = 16000, |
| max_duration_seconds: int = 30, |
| return_attention_mask: bool = True, |
| ) -> Tuple[torch.Tensor, List[torch.Tensor]]: |
| """extract audio mel features |
| |
| Args: |
| audios: |
| feature_extractor: WhisperFeatureExtractor |
| sampling_rate: |
| max_duration_seconds: |
| return_attention_mask: |
| |
| Returns: |
| (audio_features, audio_feature_lens) |
| audio_features: [batch_size, n_mels, max_frames] |
| audio_feature_lens: |
| """ |
| if isinstance(audios, np.ndarray): |
| audios_list = [[audios]] |
| elif len(audios) > 0 and isinstance(audios[0], np.ndarray): |
| audios_list = [audios] |
| else: |
| audios_list = audios |
|
|
| audio_features_all = [] |
| audio_feature_lens_list = [] |
|
|
| for batch_audios in audios_list: |
| batch_lens = [] |
|
|
| for audio in batch_audios: |
| chunks = chunk_audio(audio, max_duration_seconds, sampling_rate) |
|
|
| for chunk in chunks: |
| audio_input = feature_extractor( |
| chunk, |
| sampling_rate=sampling_rate, |
| return_tensors="pt", |
| padding="max_length", |
| return_attention_mask=return_attention_mask, |
| ) |
|
|
| audio_feature = audio_input["input_features"] |
|
|
| if return_attention_mask: |
| actual_len = audio_input["attention_mask"].sum(dim=1) |
| audio_feature = audio_feature[:, :, : actual_len[0]] |
| batch_lens.append(actual_len[0]) |
| else: |
| batch_lens.append(torch.tensor(audio_feature.shape[2])) |
|
|
| audio_features_all.append(audio_feature.squeeze(0)) |
|
|
| if len(batch_lens) > 0: |
| audio_feature_lens_list.append(torch.hstack(batch_lens)) |
| else: |
| audio_feature_lens_list.append(torch.tensor([])) |
|
|
| |
| if audio_features_all: |
| audio_features = torch.nn.utils.rnn.pad_sequence( |
| [feat.transpose(0, 1) for feat in audio_features_all], batch_first=True, padding_value=0.0 |
| ).transpose( |
| 1, 2 |
| ) |
| else: |
| audio_features = torch.tensor([]) |
|
|
| return audio_features, audio_feature_lens_list |
|
|
|
|
| def regroup_audio_features( |
| audio_features: torch.Tensor, audio_feature_lens: List[torch.Tensor], regroup_seconds: int, fps: int = 100 |
| ) -> Tuple[torch.Tensor, List[torch.Tensor]]: |
| """regroup audio features to fixed duration |
| |
| Args: |
| audio_features: [batch, n_mels, frames] |
| audio_feature_lens: each batch's actual length |
| regroup_seconds: regroup duration (seconds) |
| fps: frames per second |
| |
| Returns: |
| (regrouped_features, regrouped_lens) |
| """ |
| |
| all_lens = [] |
| for lens in audio_feature_lens: |
| if isinstance(lens, torch.Tensor): |
| all_lens.extend(lens.tolist()) |
| elif isinstance(lens, list): |
| all_lens.extend([int(x) for x in lens]) |
|
|
| if len(all_lens) == 0: |
| return torch.tensor([]), [] |
|
|
| |
| flat_slices = [audio_features[i, :, :L] for i, L in enumerate(all_lens)] |
|
|
| if len(flat_slices) == 1: |
| full_feat = flat_slices[0] |
| else: |
| full_feat = torch.cat(flat_slices, dim=1) |
|
|
| |
| frames_per_seg = int(regroup_seconds * fps) |
| segments = [] |
|
|
| for start in range(0, full_feat.size(1), frames_per_seg): |
| seg = full_feat[:, start : start + frames_per_seg] |
| if seg.size(1) > 0: |
| segments.append(seg) |
|
|
| if len(segments) == 0: |
| return torch.tensor([]), [] |
|
|
| |
| seg_lens = [s.size(1) for s in segments] |
| segs_transposed = [s.transpose(0, 1) for s in segments] |
|
|
| padded = torch.nn.utils.rnn.pad_sequence(segs_transposed, batch_first=True, padding_value=0.0) |
|
|
| padded = padded.transpose(1, 2) |
| lens_tensor = torch.tensor(seg_lens, dtype=torch.int32, device=padded.device) |
|
|
| return padded, [lens_tensor] |
|
|
|
|
| def calculate_mel_frames(audio_samples: int, n_fft: int = 400, hop_length: int = 160) -> int: |
| """calculate mel frames |
| |
| Args: |
| audio_samples: |
| n_fft: |
| hop_length: |
| |
| Returns: |
| mel frames |
| """ |
| if audio_samples < n_fft: |
| return 0 |
| return 1 + (audio_samples - n_fft) // hop_length |
|
|
|
|
| def samples_to_ms(samples: int, sample_rate: int = 16000) -> float: |
| """convert samples to milliseconds |
| |
| Args: |
| samples: |
| sample_rate: |
| |
| Returns: |
| milliseconds |
| """ |
| return samples / sample_rate * 1000 |
|
|
|
|
| def ms_to_samples(ms: float, sample_rate: int = 16000) -> int: |
| """convert milliseconds to samples |
| |
| Args: |
| ms: milliseconds |
| sample_rate: |
| |
| Returns: |
| samples |
| """ |
| return int(ms * sample_rate / 1000) |
|
|
|
|
| def validate_audio_format(audio: np.ndarray, expected_rate: int = 16000, max_duration: Optional[float] = None) -> bool: |
| """validate audio format |
| |
| Args: |
| audio: audio data |
| expected_rate: |
| max_duration: |
| |
| Returns: |
| whether valid audio format or not |
| """ |
| if not isinstance(audio, np.ndarray): |
| return False |
|
|
| if audio.ndim != 1: |
| return False |
|
|
| if audio.dtype not in [np.float32, np.float64, np.int16, np.int32]: |
| return False |
|
|
| if max_duration is not None: |
| max_samples = int(max_duration * expected_rate) |
| if len(audio) > max_samples: |
| return False |
|
|
| return True |
|
|
