250 lines
8.0 KiB
Python
250 lines
8.0 KiB
Python
"""
|
|
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
|