Mesrulla/audio_model
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add augment noise and num layer=6

Browse Source
This commit is contained in:
Mesrulla
2026-05-09 14:11:58 +06:00
parent d31233a79a
commit 96cd0a20cb
6 changed files with 155 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/dataset.py
+24 -119
View File
@@ -4,7 +4,7 @@ import torch
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import torchaudio
from torch_audiomentations import ApplyImpulseResponse, Gain, PitchShift, LowPassFilter, HighPassFilter
from torch_audiomentations import ApplyImpulseResponse, Gain, PitchShift, LowPassFilter, HighPassFilter, PolarityInversion
from torchaudio.transforms import Resample, TimeStretch
from pathlib import Path
import pandas as pd
@@ -94,6 +94,7 @@ class CommonVoiceDataset(Dataset[BatchItem]):
tsv_path: Path,
audio_dir: Path,
noise_dir: Path,
corridor_noise_dir: Path,
tokenizer: ASRTokenizer,
sample_rate: int = 16000,
max_audio_len: int = 480000, # 30秒 @ 16kHz
@@ -123,9 +124,10 @@ class CommonVoiceDataset(Dataset[BatchItem]):
self.gain_down = Gain(min_gain_in_db=-15, max_gain_in_db=-8, p=1.0, output_type='tensor')
self.pitch_up = PitchShift(min_transpose_semitones=1, max_transpose_semitones=4, p=1.0, sample_rate=self.sample_rate, output_type='tensor')
self.pitch_down = PitchShift(min_transpose_semitones=-4, max_transpose_semitones=-1, p=1.0, sample_rate=self.sample_rate, output_type='tensor')
self.lowpass = LowPassFilter(min_cutoff_freq=600, max_cutoff_freq=2000, p=1.0, output_type='tensor')
self.highpass = HighPassFilter(min_cutoff_freq=800, max_cutoff_freq=2000, p=1.0, output_type='tensor')
self.apply_ir = ApplyImpulseResponse(ir_paths=noise_dir,convolve_mode='same', p=1, output_type="tensor")
self.lowpass = LowPassFilter(min_cutoff_freq=100, max_cutoff_freq=2000, p=1.0, output_type='tensor')
self.highpass = HighPassFilter(min_cutoff_freq=1000, max_cutoff_freq=2000, p=1.0, output_type='tensor')
self.apply_ir = ApplyImpulseResponse(ir_paths=corridor_noise_dir, convolve_mode='same', p=1, output_type="tensor")
self.polarity_inversion = PolarityInversion(p=1.0, output_type="tensor")
def __len__(self):
return len(self.data)
@@ -157,7 +159,7 @@ class CommonVoiceDataset(Dataset[BatchItem]):
if random.random() < 0.6:
waveform = self._stretch_or_compress(waveform=waveform)
if random.random() < 0.4:
if random.random() < 0.7:
waveform = self.noise_augmentor.apply_real_noise(waveform)
if random.random() < 0.3:
@@ -166,29 +168,30 @@ class CommonVoiceDataset(Dataset[BatchItem]):
# torch_audiomentations: [1, time] -> [1, 1, time]
if waveform.dim() == 2:
waveform_3d = waveform.unsqueeze(0)
# 随机选择一种频谱增强
# 随机选择一种物理特性增强 (互斥区)
choice = random.random()
if choice < 0.15:
# 增益变化（上或下）
if choice < 0.25: # [0.00 - 0.25] 25% 概率：增益
if random.random() < 0.5:
waveform_3d = self.gain_up(waveform_3d, sample_rate=self.sample_rate)
else:
waveform_3d = self.gain_down(waveform_3d, sample_rate=self.sample_rate)
elif choice < 0.25:
# 音高变化（上或下）
elif choice < 0.50: # [0.25 - 0.50] 25% 概率：音高
if random.random() < 0.5:
waveform_3d = self.pitch_up(waveform_3d, sample_rate=self.sample_rate)
else:
waveform_3d = self.pitch_down(waveform_3d, sample_rate=self.sample_rate)
elif choice < 0.30:
# 低通滤波（声音发闷）
elif choice < 0.70: # [0.50 - 0.70] 20% 概率：低通
waveform_3d = self.lowpass(waveform_3d, sample_rate=self.sample_rate)
# elif choice < 0.32:
# # 低通滤波（声音发闷）
# waveform_3d = self.apply_ir(waveform_3d, sample_rate=self.sample_rate)
elif choice < 0.35:
# 高通滤波（电话效果）
elif choice < 0.85: # [0.70 - 0.85] 15% 概率：高通
waveform_3d = self.highpass(waveform_3d, sample_rate=self.sample_rate)
elif choice < 0.95: # [0.85 - 0.95] 10% 概率：走廊混响 (IR)
# 使用你测试过最好的 0.8/0.2 比例
dry = waveform_3d.clone()
wet = self.apply_ir(waveform_3d, sample_rate=self.sample_rate)
waveform_3d = 0.8 * dry + 0.2 * wet
else: # [0.95 - 1.00] 5% 概率：极性翻转
waveform_3d = self.polarity_inversion(waveform_3d, sample_rate=self.sample_rate)
# [1, 1, time] -> [1, time]
waveform = waveform_3d.squeeze(0)
@@ -201,7 +204,7 @@ class CommonVoiceDataset(Dataset[BatchItem]):
return waveform
def _stretch_or_compress(self, waveform: Tensor) -> Tensor:
speed_factor = random.uniform(0.85, 1.4) # (Speed Change: 0.85x - 1.4x)
speed_factor = random.uniform(0.80, 1.6) # (Speed Change: 0.85x - 1.4x)
spec = torch.stft(
waveform.squeeze(0),
n_fft=400,
@@ -246,7 +249,6 @@ class CommonVoiceDataset(Dataset[BatchItem]):
def __getitem__(self, index) -> BatchItem:
row: TsvFormat = self.data.iloc[index]
audio_path: Path = self.audio_dir / row['path']
# text: str = unicodedata.normalize('NFC', row['sentence'].strip())
text: str = normalize_extended_uyghur_characters(collapse_spaces(row['sentence'].strip()))
waveform = self._load_audio(audio_path=audio_path)
@@ -297,8 +299,8 @@ def collate_fn(items: List[BatchItem]) -> Batch:
)
def create_dataloader(tsv_path: Path, audio_dir: Path, noise_dir: Path, tokenizer: ASRTokenizer, batch_size: int = 8, shuffle: bool = True, augment: bool = True, augment_prob: int = 0.5) -> DataLoader:
dataset = CommonVoiceDataset(tsv_path=tsv_path, audio_dir=audio_dir, noise_dir=noise_dir, tokenizer=tokenizer, augment=augment, augment_prob=augment_prob)
def create_dataloader(tsv_path: Path, audio_dir: Path, noise_dir: Path, corridor_noise_dir: Path, tokenizer: ASRTokenizer, batch_size: int = 8, shuffle: bool = True, augment: bool = True, augment_prob: int = 0.5) -> DataLoader:
dataset = CommonVoiceDataset(tsv_path=tsv_path, audio_dir=audio_dir, noise_dir=noise_dir, corridor_noise_dir=corridor_noise_dir, tokenizer=tokenizer, augment=augment, augment_prob=augment_prob)
return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn, pin_memory=True, num_workers=8, prefetch_factor=8, persistent_workers=True)
@@ -331,101 +333,4 @@ if __name__ == "__main__":
print(f"Target lengths: {batch['target_lengths']}")
print(f"Target texts: {batch['target_texts']}")
print(f"Audio paths: {batch['audio_paths']}")
break
import torch
import torchaudio
from silero_vad import load_silero_vad, read_audio, get_speech_timestamps
# --- 1. 加载 VAD 模型 ---
# 推荐使用 GPU (cuda) 如果可用，否则使用 CPU (cpu)
vad_model = load_silero_vad(source="local", force_onnx=False)
def process_audio_with_vad_and_asr(audio_path, inference_module):
"""
使用 Silero VAD 分割音频，然后对每个语音片段进行 ASR 转录。
Args:
audio_path (str): 输入音频文件的路径。
inference_module: 您封装了 transcribe 方法的模块或对象。
需要确保其 _load_audio 和 transcribe 方法可用。
"""
print(f"正在加载音频: {audio_path}")
# --- 2. 加载音频用于 VAD 检测 ---
# Silero VAD 推荐使用 16kHz 采样率
waveform_vad, sample_rate_vad = torchaudio.load(audio_path)
if sample_rate_vad != 16000:
# 如果采样率不是16kHz，需要重采样以供VAD使用
resampler = torchaudio.transforms.Resample(orig_freq=sample_rate_vad, new_freq=16000)
waveform_vad = resampler(waveform_vad)
sample_rate_vad = 16000
# --- 3. 获取语音时间段 ---
# get_speech_timestamps 返回一个列表，每个元素是 {'start': start_sample, 'end': end_sample} 的字典
# 采样率是16000，所以时间戳单位是 1/16000 秒
speech_timestamps = get_speech_timestamps(
waveform_vad,
vad_model,
sampling_rate=sample_rate_vad,
threshold=0.5, # 可以根据需要调整阈值
min_speech_duration_ms=250, # 最小语音持续时间，防止短噪音被误判
max_speech_duration_s=float('inf'), # 最大语音持续时间，float('inf') 表示不限制
min_silence_duration_ms=100, # 最小静音间隔，用于分割语音块
window_size_samples=1536, # VAD窗口大小
speech_pad_ms=30 # 在语音块前后添加的填充时间
)
print(f"检测到 {len(speech_timestamps)} 个语音片段。")
full_text = ""
for i, ts in enumerate(speech_timestamps):
start_sample = int(ts['start'])
end_sample = int(ts['end'])
# 计算时间戳（秒）
start_time = start_sample / sample_rate_vad
end_time = end_sample / sample_rate_vad
print(f"\n处理第 {i+1} 个片段: 时间范围 [{start_time:.2f}s - {end_time:.2f}s]")
# --- 4. 从原始音频中提取此片段 ---
# 注意：这里假设您的 transcribe 函数可以接受 waveform 张量。
# 我们需要从原始可能不同采样率的音频中提取片段，或者用VAD处理过的waveform。
# 为了匹配您原来的 _load_audio 方式，我们用 torchaudio 再次精确加载片段。
# 计算原始音频中的样本索引（如果原始音频采样率与VAD不同）
original_waveform, original_sr = torchaudio.load(audio_path)
if sample_rate_vad != original_sr:
# 如果VAD和原始音频采样率不同，需要重新映射索引
start_idx_orig = int(start_sample * (original_sr / sample_rate_vad))
end_idx_orig = int(end_sample * (original_sr / sample_rate_vad))
else:
start_idx_orig = start_sample
end_idx_orig = end_sample
segment_waveform = original_waveform[:, start_idx_orig:end_idx_orig]
# --- 5. 对该片段进行 ASR 转录 ---
# 这里调用您原有的 transcribe 方法
try:
# 假设 transcribe 方法接受一个 waveform tensor
segment_text = inference_module.transcribe(waveform=segment_waveform)
print(f" -> 转录结果: {segment_text}")
full_text += f"[{start_time:.2f}-{end_time:.2f}s] {segment_text}\n"
except Exception as e:
print(f" -> 转录第 {i+1} 个片段时出错: {e}")
print("\n--- 完整转录结果 ---")
print(full_text)
# --- 使用示例 ---
# 假设您有一个名为 'inference' 的模块对象，它有 _load_audio 和 transcribe 方法
# audio_path = "your_audio_file.wav"
# process_audio_with_vad_and_asr(audio_path, inference)
break