diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..200b914 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +# PyCharm IDE +.idea/ +*.iml + +# Byte-compiled / optimized / DLL files +__pycache__/ + +# Environments +.env +.venv +env/ +venv/ +ENV/ \ No newline at end of file diff --git a/2023-09-20_transcripts_Video-SRS.py b/2023-09-20_transcripts_Video-SRS.py deleted file mode 100644 index 31f9a15..0000000 --- a/2023-09-20_transcripts_Video-SRS.py +++ /dev/null @@ -1,14 +0,0 @@ -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/README.md b/README.md index 1db68e4..8bf4ba0 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,5 @@ # Project_whisperTrans01_Code - - ## Content Transcriber code @@ -13,4 +11,34 @@ Currently very unorganized - clean the code - organize the stuff - add functions for other data sources -- seperate data production form code \ No newline at end of file +- seperate data production form code + +``` +project_root/ +├── src/ +│ ├── file_manager.py # Contains FileManager class +│ ├── transcriber.py # Contains Transcriber class +│ ├── transcript_generator.py # Contains TranscriptGenerator class +│ ├── data_handling.py # Contains download_audio and convert_to_wav methods +│ └── main.py # Main script that runs the process +├── tests/ +│ ├── test_file_manager.py +│ ├── test_transcriber.py +│ └── test_transcript_generator.py +├── data/ +│ ├── audio/ # downloaded mp4 and output wav files +│ └── transcripts/ # JSON and text transcript files +└── README.md +``` + +- **src/**: This directory will contain all your source code. Inside the src folder, I added: + - **main.py**: This will be the main entry point for the program. It will use the `TranscriptGenerator` class which + orchestrates the use of `FileManager` and `Transcriber`. + - **file_manager.py**: Contains the `FileManager` class implementation. + - **transcriber.py**: Contains the `Transcriber` class implementation. + - **transcript_generator.py**: Contains the `TranscriptGenerator` class implementation. + +- **tests/**: This directory contains all your test scripts corresponding to your source files. + +- **data/**: This directory contains all the data needed for your code. For simplicity, I categorize the files into + separate directories for audio and transcript files. \ No newline at end of file diff --git a/src/data_handling.py b/src/data_handling.py deleted file mode 100644 index 9c6da39..0000000 --- a/src/data_handling.py +++ /dev/null @@ -1,13 +0,0 @@ -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/file_manager.py b/src/file_manager.py new file mode 100644 index 0000000..40755eb --- /dev/null +++ b/src/file_manager.py @@ -0,0 +1,41 @@ +import os + +from pydub import AudioSegment +from pytube import YouTube + + +class FileManager: + def __init__(self): + pass + + def download_youtube_audio(self, youtube_link, filename): + yt = YouTube(youtube_link) + + # Get the highest quality audio stream + audio_stream = yt.streams.get_audio_only() + + # Download if the file doesn't already exist + if not os.path.exists(filename): + audio_stream.download(filename=filename) + + def convert_to_wav(self, mp4_file, wav_file): + sound = AudioSegment.from_file(mp4_file, format="mp4") + sound.export(wav_file, format="wav") + + def load_local_file(self, local_file_path): + # Check the file type + ext = os.path.splitext(local_file_path)[-1].lower() + + if ext == ".mp4": + wav_file = local_file_path.replace(".mp4", ".wav") + self.convert_to_wav(local_file_path, wav_file) + return wav_file + elif ext == ".mp3": + wav_file = local_file_path.replace(".mp3", ".wav") + audio = AudioSegment.from_file(local_file_path, format="mp3") + audio.export(wav_file, format="wav") + return wav_file + elif ext == ".wav": + return local_file_path + else: + raise ValueError("Unsupported file format.") diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..9a3a7b0 --- /dev/null +++ b/src/main.py @@ -0,0 +1,16 @@ +from src.transcript_generator import TranscriptGenerator + + +def main(): + # input could be a YouTube link or a local file + input_sources = [ + 'https://www.youtube.com/watch?v=BVVo1BaL6JA', + '/path/to/your/local/file.mp3' + ] + + transcript_generator = TranscriptGenerator() + transcript_generator.generate(input_sources) + + +if __name__ == "__main__": + main() diff --git a/src/transcriber.py b/src/transcriber.py index d4d9d54..b13232f 100644 --- a/src/transcriber.py +++ b/src/transcriber.py @@ -1,40 +1,63 @@ -google_colab_gpu = "cuda" if torch.cuda.is_available() else "cpu" +import gc +import json -def transcribe_wav_file(wav_filename, index): - batch_size = 16 - compute_type = "float16" +import torch +import whisperx - 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) +class Transcriber: - 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) + def __init__(self, language_code="en", device="cuda", compute_type="float16"): + self.device = device if torch.cuda.is_available() else "cpu" - # clear GPU memory - del model - del model_a - del metadata - gc.collect() - torch.cuda.empty_cache() + # Initialize the transcription model + self.transcription_model = whisperx.load_model( + "large-v2", device=self.device, compute_type=compute_type) - diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', device=google_colab_gpu) + # Initialize the alignment model + self.alignment_model, self.metadata = whisperx.load_align_model( + language_code=language_code, device=self.device) + self.transcript_output = None - diarize_segments = diarize_model(audio) - result = whisperx.assign_word_speakers(diarize_segments, result) + def transcribe(self, wav_filename, batch_size=16): + # Load the audio file + audio = whisperx.load_audio(wav_filename) - # clear GPU memory - del diarize_model - del diarize_segments - gc.collect() - torch.cuda.empty_cache() + # Transcribe the audio file + result = self.transcription_model.transcribe(audio, batch_size=batch_size) - # Save full results to a json file - with open(f'result_{index+1}.json', 'w') as f: - json.dump(result, f) + # Use the alignment model + result = whisperx.align( + result["segments"], self.alignment_model, self.metadata, audio, self.device, + return_char_alignments=False) - # 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 + # Clear GPU memory + self._clear_memory() + + # Run Diarization + diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', + device=self.device) + diarize_segments = diarize_model(audio) + result = whisperx.assign_word_speakers(diarize_segments, result) + + # Clear GPU memory again + self._clear_memory() + self.transcript_output = result + return result + + def save_transcript(self, output_json_file, output_txt_file): + # Save full results to a json file + with open(output_json_file, 'w') as f: + json.dump(self.transcript_output, f) + + # Save just the transcripts to a .txt file + transcripts = [seg['value'] for seg in self.transcript_output['segments']] + with open(output_txt_file, 'w') as f: + f.write('\n'.join(transcripts)) + + def _clear_memory(self): + del self.transcription_model + del self.alignment_model + del self.metadata + gc.collect() + torch.cuda.empty_cache() diff --git a/src/transcript_generator.py b/src/transcript_generator.py new file mode 100644 index 0000000..79343cb --- /dev/null +++ b/src/transcript_generator.py @@ -0,0 +1,27 @@ +from src.file_manager import FileManager +from src.transcriber import Transcriber + + +class TranscriptGenerator: + def __init__(self): + self.file_manager = FileManager() + self.transcriber = Transcriber() + + def generate(self, input_sources): + for i, source in enumerate(input_sources): + if source.startswith('http'): + mp4_filename = f"audio_{i + 1}.mp4" + wav_filename = f"audio_{i + 1}.wav" + + self.file_manager.download_youtube_audio(source, mp4_filename) + self.file_manager.convert_to_wav(mp4_filename, wav_filename) + else: + wav_filename = self.file_manager.load_local_file(source) + + transcription = self.transcriber.transcribe(wav_filename) + + json_output_file = f'result_{i + 1}.json' + txt_output_file = f'transcripts_{i + 1}.txt' + self.transcriber.save_transcript(json_output_file, txt_output_file) + + print('Transcription generation completed.') diff --git a/tests/test_file_manager.py b/tests/test_file_manager.py new file mode 100644 index 0000000..935e5cd --- /dev/null +++ b/tests/test_file_manager.py @@ -0,0 +1,33 @@ +import os +import unittest + +from src.file_manager import FileManager + + +class FileManagerTest(unittest.TestCase): + def setUp(self): + self.file_manager = FileManager() + self.dir = '' # replace with your directory + + def test_download_youtube_audio(self): + youtube_link = '' # replace with a YouTube link + filename = f'{self.dir}/audio.mp4' + self.file_manager.download_youtube_audio(youtube_link, filename) + self.assertTrue(os.path.exists(filename)) + os.remove(filename) # clean up after test + + def test_convert_to_wav(self): + mp4_file = '' # replace with an MP4 file in your directory + wav_file = mp4_file.replace('.mp4', '.wav') + self.file_manager.convert_to_wav(mp4_file, wav_file) + self.assertTrue(os.path.exists(wav_file)) + os.remove(wav_file) # clean up after test + + def test_load_local_file(self): + local_file_path = '' # replace with a file in your directory + result_path = self.file_manager.load_local_file(local_file_path) + self.assertEqual(result_path, local_file_path.replace('.mp4', '.wav')) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py new file mode 100644 index 0000000..f104893 --- /dev/null +++ b/tests/test_transcriber.py @@ -0,0 +1,32 @@ +import os +import unittest + +from src.transcriber import Transcriber + + +class TranscriberTest(unittest.TestCase): + def setUp(self): + self.transcriber = Transcriber() + + def test_transcribe(self): + wav_file = '' # Replace with path to a .wav file on your machine + result = self.transcriber.transcribe(wav_file) + self.assertIsNotNone(result) + + def test_save_transcript(self): + wav_file = '' # Replace with path to a .wav file on your machine + result = self.transcriber.transcribe(wav_file) + + json_output_file = '' # Temporary output file for testing + txt_output_file = '' # Temporary output file for testing + self.transcriber.save_transcript(json_output_file, txt_output_file) + + self.assertTrue(os.path.exists(json_output_file)) + self.assertTrue(os.path.exists(txt_output_file)) + + os.remove(json_output_file) # clean up after test + os.remove(txt_output_file) # clean up after test + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_transcript_generator.py b/tests/test_transcript_generator.py new file mode 100644 index 0000000..ffebfbf --- /dev/null +++ b/tests/test_transcript_generator.py @@ -0,0 +1,19 @@ +import unittest + +from src.transcript_generator import TranscriptGenerator + + +class TranscriptGeneratorTest(unittest.TestCase): + def setUp(self): + self.transcript_generator = TranscriptGenerator() + + def test_generate(self): + input_sources = [ + 'https://www.youtube.com/watch?v=BVVo1BaL6JA', + '/path/to/your/local/file.mp3' + ] + self.transcript_generator.generate(input_sources) + + +if __name__ == "__main__": + unittest.main()