Upload playground_yeet_the_unclean.py

playground_yeet_the_unclean.py

@@ -0,0 +1,212 @@
+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()