feat: 静画に対応

server/main.py

@@ -132,6 +132,25 @@ class BakeRequest(BaseModel):
     format: str = "mp4"
 
 
+class GenerateImagesRequest(BaseModel):
+    image_dir: str
+    filenames: List[str]
+    output_dir: str
+    start_index: int = 0
+    end_index: int = -1
+    conf_threshold: float = 0.5
+    iou_threshold: float = 0.45
+
+
+class BakeImagesRequest(BaseModel):
+    image_dir: str
+    filenames: List[str]
+    output_dir: str
+    detections_path: str
+    blur_size: int = 50
+    display_scale: float = 1.0
+
+
 class _FFmpegPipeWriter:
     """Write BGR frames to ffmpeg stdin."""
 
@@ -302,6 +321,267 @@ def _scale_bbox(
     return [x1, y1, out_w, out_h]
 
 
+def _apply_face_blur_inplace(
+    frame: np.ndarray,
+    frame_boxes: list,
+    src_width: int,
+    src_height: int,
+    blur_size: int,
+    display_scale: float,
+    blur_margin: int,
+) -> None:
+    """検出済み顔領域にガウスぼかしを適用する（in-place）。"""
+    if not frame_boxes:
+        return
+
+    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
+
+        cx = x + w / 2
+        cy = y + h / 2
+        dw = max(1, int(w * display_scale))
+        dh = max(1, int(h * display_scale))
+        dx = int(cx - dw / 2)
+        dy = int(cy - dh / 2)
+
+        roi_x1 = max(0, dx - blur_margin)
+        roi_y1 = max(0, dy - blur_margin)
+        roi_x2 = min(src_width, dx + dw + blur_margin)
+        roi_y2 = min(src_height, dy + dh + blur_margin)
+        roi_width = roi_x2 - roi_x1
+        roi_height = roi_y2 - roi_y1
+        if roi_width <= 0 or roi_height <= 0:
+            continue
+
+        roi_src = frame[roi_y1:roi_y2, roi_x1:roi_x2]
+        small_w = max(1, roi_width // 2)
+        small_h = max(1, roi_height // 2)
+        roi_small = cv2.resize(roi_src, (small_w, small_h), interpolation=cv2.INTER_LINEAR)
+        small_blur_size = max(3, (blur_size // 2) | 1)
+        roi_small_blurred = cv2.GaussianBlur(roi_small, (small_blur_size, small_blur_size), 0)
+        roi_blurred = cv2.resize(roi_small_blurred, (roi_width, roi_height), interpolation=cv2.INTER_LINEAR)
+
+        roi_mask = np.zeros((roi_height, roi_width), dtype=np.uint8)
+        center = (int(cx) - roi_x1, int(cy) - roi_y1)
+        axes = (max(1, dw // 2), max(1, dh // 2))
+        cv2.ellipse(roi_mask, center, axes, 0, 0, 360, 255, -1)
+
+        result = roi_src.copy()
+        cv2.copyTo(roi_blurred, roi_mask, result)
+        frame[roi_y1:roi_y2, roi_x1:roi_x2] = result
+
+
+def process_images_task(task_id: str, req: GenerateImagesRequest):
+    """画像シーケンスから顔を検出して msgpack キャッシュを保存する。"""
+    try:
+        tasks[task_id].status = TaskStatus.PROCESSING
+        cancel_event = cancel_events.get(task_id)
+
+        if not os.path.exists(req.image_dir):
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = f"Image directory not found: {req.image_dir}"
+            return
+        if not req.filenames:
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = "No filenames provided"
+            return
+
+        detector = get_detector(
+            conf_threshold=req.conf_threshold,
+            iou_threshold=req.iou_threshold,
+        )
+        _ = detector.model
+
+        total_files = len(req.filenames)
+        start_idx = max(0, req.start_index)
+        end_idx = req.end_index if req.end_index >= 0 else total_files - 1
+        end_idx = min(end_idx, total_files - 1)
+
+        if start_idx > end_idx:
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = "Invalid index range"
+            return
+
+        indices = list(range(start_idx, end_idx + 1))
+        tasks[task_id].total = len(indices)
+        os.makedirs(req.output_dir, exist_ok=True)
+        output_msgpack_path = os.path.join(req.output_dir, "detections.msgpack")
+
+        # 画像サイズを最初のファイルから取得
+        first_path = os.path.join(req.image_dir, req.filenames[start_idx])
+        first_img = cv2.imread(first_path)
+        if first_img is None:
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = f"Cannot read image: {first_path}"
+            return
+        height, width = first_img.shape[:2]
+
+        frame_buffer: List[np.ndarray] = []
+        frame_detections: List[List[List[float]]] = []
+        batch_size = 5
+        current_count = 0
+
+        def process_batch():
+            nonlocal current_count
+            if not frame_buffer:
+                return
+            batch_det = detector.detect_batch(frame_buffer)
+            for detections in batch_det:
+                packed: List[List[float]] = []
+                for x, y, w, h, conf in detections:
+                    bx, by, bw, bh = int(x), int(y), int(w), int(h)
+                    bx = max(0, bx)
+                    by = max(0, by)
+                    bw = min(width - bx, bw)
+                    bh = min(height - by, bh)
+                    if bw <= 0 or bh <= 0:
+                        continue
+                    packed.append([bx, by, bw, bh, float(conf)])
+                frame_detections.append(packed)
+                current_count += 1
+                tasks[task_id].progress = current_count
+            frame_buffer.clear()
+
+        print(
+            f"[FaceMask] Starting image detection: {req.image_dir} "
+            f"({len(indices)} images) -> {output_msgpack_path}"
+        )
+
+        for file_idx in indices:
+            if cancel_event and cancel_event.is_set():
+                tasks[task_id].status = TaskStatus.CANCELLED
+                tasks[task_id].message = "Cancelled by user"
+                break
+
+            img_path = os.path.join(req.image_dir, req.filenames[file_idx])
+            frame = cv2.imread(img_path)
+            if frame is None:
+                frame_detections.append([])
+                current_count += 1
+                tasks[task_id].progress = current_count
+                continue
+
+            frame_buffer.append(frame)
+            if len(frame_buffer) >= batch_size:
+                process_batch()
+
+        if frame_buffer:
+            process_batch()
+
+        if tasks[task_id].status == TaskStatus.PROCESSING:
+            payload = {
+                "version": 1,
+                "image_dir": req.image_dir,
+                "filenames": req.filenames,
+                "start_frame": start_idx,
+                "end_frame": start_idx + len(frame_detections) - 1,
+                "width": width,
+                "height": height,
+                "fps": 0.0,
+                "mask_scale": 1.0,
+                "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_msgpack_path
+            tasks[task_id].message = "Image detection cache completed"
+            print(f"[FaceMask] Image detection done: {output_msgpack_path}")
+
+    except Exception as e:
+        tasks[task_id].status = TaskStatus.FAILED
+        tasks[task_id].message = str(e)
+        traceback.print_exc()
+    finally:
+        if task_id in cancel_events:
+            del cancel_events[task_id]
+
+
+def process_bake_images_task(task_id: str, req: BakeImagesRequest):
+    """画像シーケンスに顔ぼかしを適用して新ディレクトリへ書き出す。"""
+    try:
+        tasks[task_id].status = TaskStatus.PROCESSING
+        cancel_event = cancel_events.get(task_id)
+
+        if not os.path.exists(req.image_dir):
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = f"Image directory not found: {req.image_dir}"
+            return
+        if not os.path.exists(req.detections_path):
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = f"Detections file not found: {req.detections_path}"
+            return
+
+        with open(req.detections_path, "rb") as f:
+            payload = msgpack.unpackb(f.read(), raw=False)
+        frames_detections = payload.get("frames")
+        if not isinstance(frames_detections, list):
+            tasks[task_id].status = TaskStatus.FAILED
+            tasks[task_id].message = "Invalid detections format: 'frames' is missing"
+            return
+
+        det_start_frame = int(payload.get("start_frame", 0))
+
+        blur_size = max(1, int(req.blur_size))
+        if blur_size % 2 == 0:
+            blur_size += 1
+        display_scale = max(0.1, float(req.display_scale))
+        blur_margin = blur_size // 2
+
+        os.makedirs(req.output_dir, exist_ok=True)
+        total = len(req.filenames)
+        tasks[task_id].total = total
+
+        print(
+            f"[FaceMask] Starting image bake: {req.image_dir} "
+            f"({total} images) -> {req.output_dir}"
+        )
+
+        for i, filename in enumerate(req.filenames):
+            if cancel_event and cancel_event.is_set():
+                tasks[task_id].status = TaskStatus.CANCELLED
+                tasks[task_id].message = "Cancelled by user"
+                return
+
+            src_path = os.path.join(req.image_dir, filename)
+            frame = cv2.imread(src_path)
+            if frame is None:
+                tasks[task_id].progress = i + 1
+                continue
+
+            h, w = frame.shape[:2]
+            det_idx = i - det_start_frame
+            frame_boxes = (
+                frames_detections[det_idx]
+                if 0 <= det_idx < len(frames_detections)
+                else []
+            )
+            _apply_face_blur_inplace(frame, frame_boxes, w, h, blur_size, display_scale, blur_margin)
+
+            out_path = os.path.join(req.output_dir, filename)
+            cv2.imwrite(out_path, frame)
+            tasks[task_id].progress = i + 1
+
+        if tasks[task_id].status == TaskStatus.PROCESSING:
+            tasks[task_id].status = TaskStatus.COMPLETED
+            tasks[task_id].result_path = req.output_dir
+            tasks[task_id].message = "Image blur bake completed"
+            print(f"[FaceMask] Image bake completed: {req.output_dir}")
+
+    except Exception as e:
+        tasks[task_id].status = TaskStatus.FAILED
+        tasks[task_id].message = str(e)
+        traceback.print_exc()
+    finally:
+        if task_id in cancel_events:
+            del cancel_events[task_id]
+
+
 def process_video_task(task_id: str, req: GenerateRequest):
     """Background task to detect faces and save bbox cache as msgpack."""
     cap = None
@@ -971,6 +1251,26 @@ def bake_blur_endpoint(req: BakeRequest, background_tasks: BackgroundTasks):
     background_tasks.add_task(process_bake_task, task_id, req)
     return task
 
+@app.post("/generate_images", response_model=Task)
+def generate_images_endpoint(req: GenerateImagesRequest, 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_images_task, task_id, req)
+    return task
+
+
+@app.post("/bake_image_blur", response_model=Task)
+def bake_image_blur_endpoint(req: BakeImagesRequest, 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_images_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: