diff --git a/amt/audio.py b/amt/audio.py index f2de66e..5fa796a 100644 --- a/amt/audio.py +++ b/amt/audio.py @@ -235,11 +235,13 @@ def _get_noise(self, noise_paths: list): def apply_reverb(self, wav: torch.Tensor): # wav: (bz, L) if self.num_impulse == 0: - logging.warning("Warning: No files in amt/assets/impulse - skipping apply_reverb") + logging.warning( + "Warning: No files in amt/assets/impulse - skipping apply_reverb" + ) return wav elif self.reverb_factor == 0: return wav - + batch_size, _ = wav.shape reverb_strength = ( @@ -265,7 +267,9 @@ def apply_reverb(self, wav: torch.Tensor): def apply_noise(self, wav: torch.tensor): if self.num_noise == 0: - logging.warning("Warning: No files in amt/assets/noise - skipping apply_noise") + logging.warning( + "Warning: No files in amt/assets/noise - skipping apply_noise" + ) return wav batch_size, _ = wav.shape @@ -283,7 +287,9 @@ def apply_noise(self, wav: torch.tensor): def apply_applause(self, wav: torch.tensor): if self.num_applause == 0: - logging.warning("Warning: No files in amt/assets/applause - skipping apply_applause") + logging.warning( + "Warning: No files in amt/assets/applause - skipping apply_applause" + ) return wav batch_size, _ = wav.shape diff --git a/amt/data.py b/amt/data.py index 688af93..2638636 100644 --- a/amt/data.py +++ b/amt/data.py @@ -64,6 +64,33 @@ def get_mid_segments( start_ms += stride_len_ms +def _get_single_wav_segment( + audio_path: str, + sample_rate: int, + orig_sample_rate: int, + segment_samples: tuple[int] | None = None, +): + if segment_samples is not None: + start_sample, end_sample = segment_samples + orig_start_sample = int(start_sample * (orig_sample_rate / sample_rate)) + orig_end_sample = int(end_sample * (orig_sample_rate / sample_rate)) + + wav, _ = torchaudio.load( + audio_path, + frame_offset=orig_start_sample, + num_frames=orig_end_sample - orig_start_sample, + ) + wav = wav.mean(0) + else: + wav, _ = torchaudio.load(audio_path) + wav = wav.mean(0) + + wav = torchaudio.functional.resample( + wav, orig_freq=orig_sample_rate, new_freq=sample_rate + ) + return wav + + def get_wav_segments( audio_path: str, stride_factor: int | None = None, @@ -95,48 +122,72 @@ def get_wav_segments( else: start_sample, end_sample = 0, None - stream.add_basic_audio_stream( - frames_per_chunk=stride_samples, - stream_index=0, - sample_rate=sample_rate, + wav_info = torchaudio.info(audio_path) + seg_len_s = ( + end_time_s - start_time_s + if segment is not None + else wav_info.num_frames / wav_info.sample_rate ) - buffer = torch.tensor([], dtype=torch.float32) - total_samples = start_sample - for stride_seg in stream.stream(): - seg_chunk = stride_seg[0].mean(1) + if seg_len_s <= chunk_len: + # Short audio segment - yield only one segment + wav = _get_single_wav_segment( + audio_path=audio_path, + sample_rate=sample_rate, + orig_sample_rate=wav_info.sample_rate, + segment_samples=( + (start_sample, end_sample) if segment is not None else None + ), + ) + if pad_last is True: + yield torch.nn.functional.pad(wav, (0, chunk_samples - len(wav))) + else: + yield wav - if end_sample and total_samples + seg_chunk.shape[0] > end_sample: - samples_to_use = end_sample - total_samples - seg_chunk = seg_chunk[:samples_to_use] + else: + # Yield segments in order + stream.add_basic_audio_stream( + frames_per_chunk=stride_samples, + stream_index=0, + sample_rate=sample_rate, + ) - total_samples += seg_chunk.shape[0] + buffer = torch.tensor([], dtype=torch.float32) + total_samples = start_sample + for stride_seg in stream.stream(): + seg_chunk = stride_seg[0].mean(1) - # Pad seg_chunk if required - if seg_chunk.shape[0] < stride_samples: - seg_chunk = F.pad( - seg_chunk, - (0, stride_samples - seg_chunk.shape[0]), - mode="constant", - value=0.0, - ) + if end_sample and total_samples + seg_chunk.shape[0] > end_sample: + samples_to_use = end_sample - total_samples + seg_chunk = seg_chunk[:samples_to_use] - if buffer.shape[0] < chunk_samples: - buffer = torch.cat((buffer, seg_chunk), dim=0) - else: - buffer = torch.cat((buffer[stride_samples:], seg_chunk), dim=0) + total_samples += seg_chunk.shape[0] + + # Pad seg_chunk if required + if seg_chunk.shape[0] < stride_samples: + seg_chunk = F.pad( + seg_chunk, + (0, stride_samples - seg_chunk.shape[0]), + mode="constant", + value=0.0, + ) - if buffer.shape[0] == chunk_samples: - yield buffer + if buffer.shape[0] < chunk_samples: + buffer = torch.cat((buffer, seg_chunk), dim=0) + else: + buffer = torch.cat((buffer[stride_samples:], seg_chunk), dim=0) - if end_sample and total_samples >= end_sample: - break + if buffer.shape[0] == chunk_samples: + yield buffer - if pad_last and buffer.shape[0] > stride_samples: - yield torch.nn.functional.pad( - buffer[stride_samples:], - (0, chunk_samples - len(buffer[stride_samples:])), - ) + if end_sample and total_samples >= end_sample: + break + + if pad_last and buffer.shape[0] > stride_samples: + yield torch.nn.functional.pad( + buffer[stride_samples:], + (0, chunk_samples - len(buffer[stride_samples:])), + ) def get_paired_wav_mid_segments( diff --git a/amt/inference/transcribe.py b/amt/inference/transcribe.py index adc072c..fd000b8 100644 --- a/amt/inference/transcribe.py +++ b/amt/inference/transcribe.py @@ -757,16 +757,17 @@ def transcribe_file( logger.info(f"Seen eos_tok in audio chunk {idx}: {file_path}") seq = seq[:-1] + concat_seq += _shift_onset( + seq[init_idx:], + idx * CHUNK_LEN_MS, + ) + if len(next_seq) == 1: logger.info( f"Skipping audio chunk {idx} (silence): {file_path}" ) seq = [tokenizer.bos_tok] else: - concat_seq += _shift_onset( - seq[init_idx:], - idx * CHUNK_LEN_MS, - ) seq = next_seq idx += 1 @@ -879,7 +880,7 @@ def remove_failures_from_queue_(_queue: Queue, _pid: int): logger.info( f"Finished file: {file_path} (segment: {idx if idx is not None else 'full'})" ) - if len(seq) < 500: + if len(seq) < 10: logger.info( f"Skipping seq - too short (segment {idx if idx is not None else 'full'})" )