AR/exps/get_beats_librilight.py

# 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()