AzALlN/whisperX

1import whisperx as whisper2 3from deep_translator import GoogleTranslator4import os5from whisperx.utils import write_vtt, write_srt, write_ass, write_tsv, write_txt6 7 8def detect_language(filename, model):9    # load audio and pad/trim it to fit 30 seconds10    audio = whisper.load_audio(file=filename)11    audio = whisper.pad_or_trim(audio)12    # make log-Mel spectrogram and move to the same device as the model13    mel = whisper.log_mel_spectrogram(audio).to(model.device)14    _, probs = model.detect_language(mel)15    print(f"Detected language: {max(probs, key=probs.get)}")16    return {"detected_language": max(probs, key=probs.get)}17 18 19def translate_to_english(transcription, json=False):20    if json:21        for text in transcription:22            text["text"] = GoogleTranslator(source="auto", target="en").translate(23                text["text"]24            )25    else:26 27        for text in transcription["segments"]:28            text["text"] = GoogleTranslator(source="auto", target="en").translate(29                text["text"]30            )31    return transcription32 33 34def write(filename, dtype, result_aligned):35 36    if dtype == "vtt":37        with open(38            os.path.join(".", os.path.splitext(filename)[0] + ".vtt"),39            "w",40            encoding="utf-8",41        ) as vtt:42            write_vtt(result_aligned["segments"], file=vtt)43    if dtype == "srt":44        with open(45            os.path.join(".", os.path.splitext(filename)[0] + ".srt"),46            "w",47            encoding="utf-8",48        ) as srt:49            write_srt(result_aligned["segments"], file=srt)50    if dtype == "ass":51        with open(52            os.path.join(".", os.path.splitext(filename)[0] + ".ass"),53            "w",54            encoding="utf-8",55        ) as ass:56            write_ass(result_aligned["segments"], file=ass)57    if dtype == "tsv":58        with open(59            os.path.join(".", os.path.splitext(filename)[0] + ".tsv"),60            "w",61            encoding="utf-8",62        ) as tsv:63            write_tsv(result_aligned["segments"], file=tsv)64    if dtype == "plain text":65        print("here")66        print(filename)67        with open(68            os.path.join(".", os.path.splitext(filename)[0] + ".txt"),69            "w",70            encoding="utf-8",71        ) as txt:72            write_txt(result_aligned["segments"], file=txt)73 74 75def read(filename, transc):76    if transc == "plain text":77        transc = "txt"78    filename = filename.split(".")[0]79    print(filename)80    with open(f"{filename}.{transc}", encoding="utf-8") as f:81        content = f.readlines()82    content = " ".join(z for z in content)83    return content84 85 86from constants import language_dict87 88 89def get_key(val):90    for key, value in language_dict.items():91        if val == value:92            return key93    return "Key not found"94 95 96