Blur Bake

server/main.py

@@ -83,6 +83,45 @@ class GenerateRequest(BaseModel):
     iou_threshold: float = 0.45
     mask_scale: float = 1.5
 
+
+class BakeRequest(BaseModel):
+    video_path: str
+    mask_path: str
+    output_path: str
+    blur_size: int = 50
+    format: str = "mp4"
+
+
+def _build_video_writer(
+    output_path: str,
+    fmt: str,
+    fps: float,
+    width: int,
+    height: int,
+) -> cv2.VideoWriter:
+    """Create VideoWriter with codec fallback per format."""
+    format_key = fmt.lower()
+    codec_candidates = {
+        "mp4": ["avc1", "mp4v"],
+        "mov": ["avc1", "mp4v"],
+        "avi": ["MJPG", "XVID"],
+    }.get(format_key, ["mp4v"])
+
+    for codec in codec_candidates:
+        writer = cv2.VideoWriter(
+            output_path,
+            cv2.VideoWriter_fourcc(*codec),
+            fps,
+            (width, height),
+            isColor=True,
+        )
+        if writer.isOpened():
+            print(f"[FaceMask] Using output codec: {codec}")
+            return writer
+        writer.release()
+
+    raise RuntimeError(f"Failed to create video writer for format='{fmt}'")
+
 def process_video_task(task_id: str, req: GenerateRequest):
     """Background task to process video with async MP4 output."""
     writer = None
@@ -162,6 +201,72 @@ def process_video_task(task_id: str, req: GenerateRequest):
         # 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."""
@@ -182,8 +287,8 @@ def process_video_task(task_id: str, req: GenerateRequest):
                     mask_scale=req.mask_scale
                 )
 
-                # Async write to queue
-                write_queue.put(mask)
+                # Temporal blend before async write
+                push_mask_temporal(mask)
 
             # Clear buffer
             frame_buffer.clear()
@@ -231,6 +336,7 @@ def process_video_task(task_id: str, req: GenerateRequest):
         # Process remaining frames in buffer
         if frame_buffer:
             process_batch()
+        flush_temporal_tail()
 
         # Cleanup
         writer_running.clear()
@@ -258,6 +364,128 @@ def process_video_task(task_id: str, req: GenerateRequest):
         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."""
+    src_cap = None
+    mask_cap = None
+    writer = None
+
+    try:
+        tasks[task_id].status = TaskStatus.PROCESSING
+        cancel_event = cancel_events.get(task_id)
+
+        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
+
+        if not os.path.exists(req.mask_path):
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = f"Mask video not found: {req.mask_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():
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = "Failed to open mask video"
+            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
+        if total <= 0:
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = "Source/mask 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
+
+        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"
+            )
+
+        for idx in range(total):
+            if cancel_event and cancel_event.is_set():
+                tasks[task_id].status = TaskStatus.CANCELLED
+                tasks[task_id].message = "Cancelled by user"
+                break
+
+            src_ok, src_frame = src_cap.read()
+            mask_ok, mask_frame = mask_cap.read()
+            if not src_ok or not mask_ok:
+                break
+
+            if mask_frame.ndim == 3:
+                mask_gray = cv2.cvtColor(mask_frame, cv2.COLOR_BGR2GRAY)
+            else:
+                mask_gray = mask_frame
+
+            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,
+                )
+
+            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
+            )
+            writer.write(np.clip(composed, 0, 255).astype(np.uint8))
+
+            tasks[task_id].progress = idx + 1
+
+        if tasks[task_id].status == TaskStatus.PROCESSING:
+            tasks[task_id].status = TaskStatus.COMPLETED
+            tasks[task_id].result_path = req.output_path
+            tasks[task_id].message = "Blur bake completed"
+            print(f"[FaceMask] Bake completed: {req.output_path}")
+
+    except Exception as e:
+        tasks[task_id].status = TaskStatus.FAILED
+        tasks[task_id].message = str(e)
+        print(f"Error in bake task {task_id}: {e}")
+        traceback.print_exc()
+    finally:
+        if src_cap:
+            src_cap.release()
+        if mask_cap:
+            mask_cap.release()
+        if writer:
+            writer.release()
+        if task_id in cancel_events:
+            del cancel_events[task_id]
+
 def check_gpu_available() -> dict:
     """
     Check if GPU is available for inference.
@@ -418,6 +646,17 @@ def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTas
     background_tasks.add_task(process_video_task, task_id, req)
     return task
 
+
+@app.post("/bake_blur", response_model=Task)
+def bake_blur_endpoint(req: BakeRequest, background_tasks: BackgroundTasks):
+    task_id = str(uuid.uuid4())
+    task = Task(id=task_id, status=TaskStatus.PENDING)
+    tasks[task_id] = task
+    cancel_events[task_id] = threading.Event()
+
+    background_tasks.add_task(process_bake_task, task_id, req)
+    return task
+
 @app.get("/tasks/{task_id}", response_model=Task)
 def get_task(task_id: str):
     if task_id not in tasks: