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

feat: change waveform

Browse Source
This commit is contained in:
Mesrulla
2026-05-07 11:29:21 +06:00
commit d31233a79a
21 changed files with 5330 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/dataset.py
+431
View File
@@ -0,0 +1,431 @@
import random
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 torchaudio.transforms import Resample, TimeStretch
from pathlib import Path
import pandas as pd
from typing import List, TypedDict
from handle.text_normalizer import collapse_spaces, normalize_extended_uyghur_characters
from tokenizer import ASRTokenizer
# 单个样本的数据结构（Dataset.__getitem__ 返回）
class BatchItem(TypedDict):
waveform: Tensor # [time]
target_ids: Tensor # [seq_len] 目标文本的token IDs
target_text: str # 原始文本
audio_path: str # 音频文件路径
# 批量数据的数据结构（collate_fn 返回，DataLoader 输出）
class Batch(TypedDict):
waveforms: Tensor # [batch, time]
targets: Tensor # [batch, max_len] padding后的目标IDs
waveform_lengths: Tensor # [batch] 每个样本的实际Waveform长度
target_lengths: Tensor # [batch] 每个样本的实际目标长度
target_texts: List[str] # [batch] 原始文本列表
audio_paths: List[str] # [batch] 音频路径列表
class TsvFormat(TypedDict):
client_id: str
path: str
sentence: str
up_votes: int
down_votes: int
age: str
gender: str
locale: str
class NoiseAugmentor:
def __init__(self, noise_root: Path, sample_rate: int=16000):
self.sample_rate = sample_rate
self.noise_files = list(Path(noise_root).rglob("*.wav"))
def apply_real_noise(self, waveform: Tensor):
# 1. 随机选一个噪音文件
noise_path = random.choice(self.noise_files)
noise_waveform, sr = torchaudio.load_with_torchcodec(noise_path)
# Resample to target sample rate.
if sr != self.sample_rate:
noise_waveform = Resample(sr, self.sample_rate)(noise_waveform)
# Convert to mono if it is setro.
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0, keepdim=True)
# 3. 截取或填充，使其长度与语音一致
sig_len = waveform.shape[1]
noise_len = noise_waveform.shape[1]
if noise_len >= sig_len:
# 随机截取一段
start = random.randint(0, noise_len - sig_len)
noise_waveform = noise_waveform[:, start:start + sig_len]
else:
full_noise = torch.zeros_like(waveform)
start = random.randint(0, sig_len - noise_len)
full_noise[:, start : start + noise_len] = noise_waveform
noise_waveform = full_noise
# 4. 设定随机信噪比 SNR (5dB 到 20dB)
snr_db = random.uniform(5, 20)
# 5. 混合
return self._mix_at_snr(waveform, noise_waveform, snr_db)
def _mix_at_snr(self, signal: Tensor, noise: Tensor, snr_db: float):
s_p = signal.pow(2).mean()
n_p = noise.pow(2).mean()
snr_linear = 10**(snr_db / 10)
scale = torch.sqrt(s_p / (n_p * snr_linear + 1e-8))
noisy = signal + scale * noise
# 归一化，防止溢出
return noisy / (noisy.abs().max() + 1e-8)
class CommonVoiceDataset(Dataset[BatchItem]):
def __init__(
self,
tsv_path: Path,
audio_dir: Path,
noise_dir: Path,
tokenizer: ASRTokenizer,
sample_rate: int = 16000,
max_audio_len: int = 480000, # 30秒 @ 16kHz
augment: bool = True,
augment_prob: float = 0.5,
) -> None:
super().__init__()
self.noise_augmentor = NoiseAugmentor(noise_root=noise_dir, sample_rate=sample_rate)
self.audio_dir = Path(audio_dir)
self.tokenizer = tokenizer
self.sample_rate = sample_rate
self.max_audio_len = max_audio_len
self.augment = augment
self.augment_prob = augment_prob
self.data: pd.DataFrame = pd.read_csv(tsv_path, sep='\t')
valid_indices = []
for index, row in self.data.iterrows():
audio_path: Path = self.audio_dir / row['path']
if audio_path.exists():
valid_indices.append(index)
self.data = self.data.loc[valid_indices].reset_index(drop=True)
self.gain_up = Gain(min_gain_in_db=4, max_gain_in_db=8, p=1.0, output_type='tensor')
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")
def __len__(self):
return len(self.data)
def _load_audio(self, audio_path: Path) -> Tensor:
waveform, sample_rate = torchaudio.load_with_torchcodec(audio_path)
# Resample to target sample rate.
if sample_rate != self.sample_rate:
waveform = Resample(sample_rate, self.sample_rate)(waveform)
# Convert to mono if it is setro.
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0, keepdim=True)
# Normalization
waveform = waveform / waveform.abs().max()
# Clip waveform exceeds from max length.
if waveform.shape[1] > self.max_audio_len:
waveform = waveform[:, :self.max_audio_len]
return waveform
def _augment_waveform(self, waveform: Tensor) -> Tensor:
if not self.augment or random.random() > self.augment_prob:
return waveform
# 1. voice Stretch/Compress
if random.random() < 0.6:
waveform = self._stretch_or_compress(waveform=waveform)
if random.random() < 0.4:
waveform = self.noise_augmentor.apply_real_noise(waveform)
if random.random() < 0.3:
waveform = self._time_mask_waveform(waveform=waveform)
# torch_audiomentations: [1, time] -> [1, 1, time]
if waveform.dim() == 2:
waveform_3d = waveform.unsqueeze(0)
# 随机选择一种频谱增强
choice = random.random()
if choice < 0.15:
# 增益变化（上或下）
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:
# 音高变化（上或下）
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:
# 低通滤波（声音发闷）
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:
# 高通滤波（电话效果）
waveform_3d = self.highpass(waveform_3d, sample_rate=self.sample_rate)
# [1, 1, time] -> [1, time]
waveform = waveform_3d.squeeze(0)
# 防止多次 augment 后振幅溢出，最后归一化
max_amp = waveform.abs().max()
if max_amp > 1.0:
waveform = waveform / max_amp
return waveform
def _stretch_or_compress(self, waveform: Tensor) -> Tensor:
speed_factor = random.uniform(0.85, 1.4) # (Speed Change: 0.85x - 1.4x)
spec = torch.stft(
waveform.squeeze(0),
n_fft=400,
hop_length=160,
window=torch.hann_window(400).to(waveform.device),
return_complex=True
)
# 时间拉伸（不改变音高）
stretched_spec = TimeStretch(hop_length=160, n_freq=spec.shape[-2], fixed_rate=speed_factor)(spec)
# 转回波形
waveform_stretched = torch.istft(stretched_spec, n_fft=400, hop_length=160, window=torch.hann_window(400).to(waveform.device)).unsqueeze(0)
return waveform_stretched
def _time_mask_waveform(self, waveform: Tensor) -> Tensor:
audio_len = waveform.shape[1]
sr = self.sample_rate # 16000
# 设置参数：单次遮盖最长 0.4 秒 (6400个点)
max_mask_time = 0.4
max_mask_samples = int(sr * max_mask_time)
# 根据音频长度决定遮盖次数：
# 比如每 3 秒钟允许遮盖 1 次
num_masks = max(1, audio_len // (sr * 3))
for _ in range(num_masks):
# 每次随机遮盖 0.1s 到 0.4s
current_mask_len = random.randint(int(sr * 0.1), max_mask_samples)
if audio_len > current_mask_len:
start_pos = random.randint(0, audio_len - current_mask_len)
# 填充微小噪音（模拟环境底噪）
noise = torch.randn(1, current_mask_len).to(waveform.device) * 0.002
waveform[:, start_pos : start_pos + current_mask_len] = noise
return waveform
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)
waveform = self._augment_waveform(waveform)
waveform = waveform.squeeze(0)
return BatchItem(
waveform=waveform,
target_ids=torch.tensor(self.tokenizer.encode(text), dtype=torch.long),
target_text=text,
audio_path=str(audio_path)
)
def collate_fn(items: List[BatchItem]) -> Batch:
max_waveform_len = max(item['waveform'].shape[0] for item in items)
max_target_len = max(len(item['target_ids']) for item in items)
batch_size = len(items)
waveforms = torch.zeros(batch_size, max_waveform_len)
targets = torch.zeros(batch_size, max_target_len, dtype=torch.long)
waveform_lengths = torch.zeros(batch_size, dtype=torch.long)
target_lengths = torch.zeros(batch_size, dtype=torch.long)
target_texts = []
audio_paths = []
for i, item in enumerate(items):
waveform_len = item['waveform'].shape[0]
target_len = len(item['target_ids'])
waveforms[i, :waveform_len] = item['waveform']
targets[i, :target_len] = item['target_ids']
waveform_lengths[i] = waveform_len
target_lengths[i] = target_len
target_texts.append(item['target_text'])
audio_paths.append(item['audio_path'])
return Batch(
waveforms=waveforms,
targets=targets,
waveform_lengths=waveform_lengths,
target_lengths=target_lengths,
target_texts=target_texts,
audio_paths=audio_paths
)
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)
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)
# ============ 测试代码 ============
if __name__ == "__main__":
from pathlib import Path
workspace_dir = Path(__file__).parent.parent
# 初始化tokenizer
tokenizer = ASRTokenizer(workspace_dir / 'config' / 'asr_vocab.json')
# 创建数据加载器
dataloader = create_dataloader(
tsv_path=workspace_dir / 'data' / 'ug' / 'train.tsv',
audio_dir=workspace_dir / 'data' / 'ug' / 'clips',
tokenizer=tokenizer,
batch_size=2,
shuffle=True,
)
# 测试加载一个batch
print("测试数据加载:")
for batch in dataloader:
batch: Batch
print(f"Mel specs shape: {batch['mel_specs'].shape}")
print(f"Targets shape: {batch['targets'].shape}")
print(f"Mel lengths: {batch['mel_lengths']}")
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)