GPT-SoVITS-experiment / AR /exps /get_beats_librilight.py
Ricecake123's picture
first commit
e79b770
raw
history blame
11.8 kB
# Use AudioTag tool BEATs to filter out audios who's top1 tag is not 'speech'
# non_speech.npy, 存储一个 python dict 表示非 speech 类型的音频的 tag, 更小,加载和搜索速度更快
# audio_tag 目录存储 {utt_id}.txt, 第一行是小写的 top1 tag
import argparse
import os
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import librosa
import numpy as np
import torch
import tqdm
from AR.exps.beats.BEATs import BEATs
from AR.exps.beats.BEATs import BEATsConfig
from AR.exps.beats.config import id_name_dict
from soundstorm.s2.exps.hubert.feature_utils import get_shard_range
from soundstorm.utils import check_txt_file
def get_BEATs_top1(wav,
BEATs_model,
BEATs_label_dict,
device: str='cpu',
topk: int=1):
wav = torch.tensor(wav).unsqueeze(0).to(device)
padding_mask = torch.zeros(wav.shape).bool().to(device)
probs = BEATs_model.extract_features(wav, padding_mask=padding_mask)[0]
# 单条推理
probs = probs[0]
topk_label_prob, topk_label_idx = probs.topk(k=topk)
topk_label = [
BEATs_label_dict[label_idx.item()] for label_idx in topk_label_idx
]
topk_label_name = [id_name_dict[label] for label in topk_label]
top1_label = topk_label_name[0]
return top1_label
def process_sentence(args,
fp: Path,
train_dump_dir: Path,
dev_dump_dir: Path,
test_dump_dir: Path,
VAD_dict,
BEATs_model,
BEATs_label_dict,
device: str='cpu'):
utt_id = fp.stem
sr = args.sr
record = []
train_audio_tag_dir = train_dump_dir / "audio_tag"
train_audio_tag_dir.mkdir(parents=True, exist_ok=True)
dev_audio_tag_dir = dev_dump_dir / "audio_tag"
dev_audio_tag_dir.mkdir(parents=True, exist_ok=True)
test_audio_tag_dir = test_dump_dir / "audio_tag"
test_audio_tag_dir.mkdir(parents=True, exist_ok=True)
try:
# get info for path
wav_path_list = str(fp).strip().split('/')
sub_dataset, spk_id, book_name = wav_path_list[-4], wav_path_list[
-3], wav_path_list[-2]
wav_name = wav_path_list[-1][:-5]
assert wav_name == utt_id
# key_name for big wav
key_name = f'{wav_name}#{sub_dataset}#{spk_id}#{book_name}'
# 判断 VAD 字典中不存在该条音频信息的情况
if key_name not in VAD_dict.keys():
print(key_name, 'not in VAD_dict !')
return record
wav = None
sorted_split_VAD_dict = sorted(VAD_dict[key_name].items())
len_dict = len(sorted_split_VAD_dict)
for index, item in enumerate(sorted_split_VAD_dict):
split_name, value = item
start, end = value
# train | dev | test
if index == len_dict - 1:
subset = 'test'
audio_tag_path = test_audio_tag_dir / (split_name + ".txt")
elif index == len_dict - 2:
subset = 'dev'
audio_tag_path = dev_audio_tag_dir / (split_name + ".txt")
else:
subset = 'train'
audio_tag_path = train_audio_tag_dir / (split_name + ".txt")
if os.path.exists(audio_tag_path) and check_txt_file(
audio_tag_path):
# print(audio_tag_path, 'exits!')
pass
else:
# 这里加判断保证在 sub wav 的循环中只 load 一次
if wav is None:
# load big wav
# 在最外层 load 如果 sub wav 的特征都存在了就会白白消耗 load 的时间
wav, _ = librosa.load(str(fp), sr=sr)
sub_wav = wav[int(start * sr):int(end * sr)]
audio_tag_top1 = get_BEATs_top1(
wav=sub_wav,
BEATs_model=BEATs_model,
BEATs_label_dict=BEATs_label_dict,
device=device)
with open(audio_tag_path, 'w') as f:
f.write(audio_tag_top1)
sub_record = {
"utt_id": split_name,
"audio_tag_path": audio_tag_path,
"subset": subset
}
# recodrd 变成 List of Dict
record.append(sub_record)
except Exception:
print("occur Exception")
traceback.print_exc()
# record 有可能是一个不完整的 List
return record
return record
def process_sentences(args,
fps: Path,
train_dump_dir: Path,
dev_dump_dir: Path,
test_dump_dir: Path,
VAD_dict,
BEATs_model,
BEATs_label_dict,
device: str='cpu',
nprocs: int=1):
print("nprocs:", nprocs)
if nprocs == 1:
results = []
for fp in tqdm.tqdm(fps, total=len(fps)):
record = process_sentence(
args=args,
fp=fp,
train_dump_dir=train_dump_dir,
dev_dump_dir=dev_dump_dir,
test_dump_dir=test_dump_dir,
VAD_dict=VAD_dict,
BEATs_model=BEATs_model,
BEATs_label_dict=BEATs_label_dict,
device=device)
if record:
results.append(record)
else:
with ThreadPoolExecutor(nprocs) as pool:
futures = []
with tqdm.tqdm(total=len(fps)) as progress:
for fp in fps:
future = pool.submit(process_sentence, args, fp,
train_dump_dir, dev_dump_dir,
test_dump_dir, VAD_dict, BEATs_model,
BEATs_label_dict, device)
future.add_done_callback(lambda p: progress.update())
futures.append(future)
results = []
for ft in futures:
record = ft.result()
if record:
results.append(record)
# torch.save() to a large `.pth` file
non_speech_dict = dict()
non_speech_dict['train'] = {}
non_speech_dict['dev'] = {}
non_speech_dict['test'] = {}
# record 是 List of Dict, 一条大 wav 一个 record,一条小 wav 一个 sub_recored
print(f"start to save {args.rank}_{args.nshard}.npy ...")
save_start_time = time.time()
for record in tqdm.tqdm(results, total=len(results), colour='green'):
for sub_record in record:
# 这里加 try, 因为 txt 文件可能损坏
try:
utt_id = sub_record["utt_id"]
subset = sub_record["subset"]
audio_tag_top1 = check_txt_file(sub_record["audio_tag_path"])
if audio_tag_top1 is not False:
if 'speech' not in audio_tag_top1.lower():
non_speech_dict[subset][utt_id] = audio_tag_top1
else:
# print(f'audio tag result of {utt_id} is speech')
pass
else:
print(f'audio tag result of {utt_id} is False')
except Exception:
print(f"{utt_id} occur Exception")
traceback.print_exc()
continue
train_filename = train_dump_dir / f'non_speech_{args.rank}_{args.nshard}.npy'
dev_filename = dev_dump_dir / f'non_speech_{args.rank}_{args.nshard}.npy'
test_filename = test_dump_dir / f'non_speech_{args.rank}_{args.nshard}.npy'
np.save(train_filename, non_speech_dict['train'])
print(f"npy file '{train_filename}' write down")
np.save(dev_filename, non_speech_dict['dev'])
print(f"npy file '{dev_filename}' write down")
np.save(test_filename, non_speech_dict['test'])
print(f"npy file '{test_filename}' write down")
print('time of save stage:', time.time() - save_start_time)
def main():
# parse config and args
parser = argparse.ArgumentParser(
description="Use AudioTag tool BEATs to filter out audios who's top1 tag is not 'speech'."
)
parser.add_argument(
"--data_dir", default=None, type=str, help="directory to dataset.")
parser.add_argument(
"--dump_dir",
type=str,
required=True,
help="directory to dump feature files.")
parser.add_argument(
"--num-cpu", type=int, default=1, help="number of process.")
parser.add_argument(
'--sr', type=int, default=16000, help='sample rate of model')
# For LibriLight dataset
parser.add_argument(
"--sub_dataset",
default="small",
type=str,
help="name of sub dataset of LibriLight",
choices=['small', 'medium', 'large', 'duplicate'], )
parser.add_argument(
"--VAD_path", type=str, default='./VAD/librilight_segment_dict.npy')
parser.add_argument("--nshard", type=int, default=3)
parser.add_argument("--rank", type=int, default=0)
# for BEATs
parser.add_argument(
"--BEATs_ckpt_path",
type=str,
default='./pretrained_model/BEATs_iter1_finetuned_on_AS2M_cpt1.pt')
args = parser.parse_args()
data_dir = Path(args.data_dir).expanduser()
dump_dir = Path(args.dump_dir).expanduser()
# use absolute path
dump_dir = dump_dir.resolve()
dump_dir.mkdir(parents=True, exist_ok=True)
assert data_dir.is_dir()
# sub_dataset here
sub_dataset_dir = data_dir / args.sub_dataset
# olny spk_id in list, sort by lexicographical order
speaker_list = sorted(os.listdir(sub_dataset_dir))
start, end = get_shard_range(len(speaker_list), args.nshard, args.rank)
# speaker_list for this rank
speaker_list = speaker_list[start:end]
all_wav_files = []
for speaker in speaker_list:
wav_files = sorted(list((sub_dataset_dir / speaker).rglob("*/*.flac")))
# filter out ._*.flac
wav_files = [
file for file in wav_files if not file.name.startswith('._')
]
all_wav_files += wav_files
print(f"num of wav files in rank {args.rank}:", len(all_wav_files))
# get VAD info
VAD_dict = np.load(args.VAD_path, allow_pickle=True).item()
sub_dataset_dump_dir = dump_dir / args.sub_dataset
sub_dataset_dump_dir.mkdir(parents=True, exist_ok=True)
train_dump_dir = sub_dataset_dump_dir / "train"
train_dump_dir.mkdir(parents=True, exist_ok=True)
dev_dump_dir = sub_dataset_dump_dir / "dev"
dev_dump_dir.mkdir(parents=True, exist_ok=True)
test_dump_dir = sub_dataset_dump_dir / "test"
test_dump_dir.mkdir(parents=True, exist_ok=True)
BEATs_ckpt = torch.load(args.BEATs_ckpt_path)
BEATs_cfg = BEATsConfig(BEATs_ckpt['cfg'])
BEATs_model = BEATs(BEATs_cfg)
BEATs_model.load_state_dict(BEATs_ckpt['model'])
BEATs_model.eval()
# cpu or cuda
device = 'cpu'
BEATs_model.to(device)
BEATs_label_dict = BEATs_ckpt['label_dict']
# 每条大 wav 分出一个 dev 一个 test,比例大概是 96:2:2
if all_wav_files:
process_sentences(
args=args,
fps=all_wav_files,
train_dump_dir=train_dump_dir,
dev_dump_dir=dev_dump_dir,
test_dump_dir=test_dump_dir,
VAD_dict=VAD_dict,
BEATs_model=BEATs_model,
BEATs_label_dict=BEATs_label_dict,
device=device,
nprocs=args.num_cpu)
if __name__ == "__main__":
main()