YOLO
This commit is contained in:
@@ -1,5 +1,4 @@
|
||||
"""Core module exports."""
|
||||
|
||||
from .face_detector import FaceDetector
|
||||
from .async_generator import AsyncMaskGenerator, get_generator
|
||||
from .compositor_setup import create_mask_blur_node_tree, get_or_create_blur_node_tree
|
||||
|
||||
+50
-65
@@ -43,14 +43,14 @@ class AsyncMaskGenerator:
|
||||
start_frame: int,
|
||||
end_frame: int,
|
||||
fps: float,
|
||||
scale_factor: float = 1.1,
|
||||
min_neighbors: int = 5,
|
||||
conf_threshold: float = 0.5,
|
||||
iou_threshold: float = 0.45,
|
||||
mask_scale: float = 1.5,
|
||||
on_complete: Optional[Callable] = None,
|
||||
on_progress: Optional[Callable] = None,
|
||||
):
|
||||
"""
|
||||
Start asynchronous mask generation.
|
||||
Start asynchronous mask generation with YOLO GPU acceleration.
|
||||
|
||||
Args:
|
||||
video_path: Path to source video file
|
||||
@@ -58,8 +58,8 @@ class AsyncMaskGenerator:
|
||||
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
|
||||
conf_threshold: YOLO confidence threshold
|
||||
iou_threshold: YOLO NMS IoU threshold
|
||||
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)
|
||||
@@ -93,8 +93,8 @@ class AsyncMaskGenerator:
|
||||
start_frame,
|
||||
end_frame,
|
||||
fps,
|
||||
scale_factor,
|
||||
min_neighbors,
|
||||
conf_threshold,
|
||||
iou_threshold,
|
||||
mask_scale,
|
||||
),
|
||||
daemon=True,
|
||||
@@ -120,77 +120,62 @@ class AsyncMaskGenerator:
|
||||
start_frame: int,
|
||||
end_frame: int,
|
||||
fps: float,
|
||||
scale_factor: float,
|
||||
min_neighbors: int,
|
||||
conf_threshold: float,
|
||||
iou_threshold: float,
|
||||
mask_scale: float,
|
||||
):
|
||||
"""
|
||||
Worker thread function. Runs face detection and saves masks.
|
||||
|
||||
IMPORTANT: Do NOT use bpy in this function!
|
||||
Worker thread function. Delegates to inference server and polls status.
|
||||
"""
|
||||
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
|
||||
import time
|
||||
from .inference_client import get_client
|
||||
|
||||
try:
|
||||
# Initialize detector
|
||||
detector = FaceDetector(
|
||||
scale_factor=scale_factor,
|
||||
min_neighbors=min_neighbors,
|
||||
client = get_client()
|
||||
|
||||
# Start task on server
|
||||
print(f"[FaceMask] Requesting generation on server...")
|
||||
task_id = client.generate_mask(
|
||||
video_path=video_path,
|
||||
output_dir=output_dir,
|
||||
start_frame=start_frame,
|
||||
end_frame=end_frame,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
mask_scale=mask_scale,
|
||||
)
|
||||
print(f"[FaceMask] Task started: {task_id}")
|
||||
|
||||
# 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:
|
||||
# Poll loop
|
||||
while self.is_running:
|
||||
status = client.get_task_status(task_id)
|
||||
state = status.get("status")
|
||||
|
||||
if state == "completed":
|
||||
self.result_queue.put(("done", output_dir))
|
||||
return
|
||||
|
||||
elif state == "failed":
|
||||
error_msg = status.get("message", "Unknown server error")
|
||||
print(f"[FaceMask] Server task failed: {error_msg}")
|
||||
self.result_queue.put(("error", error_msg))
|
||||
return
|
||||
|
||||
elif state == "cancelled":
|
||||
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))
|
||||
progress = status.get("progress", 0)
|
||||
if progress > 0:
|
||||
self.progress_queue.put(("progress", progress))
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
cap.release()
|
||||
|
||||
# Report completion
|
||||
self.result_queue.put(("done", output_dir))
|
||||
# If loop exited but task not done, cancel server task
|
||||
print("[FaceMask] Cancelling server task...")
|
||||
client.cancel_task(task_id)
|
||||
self.result_queue.put(("cancelled", None))
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1,160 +0,0 @@
|
||||
"""
|
||||
Face detector using OpenCV Haar Cascades.
|
||||
|
||||
This module provides face detection functionality optimized for
|
||||
privacy blur in video editing workflows.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List, Tuple, Optional
|
||||
import numpy as np
|
||||
|
||||
|
||||
class FaceDetector:
|
||||
"""
|
||||
Face detector using OpenCV Haar Cascades.
|
||||
|
||||
Optimized for privacy blur use case:
|
||||
- Detects frontal faces
|
||||
- Configurable detection sensitivity
|
||||
- Generates feathered masks for smooth blur edges
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scale_factor: float = 1.1,
|
||||
min_neighbors: int = 5,
|
||||
min_size: Tuple[int, int] = (30, 30),
|
||||
):
|
||||
"""
|
||||
Initialize the face detector.
|
||||
|
||||
Args:
|
||||
scale_factor: Image pyramid scale factor
|
||||
min_neighbors: Minimum neighbors for detection
|
||||
min_size: Minimum face size in pixels
|
||||
"""
|
||||
self.scale_factor = scale_factor
|
||||
self.min_neighbors = min_neighbors
|
||||
self.min_size = min_size
|
||||
self._classifier = None
|
||||
|
||||
@property
|
||||
def classifier(self):
|
||||
"""Lazy-load the Haar cascade classifier."""
|
||||
if self._classifier is None:
|
||||
import cv2
|
||||
|
||||
# Use haarcascade for frontal face detection
|
||||
cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
|
||||
|
||||
if not os.path.exists(cascade_path):
|
||||
raise RuntimeError(f"Haar cascade not found: {cascade_path}")
|
||||
|
||||
self._classifier = cv2.CascadeClassifier(cascade_path)
|
||||
|
||||
return self._classifier
|
||||
|
||||
def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int]]:
|
||||
"""
|
||||
Detect faces in a frame.
|
||||
|
||||
Args:
|
||||
frame: BGR image as numpy array
|
||||
|
||||
Returns:
|
||||
List of face bounding boxes as (x, y, width, height)
|
||||
"""
|
||||
import cv2
|
||||
|
||||
# Convert to grayscale for detection
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Detect faces
|
||||
faces = self.classifier.detectMultiScale(
|
||||
gray,
|
||||
scaleFactor=self.scale_factor,
|
||||
minNeighbors=self.min_neighbors,
|
||||
minSize=self.min_size,
|
||||
flags=cv2.CASCADE_SCALE_IMAGE,
|
||||
)
|
||||
|
||||
# Convert to list of tuples
|
||||
return [tuple(face) for face in faces]
|
||||
|
||||
def generate_mask(
|
||||
self,
|
||||
frame_shape: Tuple[int, int, int],
|
||||
detections: List[Tuple[int, int, int, int]],
|
||||
mask_scale: float = 1.5,
|
||||
feather_radius: int = 20,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Generate a mask image from face detections.
|
||||
|
||||
Args:
|
||||
frame_shape: Shape of the original frame (height, width, channels)
|
||||
detections: List of face bounding boxes
|
||||
mask_scale: Scale factor for mask region (1.0 = exact bounding box)
|
||||
feather_radius: Radius for edge feathering
|
||||
|
||||
Returns:
|
||||
Grayscale mask image (white = blur, black = keep)
|
||||
"""
|
||||
import cv2
|
||||
|
||||
height, width = frame_shape[:2]
|
||||
mask = np.zeros((height, width), dtype=np.uint8)
|
||||
|
||||
for (x, y, w, h) in detections:
|
||||
# Scale the bounding box
|
||||
center_x = x + w // 2
|
||||
center_y = y + h // 2
|
||||
|
||||
scaled_w = int(w * mask_scale)
|
||||
scaled_h = int(h * mask_scale)
|
||||
|
||||
# Calculate scaled bounding box
|
||||
x1 = max(0, center_x - scaled_w // 2)
|
||||
y1 = max(0, center_y - scaled_h // 2)
|
||||
x2 = min(width, center_x + scaled_w // 2)
|
||||
y2 = min(height, center_y + scaled_h // 2)
|
||||
|
||||
# Draw ellipse for more natural face shape
|
||||
cv2.ellipse(
|
||||
mask,
|
||||
(center_x, center_y),
|
||||
(scaled_w // 2, scaled_h // 2),
|
||||
0, # angle
|
||||
0, 360, # arc
|
||||
255, # color (white)
|
||||
-1, # filled
|
||||
)
|
||||
|
||||
# Apply Gaussian blur for feathering
|
||||
if feather_radius > 0 and len(detections) > 0:
|
||||
# Ensure kernel size is odd
|
||||
kernel_size = feather_radius * 2 + 1
|
||||
mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def detect_faces_batch(
|
||||
frames: List[np.ndarray],
|
||||
detector: Optional[FaceDetector] = None,
|
||||
) -> List[List[Tuple[int, int, int, int]]]:
|
||||
"""
|
||||
Detect faces in multiple frames.
|
||||
|
||||
Args:
|
||||
frames: List of BGR images
|
||||
detector: Optional detector instance (creates one if not provided)
|
||||
|
||||
Returns:
|
||||
List of detection lists, one per frame
|
||||
"""
|
||||
if detector is None:
|
||||
detector = FaceDetector()
|
||||
|
||||
return [detector.detect(frame) for frame in frames]
|
||||
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
Client for interacting with the external inference server.
|
||||
|
||||
Manages the server process and handles HTTP communication
|
||||
using standard library (avoiding requests dependency).
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
import json
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
import threading
|
||||
import os
|
||||
import signal
|
||||
from typing import Optional, Dict, Any, Tuple
|
||||
|
||||
|
||||
class InferenceClient:
|
||||
"""Client for the YOLO inference server."""
|
||||
|
||||
SERVER_URL = "http://127.0.0.1:8181"
|
||||
|
||||
def __init__(self):
|
||||
self.server_process: Optional[subprocess.Popen] = None
|
||||
self._server_lock = threading.Lock()
|
||||
|
||||
def start_server(self):
|
||||
"""Start the inference server process."""
|
||||
with self._server_lock:
|
||||
if self.is_server_running():
|
||||
return
|
||||
|
||||
print("[FaceMask] Starting inference server...")
|
||||
|
||||
# Find project root
|
||||
# Assuming this file is in core/inference_client.py
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
server_script = os.path.join(root_dir, "server", "main.py")
|
||||
|
||||
# Use system python (assumed to have dependencies via Nix/venv)
|
||||
# In user's environment, 'python' should refer to the environment python
|
||||
python_cmd = "python"
|
||||
|
||||
# Start process
|
||||
self.server_process = subprocess.Popen(
|
||||
[python_cmd, server_script],
|
||||
cwd=root_dir,
|
||||
text=True,
|
||||
preexec_fn=os.setsid, # Create new process group
|
||||
)
|
||||
|
||||
# Wait for startup
|
||||
for _ in range(20): # Wait up to 10 seconds
|
||||
if self.is_server_running():
|
||||
print("[FaceMask] Server started successfully")
|
||||
return
|
||||
|
||||
# Check if process died
|
||||
if self.server_process.poll() is not None:
|
||||
raise RuntimeError(f"Server failed to start (rc={self.server_process.returncode})")
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
raise RuntimeError("Server startup timed out")
|
||||
|
||||
def stop_server(self):
|
||||
"""Stop the inference server."""
|
||||
with self._server_lock:
|
||||
if self.server_process:
|
||||
print("[FaceMask] Stopping server...")
|
||||
try:
|
||||
os.killpg(os.getpgid(self.server_process.pid), signal.SIGTERM)
|
||||
self.server_process.wait(timeout=3)
|
||||
except (ProcessLookupError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
finally:
|
||||
self.server_process = None
|
||||
|
||||
def is_server_running(self) -> bool:
|
||||
"""Check if server is responding."""
|
||||
try:
|
||||
with urllib.request.urlopen(f"{self.SERVER_URL}/status", timeout=1) as response:
|
||||
return response.status == 200
|
||||
except (urllib.error.URLError, ConnectionRefusedError, TimeoutError):
|
||||
return False
|
||||
|
||||
def generate_mask(
|
||||
self,
|
||||
video_path: str,
|
||||
output_dir: str,
|
||||
start_frame: int,
|
||||
end_frame: int,
|
||||
conf_threshold: float,
|
||||
iou_threshold: float,
|
||||
mask_scale: float,
|
||||
) -> str:
|
||||
"""
|
||||
Request mask generation.
|
||||
|
||||
Returns:
|
||||
task_id (str)
|
||||
"""
|
||||
if not self.is_server_running():
|
||||
self.start_server()
|
||||
|
||||
data = {
|
||||
"video_path": video_path,
|
||||
"output_dir": output_dir,
|
||||
"start_frame": start_frame,
|
||||
"end_frame": end_frame,
|
||||
"conf_threshold": conf_threshold,
|
||||
"iou_threshold": iou_threshold,
|
||||
"mask_scale": mask_scale,
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{self.SERVER_URL}/generate",
|
||||
data=json.dumps(data).encode('utf-8'),
|
||||
headers={'Content-Type': 'application/json'},
|
||||
method='POST'
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req) as response:
|
||||
result = json.loads(response.read().decode('utf-8'))
|
||||
return result['id']
|
||||
except urllib.error.HTTPError as e:
|
||||
raise RuntimeError(f"Server error: {e.read().decode('utf-8')}")
|
||||
|
||||
def get_task_status(self, task_id: str) -> Dict[str, Any]:
|
||||
"""Get status of a task."""
|
||||
try:
|
||||
with urllib.request.urlopen(f"{self.SERVER_URL}/tasks/{task_id}") as response:
|
||||
return json.loads(response.read().decode('utf-8'))
|
||||
except urllib.error.HTTPError:
|
||||
return {"status": "unknown"}
|
||||
|
||||
def cancel_task(self, task_id: str):
|
||||
"""Cancel a task."""
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
f"{self.SERVER_URL}/tasks/{task_id}/cancel",
|
||||
method='POST'
|
||||
)
|
||||
with urllib.request.urlopen(req):
|
||||
pass
|
||||
except urllib.error.HTTPError:
|
||||
pass
|
||||
|
||||
|
||||
# Singleton
|
||||
_client: Optional[InferenceClient] = None
|
||||
|
||||
def get_client() -> InferenceClient:
|
||||
global _client
|
||||
if _client is None:
|
||||
_client = InferenceClient()
|
||||
return _client
|
||||
Reference in New Issue
Block a user