feat: 中間データのmsgpack移行
This commit is contained in:
+171
-218
@@ -33,7 +33,6 @@ fix_library_path()
|
||||
|
||||
import threading
|
||||
import uuid
|
||||
import queue
|
||||
import traceback
|
||||
from typing import Dict, Optional, List
|
||||
from pathlib import Path
|
||||
@@ -43,6 +42,7 @@ from pydantic import BaseModel
|
||||
import uvicorn
|
||||
import cv2
|
||||
import numpy as np
|
||||
import msgpack
|
||||
|
||||
# Add project root to path for imports if needed
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
@@ -86,7 +86,7 @@ class GenerateRequest(BaseModel):
|
||||
|
||||
class BakeRequest(BaseModel):
|
||||
video_path: str
|
||||
mask_path: str
|
||||
detections_path: str
|
||||
output_path: str
|
||||
blur_size: int = 50
|
||||
format: str = "mp4"
|
||||
@@ -122,185 +122,79 @@ def _build_video_writer(
|
||||
|
||||
raise RuntimeError(f"Failed to create video writer for format='{fmt}'")
|
||||
|
||||
def _scale_bbox(
|
||||
x: int,
|
||||
y: int,
|
||||
w: int,
|
||||
h: int,
|
||||
scale: float,
|
||||
frame_width: int,
|
||||
frame_height: int,
|
||||
) -> Optional[List[int]]:
|
||||
"""Scale bbox around center and clamp to frame boundaries."""
|
||||
if w <= 0 or h <= 0:
|
||||
return None
|
||||
|
||||
center_x = x + (w * 0.5)
|
||||
center_y = y + (h * 0.5)
|
||||
scaled_w = max(1, int(w * scale))
|
||||
scaled_h = max(1, int(h * scale))
|
||||
|
||||
x1 = max(0, int(center_x - scaled_w * 0.5))
|
||||
y1 = max(0, int(center_y - scaled_h * 0.5))
|
||||
x2 = min(frame_width, x1 + scaled_w)
|
||||
y2 = min(frame_height, y1 + scaled_h)
|
||||
out_w = x2 - x1
|
||||
out_h = y2 - y1
|
||||
if out_w <= 0 or out_h <= 0:
|
||||
return None
|
||||
return [x1, y1, out_w, out_h]
|
||||
|
||||
|
||||
def process_video_task(task_id: str, req: GenerateRequest):
|
||||
"""Background task to process video with async MP4 output."""
|
||||
writer = None
|
||||
write_queue = None
|
||||
writer_thread = None
|
||||
"""Background task to detect faces and save bbox cache as msgpack."""
|
||||
cap = None
|
||||
|
||||
try:
|
||||
tasks[task_id].status = TaskStatus.PROCESSING
|
||||
cancel_event = cancel_events.get(task_id)
|
||||
|
||||
# Verify video exists
|
||||
if not os.path.exists(req.video_path):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = f"Video not found: {req.video_path}"
|
||||
return
|
||||
|
||||
# Initialize detector (will load model on first run)
|
||||
print(f"Loading detector for task {task_id}...")
|
||||
detector = get_detector(
|
||||
conf_threshold=req.conf_threshold,
|
||||
iou_threshold=req.iou_threshold
|
||||
iou_threshold=req.iou_threshold,
|
||||
)
|
||||
_ = detector.model
|
||||
|
||||
# Open video
|
||||
cap = cv2.VideoCapture(req.video_path)
|
||||
if not cap.isOpened():
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to open video"
|
||||
return
|
||||
|
||||
# Get video properties
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
end_frame = min(req.end_frame, total_video_frames - 1)
|
||||
frames_to_process = end_frame - req.start_frame + 1
|
||||
|
||||
tasks[task_id].total = frames_to_process
|
||||
|
||||
# Ensure output directory exists
|
||||
os.makedirs(req.output_dir, exist_ok=True)
|
||||
|
||||
# Setup MP4 writer (grayscale)
|
||||
output_video_path = os.path.join(req.output_dir, "mask.mp4")
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height), isColor=False)
|
||||
|
||||
if not writer.isOpened():
|
||||
if frames_to_process <= 0:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to create video writer"
|
||||
cap.release()
|
||||
tasks[task_id].message = "Invalid frame range"
|
||||
return
|
||||
|
||||
# Async writer setup
|
||||
write_queue = queue.Queue(maxsize=30) # Buffer up to 30 frames
|
||||
writer_running = threading.Event()
|
||||
writer_running.set()
|
||||
tasks[task_id].total = frames_to_process
|
||||
os.makedirs(req.output_dir, exist_ok=True)
|
||||
output_msgpack_path = os.path.join(req.output_dir, "detections.msgpack")
|
||||
|
||||
def async_writer():
|
||||
"""Background thread for writing frames to video."""
|
||||
while writer_running.is_set() or not write_queue.empty():
|
||||
try:
|
||||
mask = write_queue.get(timeout=0.1)
|
||||
if mask is not None:
|
||||
writer.write(mask)
|
||||
write_queue.task_done()
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
writer_thread = threading.Thread(target=async_writer, daemon=True)
|
||||
writer_thread.start()
|
||||
|
||||
print(f"Starting processing: {req.video_path} ({frames_to_process} frames) -> {output_video_path}")
|
||||
|
||||
# Batch processing configuration
|
||||
BATCH_SIZE = 5 # Optimal batch size for 4K video (72.9% improvement)
|
||||
frame_buffer = []
|
||||
TEMPORAL_SIDE_WEIGHT = 0.7
|
||||
TEMPORAL_CENTER_WEIGHT = 1.0
|
||||
|
||||
# Temporal blending state (streaming, low-memory)
|
||||
prev_mask = None
|
||||
curr_mask = None
|
||||
wrote_first_frame = False
|
||||
|
||||
def _scale_mask(mask: np.ndarray, weight: float) -> np.ndarray:
|
||||
"""Scale mask intensity for temporal blending."""
|
||||
if weight == 1.0:
|
||||
return mask
|
||||
return cv2.convertScaleAbs(mask, alpha=weight, beta=0)
|
||||
|
||||
def _blend_edge(base: np.ndarray, neighbor: np.ndarray) -> np.ndarray:
|
||||
"""Blend for first/last frame (one-sided temporal context)."""
|
||||
base_w = _scale_mask(base, TEMPORAL_CENTER_WEIGHT)
|
||||
neighbor_w = _scale_mask(neighbor, TEMPORAL_SIDE_WEIGHT)
|
||||
return cv2.max(base_w, neighbor_w)
|
||||
|
||||
def _blend_middle(prev: np.ndarray, cur: np.ndarray, nxt: np.ndarray) -> np.ndarray:
|
||||
"""Blend for middle frames (previous/current/next temporal context)."""
|
||||
prev_w = _scale_mask(prev, TEMPORAL_SIDE_WEIGHT)
|
||||
cur_w = _scale_mask(cur, TEMPORAL_CENTER_WEIGHT)
|
||||
nxt_w = _scale_mask(nxt, TEMPORAL_SIDE_WEIGHT)
|
||||
return cv2.max(cur_w, cv2.max(prev_w, nxt_w))
|
||||
|
||||
def push_mask_temporal(raw_mask: np.ndarray):
|
||||
"""Push mask and emit blended output in frame order."""
|
||||
nonlocal prev_mask, curr_mask, wrote_first_frame
|
||||
|
||||
if prev_mask is None:
|
||||
prev_mask = raw_mask
|
||||
return
|
||||
|
||||
if curr_mask is None:
|
||||
curr_mask = raw_mask
|
||||
return
|
||||
|
||||
if not wrote_first_frame:
|
||||
write_queue.put(_blend_edge(prev_mask, curr_mask))
|
||||
wrote_first_frame = True
|
||||
|
||||
# Emit blended current frame using prev/current/next
|
||||
write_queue.put(_blend_middle(prev_mask, curr_mask, raw_mask))
|
||||
|
||||
# Slide temporal window
|
||||
prev_mask = curr_mask
|
||||
curr_mask = raw_mask
|
||||
|
||||
def flush_temporal_tail():
|
||||
"""Flush remaining masks after all frames are processed."""
|
||||
if prev_mask is None:
|
||||
return
|
||||
|
||||
# Single-frame case
|
||||
if curr_mask is None:
|
||||
write_queue.put(_scale_mask(prev_mask, TEMPORAL_CENTER_WEIGHT))
|
||||
return
|
||||
|
||||
# Two-frame case
|
||||
if not wrote_first_frame:
|
||||
write_queue.put(_blend_edge(prev_mask, curr_mask))
|
||||
|
||||
# Always emit last frame with one-sided blend
|
||||
write_queue.put(_blend_edge(curr_mask, prev_mask))
|
||||
|
||||
def process_batch():
|
||||
"""Process accumulated batch of frames."""
|
||||
if not frame_buffer:
|
||||
return
|
||||
|
||||
# Batch inference at full resolution
|
||||
batch_detections = detector.detect_batch(frame_buffer)
|
||||
|
||||
# Generate masks for each frame
|
||||
for i, detections in enumerate(batch_detections):
|
||||
frame = frame_buffer[i]
|
||||
|
||||
# Generate mask at original resolution
|
||||
mask = detector.generate_mask(
|
||||
frame.shape,
|
||||
detections,
|
||||
mask_scale=req.mask_scale
|
||||
)
|
||||
|
||||
# Temporal blend before async write
|
||||
push_mask_temporal(mask)
|
||||
|
||||
# Clear buffer
|
||||
frame_buffer.clear()
|
||||
|
||||
# Seek once to the starting frame. Avoid random-access seek on every frame.
|
||||
if req.start_frame > 0:
|
||||
seek_ok = cap.set(cv2.CAP_PROP_POS_FRAMES, req.start_frame)
|
||||
if not seek_ok:
|
||||
print(
|
||||
f"[FaceMask] Warning: CAP_PROP_POS_FRAMES seek failed, "
|
||||
f"fallback to sequential skip ({req.start_frame} frames)"
|
||||
)
|
||||
for _ in range(req.start_frame):
|
||||
ret, _ = cap.read()
|
||||
if not ret:
|
||||
@@ -310,65 +204,97 @@ def process_video_task(task_id: str, req: GenerateRequest):
|
||||
)
|
||||
return
|
||||
|
||||
# Process loop with batching
|
||||
frame_buffer: List[np.ndarray] = []
|
||||
frame_detections: List[List[List[float]]] = []
|
||||
batch_size = 5
|
||||
current_count = 0
|
||||
for frame_idx in range(req.start_frame, end_frame + 1):
|
||||
|
||||
def process_batch():
|
||||
nonlocal current_count
|
||||
if not frame_buffer:
|
||||
return
|
||||
|
||||
batch_detections = detector.detect_batch(frame_buffer)
|
||||
for detections in batch_detections:
|
||||
packed_detections: List[List[float]] = []
|
||||
for x, y, w, h, conf in detections:
|
||||
scaled = _scale_bbox(
|
||||
int(x),
|
||||
int(y),
|
||||
int(w),
|
||||
int(h),
|
||||
float(req.mask_scale),
|
||||
width,
|
||||
height,
|
||||
)
|
||||
if scaled is None:
|
||||
continue
|
||||
packed_detections.append(
|
||||
[scaled[0], scaled[1], scaled[2], scaled[3], float(conf)]
|
||||
)
|
||||
frame_detections.append(packed_detections)
|
||||
current_count += 1
|
||||
tasks[task_id].progress = current_count
|
||||
|
||||
frame_buffer.clear()
|
||||
|
||||
print(
|
||||
f"Starting detection cache generation: {req.video_path} "
|
||||
f"({frames_to_process} frames) -> {output_msgpack_path}"
|
||||
)
|
||||
|
||||
for _ in range(req.start_frame, end_frame + 1):
|
||||
if cancel_event and cancel_event.is_set():
|
||||
tasks[task_id].status = TaskStatus.CANCELLED
|
||||
tasks[task_id].message = "Cancelled by user"
|
||||
break
|
||||
|
||||
# Read next frame sequentially (after one-time initial seek)
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
if ret:
|
||||
# Store frame for batch processing
|
||||
frame_buffer.append(frame)
|
||||
frame_buffer.append(frame)
|
||||
if len(frame_buffer) >= batch_size:
|
||||
process_batch()
|
||||
|
||||
# Process batch when full
|
||||
if len(frame_buffer) >= BATCH_SIZE:
|
||||
process_batch()
|
||||
|
||||
# Update progress
|
||||
current_count += 1
|
||||
tasks[task_id].progress = current_count
|
||||
|
||||
# Process remaining frames in buffer
|
||||
if frame_buffer:
|
||||
process_batch()
|
||||
flush_temporal_tail()
|
||||
|
||||
# Cleanup
|
||||
writer_running.clear()
|
||||
write_queue.join() # Wait for all frames to be written
|
||||
if writer_thread:
|
||||
writer_thread.join(timeout=5)
|
||||
|
||||
cap.release()
|
||||
if writer:
|
||||
writer.release()
|
||||
|
||||
if tasks[task_id].status == TaskStatus.PROCESSING:
|
||||
payload = {
|
||||
"version": 1,
|
||||
"video_path": req.video_path,
|
||||
"start_frame": req.start_frame,
|
||||
"end_frame": req.start_frame + len(frame_detections) - 1,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps,
|
||||
"mask_scale": float(req.mask_scale),
|
||||
"frames": frame_detections,
|
||||
}
|
||||
with open(output_msgpack_path, "wb") as f:
|
||||
f.write(msgpack.packb(payload, use_bin_type=True))
|
||||
|
||||
tasks[task_id].status = TaskStatus.COMPLETED
|
||||
tasks[task_id].result_path = output_video_path # Return video path
|
||||
tasks[task_id].message = "Processing completed successfully"
|
||||
print(f"Task {task_id} completed: {output_video_path}")
|
||||
|
||||
tasks[task_id].result_path = output_msgpack_path
|
||||
tasks[task_id].message = "Detection cache completed"
|
||||
print(f"Task {task_id} completed: {output_msgpack_path}")
|
||||
|
||||
except Exception as e:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = str(e)
|
||||
print(f"Error in task {task_id}: {e}")
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# Cleanup
|
||||
if cap:
|
||||
cap.release()
|
||||
if task_id in cancel_events:
|
||||
del cancel_events[task_id]
|
||||
|
||||
|
||||
def process_bake_task(task_id: str, req: BakeRequest):
|
||||
"""Background task to bake blur into a regular video file."""
|
||||
"""Background task to bake blur using bbox detections in msgpack."""
|
||||
src_cap = None
|
||||
mask_cap = None
|
||||
writer = None
|
||||
|
||||
try:
|
||||
@@ -379,59 +305,56 @@ def process_bake_task(task_id: str, req: BakeRequest):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = f"Video not found: {req.video_path}"
|
||||
return
|
||||
|
||||
if not os.path.exists(req.mask_path):
|
||||
if not os.path.exists(req.detections_path):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = f"Mask video not found: {req.mask_path}"
|
||||
tasks[task_id].message = f"Detections file not found: {req.detections_path}"
|
||||
return
|
||||
|
||||
src_cap = cv2.VideoCapture(req.video_path)
|
||||
mask_cap = cv2.VideoCapture(req.mask_path)
|
||||
|
||||
if not src_cap.isOpened():
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to open source video"
|
||||
return
|
||||
if not mask_cap.isOpened():
|
||||
|
||||
with open(req.detections_path, "rb") as f:
|
||||
payload = msgpack.unpackb(f.read(), raw=False)
|
||||
frames = payload.get("frames")
|
||||
if not isinstance(frames, list):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to open mask video"
|
||||
tasks[task_id].message = "Invalid detections format: 'frames' is missing"
|
||||
return
|
||||
|
||||
src_fps = src_cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
src_width = int(src_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
src_height = int(src_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
src_frames = int(src_cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
mask_frames = int(mask_cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
if src_width <= 0 or src_height <= 0:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Invalid source video dimensions"
|
||||
return
|
||||
|
||||
total = min(src_frames, mask_frames) if src_frames > 0 and mask_frames > 0 else 0
|
||||
total = min(src_frames, len(frames)) if src_frames > 0 else len(frames)
|
||||
if total <= 0:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Source/mask frame count is zero"
|
||||
tasks[task_id].message = "Source/detections frame count is zero"
|
||||
return
|
||||
tasks[task_id].total = total
|
||||
|
||||
output_dir = os.path.dirname(req.output_path)
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
writer = _build_video_writer(req.output_path, req.format, src_fps, src_width, src_height)
|
||||
|
||||
# Kernel size must be odd and >= 1
|
||||
blur_size = max(1, int(req.blur_size))
|
||||
if blur_size % 2 == 0:
|
||||
blur_size += 1
|
||||
feather_radius = max(3, min(25, blur_size // 3))
|
||||
feather_kernel = feather_radius * 2 + 1
|
||||
|
||||
print(f"[FaceMask] Starting blur bake: {req.video_path} + {req.mask_path} -> {req.output_path}")
|
||||
if src_frames != mask_frames:
|
||||
print(
|
||||
f"[FaceMask] Warning: frame count mismatch "
|
||||
f"(src={src_frames}, mask={mask_frames}), processing {total} frames"
|
||||
)
|
||||
print(
|
||||
f"[FaceMask] Starting blur bake (bbox-msgpack): {req.video_path} + "
|
||||
f"{req.detections_path} -> {req.output_path}"
|
||||
)
|
||||
|
||||
for idx in range(total):
|
||||
if cancel_event and cancel_event.is_set():
|
||||
@@ -440,29 +363,61 @@ def process_bake_task(task_id: str, req: BakeRequest):
|
||||
break
|
||||
|
||||
src_ok, src_frame = src_cap.read()
|
||||
mask_ok, mask_frame = mask_cap.read()
|
||||
if not src_ok or not mask_ok:
|
||||
if not src_ok:
|
||||
break
|
||||
|
||||
if mask_frame.ndim == 3:
|
||||
mask_gray = cv2.cvtColor(mask_frame, cv2.COLOR_BGR2GRAY)
|
||||
else:
|
||||
mask_gray = mask_frame
|
||||
frame_boxes = frames[idx] if idx < len(frames) else []
|
||||
if not frame_boxes:
|
||||
writer.write(src_frame)
|
||||
tasks[task_id].progress = idx + 1
|
||||
continue
|
||||
|
||||
if mask_gray.shape[0] != src_height or mask_gray.shape[1] != src_width:
|
||||
mask_gray = cv2.resize(
|
||||
mask_gray,
|
||||
(src_width, src_height),
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
mask_gray = np.zeros((src_height, src_width), dtype=np.uint8)
|
||||
for box in frame_boxes:
|
||||
if not isinstance(box, list) or len(box) < 4:
|
||||
continue
|
||||
x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
|
||||
if w <= 0 or h <= 0:
|
||||
continue
|
||||
center = (x + w // 2, y + h // 2)
|
||||
axes = (max(1, w // 2), max(1, h // 2))
|
||||
cv2.ellipse(mask_gray, center, axes, 0, 0, 360, 255, -1)
|
||||
|
||||
blurred = cv2.GaussianBlur(src_frame, (blur_size, blur_size), 0)
|
||||
alpha = (mask_gray.astype(np.float32) / 255.0)[..., np.newaxis]
|
||||
composed = (src_frame.astype(np.float32) * (1.0 - alpha)) + (
|
||||
blurred.astype(np.float32) * alpha
|
||||
if cv2.countNonZero(mask_gray) == 0:
|
||||
writer.write(src_frame)
|
||||
tasks[task_id].progress = idx + 1
|
||||
continue
|
||||
|
||||
mask_gray = cv2.GaussianBlur(mask_gray, (feather_kernel, feather_kernel), 0)
|
||||
_, mask_binary = cv2.threshold(mask_gray, 2, 255, cv2.THRESH_BINARY)
|
||||
non_zero_coords = cv2.findNonZero(mask_binary)
|
||||
if non_zero_coords is None:
|
||||
writer.write(src_frame)
|
||||
tasks[task_id].progress = idx + 1
|
||||
continue
|
||||
|
||||
x, y, w, h = cv2.boundingRect(non_zero_coords)
|
||||
blur_margin = max(1, (blur_size // 2) + feather_radius)
|
||||
x1 = max(0, x - blur_margin)
|
||||
y1 = max(0, y - blur_margin)
|
||||
x2 = min(src_width, x + w + blur_margin)
|
||||
y2 = min(src_height, y + h + blur_margin)
|
||||
|
||||
roi_src = src_frame[y1:y2, x1:x2]
|
||||
roi_mask = mask_gray[y1:y2, x1:x2]
|
||||
if roi_src.size == 0:
|
||||
writer.write(src_frame)
|
||||
tasks[task_id].progress = idx + 1
|
||||
continue
|
||||
|
||||
roi_blurred = cv2.GaussianBlur(roi_src, (blur_size, blur_size), 0)
|
||||
roi_alpha = (roi_mask.astype(np.float32) / 255.0)[..., np.newaxis]
|
||||
roi_composed = (roi_src.astype(np.float32) * (1.0 - roi_alpha)) + (
|
||||
roi_blurred.astype(np.float32) * roi_alpha
|
||||
)
|
||||
writer.write(np.clip(composed, 0, 255).astype(np.uint8))
|
||||
|
||||
output_frame = src_frame.copy()
|
||||
output_frame[y1:y2, x1:x2] = np.clip(roi_composed, 0, 255).astype(np.uint8)
|
||||
writer.write(output_frame)
|
||||
tasks[task_id].progress = idx + 1
|
||||
|
||||
if tasks[task_id].status == TaskStatus.PROCESSING:
|
||||
@@ -479,8 +434,6 @@ def process_bake_task(task_id: str, req: BakeRequest):
|
||||
finally:
|
||||
if src_cap:
|
||||
src_cap.release()
|
||||
if mask_cap:
|
||||
mask_cap.release()
|
||||
if writer:
|
||||
writer.release()
|
||||
if task_id in cancel_events:
|
||||
|
||||
Reference in New Issue
Block a user