File size: 5,586 Bytes
d8bfe4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import os
import glob
import torch
import torchaudio
import torchaudio.transforms as T
from tqdm import tqdm
def merge_audio_data(root_dir, output_wav_root, output_tsv_path):
"""
遍历 root_dir 下的 lst 文件,将 prompt 和 target 音频拼接。
输出音频到 output_wav_root,输出元数据到 output_tsv_path。
"""
# 1. 准备工作
sub_dirs = ['en', 'zh']
silence_duration = 0.1 # 拼接中间插入 0.3秒 静音
# 创建输出目录
os.makedirs(output_wav_root, exist_ok=True)
tsv_lines = []
print(f"Scanning root directory: {root_dir}")
# 2. 收集 .lst 文件
lst_files = []
for lang in sub_dirs:
lang_dir = os.path.join(root_dir, lang)
if not os.path.exists(lang_dir):
continue
curr_lst_files = glob.glob(os.path.join(lang_dir, "*.lst"))
lst_files.extend(curr_lst_files)
print(f"Found {len(lst_files)} .lst files.")
success_count = 0
for lst_path in tqdm(lst_files, desc="Merging Audios"):
# 获取当前是 en 还是 zh
# 假设路径结构是 .../en/xxx.lst,所以 dirname 是 .../en,再 basename 就是 en
lang_category = os.path.basename(os.path.dirname(lst_path))
# 对应的输出子目录
save_dir = os.path.join(output_wav_root, lang_category)
os.makedirs(save_dir, exist_ok=True)
base_dir = os.path.dirname(lst_path)
with open(lst_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if not line:
continue
parts = line.split('|')
if len(parts) < 4:
continue
utt_id = parts[0]
prompt_text = parts[1]
prompt_rel_path = parts[2]
target_text = parts[3]
# -------------------------------------------------
# ✅ 文本拼接逻辑修改
# -------------------------------------------------
if lang_category == 'zh':
# 中文模式:直接拼接,不加空格
full_text = f"{prompt_text}{target_text}"
else:
# 英文模式:中间加空格
full_text = f"{prompt_text} {target_text}"
# -------------------------------------------------
# 获取路径
# -------------------------------------------------
prompt_abs_path = os.path.join(base_dir, prompt_rel_path)
target_abs_path = os.path.join(base_dir, "wavs", f"{utt_id}.wav")
new_filename = f"{utt_id}_merged.wav"
new_file_path = os.path.join(save_dir, new_filename)
if not (os.path.exists(prompt_abs_path) and os.path.exists(target_abs_path)):
continue
try:
# -------------------------------------------------
# 音频处理
# -------------------------------------------------
# load 默认返回 float32 类型
wav_p, sr_p = torchaudio.load(prompt_abs_path)
wav_t, sr_t = torchaudio.load(target_abs_path)
# 统一通道数 (转为单声道)
if wav_p.shape[0] > 1: wav_p = torch.mean(wav_p, dim=0, keepdim=True)
if wav_t.shape[0] > 1: wav_t = torch.mean(wav_t, dim=0, keepdim=True)
# 统一采样率 (以 prompt 的 SR 为基准,通常是 16k, 22k 或 24k)
target_sr = sr_p
if sr_t != target_sr:
resampler = T.Resample(sr_t, target_sr)
wav_t = resampler(wav_t)
# 创建静音片段
silence_samples = int(silence_duration * target_sr)
silence_wav = torch.zeros(1, silence_samples)
# 拼接
merged_wav = torch.cat([wav_p, silence_wav, wav_t], dim=1)
# 保存
torchaudio.save(new_file_path, merged_wav, target_sr)
# 计算总时长
total_duration = merged_wav.shape[1] / target_sr
# 记录 (路径 \t 文本 \t 时长 \t 采样率)
tsv_lines.append(f"{new_file_path}\t{full_text}\t{total_duration:.4f}\t{target_sr}")
success_count += 1
except Exception as e:
print(f"[Error] Failed to process {utt_id}: {e}")
continue
# 4. 写入 TSV
print(f"Saving metadata to {output_tsv_path}...")
with open(output_tsv_path, 'w', encoding='utf-8') as f:
for line in tsv_lines:
f.write(line + "\n")
print(f"Done. Processed {success_count} files.")
if __name__ == "__main__":
# 配置你的路径
INPUT_ROOT = "/workspace/echoloc/datas/ntu_enzh_speech/"
OUTPUT_WAV_ROOT = "/workspace/echoloc/datas/ntu_enzh_speech/merged_reconstruction/wavs"
OUTPUT_TSV_PATH = "/workspace/echoloc/datas/ntu_enzh_speech/merged_reconstruction/test_merged.tsv"
merge_audio_data(INPUT_ROOT, OUTPUT_WAV_ROOT, OUTPUT_TSV_PATH)
|