refactor(buzz): group video params into _VideoParams dataclass in FfmpegFrameReader (#1543)

Co-authored-by: Raivis Dejus <raivisd@scandiweb.com>

Author: josiasdev

buzz/ffmpeg_video_player.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from dataclasses import dataclass
 import os
 import subprocess
 import sys
@@ -74,6 +75,16 @@ def probe_video(file_path: str) -> dict:
     }
 
 
+@dataclass
+class _VideoParams:
+    file_path: str
+    width: int
+    height: int
+    fps: float
+    frame_size: int
+    start_ms: float
+
+
 class FfmpegFrameReader:
     """
     Reads raw RGB24 video frames from a file via ffmpeg in a background thread.
@@ -83,15 +94,17 @@ class FfmpegFrameReader:
     MAX_BUFFER = 60
 
     def __init__(self, file_path: str, width: int, height: int, fps: float, start_ms: int = 0):
-        self._file_path = file_path
-        self._width = width
-        self._height = height
-        self._fps = fps
-        self._frame_size = width * height * 3
+        self._params = _VideoParams(
+            file_path=file_path,
+            width=width,
+            height=height,
+            fps=fps,
+            frame_size=width * height * 3,
+            start_ms=float(start_ms),
+        )
 
         self._lock = threading.Lock()
         self._frames: deque = deque()
-        self._start_ms = float(start_ms)
         self._done = False
         self._stop_event = threading.Event()
         self._proc: Optional[subprocess.Popen] = None
@@ -108,21 +121,21 @@ def _read_loop(self):
         ffmpeg = _find_ffmpeg()
         cmd = [
             ffmpeg,
-            "-ss", str(self._start_ms / 1000.0),
-            "-i", self._file_path,
+            "-ss", str(self._params.start_ms / 1000.0),
+            "-i", self._params.file_path,
             "-f", "rawvideo",
             "-pix_fmt", "rgb24",
-            "-s", f"{self._width}x{self._height}",
-            "-r", str(self._fps),
+            "-s", f"{self._params.width}x{self._params.height}",
+            "-r", str(self._params.fps),
             "-an",
             "-",
         ]
         try:
             self._proc = subprocess.Popen(
                 cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
             )
-            frame_ms = 1000.0 / self._fps
-            pos = self._start_ms
+            frame_ms = 1000.0 / self._params.fps
+            pos = self._params.start_ms
 
             while not self._stop_event.is_set():
                 # Throttle when buffer is full
@@ -135,14 +148,14 @@ def _read_loop(self):
                 if self._stop_event.is_set():
                     break
 
-                raw = self._proc.stdout.read(self._frame_size)
-                if len(raw) < self._frame_size:
+                raw = self._proc.stdout.read(self._params.frame_size)
+                if len(raw) < self._params.frame_size:
                     with self._lock:
                         self._done = True
                     break
 
                 arr = np.frombuffer(raw, dtype=np.uint8).reshape(
-                    (self._height, self._width, 3)
+                    (self._params.height, self._params.width, 3)
                 ).copy()
                 with self._lock:
                     self._frames.append((pos, arr))
@@ -167,7 +180,7 @@ def get_frame_for_position(self, position_ms: float) -> Optional[np.ndarray]:
             if not self._frames:
                 return None
 
-            frame_ms = 1000.0 / self._fps
+            frame_ms = 1000.0 / self._params.fps
             # Drain frames that are more than one frame period behind
             while len(self._frames) > 1:
                 ts, _ = self._frames[0]

buzz/file_transcriber_queue_worker.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Optional, Tuple, List, Set
 from uuid import UUID
+from dataclasses import dataclass
 
 # Fix SSL certificate verification for bundled applications (macOS, Windows)
 # This must be done before importing demucs which uses torch.hub with urllib
@@ -97,11 +98,15 @@ def callback(progress):
 from buzz.transcriber.whisper_file_transcriber import WhisperFileTranscriber
 
 
+@dataclass
+class _CurrentTranscription:
+    task: Optional[FileTranscriptionTask] = None
+    transcriber: Optional[FileTranscriber] = None
+    transcriber_thread: Optional[QThread] = None
+
+
 class FileTranscriberQueueWorker(QObject):
     tasks_queue: multiprocessing.Queue
-    current_task: Optional[FileTranscriptionTask] = None
-    current_transcriber: Optional[FileTranscriber] = None
-    current_transcriber_thread: Optional[QThread] = None
 
     task_started = pyqtSignal(FileTranscriptionTask)
     task_progress = pyqtSignal(FileTranscriptionTask, float)
@@ -115,8 +120,8 @@ class FileTranscriberQueueWorker(QObject):
     def __init__(self, parent: Optional[QObject] = None):
         super().__init__(parent)
         self.tasks_queue = queue.Queue()
+        self.current = _CurrentTranscription()
         self.canceled_tasks: Set[UUID] = set()
-        self.current_transcriber = None
         self.speech_path = None
         self.speech_extractor_process = None
         self.is_running = False
@@ -143,7 +148,7 @@ def run(self):
 
         self.is_running = True
 
-        if self.current_task.transcription_options.extract_speech:
+        if self.current.task.transcription_options.extract_speech:
             status = self._setup_speech_extraction()
             if status == "error":
                 self.is_running = False
@@ -153,22 +158,22 @@ def run(self):
             return
 
         logging.debug("Starting next transcription task")
-        self.task_progress.emit(self.current_task, 0)
+        self.task_progress.emit(self.current.task, 0)
 
         self._create_transcriber()
         self._setup_transcriber_thread()
 
     def _cleanup_previous_transcriber(self):
-        if self.current_transcriber is not None:
-            self.current_transcriber.stop()
-            self.current_transcriber = None
+        if self.current.transcriber is not None:
+            self.current.transcriber.stop()
+            self.current.transcriber = None
 
     def _get_next_task(self) -> bool:
         while True:
-            self.current_task = self.tasks_queue.get()
-            if self.current_task is None:
+            self.current.task = self.tasks_queue.get()
+            if self.current.task is None:
                 return False
-            if self.current_task.uid in self.canceled_tasks:
+            if self.current.task.uid in self.canceled_tasks:
                 continue
             return True
 
@@ -182,21 +187,21 @@ def _setup_speech_extraction(self) -> str:
             import torch
             device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        task_file_path = Path(self.current_task.file_path)
+        task_file_path = Path(self.current.task.file_path)
         speech_path = task_file_path.with_name(f"{task_file_path.stem}_speech.mp3")
 
         status = self._extract_speech(str(task_file_path), str(speech_path), device)
 
         if status == "error":
             self.task_error.emit(
-                self.current_task,
+                self.current.task,
                 _("Speech extraction failed! Check your internet connection \u2014 a model may need to be downloaded."),
             )
         elif status == "ok":
             self.speech_path = speech_path
-            if not self.current_task.original_file_path:
-                self.current_task.original_file_path = str(task_file_path)
-            self.current_task.file_path = str(speech_path)
+            if not self.current.task.original_file_path:
+                self.current.task.original_file_path = str(task_file_path)
+            self.current.task.file_path = str(speech_path)
 
         return status
 
@@ -208,14 +213,14 @@ def _run_plugins(self) -> bool:
         """
         if self.plugin_manager is not None:
             try:
-                self.plugin_manager.run_before_transcription(self.current_task)
+                self.plugin_manager.run_before_transcription(self.current.task)
             except Exception as e:
                 logging.error(f"Plugin before_transcription failed: {e}", exc_info=True)
 
-            should_skip, skip_segments = self.plugin_manager.run_check_skip(self.current_task)
+            should_skip, skip_segments = self.plugin_manager.run_check_skip(self.current.task)
             if should_skip:
                 logging.debug("Skipping transcription task (plugin signaled skip)")
-                self.current_task.status = FileTranscriptionTask.Status.SKIPPED
+                self.current.task.status = FileTranscriptionTask.Status.SKIPPED
                 self.on_task_completed(skip_segments)
                 self.is_running = False
                 self._on_task_finished()
@@ -224,49 +229,49 @@ def _run_plugins(self) -> bool:
         return True
 
     def _create_transcriber(self):
-        model_type = self.current_task.transcription_options.model.model_type
+        model_type = self.current.task.transcription_options.model.model_type
         if model_type == ModelType.OPEN_AI_WHISPER_API:
-            self.current_transcriber = OpenAIWhisperAPIFileTranscriber(
-                task=self.current_task
+            self.current.transcriber = OpenAIWhisperAPIFileTranscriber(
+                task=self.current.task
             )
         elif (
             model_type == ModelType.WHISPER_CPP
             or model_type == ModelType.HUGGING_FACE
             or model_type == ModelType.WHISPER
             or model_type == ModelType.FASTER_WHISPER
         ):
-            self.current_transcriber = WhisperFileTranscriber(task=self.current_task)
+            self.current.transcriber = WhisperFileTranscriber(task=self.current.task)
         else:
             raise Exception(f"Unknown model type: {model_type}")
 
     def _setup_transcriber_thread(self):
-        self.current_transcriber_thread = QThread(self)
+        self.current.transcriber_thread = QThread(self)
 
-        self.current_transcriber.moveToThread(self.current_transcriber_thread)
+        self.current.transcriber.moveToThread(self.current.transcriber_thread)
 
-        self.current_transcriber_thread.started.connect(self.current_transcriber.run)
-        self.current_transcriber.completed.connect(self.current_transcriber_thread.quit)
-        self.current_transcriber.error.connect(self.current_transcriber_thread.quit)
+        self.current.transcriber_thread.started.connect(self.current.transcriber.run)
+        self.current.transcriber.completed.connect(self.current.transcriber_thread.quit)
+        self.current.transcriber.error.connect(self.current.transcriber_thread.quit)
 
-        self.current_transcriber.completed.connect(self.current_transcriber.deleteLater)
-        self.current_transcriber.error.connect(self.current_transcriber.deleteLater)
-        self.current_transcriber_thread.finished.connect(
-            self.current_transcriber_thread.deleteLater
+        self.current.transcriber.completed.connect(self.current.transcriber.deleteLater)
+        self.current.transcriber.error.connect(self.current.transcriber.deleteLater)
+        self.current.transcriber_thread.finished.connect(
+            self.current.transcriber_thread.deleteLater
         )
 
-        self.current_transcriber.progress.connect(self.on_task_progress)
-        self.current_transcriber.download_progress.connect(
+        self.current.transcriber.progress.connect(self.on_task_progress)
+        self.current.transcriber.download_progress.connect(
             self.on_task_download_progress
         )
-        self.current_transcriber.error.connect(self.on_task_error)
+        self.current.transcriber.error.connect(self.on_task_error)
 
-        self.current_transcriber.completed.connect(self.on_task_completed)
+        self.current.transcriber.completed.connect(self.on_task_completed)
 
-        self.current_transcriber.error.connect(lambda: self._on_task_finished())
-        self.current_transcriber.completed.connect(lambda: self._on_task_finished())
+        self.current.transcriber.error.connect(lambda: self._on_task_finished())
+        self.current.transcriber.completed.connect(lambda: self._on_task_finished())
 
-        self.task_started.emit(self.current_task)
-        self.current_transcriber_thread.start()
+        self.task_started.emit(self.current.task)
+        self.current.transcriber_thread.start()
 
     def _extract_speech(self, file_path: str, speech_path: str, device: str) -> str:
         """Run demucs speech extraction in a separate process.
@@ -298,7 +303,7 @@ def _extract_speech(self, file_path: str, speech_path: str, device: str) -> str:
                     audio_length = int(message[2] * 100)
                     if audio_length:
                         self.task_progress.emit(
-                            self.current_task,
+                            self.current.task,
                             int(message[1] * 100) / audio_length,
                         )
                 elif kind == "done":
@@ -364,44 +369,44 @@ def add_task(self, task: FileTranscriptionTask):
     def cancel_task(self, task_id: UUID):
         self.canceled_tasks.add(task_id)
 
-        if self.current_task is not None and self.current_task.uid == task_id:
+        if self.current.task is not None and self.current.task.uid == task_id:
             self._terminate_speech_extractor_process()
 
-            if self.current_transcriber is not None:
-                self.current_transcriber.stop()
+            if self.current.transcriber is not None:
+                self.current.transcriber.stop()
 
-            if self.current_transcriber_thread is not None:
-                if not self.current_transcriber_thread.wait(5000):
+            if self.current.transcriber_thread is not None:
+                if not self.current.transcriber_thread.wait(5000):
                     logging.warning("Transcriber thread did not terminate gracefully")
-                    self.current_transcriber_thread.terminate()
+                    self.current.transcriber_thread.terminate()
 
     def on_task_error(self, error: str):
         if (
-            self.current_task is not None
-            and self.current_task.uid not in self.canceled_tasks
+            self.current.task is not None
+            and self.current.task.uid not in self.canceled_tasks
         ):
             # Check if the error indicates cancellation
             if "canceled" in error.lower() or "cancelled" in error.lower():
-                self.current_task.status = FileTranscriptionTask.Status.CANCELED
-                self.current_task.error = error
+                self.current.task.status = FileTranscriptionTask.Status.CANCELED
+                self.current.task.error = error
             else:
-                self.current_task.status = FileTranscriptionTask.Status.FAILED
-                self.current_task.error = error
-            self.task_error.emit(self.current_task, error)
+                self.current.task.status = FileTranscriptionTask.Status.FAILED
+                self.current.task.error = error
+            self.task_error.emit(self.current.task, error)
 
     @pyqtSlot(tuple)
     def on_task_progress(self, progress: Tuple[int, int]):
-        if self.current_task is not None:
-            self.task_progress.emit(self.current_task, progress[0] / progress[1])
+        if self.current.task is not None:
+            self.task_progress.emit(self.current.task, progress[0] / progress[1])
 
     def on_task_download_progress(self, fraction_downloaded: float):
-        if self.current_task is not None:
-            self.task_download_progress.emit(self.current_task, fraction_downloaded)
+        if self.current.task is not None:
+            self.task_download_progress.emit(self.current.task, fraction_downloaded)
 
     @pyqtSlot(list)
     def on_task_completed(self, segments: List[Segment]):
-        if self.current_task is not None:
-            self.task_completed.emit(self.current_task, segments)
+        if self.current.task is not None:
+            self.task_completed.emit(self.current.task, segments)
 
         if self.speech_path is not None:
             try:
@@ -412,8 +417,8 @@ def on_task_completed(self, segments: List[Segment]):
 
     def stop(self):
         self.tasks_queue.put(None)
-        if self.current_transcriber is not None:
-            self.current_transcriber.stop()
+        if self.current.transcriber is not None:
+            self.current.transcriber.stop()
 
         # Terminate the speech extraction process if one is still running.
         self._terminate_speech_extractor_process()