diff --git a/README.md b/README.md index 8f1165a..d71ed6d 100644 --- a/README.md +++ b/README.md @@ -7,4 +7,9 @@ Requires nvcc v12.8 You can create environment running ``` conda env create -f environment.yml +``` +or for CUDA/Nvidia: +``` +conda create -n notecast -c pytorch -c nvidia pytorch torchvision torchaudio transformers python=3.12 +conda install ffmpeg customtkinter -c conda-forge -c bioconda ``` \ No newline at end of file diff --git a/env_cuda.yml b/env_cuda.yml index 192b8ef..bcd6f71 100644 --- a/env_cuda.yml +++ b/env_cuda.yml @@ -3,19 +3,12 @@ channels: - pytorch - nvidia - conda-forge - - defaults + - bioconda dependencies: - - python=3.10 - - pytorch::pytorch=2.0.1 - - pytorch::torchvision=0.15.2 - - pytorch::torchaudio=2.0.2 - - pytorch::cudatoolkit=11.8 - - transformers=4.30.0 - - accelerate=0.20.0 + - pytorch + - torchvision + - torchaudio + - transformers + - python=3.12 - ffmpeg - - pyinstaller - - customtkinter - - isort - - mypy - - black - - tqdm \ No newline at end of file + - customtkinter \ No newline at end of file diff --git a/transcription/audio_transcription.py b/transcription/audio_transcription.py index aff804e..7bc0238 100644 --- a/transcription/audio_transcription.py +++ b/transcription/audio_transcription.py @@ -18,6 +18,7 @@ class AudioTranscription: filepath: str waveform: torch.Tensor sampling_rate: int + file_format: str chunks: list = [] batches: list = [] @@ -46,6 +47,9 @@ class AudioTranscription: self.filepath = filepath self.language = language self.logger = logger + + # setting file extension + self.file_format = filepath.split(".")[-1] # extracting configuration self.device = device_configuration.device @@ -57,24 +61,28 @@ class AudioTranscription: self.chunks: list = [] self.batches: list = [] self.all_transcription: list = [] - try: - logger.info("Loading model WhisperProcessor...") + + def _load_model(self) -> None: + self.logger.info("Loading model WhisperProcessor...") + try: self.processor = WhisperProcessor.from_pretrained(self.model_name) - self.model = WhisperForConditionalGeneration.from_pretrained( - self.model_name, torch_dtype=self.torch_dtype + self.model_name, torch_dtype = self.torch_dtype ).to(self.device) - - logger.info("Model loaded.") - - self.waveform, self.sampling_rate = torchaudio.load( - filepath, format="mp3", backend="ffmpeg" - ) - logger.info(f"Successfully loaded file {filepath}.") - + + self.logger.info("Model loaded.") except Exception as e: - logger.error(f"Unable to load file {self.filepath}: {e}") - raise + self.logger.error(f"Error while loading model: {e}") + + def _load_file(self) -> None: + self.logger.info(f"Loading file {self.filepath}") + try: + self.waveform, self.sampling_rate = torchaudio.load( + self.filepath, format=self.file_format, backend="ffmpeg" + ) + self.logger.info(f"Successfully loaded file {self.filepath}.") + except Exception as e: + self.logger.error(f"Unable to load file {self.filepath}: {e}") def _resample(self) -> None: self.waveform = torchaudio.functional.resample( @@ -86,15 +94,15 @@ class AudioTranscription: self.waveform = self.waveform.mean(dim=0, keepdim=True) self.waveform = self.waveform.squeeze(0) - def _split_to_chunks(self, chunk_length_s: int = 30) -> None: + def _split_to_chunks(self, shift: bool = False) -> None: self.logger.info(f"Splitting audio on chunks...") - self.chunk_size = chunk_length_s * 16000 # 16kHz after resampling + self.chunk_size = self.custom_chunk_length * 16000 # 16kHz after resampling total_samples = self.waveform.shape[0] chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size self.logger.info( - f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk." + f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {self.custom_chunk_length} seconds per chunk." ) self.chunks = [] @@ -104,12 +112,13 @@ class AudioTranscription: chunk = self.waveform[start:end].cpu().numpy().astype("float32") self.chunks.append(chunk) - def _resplit_to_batches(self) -> None: + def _resplit_chunks_to_batches(self) -> None: self.logger.info(f"Splitting chunks into batches...") self.batches = [] for i in range(0, len(self.chunks), self.custom_batch_length): batch = self.chunks[i : i + self.custom_batch_length] self.batches.append(batch) + self.logger.info(f"Total: {len(self.batches)} batches") def _process_all_batches(self) -> None: start_time = time.time() @@ -149,10 +158,11 @@ class AudioTranscription: self.logger.error(f"Errors occured while processing chunks: {e}") def transcribe_audio(self) -> str: - # TODO: maybe something else, not str? + self._load_model() + self._load_file() self._resample() self._to_mono() self._split_to_chunks() - self._resplit_to_batches() + self._resplit_chunks_to_batches() self._process_all_batches() return " ".join(self.all_transcription) diff --git a/ui/tooltip.py b/ui/tooltip.py new file mode 100644 index 0000000..8af3341 --- /dev/null +++ b/ui/tooltip.py @@ -0,0 +1,24 @@ +import customtkinter as ctk + +class ToolTip(ctk.CTkToplevel): + def __init__(self, widget, text): + super().__init__() + self.withdraw() + self.overrideredirect(True) + self.attributes("-topmost", True) + + self.label = ctk.CTkLabel(self, text=text, fg_color="gray20", corner_radius=6, padx=10, pady=5) + self.label.pack() + + self.widget = widget + widget.bind("", self.show_tooltip) + widget.bind("", self.hide_tooltip) + + def show_tooltip(self, event=None): + x = self.widget.winfo_rootx() + 20 + y = self.widget.winfo_rooty() + 20 + self.geometry(f"+{x}+{y}") + self.deiconify() + + def hide_tooltip(self, event=None): + self.withdraw() \ No newline at end of file diff --git a/ui/ui.py b/ui/ui.py index ccb16a8..69a3ece 100644 --- a/ui/ui.py +++ b/ui/ui.py @@ -11,6 +11,7 @@ from transcription.audio_transcription import AudioTranscription from transcription.device_configuration import DeviceConfiguration from transcription.torch_checker import check_torch from ui.ui_log_handler import setup_ui_logger +from ui.tooltip import ToolTip WINDOW_WIDTH = 900 WINDOW_HEIGHT = 650 @@ -101,6 +102,8 @@ class TranscriberApp(ctk.CTk): ctrl_frame, text="Stop", command=self._stop_transcription, state="disabled" ) self.stop_btn.pack(side="left", padx=10, pady=5) + + # TODO: add unload model button here # log box self.log_box = scrolledtext.ScrolledText( @@ -112,6 +115,7 @@ class TranscriberApp(ctk.CTk): self.log_box.pack(padx=20, pady=10, expand=True, fill="both") def _build_settings_tab(self): + # TODO: add tooltips here pad = 20 ### Model @@ -175,20 +179,27 @@ class TranscriberApp(ctk.CTk): path = filedialog.askopenfilename( title="Select input audio file", filetypes=[ - ("Audio files", "*.wav *.mp3 *.m4a *.flac"), + ("Media files", "*.wav *.mp3 *.m4a *.flac *.ogg *.mp4 *.mkv *.avi"), + ("Audio files", "*.wav *.mp3 *.m4a *.flac *.ogg"), + ("Video files", "*.mp4 *.mkv *.avi"), ("All files", "*.*"), - ], + ] ) if path: self.input_file_var.set(path) def _browse_output(self): - path = filedialog.asksaveasfilename( - title="Select output file", - defaultextension=".txt", - filetypes=[("Text files", "*.txt"), ("All files", "*.*")], + # TODO: add custom filename here + directory = filedialog.askdirectory( + title="Select output directory", ) - if path: + if directory: + input_name = os.path.basename(self.input_file_var.get()) + name, _ = os.path.splitext(input_name) + # TODO: redo output_name logic maybe? + output_name = f"{"".join(name.split(".").pop())}.txt" + path = os.path.join(directory, output_name) + self.output_file_var.set(path) def _check_torch(self): diff --git a/ui/ui_log_handler.py b/ui/ui_log_handler.py index ec7fb10..d617ead 100644 --- a/ui/ui_log_handler.py +++ b/ui/ui_log_handler.py @@ -13,7 +13,7 @@ class UILogHandler(logging.Handler): self.text_widget.see(tk.END) -# TODO: maybe some tqdm here, not in console? +# TODO: status should be here def setup_ui_logger(text_widget: tk.Text, level=logging.INFO): logger = logging.getLogger("UI_LOGGER") logger.setLevel(level)