playground_yeet_the_unclean.py

import multiprocessing
import os
import shutil

import librosa as lb
import numpy as np
import soundfile as sf
from deepmultilingualpunctuation import PunctuationModel
from pyannote.audio import Pipeline
from rpunct import RestorePuncts
from tqdm import tqdm


class UncleanYeeter:
    def __init__(self):
        """
        all the models and persistent stuff
        """
        self.diarizer = Pipeline.from_pretrained("pyannote/speaker-diarization@2.1")

    def create_list_of_samples_marked_for_deletion(self, list_of_audios):
        marked_for_yeeting = list()
        for audio_file in tqdm(list_of_audios):
            try:
                wav, sr = sf.read(audio_file)
            except RuntimeError:
                print(f"PROBLEMATIC FILE: {audio_file}")
                continue
            wav = to_mono(wav)

            # check duration
            if 5 < len(wav) / sr < 15:
                continue

            # check SNR
            if wada_snr(wav) < 20.0:
                continue

            # check amount of speakers
            try:
                output = self.diarizer(audio_file)
            except ValueError:
                print("Diarizer is unhappy")
                continue
            speakers = set()
            for _, _, speaker in output.itertracks(yield_label=True):
                speakers.add(speaker)
            if len(speakers) > 1:
                continue

            marked_for_yeeting.append(audio_file.split("/")[-1])

        # save list of files to be yoten to a file for later yeeting
        with open("files_to_keep.txt", "a", encoding="utf8") as file:
            file.write("\n".join(marked_for_yeeting) + "\n")
        print(marked_for_yeeting)


class Punctuator:
    def __init__(self, lang="eng"):
        if lang == "en":
            model = RestorePuncts()
            self.punctuate_transcripts = model.punctuate  # pass a string into it and you get a punctuated string returned
        else:
            model = PunctuationModel()
            self.punctuate_transcripts = model.restore_punctuation  # pass a string into it and you get a punctuated string returned


def wada_snr(wav):
    # Direct blind estimation of the SNR of a speech signal.
    #
    # Paper on WADA SNR:
    #   http://www.cs.cmu.edu/~robust/Papers/KimSternIS08.pdf
    #
    # This function was adapted from this matlab code:
    #   https://labrosa.ee.columbia.edu/projects/snreval/#9

    # init
    eps = 1e-10
    # next 2 lines define a fancy curve derived from a gamma distribution -- see paper
    db_vals = np.arange(-20, 101)
    g_vals = np.array(
        [0.40974774, 0.40986926, 0.40998566, 0.40969089, 0.40986186, 0.40999006, 0.41027138, 0.41052627, 0.41101024, 0.41143264, 0.41231718, 0.41337272, 0.41526426, 0.4178192, 0.42077252, 0.42452799, 0.42918886, 0.43510373, 0.44234195, 0.45161485, 0.46221153, 0.47491647, 0.48883809, 0.50509236, 0.52353709, 0.54372088, 0.56532427,
         0.58847532, 0.61346212, 0.63954496, 0.66750818, 0.69583724, 0.72454762, 0.75414799, 0.78323148, 0.81240985, 0.84219775, 0.87166406, 0.90030504, 0.92880418, 0.95655449, 0.9835349, 1.01047155, 1.0362095, 1.06136425, 1.08579312, 1.1094819, 1.13277995, 1.15472826, 1.17627308, 1.19703503, 1.21671694, 1.23535898, 1.25364313,
         1.27103891, 1.28718029, 1.30302865, 1.31839527, 1.33294817, 1.34700935, 1.3605727, 1.37345513, 1.38577122, 1.39733504, 1.40856397, 1.41959619, 1.42983624, 1.43958467, 1.44902176, 1.45804831, 1.46669568, 1.47486938, 1.48269965, 1.49034339, 1.49748214, 1.50435106, 1.51076426, 1.51698915, 1.5229097, 1.528578, 1.53389835, 1.5391211,
         1.5439065, 1.54858517, 1.55310776, 1.55744391, 1.56164927, 1.56566348, 1.56938671, 1.57307767, 1.57654764, 1.57980083, 1.58304129, 1.58602496, 1.58880681, 1.59162477, 1.5941969, 1.59693155, 1.599446, 1.60185011, 1.60408668, 1.60627134, 1.60826199, 1.61004547, 1.61192472, 1.61369656, 1.61534074, 1.61688905, 1.61838916, 1.61985374,
         1.62135878, 1.62268119, 1.62390423, 1.62513143, 1.62632463, 1.6274027, 1.62842767, 1.62945532, 1.6303307, 1.63128026, 1.63204102])

    # peak normalize, get magnitude, clip lower bound
    wav = np.array(wav)
    wav = wav / abs(wav).max()
    abs_wav = abs(wav)
    abs_wav[abs_wav < eps] = eps

    # calcuate statistics
    # E[|z|]
    v1 = max(eps, abs_wav.mean())
    # E[log|z|]
    v2 = np.log(abs_wav).mean()
    # log(E[|z|]) - E[log(|z|)]
    v3 = np.log(v1) - v2

    # table interpolation
    wav_snr_idx = None
    if any(g_vals < v3):
        wav_snr_idx = np.where(g_vals < v3)[0].max()
    # handle edge cases or interpolate
    if wav_snr_idx is None:
        wav_snr = db_vals[0]
    elif wav_snr_idx == len(db_vals) - 1:
        wav_snr = db_vals[-1]
    else:
        wav_snr = db_vals[wav_snr_idx] + \
                  (v3 - g_vals[wav_snr_idx]) / (g_vals[wav_snr_idx + 1] - g_vals[wav_snr_idx]) * (db_vals[wav_snr_idx + 1] - db_vals[wav_snr_idx])

    # Calculate SNR
    dEng = sum(wav ** 2)
    dFactor = 10 ** (wav_snr / 10)
    dNoiseEng = dEng / (1 + dFactor)  # Noise energy
    dSigEng = dEng * dFactor / (1 + dFactor)  # Signal energy
    snr = 10 * np.log10(dSigEng / dNoiseEng)

    return snr


def to_mono(x):
    """
    make sure we deal with a 1D array
    """
    if len(x.shape) == 2:
        return lb.to_mono(np.transpose(x))
    else:
        return x


def clean_mls_ger():
    clean_mls("mls_german", "de")


def clean_mls_fr():
    clean_mls("mls_french", "fr")


def clean_mls_it():
    clean_mls("mls_italian", "it")


def clean_mls_eng():
    clean_mls("mls_english", "en")


def clean_mls(lang_dir, lang):
    punco = Punctuator(lang=lang)
    new_file = ""
    shutil.copy(f"/mount/resources/speech/corpora/MultiLingLibriSpeech/{lang_dir}/train/transcripts.txt", f"/mount/resources/speech/corpora/MultiLingLibriSpeech/{lang_dir}/train/orig_transcripts.txt")
    with open(f"/mount/resources/speech/corpora/MultiLingLibriSpeech/{lang_dir}/train/transcripts.txt", "r", encoding="utf8") as file:
        sentence_list = file.read().split("\n")
    for sentence in tqdm(sentence_list):
        if sentence.strip() == "":
            continue
        sent_id = sentence.split()[0]
        punc_sent = punco.punctuate_transcripts(" ".join(sentence.split()[1:]))
        new_file = new_file + f"{sent_id}\t{punc_sent}\n"
    with open(f"/mount/resources/speech/corpora/MultiLingLibriSpeech/{lang_dir}/train/transcripts.txt", "w", encoding="utf8") as file:
        file.write(new_file)


def build_path_to_transcript_dict_gigaspeech():
    path_to_transcript = dict()
    root = "/mount/resources/speech/corpora/GigaSpeech/"
    with open(os.path.join(root, "transcripts.txt"), "r", encoding="utf8") as file:
        lookup = file.read()
    for line in lookup.split("\n"):
        if line.strip() != "":
            norm_transcript = line.split("\t")[1]
            wav_path = os.path.join(root, "wavs", line.split("\t")[0])
            if os.path.exists(wav_path):
                path_to_transcript[wav_path] = norm_transcript
    return path_to_transcript


def split_list(lst, n):
    if n <= 0:
        return []

    quotient, remainder = divmod(len(lst), n)
    shards = [lst[i * quotient + min(i, remainder):(i + 1) * quotient + min(i + 1, remainder)] for i in range(n)]
    return shards

def yonkus(shard):
    yeet = UncleanYeeter()
    yeet.create_list_of_samples_marked_for_deletion(shard)

if __name__ == '__main__':
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
    print(f"Making GPU {os.environ['CUDA_VISIBLE_DEVICES']} the only visible device.")

    list_of_files = os.listdir("/mount/resources/speech/corpora/GigaSpeech/wavs")
    absolute_list_of_files = list()
    for filo in list_of_files:
        absolute_list_of_files.append(f"/mount/resources/speech/corpora/GigaSpeech/wavs/{filo}")
    processes = list()
    for sublist in split_list(absolute_list_of_files, 20):
        processes.append(multiprocessing.Process(args=(sublist,), target=yonkus, daemon=True))
        processes[-1].start()
    for processo in processes:
        processo.join()
    # clean_mls_it()
    # clean_mls_fr()
    # clean_mls_ger()
    # clean_mls_eng()