diff --git a/.gitignore b/.gitignore index e95ff19..200b914 100644 --- a/.gitignore +++ b/.gitignore @@ -10,7 +10,4 @@ __pycache__/ .venv env/ venv/ -ENV/ - -# Large files -data/models # local storage of transcriber models \ No newline at end of file +ENV/ \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..29cb7ea --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Gerrit Anders + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md index 7f5e920..1c1826f 100644 --- a/README.md +++ b/README.md @@ -1,15 +1,27 @@ -# Project_whisperTrans01_Code +# Project_whisperTrans_Code -## Content +## Overview -Transcriber code +whisperTrans is a video2transcript tool built upon [Whisper](https://github.com/openai/whisper) +and [WhisperX](https://github.com/m-bain/whisperX). For automated speech recognition (ASR). +The project provides an efficient implementation for running transcriptions locally on open source models from local +files and YouTube videos. It is structured to handle large-scale audio transcription tasks with accuracy and ease using +diarization and alignment techniques. ## Setup -Before running this script, make sure to install the necessary python packages using the following command: +The project is tested under python 3.10.12. +In order to edit video material ffmpeg is required. Furthermore, machine learning models are employed in this project to +convert audio content to text. Which models are used depends on the language of the video content. The models are +downloaded on demand from the model hosting platform [huggingface](https://huggingface.co/). To do so you need to create +an account and set the environment variable HF_TOKEN to your token. This is necessary to accept the conditions of the +models used. + +Before running this script, make sure to install the necessary python packages. We recommend using a virtual environment +to manage dependencies. You can install the required packages with the following command: ```bash -pip install -r requirements.txt +pip install -r requirements.txt --index-url https://download.pytorch.org/whl/cu118 ``` If your system does not have ffmpeg installed you can install it via: @@ -18,9 +30,44 @@ If your system does not have ffmpeg installed you can install it via: sudo apt install ffmpeg ``` -## Running the script +For GPU usage you need to install the CUDA toolkit. You can install it via: -To generate transcripts create a new job with your specifications (see jobs section below). +```bash +sudo apt install nvidia-cuda-toolkit-11-8 +``` + +This might require to add the package the following way: + +```bash +wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb +sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb +sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/ +``` + +Afterward, set the environmental variables: + +```bash +echo 'export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}' >> ~/.bashrc +echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc +``` + +and reload your `~/.bashrc` file to apply these changes: + +```bash +source ~/.bashrc +``` + +## Prerequisites + +- Python 3.10 +- ffmpeg +- Access to Hugging Face for model downloads +- CUDA 11.8, and cuBLAS 11.x and cuDNN 8.x +- CUDA-compatible GPU with <8GB VRAM (recommended for faster processing) + +## Usage + +To generate transcripts create a new job with your specifications (see [Jobs secion](#jobs) below).
If you need to download the transcriber model set your huggingface token: @@ -36,65 +83,68 @@ On Windows (Command prompt): set HF_TOKEN="your-token-here" ``` -Then run the following command to generate your transcript: +After setting up the token, execute the script with: ```bash -main.py --job {name} +python main.py --job {name} ``` -where name is the name of your job. +Here, `name` is the identifier for your transcription job. +Note that you might need to go to huggingface to accept the conditions of the models used. Especially when running the +transcriber for the first time on your huggingface account. ## Jobs Jobs specify tasks for transcript creation. A job contains the following information. - `name`: string: The name of the job -- `source`: List[string]: The sources of the videos (paths, YouTube links, TibAV links) -- `modelpath`: Optional[string]: The path to the model if already saved locally (will be downloaded temporarily if not) -- `language`: Optional[string]: The language of the video (If not specified employs auto-detection) +- `source`: List[string]: The sources of the videos (paths, YouTube links) +- `language`: Optional[string]: The language of the video (If not specified employs auto-detection) as a language code + (e.g. 'en', 'de', 'fr') +- `output_format`: -## to-do +## Contributing -- seperate data production form code -- remember to accept the conditions on diarization if you use huggingface -- do a config json -- do option to have models already stored and point to them (is this necessary? > if you do not want to interact with - huggingface yes) - - also for single language models -- check warnings -- downloader for tib-av -- more optional setting -- add progress bar +Contributions are welcome! Please fork the repository and make your changes in a separate branch. Submit a pull request +with your proposed modifications. For major changes, please open an issue first to discuss what you would like to +change. + +## License + +This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details. + +## Contact + +This project is maintained by [Gerrit Anders](https://www.iwm-tuebingen.de/www/de/personen/ma.html?uid=ganders). +For any questions or support, please open an issue on this repository or reach out via email +at [g.anders@iwm-tuebingen.de](mailto:g.anders@iwm-tuebingen.de). + +## Future plans + +- **Tib-AV Downloader**: Add a downloader capability for [TibAV](https://av.tib.eu/) video links. +- **Progress Bar**: Implement a progress bar for better user feedback during processing. +- **Local Model Directory**: Allow models to be stored locally for offline use and faster processing. +- **Extend Mock Tests**: Add actual references for tests. ## Structure ``` project_root/ ├── src/ -│ ├── file_manager.py # Contains FileManager class +│ ├── video_manager.py # Contains VideoManager 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 +│ └── data_handling.py # Contains download_audio and convert_to_wav methods ├── tests/ │ ├── test_file_manager.py │ ├── test_transcriber.py │ └── test_transcript_generator.py ├── data/ -│ ├── audio/ # downloaded mp4 and output wav files -│ ├── models/ # models used for transcripts (are ignored by git and not necessarily) -│ └── transcripts/ # JSON and text transcript files +│ ├── logs/ # logs of previous jobs +│ ├── jobs/ # place to store job files +│ └── results/ # place to store results +├── main.py # Main script that runs the process +├── requirements.txt +├── LICENSE └── 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, model and transcript files. \ No newline at end of file +``` \ No newline at end of file diff --git a/data/jobs/example_job.json b/data/jobs/example_job.json new file mode 100644 index 0000000..7b3766f --- /dev/null +++ b/data/jobs/example_job.json @@ -0,0 +1,11 @@ +{ + "name": "myjob1", + "source": [ + "https://www.youtube.com/watch?v=PrGl3w2sAiw" + ], + "transcript_version": { + "get_json_version": true, + "get_txt_version": true, + "get_sentence_level_json_version": true + } +} \ No newline at end of file diff --git a/data/logs/HERE_logfiles_are_saved.txt b/data/logs/HERE_logfiles_are_saved.txt new file mode 100644 index 0000000..e69de29 diff --git a/main.py b/main.py index 2c5c6f1..8f09525 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,49 @@ +import json +import logging import os +from argparse import ArgumentParser, Namespace +from pathlib import Path +from src.logger import setup_logger 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"] + parser: ArgumentParser = ArgumentParser( + description="Transcript Generator for Videos" + ) + + parser.add_argument("--job", type=str, help="The name of the job", required=True) + + arguments: Namespace = parser.parse_args() + + with open(f"data/jobs/{arguments.job}", "r") as f: + job_config = json.load(f) + + log_directory: Path = Path("data/logs") + os.makedirs(log_directory, exist_ok=True) + log_file_path: Path = log_directory / f"{arguments.job}.log" + setup_logger(__name__, log_file=log_file_path) + logger: logging.Logger = logging.getLogger(__name__) + + logger.info("Starting the transcript generation.") hf_token = os.environ.get("HF_TOKEN") if not hf_token: hf_token = input("Please enter your HF_TOKEN: ") os.environ["HF_TOKEN"] = hf_token - transcript_generator = TranscriptGenerator() - transcript_generator.generate(input_sources) + transcript_generator: TranscriptGenerator = TranscriptGenerator( + job_name=job_config["name"], + language_code=job_config.get("language_code", "en"), + get_txt_version=job_config["transcript_version"]["get_txt_version"], + get_json_version=job_config["transcript_version"]["get_json_version"], + get_sentence_level_json=job_config["transcript_version"][ + "get_sentence_level_json_version" + ], + ) + + transcript_generator.generate(job_config["source"]) if __name__ == "__main__": diff --git a/requirements.txt b/requirements.txt index 9ec58e7..7e99cb6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,9 @@ -youtube-dl -pytube -pydub -ffmpeg-python -ffprobe -torch -git+https://github.com/m-bain/whisperx.git \ No newline at end of file +youtube-dl==2021.12.17 +pytube==15.0.0 +pydub==0.25.1 +ffmpeg-python==0.2.0 +ffprobe==0.5 +torch==2.0.0+cu118 +torchaudio==2.0.1+cu118 +pytorch-cuda=11.8 +git+https://github.com/m-bain/whisperx.git@d8a2b4f diff --git a/src/file_manager.py b/src/file_manager.py deleted file mode 100644 index c94f353..0000000 --- a/src/file_manager.py +++ /dev/null @@ -1,79 +0,0 @@ -import logging -import os - -from pydub import AudioSegment -from pytube import YouTube, exceptions as pytube_exceptions - -logging.basicConfig(level=logging.INFO) - - -class FileManager: - """ - FileManager class that handles operations such as downloading and converting audio. - """ - - def __init__(self): - pass - - @staticmethod - def download_youtube_audio(youtube_link: str, filename: str): - """ - Downloads audio from a YouTube link and saves it in the location specified by filename. - """ - try: - yt = YouTube(youtube_link) - except pytube_exceptions.PytubeError as e: - logging.error( - "An Error occurred while fetching the YouTube video: " + str(e) - ) - return - - audio_stream = yt.streams.get_audio_only() - - if not os.path.exists(filename): - try: - audio_stream.download(filename=filename) - logging.info(f"Downloaded {youtube_link} successfully.") - except Exception as e: - logging.error( - "An Error occurred while downloading the YouTube video: " + str(e) - ) - return - - @staticmethod - def convert_to_wav(mp4_file: str, wav_file: str): - """ - Converts a mp4 audio file into wav format. - """ - try: - sound = AudioSegment.from_file(mp4_file, format="mp4") - sound.export(wav_file, format="wav") - except Exception as e: - logging.error(f"An error occurred while converting {mp4_file} to wav: {e}") - - @staticmethod - def load_local_file(local_file_path: str) -> str: - """ - Loads a local file and converts it to wav format. - Return wav file path. - """ - ext = os.path.splitext(local_file_path)[-1].lower() - - if ext in [".mp4", ".mp3"]: - wav_file = local_file_path.replace(ext, ".wav") - try: - audio = AudioSegment.from_file(local_file_path, format=ext[1:]) - audio.export(wav_file, format="wav") - except Exception as e: - logging.error( - f"An error occurred while converting {local_file_path} to wav: {e}" - ) - - return wav_file - - elif ext == ".wav": - return local_file_path - - else: - logging.error(f"Unsupported file format: {ext}.") - raise ValueError("Unsupported file format.") diff --git a/src/logger.py b/src/logger.py new file mode 100644 index 0000000..6c41364 --- /dev/null +++ b/src/logger.py @@ -0,0 +1,47 @@ +import logging +from pathlib import Path +from typing import Optional + + +def setup_logger( + name: str, level: int = logging.INFO, log_file: Optional[Path] = None +) -> logging.Logger: + """ + Configure and return a logger with the specified name and log level. + + Parameters + ---------- + name : str + The name of the logger. + level : int (default: logging.INFO) + The log level for the logger. + log_file : Path (default: None) + The path to the log file. + """ + logger: logging.Logger = logging.getLogger(name) + logger.setLevel(level) + + console_handler: logging.StreamHandler = logging.StreamHandler() + console_handler.setLevel(level) + + if log_file: + file_handler: logging.FileHandler = logging.FileHandler(log_file) + file_handler.setLevel(level) + else: + file_handler = None + + formatter: logging.Formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + + console_handler.setFormatter(formatter) + + if file_handler: + file_handler.setFormatter(formatter) + + logger.addHandler(console_handler) + + if file_handler: + logger.addHandler(file_handler) + + return logger diff --git a/src/output_manager.py b/src/output_manager.py new file mode 100644 index 0000000..550c0c0 --- /dev/null +++ b/src/output_manager.py @@ -0,0 +1,128 @@ +import json +import logging + + +class OutputManager: + """ + A class to save the transcription result to a json file, a text file, and/or a sentence-level json file. + + Attributes: + ---------- + get_txt_version : bool + A flag to determine if the transcript should be saved in a text file. + get_json_version : bool + A flag to determine if the transcript should be saved in a JSON file. + get_sentence_level_json : bool + A flag to determine if the sentence-level transcript should be saved in a JSON file. + + Methods + ------- + save_transcript(transcript_output: dict, job_name: str, source_number: int) -> None: + Save the transcription result to a json file a text file and/or a sentence level json file. + """ + + def __init__( + self, + get_txt_version: bool, + get_json_version: bool, + get_sentence_level_json: bool, + ): + self.get_txt_version = get_txt_version + self.get_json_version = get_json_version + self.get_sentence_level_json = get_sentence_level_json + + self.logger = logging.getLogger(__name__) + + def save_transcript( + self, transcript_output: dict, job_name: str, source_number: int + ): + """ + Save the transcription result to a json file a text file and/or a sentence level json file. + + Parameters + ---------- + transcript_output : dict + The transcription result. + job_name : str + The name of the job. + source_number : int + The enumeration of the source in the job. + """ + + json_output_file_name = f"result_{job_name}_{source_number + 1}.json" + txt_output_file_name = f"transcripts_{job_name}_{source_number + 1}.txt" + sentence_level_output_file_name = ( + f"sentence_level_{job_name}_{source_number + 1}.json" + ) + if self.get_json_version: + self._save_transcript_full(json_output_file_name, transcript_output) + if self.get_txt_version: + self._save_transcript_txt(txt_output_file_name, transcript_output) + if self.get_sentence_level_json: + self._save_transcript_sentence_level( + sentence_level_output_file_name, transcript_output + ) + else: + self.logger.error("No output version specified. Not saving any output.") + + def _save_transcript_full(self, output_json_file: str, transcript_output: dict): + """ + Save the full transcription result to a json file. + + Parameters + ---------- + output_json_file : str + Path to the output json file. + transcript_output : dict + The transcription result to be saved. + """ + try: + with open(output_json_file, "w") as f: + json.dump(transcript_output, f) + except Exception as e: + self.logger.exception("Failed to save transcript to json: %s", e) + + def _save_transcript_txt(self, output_txt_file: str, transcript_output: dict): + """ + Save the transcript to a text file. + + Parameters + ---------- + output_txt_file : str + Path to the output txt file. + transcript_output : dict + The transcription result to be saved. + """ + try: + transcripts = [seg["text"] for seg in transcript_output["segments"]] + with open(output_txt_file, "w") as f: + f.write("\n".join(transcripts)) + except Exception as e: + self.logger.exception("Failed to save transcript to txt: %s", e) + + def _save_transcript_sentence_level( + self, output_json_file: str, transcript_output: dict + ): + """ + Save the transcript to a sentence-level json file. + + Parameters + ---------- + output_json_file : str + Path to the output json file. + transcript_output : dict + The transcription result to be saved. + """ + try: + sentence_level_transcript = [] + for segment in transcript_output.get("segments", []): + sentence_info = { + "start": segment.get("start"), + "end": segment.get("end"), + "text": segment.get("text"), + } + sentence_level_transcript.append(sentence_info) + with open(output_json_file, "w") as f: + json.dump(sentence_level_transcript, f, ensure_ascii=False, indent=4) + except Exception as e: + self.logger.exception("Failed to save sentence-level transcript: %s", e) diff --git a/src/transcriber.py b/src/transcriber.py index 4842b7a..01f460a 100644 --- a/src/transcriber.py +++ b/src/transcriber.py @@ -1,42 +1,61 @@ import gc -import json import logging import os import torch import whisperx -# Set up logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - class Transcriber: - """A class to transcribe audio files using whisperx.""" + """ + A class to transcribe audio files using whisperx. + + Attributes: + ---------- + language_code : str (default: "en") + Language code for whisperx transcription. Default is english. + device : str (default: "cuda") + Device to use for transcription. Default is "cuda" if available, else "cpu". + compute_type : str (default: "float16") + Precision type for whisperx transcription. + + Methods + ------- + transcribe(wav_filename: str, batch_size: int = 16) -> dict: + Transcribe the wav file. + """ def __init__( self, language_code: str = "en", device: str = "cuda", compute_type: str = "float16", + asr_options=None, ): - """ - Init method. - :param language_code: str: Language code for whisperx transcription (default is "en"). - :param device: str: Device to use for transcription (default is "cuda" if available, else "cpu"). - :param compute_type: str: Compute type for whisperx transcription (default is "float16"). - """ + if asr_options is None: + asr_options = { + "repetition_penalty": 1, + "prompt_reset_on_temperature": 0.5, + "no_repeat_ngram_size": 0, + } + self.logger = logging.getLogger(__name__) + self.device = device if torch.cuda.is_available() else "cpu" - # Initialize the transcription and alignment models + self.logger.info("Using device: %s", self.device) + try: self.transcription_model = whisperx.load_model( - "large-v2", device=self.device, compute_type=compute_type + "large-v2", + device=self.device, + compute_type=compute_type, + language=language_code, + asr_options=asr_options, ) self.alignment_model, self.metadata = whisperx.load_align_model( language_code=language_code, device=self.device ) except Exception as e: - logger.exception("Failed to initialize models: %s", e) + self.logger.exception("Failed to initialize models: %s", e) raise self.transcript_output = None self.hf_token = os.environ.get("HF_TOKEN") @@ -44,18 +63,24 @@ class Transcriber: def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict: """ Transcribe the wav file. - :param wav_filename: str: Path to the wav file. - :param batch_size: int: Batch size for whisperx transcription (default is 16). - :return: dict: Transcribed and aligned result. + + Parameters: + ---------- + wav_filename : str + Path to the wav file. + batch_size : int (default: 16) + Batch size for whisperx transcription. + + Returns: + ------- + dict + Transcription result. """ try: - # Load the audio file audio = whisperx.load_audio(wav_filename) - # Transcribe the audio file result = self.transcription_model.transcribe(audio, batch_size=batch_size) - # Use the alignment model result = whisperx.align( result["segments"], self.alignment_model, @@ -65,45 +90,22 @@ class Transcriber: return_char_alignments=False, ) - # Clear GPU memory self._clear_memory() - # Run Diarization diarize_model = whisperx.DiarizationPipeline( use_auth_token=self.hf_token, 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 except Exception as e: - logger.exception("Failed to transcribe: %s", e) + self.logger.exception("Failed to transcribe: %s", e) raise - def save_transcript(self, output_json_file: str, output_txt_file: str): - """ - Save the transcription result to a json file and a text file. - :param output_json_file: str: Path to the output json file. - :param output_txt_file: str: Path to the output text file. - """ - try: - # 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["text"] for seg in self.transcript_output["segments"]] - with open(output_txt_file, "w") as f: - f.write("\n".join(transcripts)) - except Exception as e: - logger.exception("Failed to save transcript: %s", e) - raise - - @staticmethod - def _clear_memory(): + def _clear_memory(self): """ Clear GPU memory. """ @@ -111,5 +113,5 @@ class Transcriber: gc.collect() torch.cuda.empty_cache() except Exception as e: - logger.exception("Failed to clear GPU memory: %s", e) + self.logger.exception("Failed to clear GPU memory: %s", e) raise diff --git a/src/transcript_generator.py b/src/transcript_generator.py index 51cd113..21cf03d 100644 --- a/src/transcript_generator.py +++ b/src/transcript_generator.py @@ -1,48 +1,82 @@ import logging -from typing import List -from src.file_manager import FileManager +from src.output_manager import OutputManager from src.transcriber import Transcriber - -# Set up logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) +from src.video_manager import VideoManager class TranscriptGenerator: """ TranscriptGenerator class that coordinates FileManager and Transcriber classes - to generate transcriptions from given input sources. + to generate transcriptions in desired formats from given input sources. + + Attributes: + ---------- + file_manager: FileManager + An instance of the FileManager class. + transcriber: Transcriber + job_name: str + The name of the job. + get_txt_version: bool (default: True) + A flag to determine if the transcript should be saved in a text file. + get_json_version: bool (default: True) + A flag to determine if the transcript should be saved in a JSON file. + get_sentence_level_json: bool (default: True) + A flag to determine if the sentence-level transcript should be saved in a JSON file. + + Methods + ------- + generate(input_sources: list[str]) -> None: + Generates transcriptions for the given list of input sources. """ - def __init__(self): - self.file_manager = FileManager() - self.transcriber = Transcriber() + def __init__( + self, + job_name: str, + language_code: str = "en", + get_txt_version: bool = True, + get_json_version: bool = True, + get_sentence_level_json: bool = True, + ): + self.video_manager: VideoManager = VideoManager() + self.output_manager: OutputManager = OutputManager( + get_txt_version, get_json_version, get_sentence_level_json + ) + self.transcriber: Transcriber = Transcriber(language_code) + self.job_name = job_name - def generate(self, input_sources: List[str]) -> None: + self.logger = logging.getLogger(__name__) + + def generate(self, input_sources: list[str]) -> None: """ Generates transcriptions for the given list of input sources. - Args: - input_sources (List[str]): List of input sources which could be URLs or local file paths. + + Parameters + ---------- + input_sources : list[str] + A list of input sources (local file paths or YouTube URLs) for which transcriptions should be generated. """ - for i, source in enumerate(input_sources): + for source_number, source in enumerate(input_sources): try: 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) + mp4_filename = f"audio_{self.job_name}_{source_number + 1}.mp4" + wav_filename = f"audio_{self.job_name}_{source_number + 1}.wav" + self.video_manager.download_youtube_audio(source, mp4_filename) + self.video_manager.convert_to_wav(mp4_filename, wav_filename) else: - wav_filename = self.file_manager.load_local_file(source) - 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) + wav_filename = self.video_manager.load_local_file(source) + transcript_data = self.transcriber.transcribe(wav_filename) + self.output_manager.save_transcript( + transcript_data, self.job_name, source_number + ) + + self.logger.info(f"Transcription generated for source {source}.") + except Exception as e: - logger.exception( + self.logger.exception( f"Failed to generate transcript for source {source}. Error: {str(e)}" ) continue - logger.info("Transcription generation completed.") + self.logger.info("Transcription generation completed.") diff --git a/src/video_manager.py b/src/video_manager.py new file mode 100644 index 0000000..ff997d9 --- /dev/null +++ b/src/video_manager.py @@ -0,0 +1,111 @@ +import logging +import os + +from pydub import AudioSegment +from pytube import YouTube, exceptions as pytube_exceptions + + +class VideoManager: + """ + FileManager class that handles operations on input videos such as downloading and converting to audio. + To arrive at the correct input format for transcript generation. + + Attributes: + ---------- + None + + Methods + ------- + download_youtube_audio(youtube_link: str, filename: str) -> None: + Downloads audio from a YouTube link and saves it in the location specified by filename. + convert_to_wav(mp4_file: str, wav_file: str) -> None: + Converts a mp4 audio file into wav format. + load_local_file(local_file_path: str) -> str: + Loads a local file and converts it to wav format + """ + + def __init__(self): + self.logger = logging.getLogger(__name__) + + def download_youtube_audio(self, youtube_video_link: str, filename: str) -> None: + """ + Downloads audio from a YouTube link and saves it in the location specified by filename. + + Parameters + ---------- + youtube_video_link : str + The YouTube video link from which the audio should be downloaded. + filename : str + The location where the audio should be saved + """ + try: + yt = YouTube(youtube_video_link) + except pytube_exceptions.PytubeError as e: + self.logger.error( + "An Error occurred while fetching the YouTube video: " + str(e) + ) + return None + + audio_stream = yt.streams.get_audio_only() + + if not os.path.exists(filename): + try: + audio_stream.download(filename=filename) + self.logger.info(f"Downloaded {youtube_video_link} successfully.") + except Exception as e: + self.logger.error( + "An Error occurred while downloading the YouTube video: " + str(e) + ) + return None + + def convert_to_wav(self, mp4_file: str, wav_file: str) -> None: + """ + Converts a mp4 audio file into wav format. + + Parameters + ---------- + mp4_file : str + The path to the mp4 audio file. + wav_file : str + The path where the wav file should be saved + """ + try: + sound = AudioSegment.from_file(mp4_file, format="mp4") + sound.export(wav_file, format="wav") + except Exception as e: + self.logger.error( + f"An error occurred while converting {mp4_file} to wav: {e}" + ) + + def load_local_file(self, local_file_path: str) -> str: + """ + Loads a local file and converts it to wav format. + Return wav file path. + + Parameters + ---------- + local_file_path : str + The path to the local file. + """ + file_extension = os.path.splitext(local_file_path)[-1].lower() + + if file_extension in [".mp4", ".mp3"]: + wav_file_path = local_file_path.replace(file_extension, ".wav") + try: + audio = AudioSegment.from_file( + local_file_path, format=file_extension[1:] + ) + audio.export(wav_file_path, format="wav") + except Exception as e: + self.logger.error( + f"An error occurred while converting {local_file_path} to wav: {e}" + ) + + return wav_file_path + + elif file_extension == ".wav": + return local_file_path + + else: + self.logger.error(f"Unsupported file format: {file_extension}.") + raise ValueError("Unsupported file format.") diff --git a/tests/test_file_manager.py b/tests/test_file_manager.py index 935e5cd..0465272 100644 --- a/tests/test_file_manager.py +++ b/tests/test_file_manager.py @@ -1,32 +1,32 @@ import os import unittest -from src.file_manager import FileManager +from src.video_manager import VideoManager -class FileManagerTest(unittest.TestCase): +class VideoManagerTest(unittest.TestCase): def setUp(self): - self.file_manager = FileManager() - self.dir = '' # replace with your directory + self.file_manager = VideoManager() + self.dir = "" def test_download_youtube_audio(self): - youtube_link = '' # replace with a YouTube link - filename = f'{self.dir}/audio.mp4' + 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 + os.remove(filename) def test_convert_to_wav(self): - mp4_file = '' # replace with an MP4 file in your directory - wav_file = mp4_file.replace('.mp4', '.wav') + mp4_file = "" + 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 + os.remove(wav_file) def test_load_local_file(self): - local_file_path = '' # replace with a file in your directory + local_file_path = "" result_path = self.file_manager.load_local_file(local_file_path) - self.assertEqual(result_path, local_file_path.replace('.mp4', '.wav')) + self.assertEqual(result_path, local_file_path.replace(".mp4", ".wav")) if __name__ == "__main__": diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py index f104893..28fab77 100644 --- a/tests/test_transcriber.py +++ b/tests/test_transcriber.py @@ -1,4 +1,3 @@ -import os import unittest from src.transcriber import Transcriber @@ -9,24 +8,10 @@ class TranscriberTest(unittest.TestCase): self.transcriber = Transcriber() def test_transcribe(self): - wav_file = '' # Replace with path to a .wav file on your machine + wav_file = "" 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()