""" Async Mask Generator using Thread + Queue + Timer pattern. This module provides non-blocking face mask generation for Blender. Heavy processing (face detection) runs in a worker thread while Blender's UI remains responsive via bpy.app.timers. """ import os import threading import queue from functools import partial from typing import Optional, Callable, Tuple from pathlib import Path # Will be imported when running inside Blender bpy = None class AsyncMaskGenerator: """ Asynchronous mask generator that doesn't block Blender's UI. Uses Thread + Queue + Timer pattern: - Worker thread: Face detection (can use bpy-unsafe operations) - Main thread timer: UI updates and bpy operations """ def __init__(self): self.result_queue: queue.Queue = queue.Queue() self.progress_queue: queue.Queue = queue.Queue() self.worker_thread: Optional[threading.Thread] = None self.is_running: bool = False self.total_frames: int = 0 self.current_frame: int = 0 self._on_complete: Optional[Callable] = None self._on_progress: Optional[Callable] = None def start( self, video_path: str, output_dir: str, start_frame: int, end_frame: int, fps: float, scale_factor: float = 1.1, min_neighbors: int = 5, mask_scale: float = 1.5, on_complete: Optional[Callable] = None, on_progress: Optional[Callable] = None, ): """ Start asynchronous mask generation. Args: video_path: Path to source video file output_dir: Directory to save mask images start_frame: First frame to process end_frame: Last frame to process fps: Video frame rate (for seeking) scale_factor: Face detection scale factor min_neighbors: Face detection min neighbors mask_scale: Mask region scale factor on_complete: Callback when processing completes (called from main thread) on_progress: Callback for progress updates (called from main thread) """ global bpy import bpy as _bpy bpy = _bpy if self.is_running: raise RuntimeError("Mask generation already in progress") print(f"[FaceMask] Starting mask generation: {video_path}") print(f"[FaceMask] Output directory: {output_dir}") print(f"[FaceMask] Frame range: {start_frame} - {end_frame}") self.is_running = True self.total_frames = end_frame - start_frame + 1 self.current_frame = 0 self._on_complete = on_complete self._on_progress = on_progress # Ensure output directory exists Path(output_dir).mkdir(parents=True, exist_ok=True) # Start worker thread self.worker_thread = threading.Thread( target=self._worker, args=( video_path, output_dir, start_frame, end_frame, fps, scale_factor, min_neighbors, mask_scale, ), daemon=True, ) self.worker_thread.start() # Register timer for main thread callbacks bpy.app.timers.register( self._check_progress, first_interval=0.1, ) def cancel(self): """Cancel the current processing.""" self.is_running = False if self.worker_thread and self.worker_thread.is_alive(): self.worker_thread.join(timeout=2.0) def _worker( self, video_path: str, output_dir: str, start_frame: int, end_frame: int, fps: float, scale_factor: float, min_neighbors: int, mask_scale: float, ): """ Worker thread function. Runs face detection and saves masks. IMPORTANT: Do NOT use bpy in this function! """ try: import cv2 print(f"[FaceMask] OpenCV loaded: {cv2.__version__}") from .face_detector import FaceDetector except ImportError as e: print(f"[FaceMask] Import error: {e}") self.result_queue.put(("error", str(e))) return try: # Initialize detector detector = FaceDetector( scale_factor=scale_factor, min_neighbors=min_neighbors, ) # Open video cap = cv2.VideoCapture(video_path) if not cap.isOpened(): print(f"[FaceMask] Failed to open video: {video_path}") self.result_queue.put(("error", f"Failed to open video: {video_path}")) return total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) print(f"[FaceMask] Video opened, total frames: {total_video_frames}") # Process frames for frame_idx in range(start_frame, end_frame + 1): if not self.is_running: self.result_queue.put(("cancelled", None)) return # Seek to frame cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) ret, frame = cap.read() if not ret: # Skip unreadable frames continue # Detect faces detections = detector.detect(frame) # Generate mask mask = detector.generate_mask( frame.shape, detections, mask_scale=mask_scale, ) # Save mask mask_filename = f"mask_{frame_idx:06d}.png" mask_path = os.path.join(output_dir, mask_filename) cv2.imwrite(mask_path, mask) # Report progress self.progress_queue.put(("progress", frame_idx - start_frame + 1)) cap.release() # Report completion self.result_queue.put(("done", output_dir)) except Exception as e: import traceback print(f"[FaceMask] Error: {e}") traceback.print_exc() self.result_queue.put(("error", str(e))) def _check_progress(self) -> Optional[float]: """ Timer callback for checking progress from main thread. Returns: Time until next call, or None to unregister. """ # Process all pending progress updates while not self.progress_queue.empty(): try: msg_type, data = self.progress_queue.get_nowait() if msg_type == "progress": self.current_frame = data if self._on_progress: self._on_progress(self.current_frame, self.total_frames) except queue.Empty: break # Check for completion if not self.result_queue.empty(): try: msg_type, data = self.result_queue.get_nowait() self.is_running = False if self._on_complete: self._on_complete(msg_type, data) return None # Unregister timer except queue.Empty: pass # Continue checking if still running if self.is_running: return 0.1 # Check again in 100ms return None # Unregister timer # Global instance for easy access from operators _generator: Optional[AsyncMaskGenerator] = None def get_generator() -> AsyncMaskGenerator: """Get or create the global mask generator instance.""" global _generator if _generator is None: _generator = AsyncMaskGenerator() return _generator