From 12065e3124d043cc1ece87026b6e0de0f6193640 Mon Sep 17 00:00:00 2001
From: ganders <g.anders@iwm-tuebingen.de>
Date: Wed, 20 Sep 2023 15:07:25 +0200
Subject: [PATCH] Update: added script to handle data and transcribe.

---
 2023-09-20_transcripts_Video-SRS.py | 14 ++++++++++
 requirements.txt                    |  5 ++++
 src/__init__.py                     |  1 +
 src/data_handling.py                | 13 ++++++++++
 src/transcriber.py                  | 40 +++++++++++++++++++++++++++++
 5 files changed, 73 insertions(+)
 create mode 100644 2023-09-20_transcripts_Video-SRS.py
 create mode 100644 requirements.txt
 create mode 100644 src/__init__.py
 create mode 100644 src/data_handling.py
 create mode 100644 src/transcriber.py

diff --git a/2023-09-20_transcripts_Video-SRS.py b/2023-09-20_transcripts_Video-SRS.py
new file mode 100644
index 0000000..31f9a15
--- /dev/null
+++ b/2023-09-20_transcripts_Video-SRS.py
@@ -0,0 +1,14 @@
+from src.data_handling import download_audio, convert_to_wav
+from src.transcriber import transcribe_wav_file
+
+youtube_links = ['https://www.youtube.com/watch?v=BVVo1BaL6JA', 'https://www.youtube.com/watch?v=m0gr3snccVY&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=2', 'https://www.youtube.com/watch?v=ThjgUfck8Sc&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=3']
+#youtube_links = ['https://www.youtube.com/watch?v=6A5ZSiosgAk']
+
+for i, link in enumerate(youtube_links):
+    mp4_filename = f"audio_{i+1}.mp4"
+    wav_filename = f"audio_{i+1}.wav"
+
+    download_audio(link, mp4_filename)
+    convert_to_wav(mp4_filename, wav_filename)
+
+    transcribe_wav_file(wav_filename, i)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..2300225
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+youtube-dl
+pfy
+pytube
+pydub
+git+https://github.com/m-bain/whisperx.git
\ No newline at end of file
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..3dd306f
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1 @@
+"""Functions to get video data and create transcripts as well as detailed information"""
\ No newline at end of file
diff --git a/src/data_handling.py b/src/data_handling.py
new file mode 100644
index 0000000..9c6da39
--- /dev/null
+++ b/src/data_handling.py
@@ -0,0 +1,13 @@
+def download_audio(youtube_link, filename):
+    yt = YouTube(youtube_link)
+
+    # get highest quality audio stream
+    audio_stream = yt.streams.get_audio_only()
+
+    # download
+    if not os.path.exists(filename):
+        audio_stream.download(filename=filename)
+
+def convert_to_wav(mp4_file, wav_file):
+    sound = AudioSegment.from_file(mp4_file, format="mp4")
+    sound.export(wav_file, format="wav")
\ No newline at end of file
diff --git a/src/transcriber.py b/src/transcriber.py
new file mode 100644
index 0000000..d4d9d54
--- /dev/null
+++ b/src/transcriber.py
@@ -0,0 +1,40 @@
+google_colab_gpu = "cuda" if torch.cuda.is_available() else "cpu"
+
+def transcribe_wav_file(wav_filename, index):
+    batch_size = 16
+    compute_type = "float16"
+
+    model = whisperx.load_model("large-v2", device=google_colab_gpu, compute_type=compute_type)
+
+    audio = whisperx.load_audio(wav_filename)
+    result = model.transcribe(audio, batch_size=batch_size)
+
+    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=google_colab_gpu)
+    result = whisperx.align(result["segments"], model_a, metadata, audio, google_colab_gpu, return_char_alignments=False)
+
+    # clear GPU memory
+    del model
+    del model_a
+    del metadata
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', device=google_colab_gpu)
+
+    diarize_segments = diarize_model(audio)
+    result = whisperx.assign_word_speakers(diarize_segments, result)
+
+    # clear GPU memory
+    del diarize_model
+    del diarize_segments
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # Save full results to a json file
+    with open(f'result_{index+1}.json', 'w') as f:
+        json.dump(result, f)
+
+    # Save just the transcripts to a .txt file
+    transcripts = [seg['value'] for seg in result['segments']]
+    with open(f'transcripts_{index+1}.txt', 'w') as f:
+        f.write('\n'.join(transcripts))
\ No newline at end of file