Working version

- audio/video files support added

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
 ```

env_cuda.yml

@@ -3,19 +3,12 @@ channels:
   - pytorch
   - nvidia
   - conda-forge
-  - defaults
+  - bioconda
 dependencies:
-  - python=3.10
+  - pytorch
-  - pytorch::pytorch=2.0.1
+  - torchvision
-  - pytorch::torchvision=0.15.2
+  - torchaudio
-  - pytorch::torchaudio=2.0.2
+  - transformers
-  - pytorch::cudatoolkit=11.8
+  - python=3.12
-  - transformers=4.30.0
-  - accelerate=0.20.0
   - ffmpeg
-  - pyinstaller
+  - customtkinter
-  - customtkinter
-  - isort
-  - mypy
-  - black
-  - tqdm

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.logger.info("Model loaded.")
-
-            self.waveform, self.sampling_rate = torchaudio.load(
-                filepath, format="mp3", backend="ffmpeg"
-            )
-            logger.info(f"Successfully loaded file {filepath}.")
-
         except Exception as e:
-            logger.error(f"Unable to load file {self.filepath}: {e}")
+            self.logger.error(f"Error while loading model: {e}")
-            raise
+    
+    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)

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("<Enter>", self.show_tooltip)
+        widget.bind("<Leave>", 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()

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(
+        # TODO: add custom filename here
-            title="Select output file",
+        directory = filedialog.askdirectory(
-            defaultextension=".txt",
+            title="Select output directory",
-            filetypes=[("Text files", "*.txt"), ("All files", "*.*")],
         )
-        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):

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)