Compare commits

..
11 Commits
15 changed files with 2381 additions and 434 deletions
+10 -10
View File
@@ -40,15 +40,6 @@ def register():
step=0.01, step=0.01,
) )
bpy.types.Scene.facemask_mask_scale = FloatProperty(
name="Mask Scale",
description="Scale factor for mask region (1.0 = exact face size)",
default=1.5,
min=1.0,
max=3.0,
step=0.1,
)
bpy.types.Scene.facemask_cache_dir = StringProperty( bpy.types.Scene.facemask_cache_dir = StringProperty(
name="Cache Directory", name="Cache Directory",
description="Optional cache root directory (empty = default .mask_cache)", description="Optional cache root directory (empty = default .mask_cache)",
@@ -64,6 +55,15 @@ def register():
max=501, max=501,
) )
bpy.types.Scene.facemask_bake_display_scale = FloatProperty(
name="Mask Scale",
description="Scale factor for the blur mask ellipse at bake time (1.0 = raw detection size)",
default=1.3,
min=0.5,
max=3.0,
step=0.1,
)
bpy.types.Scene.facemask_bake_format = EnumProperty( bpy.types.Scene.facemask_bake_format = EnumProperty(
name="Bake Format", name="Bake Format",
description="Output format for baked blur video", description="Output format for baked blur video",
@@ -91,9 +91,9 @@ def unregister():
# Unregister scene properties # Unregister scene properties
del bpy.types.Scene.facemask_conf_threshold del bpy.types.Scene.facemask_conf_threshold
del bpy.types.Scene.facemask_iou_threshold del bpy.types.Scene.facemask_iou_threshold
del bpy.types.Scene.facemask_mask_scale
del bpy.types.Scene.facemask_cache_dir del bpy.types.Scene.facemask_cache_dir
del bpy.types.Scene.facemask_bake_blur_size del bpy.types.Scene.facemask_bake_blur_size
del bpy.types.Scene.facemask_bake_display_scale
del bpy.types.Scene.facemask_bake_format del bpy.types.Scene.facemask_bake_format
+3 -3
View File
@@ -1,5 +1,5 @@
"""Core module exports.""" """Core module exports."""
from .async_bake_generator import AsyncBakeGenerator, get_bake_generator from .async_bake_generator import AsyncBakeGenerator as AsyncBakeGenerator, get_bake_generator as get_bake_generator
from .async_generator import AsyncMaskGenerator, get_generator from .async_generator import AsyncMaskGenerator as AsyncMaskGenerator, get_generator as get_generator
from .compositor_setup import create_mask_blur_node_tree, get_or_create_blur_node_tree from .compositor_setup import create_mask_blur_node_tree as create_mask_blur_node_tree, get_or_create_blur_node_tree as get_or_create_blur_node_tree
+87 -1
View File
@@ -32,6 +32,7 @@ class AsyncBakeGenerator:
detections_path: str, detections_path: str,
output_path: str, output_path: str,
blur_size: int, blur_size: int,
display_scale: float,
fmt: str, fmt: str,
on_complete: Optional[Callable] = None, on_complete: Optional[Callable] = None,
on_progress: Optional[Callable] = None, on_progress: Optional[Callable] = None,
@@ -53,7 +54,7 @@ class AsyncBakeGenerator:
self.worker_thread = threading.Thread( self.worker_thread = threading.Thread(
target=self._worker, target=self._worker,
args=(video_path, detections_path, output_path, blur_size, fmt), args=(video_path, detections_path, output_path, blur_size, display_scale, fmt),
daemon=True, daemon=True,
) )
self.worker_thread.start() self.worker_thread.start()
@@ -63,18 +64,102 @@ class AsyncBakeGenerator:
first_interval=0.1, first_interval=0.1,
) )
def start_images(
self,
image_dir: str,
filenames: list,
output_dir: str,
detections_path: str,
blur_size: int,
display_scale: float,
on_complete=None,
on_progress=None,
):
"""画像シーケンスのぼかしBakeを非同期で開始する。"""
global bpy
import bpy as _bpy
bpy = _bpy
if self.is_running:
raise RuntimeError("Blur bake already in progress")
self.is_running = True
self.total_frames = len(filenames)
self.current_frame = 0
self._on_complete = on_complete
self._on_progress = on_progress
self.worker_thread = threading.Thread(
target=self._worker_images,
args=(image_dir, filenames, output_dir, detections_path, blur_size, display_scale),
daemon=True,
)
self.worker_thread.start()
bpy.app.timers.register(self._check_progress, first_interval=0.1)
def cancel(self): def cancel(self):
"""Cancel the current bake processing.""" """Cancel the current bake processing."""
self.is_running = False self.is_running = False
if self.worker_thread and self.worker_thread.is_alive(): if self.worker_thread and self.worker_thread.is_alive():
self.worker_thread.join(timeout=2.0) self.worker_thread.join(timeout=2.0)
def _worker_images(
self,
image_dir: str,
filenames: list,
output_dir: str,
detections_path: str,
blur_size: int,
display_scale: float,
):
import time
from .inference_client import get_client
task_id = None
try:
client = get_client()
task_id = client.bake_image_blur(
image_dir=image_dir,
filenames=filenames,
output_dir=output_dir,
detections_path=detections_path,
blur_size=blur_size,
display_scale=display_scale,
)
while self.is_running:
status = client.get_task_status(task_id)
state = status.get("status")
total = status.get("total", 0)
if total > 0:
self.total_frames = total
progress = status.get("progress", 0)
if progress >= 0:
self.progress_queue.put(("progress", progress))
if state == "completed":
result_path = status.get("result_path", output_dir)
self.result_queue.put(("done", result_path))
return
if state == "failed":
self.result_queue.put(("error", status.get("message", "Unknown error")))
return
if state == "cancelled":
self.result_queue.put(("cancelled", None))
return
time.sleep(0.5)
if task_id:
client.cancel_task(task_id)
self.result_queue.put(("cancelled", None))
except Exception as e:
self.result_queue.put(("error", str(e)))
def _worker( def _worker(
self, self,
video_path: str, video_path: str,
detections_path: str, detections_path: str,
output_path: str, output_path: str,
blur_size: int, blur_size: int,
display_scale: float,
fmt: str, fmt: str,
): ):
import time import time
@@ -88,6 +173,7 @@ class AsyncBakeGenerator:
detections_path=detections_path, detections_path=detections_path,
output_path=output_path, output_path=output_path,
blur_size=blur_size, blur_size=blur_size,
display_scale=display_scale,
fmt=fmt, fmt=fmt,
) )
+177 -4
View File
@@ -44,7 +44,6 @@ class AsyncMaskGenerator:
fps: float, fps: float,
conf_threshold: float = 0.5, conf_threshold: float = 0.5,
iou_threshold: float = 0.45, iou_threshold: float = 0.45,
mask_scale: float = 1.5,
on_complete: Optional[Callable] = None, on_complete: Optional[Callable] = None,
on_progress: Optional[Callable] = None, on_progress: Optional[Callable] = None,
): ):
@@ -94,7 +93,6 @@ class AsyncMaskGenerator:
fps, fps,
conf_threshold, conf_threshold,
iou_threshold, iou_threshold,
mask_scale,
), ),
daemon=True, daemon=True,
) )
@@ -106,12 +104,189 @@ class AsyncMaskGenerator:
first_interval=0.1, first_interval=0.1,
) )
def start_images(
self,
image_dir: str,
filenames: list,
output_dir: str,
start_index: int,
end_index: int,
conf_threshold: float = 0.5,
iou_threshold: float = 0.45,
on_complete=None,
on_progress=None,
):
"""画像シーケンスの顔検出を非同期で開始する。"""
global bpy
import bpy as _bpy
bpy = _bpy
if self.is_running:
raise RuntimeError("Mask generation already in progress")
self.is_running = True
self.total_frames = end_index - start_index + 1
self.current_frame = 0
self._on_complete = on_complete
self._on_progress = on_progress
os.makedirs(output_dir, exist_ok=True)
self.worker_thread = threading.Thread(
target=self._worker_images,
args=(image_dir, filenames, output_dir, start_index, end_index,
conf_threshold, iou_threshold),
daemon=True,
)
self.worker_thread.start()
bpy.app.timers.register(self._check_progress, first_interval=0.1)
def start_augment_pose(
self,
detections_path: str,
total_frames: int,
conf_threshold: float = 0.5,
iou_threshold: float = 0.45,
on_complete=None,
on_progress=None,
):
"""既存キャッシュへの pose 補完を非同期で開始する。"""
global bpy
import bpy as _bpy
bpy = _bpy
if self.is_running:
raise RuntimeError("Mask generation already in progress")
self.is_running = True
self.total_frames = total_frames
self.current_frame = 0
self._on_complete = on_complete
self._on_progress = on_progress
self.worker_thread = threading.Thread(
target=self._worker_augment_pose,
args=(detections_path, conf_threshold, iou_threshold),
daemon=True,
)
self.worker_thread.start()
bpy.app.timers.register(self._check_progress, first_interval=0.1)
def _worker_augment_pose(
self,
detections_path: str,
conf_threshold: float,
iou_threshold: float,
):
"""client.augment_pose() を呼んで task_id でポーリング。"""
import time
from .inference_client import get_client
try:
client = get_client()
task_id = client.augment_pose(
detections_path=detections_path,
conf_threshold=conf_threshold,
iou_threshold=iou_threshold,
)
while self.is_running:
status = client.get_task_status(task_id)
state = status.get("status")
total = status.get("total", 0)
if total > 0:
self.total_frames = total
if state == "completed":
progress = status.get("progress", self.total_frames)
if progress >= 0:
self.progress_queue.put(("progress", progress))
result_path = status.get("result_path", detections_path)
self.result_queue.put(("done", result_path))
return
elif state == "failed":
self.result_queue.put(("error", status.get("message", "Unknown error")))
return
elif state == "cancelled":
self.result_queue.put(("cancelled", None))
return
progress = status.get("progress", 0)
if progress >= 0:
self.progress_queue.put(("progress", progress))
time.sleep(0.5)
client.cancel_task(task_id)
self.result_queue.put(("cancelled", None))
except Exception as e:
import traceback
print(f"[FaceMask] Error: {e}")
traceback.print_exc()
self.result_queue.put(("error", str(e)))
def cancel(self): def cancel(self):
"""Cancel the current processing.""" """Cancel the current processing."""
self.is_running = False self.is_running = False
if self.worker_thread and self.worker_thread.is_alive(): if self.worker_thread and self.worker_thread.is_alive():
self.worker_thread.join(timeout=2.0) self.worker_thread.join(timeout=2.0)
def _worker_images(
self,
image_dir: str,
filenames: list,
output_dir: str,
start_index: int,
end_index: int,
conf_threshold: float,
iou_threshold: float,
):
import time
from .inference_client import get_client
try:
client = get_client()
task_id = client.generate_mask_images(
image_dir=image_dir,
filenames=filenames,
output_dir=output_dir,
start_index=start_index,
end_index=end_index,
conf_threshold=conf_threshold,
iou_threshold=iou_threshold,
)
while self.is_running:
status = client.get_task_status(task_id)
state = status.get("status")
total = status.get("total", 0)
if total > 0:
self.total_frames = total
if state == "completed":
progress = status.get("progress", self.total_frames)
if progress >= 0:
self.progress_queue.put(("progress", progress))
result_path = status.get(
"result_path",
os.path.join(output_dir, "detections.msgpack"),
)
self.result_queue.put(("done", result_path))
return
elif state == "failed":
self.result_queue.put(("error", status.get("message", "Unknown error")))
return
elif state == "cancelled":
self.result_queue.put(("cancelled", None))
return
progress = status.get("progress", 0)
if progress >= 0:
self.progress_queue.put(("progress", progress))
time.sleep(0.5)
client.cancel_task(task_id)
self.result_queue.put(("cancelled", None))
except Exception as e:
self.result_queue.put(("error", str(e)))
def _worker( def _worker(
self, self,
video_path: str, video_path: str,
@@ -121,7 +296,6 @@ class AsyncMaskGenerator:
fps: float, fps: float,
conf_threshold: float, conf_threshold: float,
iou_threshold: float, iou_threshold: float,
mask_scale: float,
): ):
""" """
Worker thread function. Delegates to inference server and polls status. Worker thread function. Delegates to inference server and polls status.
@@ -141,7 +315,6 @@ class AsyncMaskGenerator:
end_frame=end_frame, end_frame=end_frame,
conf_threshold=conf_threshold, conf_threshold=conf_threshold,
iou_threshold=iou_threshold, iou_threshold=iou_threshold,
mask_scale=mask_scale,
) )
print(f"[FaceMask] Task started: {task_id}") print(f"[FaceMask] Task started: {task_id}")
+280
View File
@@ -0,0 +1,280 @@
"""
Batch processor for sequential Generate+Bake across multiple VSE strips.
Uses timer-based async chaining so Blender's UI stays responsive.
"""
from typing import List, Optional, Callable, Any
# Lazy-imported inside Blender
bpy = None
class _DummyOperator:
"""Dummy operator object for _start_bake_impl calls."""
def report(self, level, msg):
print(f"[FaceMask] Batch: {msg}")
class BatchProcessor:
"""Manages sequential Generate Detection Cache → Bake across a list of strips."""
def __init__(self):
self.is_running: bool = False
self._mode: str = "full" # "full" or "mask_only"
self._strip_names: List[str] = []
self._current_idx: int = 0
self._context: Any = None
self._cancelled: bool = False
self._results: List[dict] = []
self._on_item_complete: Optional[Callable] = None # (idx, total, name, status)
self._on_all_complete: Optional[Callable] = None # (results)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def start(self, context, strips, on_item_complete=None, on_all_complete=None, mode="full"):
"""Start batch processing for the given strips.
mode:
"full" - マスク生成(キャッシュなければ)→ Bake
"mask_only" - キャッシュを無視してマスク生成のみ(Bakeしない)
"""
global bpy
import bpy as _bpy
bpy = _bpy
if self.is_running:
raise RuntimeError("Batch already running")
self.is_running = True
self._mode = mode
self._strip_names = [s.name for s in strips]
self._current_idx = 0
self._context = context
self._cancelled = False
self._results = []
self._on_item_complete = on_item_complete
self._on_all_complete = on_all_complete
wm = context.window_manager
wm.batch_current = 0
wm.batch_total = len(self._strip_names)
wm.batch_current_name = ""
bpy.app.timers.register(self._process_next, first_interval=0.0)
def cancel(self):
"""Cancel batch. Stops currently running mask gen / bake."""
self._cancelled = True
from .async_generator import get_generator
from .async_bake_generator import get_bake_generator
gen = get_generator()
bake_gen = get_bake_generator()
if gen.is_running:
gen.cancel()
if bake_gen.is_running:
bake_gen.cancel()
# ------------------------------------------------------------------
# Internal: queue stepping
# ------------------------------------------------------------------
def _process_next(self):
"""Process the next strip in the queue (called via timer)."""
if self._cancelled:
self._finish()
return None
if self._current_idx >= len(self._strip_names):
self._finish()
return None
strip_name = self._strip_names[self._current_idx]
seq_editor = self._context.scene.sequence_editor
strip = seq_editor.strips.get(strip_name)
if strip is None:
print(f"[FaceMask] Batch: strip not found, skipping: {strip_name}")
self._results.append({"strip": strip_name, "status": "skipped"})
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, "skipped")
self._current_idx += 1
bpy.app.timers.register(self._process_next, first_interval=0.0)
return None
# Update wm progress labels
wm = self._context.window_manager
wm.batch_current = self._current_idx + 1
wm.batch_current_name = strip_name
for area in self._context.screen.areas:
if area.type == "SEQUENCE_EDITOR":
area.tag_redraw()
if self._mode == "mask_only":
# キャッシュを無視して常にマスク生成(Bakeしない)
self._start_mask_gen(strip)
else:
from .utils import check_detection_cache
if not check_detection_cache(strip.name):
self._start_mask_gen(strip)
else:
self._start_bake(strip)
return None # one-shot timer
def _schedule_next(self):
bpy.app.timers.register(self._process_next, first_interval=0.0)
# ------------------------------------------------------------------
# Mask generation
# ------------------------------------------------------------------
def _start_mask_gen(self, strip):
from ..operators.generate_mask import start_mask_gen_for_strip
strip_name = strip.name
def on_complete(status, data):
self._on_mask_done(strip_name, status, data)
def on_progress(current, total):
wm = self._context.window_manager
wm.mask_progress = current
wm.mask_total = max(total, 1)
for area in self._context.screen.areas:
if area.type == "SEQUENCE_EDITOR":
area.tag_redraw()
try:
start_mask_gen_for_strip(self._context, strip, on_complete, on_progress)
print(f"[FaceMask] Batch: started mask gen for {strip_name}")
except Exception as e:
print(f"[FaceMask] Batch: failed to start mask gen for {strip_name}: {e}")
self._on_mask_done(strip_name, "error", str(e))
def _on_mask_done(self, strip_name, status, data):
if self._cancelled or status == "cancelled":
self._results.append({"strip": strip_name, "status": "cancelled"})
self._finish()
return
if status == "error":
print(f"[FaceMask] Batch: mask gen failed for {strip_name}: {data}")
self._results.append({"strip": strip_name, "status": "error", "reason": str(data)})
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, "error")
self._current_idx += 1
self._schedule_next()
return
# Mask gen succeeded
if self._mode == "mask_only":
# Bakeしない:結果を記録して次へ
self._results.append({"strip": strip_name, "status": "done"})
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, "done")
self._current_idx += 1
self._schedule_next()
return
# full mode: proceed to bake
seq_editor = self._context.scene.sequence_editor
strip = seq_editor.strips.get(strip_name)
if strip is None:
print(f"[FaceMask] Batch: strip removed after mask gen: {strip_name}")
self._results.append({"strip": strip_name, "status": "skipped"})
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, "skipped")
self._current_idx += 1
self._schedule_next()
return
self._start_bake(strip)
# ------------------------------------------------------------------
# Bake
# ------------------------------------------------------------------
def _start_bake(self, strip):
from .async_bake_generator import get_bake_generator
from ..operators.apply_blur import _start_bake_impl
strip_name = strip.name
def on_complete_extra(status, data):
self._on_bake_done(strip_name, status, data)
bake_gen = get_bake_generator()
result = _start_bake_impl(
_DummyOperator(),
self._context,
force=False,
strip=strip,
on_complete_extra=on_complete_extra,
)
if result == {"CANCELLED"}:
# Error starting bake
print(f"[FaceMask] Batch: bake failed to start for {strip_name}")
self._results.append({"strip": strip_name, "status": "error", "reason": "bake failed to start"})
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, "error")
self._current_idx += 1
self._schedule_next()
elif not bake_gen.is_running:
# Cache hit: on_complete_extra was NOT called by _start_bake_impl
print(f"[FaceMask] Batch: bake cache hit for {strip_name}")
self._on_bake_done(strip_name, "done", None)
def _on_bake_done(self, strip_name, status, data):
if self._cancelled or status == "cancelled":
self._results.append({"strip": strip_name, "status": "cancelled"})
self._finish()
return
if status == "error":
print(f"[FaceMask] Batch: bake failed for {strip_name}: {data}")
self._results.append({"strip": strip_name, "status": "error", "reason": str(data)})
else:
self._results.append({"strip": strip_name, "status": "done"})
print(f"[FaceMask] Batch: completed {strip_name}")
if self._on_item_complete:
self._on_item_complete(self._current_idx, len(self._strip_names), strip_name, status)
self._current_idx += 1
self._schedule_next()
# ------------------------------------------------------------------
# Finish
# ------------------------------------------------------------------
def _finish(self):
self.is_running = False
wm = self._context.window_manager
wm.batch_current = 0
wm.batch_total = 0
wm.batch_current_name = ""
print(f"[FaceMask] Batch: all done. Results: {self._results}")
if self._on_all_complete:
self._on_all_complete(self._results)
for area in self._context.screen.areas:
if area.type == "SEQUENCE_EDITOR":
area.tag_redraw()
# Singleton
_batch_processor: Optional[BatchProcessor] = None
def get_batch_processor() -> BatchProcessor:
global _batch_processor
if _batch_processor is None:
_batch_processor = BatchProcessor()
return _batch_processor
+6 -6
View File
@@ -8,7 +8,7 @@ only to masked regions of a video strip.
def create_mask_blur_node_tree( def create_mask_blur_node_tree(
name: str = "FaceMaskBlur", name: str = "FaceMaskBlur",
blur_size: int = 50, blur_size: int = 50,
) -> "bpy.types.NodeTree": ) -> "bpy.types.NodeTree": # noqa: F821
""" """
Create a compositing node tree for mask-based blur. Create a compositing node tree for mask-based blur.
@@ -107,10 +107,10 @@ def create_mask_blur_node_tree(
def setup_strip_compositor_modifier( def setup_strip_compositor_modifier(
strip: "bpy.types.Strip", strip: "bpy.types.Strip", # noqa: F821
mask_strip: "bpy.types.Strip", mask_strip: "bpy.types.Strip", # noqa: F821
node_tree: "bpy.types.NodeTree", node_tree: "bpy.types.NodeTree", # noqa: F821
) -> "bpy.types.SequenceModifier": ) -> "bpy.types.SequenceModifier": # noqa: F821
""" """
Add a Compositor modifier to a strip using the mask-blur node tree. Add a Compositor modifier to a strip using the mask-blur node tree.
@@ -148,7 +148,7 @@ def setup_strip_compositor_modifier(
return modifier return modifier
def get_or_create_blur_node_tree(blur_size: int = 50) -> "bpy.types.NodeTree": def get_or_create_blur_node_tree(blur_size: int = 50) -> "bpy.types.NodeTree": # noqa: F821
""" """
Get existing or create new blur node tree with specified blur size. Get existing or create new blur node tree with specified blur size.
+119 -2
View File
@@ -204,7 +204,6 @@ class InferenceClient:
end_frame: int, end_frame: int,
conf_threshold: float, conf_threshold: float,
iou_threshold: float, iou_threshold: float,
mask_scale: float,
) -> str: ) -> str:
""" """
Request mask generation. Request mask generation.
@@ -222,7 +221,6 @@ class InferenceClient:
"end_frame": end_frame, "end_frame": end_frame,
"conf_threshold": conf_threshold, "conf_threshold": conf_threshold,
"iou_threshold": iou_threshold, "iou_threshold": iou_threshold,
"mask_scale": mask_scale,
} }
req = urllib.request.Request( req = urllib.request.Request(
@@ -239,6 +237,36 @@ class InferenceClient:
except urllib.error.HTTPError as e: except urllib.error.HTTPError as e:
raise RuntimeError(f"Server error: {e.read().decode('utf-8')}") raise RuntimeError(f"Server error: {e.read().decode('utf-8')}")
def augment_pose(
self,
detections_path: str,
conf_threshold: float,
iou_threshold: float,
) -> str:
"""既存キャッシュに pose 推定結果を追加合成する。task_id を返す。"""
if not self.is_server_running():
self.start_server()
data = {
"detections_path": detections_path,
"conf_threshold": conf_threshold,
"iou_threshold": iou_threshold,
}
req = urllib.request.Request(
f"{self.SERVER_URL}/augment_pose",
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]: def get_task_status(self, task_id: str) -> Dict[str, Any]:
"""Get status of a task.""" """Get status of a task."""
try: try:
@@ -249,12 +277,30 @@ class InferenceClient:
except urllib.error.HTTPError: except urllib.error.HTTPError:
return {"status": "unknown"} return {"status": "unknown"}
def get_video_info(self, video_path: str) -> Dict[str, Any]:
"""Get video metadata from the inference server."""
if not self.is_server_running():
self.start_server()
req = urllib.request.Request(
f"{self.SERVER_URL}/video_info",
data=json.dumps({"video_path": video_path}).encode("utf-8"),
headers={"Content-Type": "application/json"},
method="POST",
)
try:
with urllib.request.urlopen(req) as response:
return json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as e:
raise RuntimeError(f"Server error: {e.read().decode('utf-8')}")
def bake_blur( def bake_blur(
self, self,
video_path: str, video_path: str,
detections_path: str, detections_path: str,
output_path: str, output_path: str,
blur_size: int, blur_size: int,
display_scale: float,
fmt: str, fmt: str,
) -> str: ) -> str:
""" """
@@ -271,6 +317,7 @@ class InferenceClient:
"detections_path": detections_path, "detections_path": detections_path,
"output_path": output_path, "output_path": output_path,
"blur_size": blur_size, "blur_size": blur_size,
"display_scale": display_scale,
"format": fmt, "format": fmt,
} }
@@ -288,6 +335,76 @@ class InferenceClient:
except urllib.error.HTTPError as e: except urllib.error.HTTPError as e:
raise RuntimeError(f"Server error: {e.read().decode('utf-8')}") raise RuntimeError(f"Server error: {e.read().decode('utf-8')}")
def generate_mask_images(
self,
image_dir: str,
filenames: list,
output_dir: str,
start_index: int,
end_index: int,
conf_threshold: float,
iou_threshold: float,
) -> str:
"""画像シーケンスの顔検出タスクを開始して task_id を返す。"""
if not self.is_server_running():
self.start_server()
data = {
"image_dir": image_dir,
"filenames": filenames,
"output_dir": output_dir,
"start_index": start_index,
"end_index": end_index,
"conf_threshold": conf_threshold,
"iou_threshold": iou_threshold,
}
req = urllib.request.Request(
f"{self.SERVER_URL}/generate_images",
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 bake_image_blur(
self,
image_dir: str,
filenames: list,
output_dir: str,
detections_path: str,
blur_size: int,
display_scale: float,
) -> str:
"""画像シーケンスのぼかしBakeタスクを開始して task_id を返す。"""
if not self.is_server_running():
self.start_server()
data = {
"image_dir": image_dir,
"filenames": filenames,
"output_dir": output_dir,
"detections_path": detections_path,
"blur_size": blur_size,
"display_scale": display_scale,
}
req = urllib.request.Request(
f"{self.SERVER_URL}/bake_image_blur",
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 cancel_task(self, task_id: str): def cancel_task(self, task_id: str):
"""Cancel a task.""" """Cancel a task."""
try: try:
+9
View File
@@ -83,6 +83,15 @@ def get_detections_path_for_strip(strip_name: str) -> str:
return os.path.join(get_cache_dir_for_strip(strip_name), "detections.msgpack") return os.path.join(get_cache_dir_for_strip(strip_name), "detections.msgpack")
def check_detection_cache(strip_name: str) -> bool:
"""Detection cache ファイルが存在し有効かどうか確認する。"""
path = get_detections_path_for_strip(strip_name)
try:
return os.path.exists(path) and os.path.getsize(path) > 0
except OSError:
return False
def get_cache_info(strip_name: Optional[str] = None) -> Tuple[str, int, int]: def get_cache_info(strip_name: Optional[str] = None) -> Tuple[str, int, int]:
""" """
Get cache directory information. Get cache directory information.
+3
View File
@@ -3,15 +3,18 @@
from . import generate_mask from . import generate_mask
from . import apply_blur from . import apply_blur
from . import clear_cache from . import clear_cache
from . import batch_bake
def register(): def register():
generate_mask.register() generate_mask.register()
apply_blur.register() apply_blur.register()
clear_cache.register() clear_cache.register()
batch_bake.register()
def unregister(): def unregister():
batch_bake.unregister()
clear_cache.unregister() clear_cache.unregister()
apply_blur.unregister() apply_blur.unregister()
generate_mask.unregister() generate_mask.unregister()
+166 -50
View File
@@ -20,6 +20,7 @@ KEY_BAKED = "facemask_baked_filepath"
KEY_MODE = "facemask_source_mode" KEY_MODE = "facemask_source_mode"
KEY_FORMAT = "facemask_bake_format" KEY_FORMAT = "facemask_bake_format"
KEY_BLUR_SIZE = "facemask_bake_blur_size" KEY_BLUR_SIZE = "facemask_bake_blur_size"
KEY_DISPLAY_SCALE = "facemask_bake_display_scale"
FORMAT_EXT = { FORMAT_EXT = {
@@ -36,6 +37,12 @@ def _output_path(video_strip, detections_path: str, fmt: str) -> str:
return os.path.join(out_dir, f"{safe_name}_blurred.{ext}") return os.path.join(out_dir, f"{safe_name}_blurred.{ext}")
def _output_dir_for_images(strip, detections_path: str) -> str:
out_dir = os.path.dirname(detections_path)
safe_name = strip.name.replace("/", "_").replace("\\", "_")
return os.path.join(out_dir, f"{safe_name}_blurred")
def _reload_movie_strip(strip): def _reload_movie_strip(strip):
if hasattr(strip, "reload"): if hasattr(strip, "reload"):
try: try:
@@ -44,67 +51,80 @@ def _reload_movie_strip(strip):
pass pass
def _set_strip_source(strip, filepath: str): def _set_strip_source(strip, path: str):
strip.filepath = filepath if strip.type == "IMAGE":
strip.directory = path
else:
strip.filepath = path
_reload_movie_strip(strip) _reload_movie_strip(strip)
class SEQUENCER_OT_bake_and_swap_blur_source(Operator): def _start_bake_impl(operator, context, force: bool = False, strip=None, on_complete_extra=None):
"""Bake masked blur and replace active strip source with baked video.""" """Bakeの共通実装。force=True でキャッシュを無視して再Bakeする。
bl_idname = "sequencer.bake_and_swap_blur_source" strip: 処理対象のstrip。None の場合は active_strip を使用。
bl_label = "Bake & Swap Source" on_complete_extra: 非同期Bake完了時に追加で呼ばれるコールバック (status, data)。
bl_description = "Bake masked blur to video and swap active strip source" キャッシュヒット即時完了の場合は呼ばれない。
bl_options = {"REGISTER", "UNDO"} MOVIE / IMAGE 両対応。
"""
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
# Prevent overlapping heavy tasks
if get_mask_generator().is_running:
return False
if get_bake_generator().is_running:
return False
strip = context.scene.sequence_editor.active_strip
return bool(strip and strip.type == "MOVIE")
def execute(self, context):
seq_editor = context.scene.sequence_editor seq_editor = context.scene.sequence_editor
scene = context.scene scene = context.scene
video_strip = seq_editor.active_strip video_strip = strip if strip is not None else seq_editor.active_strip
is_image = video_strip.type == "IMAGE"
video_path = bpy.path.abspath(video_strip.filepath)
detections_path = get_detections_path_for_strip(video_strip.name) detections_path = get_detections_path_for_strip(video_strip.name)
if not os.path.exists(video_path):
self.report({"ERROR"}, f"Source video not found: {video_path}")
return {"CANCELLED"}
if not os.path.exists(detections_path): if not os.path.exists(detections_path):
self.report({"ERROR"}, f"Detection cache not found: {detections_path}") operator.report({"ERROR"}, f"Detection cache not found: {detections_path}")
return {"CANCELLED"} return {"CANCELLED"}
blur_size = int(scene.facemask_bake_blur_size)
display_scale = float(scene.facemask_bake_display_scale)
if is_image:
image_dir = bpy.path.abspath(video_strip.directory)
filenames = [elem.filename for elem in video_strip.elements]
if not os.path.isdir(image_dir):
operator.report({"ERROR"}, f"Image directory not found: {image_dir}")
return {"CANCELLED"}
output_dir = _output_dir_for_images(video_strip, detections_path)
original_source = image_dir
bake_format = None # IMAGE strips don't use format
else:
video_path = bpy.path.abspath(video_strip.filepath)
if not os.path.exists(video_path):
operator.report({"ERROR"}, f"Source video not found: {video_path}")
return {"CANCELLED"}
bake_format = scene.facemask_bake_format bake_format = scene.facemask_bake_format
output_path = _output_path(video_strip, detections_path, bake_format) output_path = _output_path(video_strip, detections_path, bake_format)
blur_size = int(scene.facemask_bake_blur_size) original_source = video_path
# Reuse baked cache when parameters match and file still exists. if not force:
# パラメータが一致するキャッシュがあればswapのみ
cached_baked_path = video_strip.get(KEY_BAKED) cached_baked_path = video_strip.get(KEY_BAKED)
cached_format = video_strip.get(KEY_FORMAT)
cached_blur_size = video_strip.get(KEY_BLUR_SIZE) cached_blur_size = video_strip.get(KEY_BLUR_SIZE)
cached_display_scale = video_strip.get(KEY_DISPLAY_SCALE)
try: try:
cached_blur_size_int = int(cached_blur_size) cached_blur_size_int = int(cached_blur_size)
except (TypeError, ValueError): except (TypeError, ValueError):
cached_blur_size_int = None cached_blur_size_int = None
if ( try:
cached_baked_path cached_display_scale_f = float(cached_display_scale)
and os.path.exists(cached_baked_path) except (TypeError, ValueError):
and cached_format == bake_format cached_display_scale_f = None
cache_exists = (
cached_baked_path and os.path.exists(cached_baked_path)
and cached_blur_size_int == blur_size and cached_blur_size_int == blur_size
): and cached_display_scale_f == display_scale
)
if not is_image:
cache_exists = cache_exists and video_strip.get(KEY_FORMAT) == bake_format
if cache_exists:
if video_strip.get(KEY_MODE) != "baked": if video_strip.get(KEY_MODE) != "baked":
video_strip[KEY_MODE] = "baked" video_strip[KEY_MODE] = "baked"
_set_strip_source(video_strip, cached_baked_path) _set_strip_source(video_strip, cached_baked_path)
self.report({"INFO"}, "Using cached baked blur") operator.report({"INFO"}, "Using cached baked blur")
return {"FINISHED"} return {"FINISHED"}
bake_generator = get_bake_generator() bake_generator = get_bake_generator()
@@ -117,17 +137,18 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
return return
if status == "done": if status == "done":
result_path = data or output_path result = data or (output_dir if is_image else output_path)
original_path = strip.get(KEY_ORIGINAL)
current_mode = strip.get(KEY_MODE, "original") current_mode = strip.get(KEY_MODE, "original")
if not original_path or current_mode != "baked": if not strip.get(KEY_ORIGINAL) or current_mode != "baked":
strip[KEY_ORIGINAL] = video_path strip[KEY_ORIGINAL] = original_source
strip[KEY_BAKED] = result_path strip[KEY_BAKED] = result
strip[KEY_MODE] = "baked" strip[KEY_MODE] = "baked"
strip[KEY_FORMAT] = bake_format
strip[KEY_BLUR_SIZE] = blur_size strip[KEY_BLUR_SIZE] = blur_size
_set_strip_source(strip, result_path) strip[KEY_DISPLAY_SCALE] = display_scale
print(f"[FaceMask] Bake completed and source swapped: {result_path}") if not is_image:
strip[KEY_FORMAT] = bake_format
_set_strip_source(strip, result)
print(f"[FaceMask] Bake completed and source swapped: {result}")
elif status == "error": elif status == "error":
print(f"[FaceMask] Bake failed: {data}") print(f"[FaceMask] Bake failed: {data}")
elif status == "cancelled": elif status == "cancelled":
@@ -137,6 +158,9 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
if area.type == "SEQUENCE_EDITOR": if area.type == "SEQUENCE_EDITOR":
area.tag_redraw() area.tag_redraw()
if on_complete_extra:
on_complete_extra(status, data)
def on_progress(current, total): def on_progress(current, total):
wm.bake_progress = current wm.bake_progress = current
wm.bake_total = max(total, 1) wm.bake_total = max(total, 1)
@@ -148,20 +172,108 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
wm.bake_total = 1 wm.bake_total = 1
try: try:
if is_image:
bake_generator.start_images(
image_dir=image_dir,
filenames=filenames,
output_dir=output_dir,
detections_path=detections_path,
blur_size=blur_size,
display_scale=display_scale,
on_complete=on_complete,
on_progress=on_progress,
)
else:
bake_generator.start( bake_generator.start(
video_path=video_path, video_path=video_path,
detections_path=detections_path, detections_path=detections_path,
output_path=output_path, output_path=output_path,
blur_size=blur_size, blur_size=blur_size,
display_scale=display_scale,
fmt=bake_format.lower(), fmt=bake_format.lower(),
on_complete=on_complete, on_complete=on_complete,
on_progress=on_progress, on_progress=on_progress,
) )
except Exception as e: except Exception as e:
self.report({"ERROR"}, f"Failed to start bake: {e}") operator.report({"ERROR"}, f"Failed to start bake: {e}")
return {"CANCELLED"} return {"CANCELLED"}
self.report({"INFO"}, "Started blur bake in background") operator.report({"INFO"}, "Started blur bake in background")
return {"FINISHED"}
class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
"""Bake masked blur (reuse cache if parameters match)."""
bl_idname = "sequencer.bake_and_swap_blur_source"
bl_label = "Bake"
bl_description = "Bake masked blur to video and swap active strip source"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_mask_generator().is_running:
return False
if get_bake_generator().is_running:
return False
strip = context.scene.sequence_editor.active_strip
return bool(strip and strip.type in {"MOVIE", "IMAGE"})
def execute(self, context):
return _start_bake_impl(self, context, force=False)
class SEQUENCER_OT_force_rebake_blur(Operator):
"""Force re-bake, ignoring any existing cached result."""
bl_idname = "sequencer.force_rebake_blur"
bl_label = "Re-bake"
bl_description = "Discard cached bake and re-bake from scratch"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_mask_generator().is_running:
return False
if get_bake_generator().is_running:
return False
strip = context.scene.sequence_editor.active_strip
return bool(strip and strip.type in {"MOVIE", "IMAGE"})
def execute(self, context):
return _start_bake_impl(self, context, force=True)
class SEQUENCER_OT_swap_to_baked_blur(Operator):
"""Swap active strip source to already-baked video (no re-bake)."""
bl_idname = "sequencer.swap_to_baked_blur"
bl_label = "Swap to Baked"
bl_description = "Switch active strip source to the baked video without re-baking"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_bake_generator().is_running:
return False
strip = context.scene.sequence_editor.active_strip
if not strip or strip.type not in {"MOVIE", "IMAGE"}:
return False
baked_path = strip.get(KEY_BAKED)
return bool(baked_path and os.path.exists(baked_path))
def execute(self, context):
strip = context.scene.sequence_editor.active_strip
baked_path = strip.get(KEY_BAKED)
_set_strip_source(strip, baked_path)
strip[KEY_MODE] = "baked"
self.report({"INFO"}, "Swapped to baked source")
return {"FINISHED"} return {"FINISHED"}
@@ -169,7 +281,7 @@ class SEQUENCER_OT_restore_original_source(Operator):
"""Restore active strip source filepath to original video.""" """Restore active strip source filepath to original video."""
bl_idname = "sequencer.restore_original_source" bl_idname = "sequencer.restore_original_source"
bl_label = "Restore Original Source" bl_label = "Restore Original"
bl_description = "Restore active strip to original source filepath" bl_description = "Restore active strip to original source filepath"
bl_options = {"REGISTER", "UNDO"} bl_options = {"REGISTER", "UNDO"}
@@ -180,7 +292,9 @@ class SEQUENCER_OT_restore_original_source(Operator):
if get_bake_generator().is_running: if get_bake_generator().is_running:
return False return False
strip = context.scene.sequence_editor.active_strip strip = context.scene.sequence_editor.active_strip
if not strip or strip.type != "MOVIE": if not strip or strip.type not in {"MOVIE", "IMAGE"}:
return False
if strip.get(KEY_MODE, "original") == "original":
return False return False
return bool(strip.get(KEY_ORIGINAL)) return bool(strip.get(KEY_ORIGINAL))
@@ -205,7 +319,7 @@ class SEQUENCER_OT_apply_mask_blur(Operator):
bl_idname = "sequencer.apply_mask_blur" bl_idname = "sequencer.apply_mask_blur"
bl_label = "Apply Mask Blur" bl_label = "Apply Mask Blur"
bl_description = "Compatibility alias for Bake & Swap Source" bl_description = "Compatibility alias for Bake"
bl_options = {"REGISTER", "UNDO"} bl_options = {"REGISTER", "UNDO"}
@classmethod @classmethod
@@ -236,6 +350,8 @@ class SEQUENCER_OT_cancel_bake_blur(Operator):
classes = [ classes = [
SEQUENCER_OT_bake_and_swap_blur_source, SEQUENCER_OT_bake_and_swap_blur_source,
SEQUENCER_OT_force_rebake_blur,
SEQUENCER_OT_swap_to_baked_blur,
SEQUENCER_OT_restore_original_source, SEQUENCER_OT_restore_original_source,
SEQUENCER_OT_cancel_bake_blur, SEQUENCER_OT_cancel_bake_blur,
SEQUENCER_OT_apply_mask_blur, SEQUENCER_OT_apply_mask_blur,
+191
View File
@@ -0,0 +1,191 @@
"""
Batch Bake operator: sequentially Generate Detection Cache → Bake
for all selected MOVIE strips in the VSE.
"""
import os
import bpy
from bpy.props import IntProperty, StringProperty
from bpy.types import Operator
from ..core.batch_processor import get_batch_processor
from ..core.async_generator import get_generator as get_mask_generator
from ..core.async_bake_generator import get_bake_generator
from .apply_blur import KEY_ORIGINAL, KEY_MODE, _set_strip_source
class SEQUENCER_OT_batch_bake_selected(Operator):
"""Generate detection cache and bake blur for all selected MOVIE/IMAGE strips."""
bl_idname = "sequencer.batch_bake_selected"
bl_label = "Batch Bake Selected"
bl_description = "Generate detection cache and bake blur for all selected MOVIE/IMAGE strips"
bl_options = {"REGISTER"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_batch_processor().is_running:
return False
if get_mask_generator().is_running:
return False
if get_bake_generator().is_running:
return False
seq_editor = context.scene.sequence_editor
return any(s.select and s.type in {"MOVIE", "IMAGE"} for s in seq_editor.strips)
def execute(self, context):
seq_editor = context.scene.sequence_editor
strips = [s for s in seq_editor.strips if s.select and s.type in {"MOVIE", "IMAGE"}]
if not strips:
self.report({"WARNING"}, "No MOVIE or IMAGE strips selected")
return {"CANCELLED"}
batch = get_batch_processor()
def on_item_complete(idx, total, strip_name, status):
pass # wm properties already updated by BatchProcessor
def on_all_complete(results):
done = sum(1 for r in results if r["status"] == "done")
total = len(results)
print(f"[FaceMask] Batch finished: {done}/{total} strips completed")
wm = context.window_manager
wm.batch_current = 0
wm.batch_total = len(strips)
wm.batch_current_name = ""
batch.start(context, strips, on_item_complete=on_item_complete, on_all_complete=on_all_complete)
self.report({"INFO"}, f"Batch bake started for {len(strips)} strips")
return {"FINISHED"}
class SEQUENCER_OT_batch_regenerate_cache(Operator):
"""Regenerate detection cache for all selected MOVIE/IMAGE strips (ignore existing cache)."""
bl_idname = "sequencer.batch_regenerate_cache"
bl_label = "Batch Regenerate Cache"
bl_description = "Regenerate detection cache for all selected MOVIE/IMAGE strips"
bl_options = {"REGISTER"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_batch_processor().is_running:
return False
if get_mask_generator().is_running:
return False
if get_bake_generator().is_running:
return False
seq_editor = context.scene.sequence_editor
return any(s.select and s.type in {"MOVIE", "IMAGE"} for s in seq_editor.strips)
def execute(self, context):
seq_editor = context.scene.sequence_editor
strips = [s for s in seq_editor.strips if s.select and s.type in {"MOVIE", "IMAGE"}]
if not strips:
self.report({"WARNING"}, "No MOVIE or IMAGE strips selected")
return {"CANCELLED"}
batch = get_batch_processor()
def on_all_complete(results):
done = sum(1 for r in results if r["status"] == "done")
print(f"[FaceMask] Batch regenerate finished: {done}/{len(results)} strips")
batch.start(
context,
strips,
on_all_complete=on_all_complete,
mode="mask_only",
)
self.report({"INFO"}, f"Batch regenerate cache started for {len(strips)} strips")
return {"FINISHED"}
class SEQUENCER_OT_batch_restore_original(Operator):
"""Restore original source for all selected MOVIE/IMAGE strips."""
bl_idname = "sequencer.batch_restore_original"
bl_label = "Batch Restore Original"
bl_description = "Restore original source filepath for all selected MOVIE/IMAGE strips"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
if get_batch_processor().is_running:
return False
seq_editor = context.scene.sequence_editor
return any(s.select and s.type in {"MOVIE", "IMAGE"} for s in seq_editor.strips)
def execute(self, context):
seq_editor = context.scene.sequence_editor
strips = [s for s in seq_editor.strips if s.select and s.type in {"MOVIE", "IMAGE"}]
restored = 0
skipped = 0
for strip in strips:
original_path = strip.get(KEY_ORIGINAL)
if not original_path or not os.path.exists(original_path):
skipped += 1
continue
if strip.get(KEY_MODE, "original") != "original":
_set_strip_source(strip, original_path)
strip[KEY_MODE] = "original"
restored += 1
self.report(
{"INFO"},
f"Restored {restored} strip(s)"
+ (f", skipped {skipped} (no original stored)" if skipped else ""),
)
return {"FINISHED"}
class SEQUENCER_OT_cancel_batch_bake(Operator):
"""Cancel ongoing batch bake."""
bl_idname = "sequencer.cancel_batch_bake"
bl_label = "Cancel Batch Bake"
bl_description = "Cancel the current batch bake process"
bl_options = {"REGISTER"}
def execute(self, context):
batch = get_batch_processor()
if batch.is_running:
batch.cancel()
self.report({"INFO"}, "Batch bake cancelled")
else:
self.report({"WARNING"}, "No batch bake in progress")
return {"FINISHED"}
classes = [
SEQUENCER_OT_batch_bake_selected,
SEQUENCER_OT_batch_regenerate_cache,
SEQUENCER_OT_batch_restore_original,
SEQUENCER_OT_cancel_batch_bake,
]
def register():
for cls in classes:
bpy.utils.register_class(cls)
bpy.types.WindowManager.batch_current = IntProperty(default=0)
bpy.types.WindowManager.batch_total = IntProperty(default=0)
bpy.types.WindowManager.batch_current_name = StringProperty(default="")
def unregister():
del bpy.types.WindowManager.batch_current_name
del bpy.types.WindowManager.batch_total
del bpy.types.WindowManager.batch_current
for cls in reversed(classes):
bpy.utils.unregister_class(cls)
+186 -96
View File
@@ -7,11 +7,99 @@ from video strips in the Video Sequence Editor.
import os import os
import bpy import bpy
from bpy.props import IntProperty from bpy.props import IntProperty, BoolProperty
from bpy.types import Operator from bpy.types import Operator
from ..core.async_generator import get_generator from ..core.async_generator import get_generator
from ..core.utils import get_cache_dir_for_strip from ..core.inference_client import get_client
from ..core.utils import get_cache_dir_for_strip, check_detection_cache
def compute_strip_frame_range(strip, scene, client) -> tuple:
"""(start_frame, end_frame, source_fps) を返す。失敗時は例外を送出。"""
video_path = bpy.path.abspath(strip.filepath)
video_info = client.get_video_info(video_path)
total_video_frames = int(video_info.get("frame_count", 0))
source_fps = float(video_info.get("fps", 0.0))
if total_video_frames <= 0:
raise ValueError(f"Could not read frame count from video: {video_path}")
if source_fps <= 0:
source_fps = scene.render.fps / scene.render.fps_base
project_fps = scene.render.fps / scene.render.fps_base
fps_ratio = source_fps / project_fps
start_frame = int(round(strip.frame_offset_start * fps_ratio))
end_frame = start_frame + int(round(strip.frame_final_duration * fps_ratio)) - 1
start_frame = max(0, min(start_frame, total_video_frames - 1))
end_frame = max(start_frame, min(end_frame, total_video_frames - 1))
return start_frame, end_frame, source_fps
def get_image_strip_files(strip) -> tuple:
"""IMAGE strip の (abs_image_dir, filenames_list) を返す。"""
image_dir = bpy.path.abspath(strip.directory)
filenames = [elem.filename for elem in strip.elements]
return image_dir, filenames
def compute_image_strip_range(strip) -> tuple:
"""IMAGE strip のアクティブ範囲 (start_index, end_index) を返す。"""
total_elements = len(strip.elements)
start_idx = max(0, int(strip.frame_offset_start))
end_idx = start_idx + int(strip.frame_final_duration) - 1
start_idx = min(start_idx, total_elements - 1)
end_idx = max(start_idx, min(end_idx, total_elements - 1))
return start_idx, end_idx
def start_mask_gen_for_strip(context, strip, on_complete, on_progress):
"""Strip のマスク生成を開始する共通処理(MOVIE / IMAGE 両対応)。
generator.is_running 等のエラー時は例外を送出する。
wm.mask_progress / mask_total を初期化してから generator.start*() を呼ぶ。
"""
scene = context.scene
wm = context.window_manager
generator = get_generator()
if generator.is_running:
raise RuntimeError("Mask generation already in progress")
output_dir = get_cache_dir_for_strip(strip.name)
os.makedirs(output_dir, exist_ok=True)
wm.mask_progress = 0
if strip.type == "IMAGE":
image_dir, filenames = get_image_strip_files(strip)
if not filenames:
raise ValueError("Image strip has no elements")
start_idx, end_idx = compute_image_strip_range(strip)
wm.mask_total = end_idx - start_idx + 1
generator.start_images(
image_dir=image_dir,
filenames=filenames,
output_dir=output_dir,
start_index=start_idx,
end_index=end_idx,
conf_threshold=scene.facemask_conf_threshold,
iou_threshold=scene.facemask_iou_threshold,
on_complete=on_complete,
on_progress=on_progress,
)
else:
client = get_client()
start_frame, end_frame, source_fps = compute_strip_frame_range(strip, scene, client)
wm.mask_total = end_frame - start_frame + 1
generator.start(
video_path=bpy.path.abspath(strip.filepath),
output_dir=output_dir,
start_frame=start_frame,
end_frame=end_frame,
fps=source_fps,
conf_threshold=scene.facemask_conf_threshold,
iou_threshold=scene.facemask_iou_threshold,
on_complete=on_complete,
on_progress=on_progress,
)
class SEQUENCER_OT_generate_face_mask(Operator): class SEQUENCER_OT_generate_face_mask(Operator):
@@ -22,56 +110,42 @@ class SEQUENCER_OT_generate_face_mask(Operator):
bl_description = "Detect faces and generate mask image sequence" bl_description = "Detect faces and generate mask image sequence"
bl_options = {'REGISTER', 'UNDO'} bl_options = {'REGISTER', 'UNDO'}
force: BoolProperty(
name="Force Regenerate",
description="既存のキャッシュを無視して再生成する",
default=False,
)
@classmethod @classmethod
def poll(cls, context): def poll(cls, context):
"""Check if operator can run."""
if not context.scene.sequence_editor: if not context.scene.sequence_editor:
return False return False
strip = context.scene.sequence_editor.active_strip strip = context.scene.sequence_editor.active_strip
if not strip: if not strip:
return False return False
return strip.type in {'MOVIE', 'IMAGE'} return strip.type in {'MOVIE', 'IMAGE'}
def execute(self, context): def execute(self, context):
strip = context.scene.sequence_editor.active_strip strip = context.scene.sequence_editor.active_strip
scene = context.scene
# Get video path # ファイル存在確認
if strip.type == 'MOVIE': if strip.type == 'MOVIE':
video_path = bpy.path.abspath(strip.filepath) video_path = bpy.path.abspath(strip.filepath)
else: else:
# Image sequence - get directory
video_path = bpy.path.abspath(strip.directory) video_path = bpy.path.abspath(strip.directory)
if not os.path.exists(video_path): if not os.path.exists(video_path):
self.report({'ERROR'}, f"Video file not found: {video_path}") self.report({'ERROR'}, f"Video file not found: {video_path}")
return {'CANCELLED'} return {'CANCELLED'}
# Determine output directory # キャッシュ確認(force=True の場合はスキップ)
output_dir = self._get_cache_dir(context, strip) if not self.force and check_detection_cache(strip.name):
self.report({'INFO'}, f"Using cached detections for {strip.name}")
# Check cache - if masks already exist, use them
expected_frame_count = strip.frame_final_end - strip.frame_final_start + 1
if self._check_cache(output_dir, expected_frame_count):
self.report({'INFO'}, f"Using cached detections from {output_dir}")
return {'FINISHED'} return {'FINISHED'}
# Get frame range
start_frame = strip.frame_final_start
end_frame = strip.frame_final_end
fps = scene.render.fps / scene.render.fps_base
# Start async generation
generator = get_generator() generator = get_generator()
if generator.is_running:
self.report({'WARNING'}, "Mask generation already in progress")
return {'CANCELLED'}
def on_complete(status, data): def on_complete(status, data):
"""Called when mask generation completes."""
wm = context.window_manager wm = context.window_manager
wm.mask_total = max(wm.mask_total, generator.total_frames) wm.mask_total = max(wm.mask_total, generator.total_frames)
if status == "done": if status == "done":
@@ -91,85 +165,25 @@ class SEQUENCER_OT_generate_face_mask(Operator):
area.tag_redraw() area.tag_redraw()
def on_progress(current, total): def on_progress(current, total):
"""Called on progress updates."""
# Update window manager properties for UI
wm = context.window_manager wm = context.window_manager
wm.mask_progress = current wm.mask_progress = current
wm.mask_total = total wm.mask_total = total
# Force UI redraw
for area in context.screen.areas: for area in context.screen.areas:
if area.type == 'SEQUENCE_EDITOR': if area.type == 'SEQUENCE_EDITOR':
area.tag_redraw() area.tag_redraw()
# Initialize progress try:
wm = context.window_manager start_mask_gen_for_strip(context, strip, on_complete, on_progress)
wm.mask_progress = 0 except RuntimeError as e:
wm.mask_total = end_frame - start_frame + 1 self.report({'WARNING'}, str(e))
return {'CANCELLED'}
# Get parameters from scene properties except Exception as e:
conf_threshold = scene.facemask_conf_threshold self.report({'ERROR'}, f"Failed to start mask generation: {e}")
iou_threshold = scene.facemask_iou_threshold return {'CANCELLED'}
mask_scale = scene.facemask_mask_scale
# Start generation
generator.start(
video_path=video_path,
output_dir=output_dir,
start_frame=0, # Frame indices in video
end_frame=end_frame - start_frame,
fps=fps,
conf_threshold=conf_threshold,
iou_threshold=iou_threshold,
mask_scale=mask_scale,
on_complete=on_complete,
on_progress=on_progress,
)
self.report({'INFO'}, f"Started mask generation for {strip.name}") self.report({'INFO'}, f"Started mask generation for {strip.name}")
return {'FINISHED'} return {'FINISHED'}
def _get_cache_dir(self, context, strip) -> str:
"""Get or create cache directory for mask images."""
cache_dir = get_cache_dir_for_strip(strip.name)
os.makedirs(cache_dir, exist_ok=True)
return cache_dir
def _check_cache(self, cache_dir: str, expected_frames: int) -> bool:
"""Check if cached masks exist and are complete.
Args:
cache_dir: Path to cache directory
expected_frames: Number of frames expected
Returns:
True if cache exists and is valid
"""
if not os.path.exists(cache_dir):
return False
detections_path = os.path.join(cache_dir, "detections.msgpack")
if not os.path.exists(detections_path):
return False
# Quick sanity check: non-empty file
try:
if os.path.getsize(detections_path) <= 0:
return False
except OSError:
return False
# Optional frame count verification if msgpack is available
try:
import msgpack
with open(detections_path, "rb") as f:
payload = msgpack.unpackb(f.read(), raw=False)
frames = payload.get("frames", [])
return len(frames) >= expected_frames * 0.9
except Exception:
return True
class SEQUENCER_OT_cancel_mask_generation(Operator): class SEQUENCER_OT_cancel_mask_generation(Operator):
"""Cancel ongoing mask generation.""" """Cancel ongoing mask generation."""
@@ -191,10 +205,88 @@ class SEQUENCER_OT_cancel_mask_generation(Operator):
return {'FINISHED'} return {'FINISHED'}
class SEQUENCER_OT_augment_pose_mask(Operator):
"""Add pose-based head detections to existing detection cache."""
bl_idname = "sequencer.augment_pose_mask"
bl_label = "Augment with Pose"
bl_description = "Run pose estimation and merge results into existing detection cache"
bl_options = {'REGISTER'}
@classmethod
def poll(cls, context):
if not context.scene.sequence_editor:
return False
strip = context.scene.sequence_editor.active_strip
if not strip or strip.type != 'MOVIE':
return False
return check_detection_cache(strip.name)
def execute(self, context):
strip = context.scene.sequence_editor.active_strip
output_dir = get_cache_dir_for_strip(strip.name)
detections_path = os.path.join(output_dir, "detections.msgpack")
if not os.path.exists(detections_path):
self.report({'ERROR'}, f"Detection cache not found: {detections_path}")
return {'CANCELLED'}
generator = get_generator()
scene = context.scene
wm = context.window_manager
wm.mask_progress = 0
wm.mask_total = 0 # サーバー側から実際の値に更新される
def on_complete(status, data):
wm.mask_total = max(wm.mask_total, generator.total_frames)
if status == "done":
wm.mask_progress = wm.mask_total
elif status in {"error", "cancelled"}:
wm.mask_progress = min(wm.mask_progress, wm.mask_total)
if status == "done":
print(f"[FaceMask] Pose augmentation completed: {data}")
elif status == "error":
print(f"[FaceMask] Error: {data}")
elif status == "cancelled":
print("[FaceMask] Pose augmentation cancelled")
for area in context.screen.areas:
if area.type == 'SEQUENCE_EDITOR':
area.tag_redraw()
def on_progress(current, total_f):
wm.mask_progress = current
wm.mask_total = total_f
for area in context.screen.areas:
if area.type == 'SEQUENCE_EDITOR':
area.tag_redraw()
try:
generator.start_augment_pose(
detections_path=detections_path,
total_frames=0,
conf_threshold=scene.facemask_conf_threshold,
iou_threshold=scene.facemask_iou_threshold,
on_complete=on_complete,
on_progress=on_progress,
)
except RuntimeError as e:
self.report({'WARNING'}, str(e))
return {'CANCELLED'}
except Exception as e:
self.report({'ERROR'}, f"Failed to start pose augmentation: {e}")
return {'CANCELLED'}
self.report({'INFO'}, f"Started pose augmentation for {strip.name}")
return {'FINISHED'}
# Registration # Registration
classes = [ classes = [
SEQUENCER_OT_generate_face_mask, SEQUENCER_OT_generate_face_mask,
SEQUENCER_OT_cancel_mask_generation, SEQUENCER_OT_cancel_mask_generation,
SEQUENCER_OT_augment_pose_mask,
] ]
@@ -202,13 +294,11 @@ def register():
for cls in classes: for cls in classes:
bpy.utils.register_class(cls) bpy.utils.register_class(cls)
# Add progress properties to window manager
bpy.types.WindowManager.mask_progress = IntProperty(default=0) bpy.types.WindowManager.mask_progress = IntProperty(default=0)
bpy.types.WindowManager.mask_total = IntProperty(default=0) bpy.types.WindowManager.mask_total = IntProperty(default=0)
def unregister(): def unregister():
# Remove properties
del bpy.types.WindowManager.mask_progress del bpy.types.WindowManager.mask_progress
del bpy.types.WindowManager.mask_total del bpy.types.WindowManager.mask_total
+124 -22
View File
@@ -11,11 +11,12 @@ from bpy.types import Panel
from ..core.async_bake_generator import get_bake_generator from ..core.async_bake_generator import get_bake_generator
from ..core.async_generator import get_generator from ..core.async_generator import get_generator
from ..core.batch_processor import get_batch_processor
from ..core.utils import ( from ..core.utils import (
get_server_status, get_server_status,
get_cache_info, get_cache_info,
format_size, format_size,
get_detections_path_for_strip, check_detection_cache,
) )
@@ -35,9 +36,15 @@ class SEQUENCER_PT_face_mask(Panel):
seq_editor = context.scene.sequence_editor seq_editor = context.scene.sequence_editor
# Note: Blender 5.0 uses 'strips' instead of 'sequences' # Note: Blender 5.0 uses 'strips' instead of 'sequences'
batch = get_batch_processor()
generator = get_generator() generator = get_generator()
bake_generator = get_bake_generator() bake_generator = get_bake_generator()
# Batch progress (highest priority)
if batch.is_running:
self._draw_batch_progress(layout, wm, batch, generator, bake_generator)
return
# Show progress if generating masks # Show progress if generating masks
if generator.is_running: if generator.is_running:
self._draw_progress(layout, wm, generator) self._draw_progress(layout, wm, generator)
@@ -65,6 +72,7 @@ class SEQUENCER_PT_face_mask(Panel):
self._draw_parameters(layout, scene) self._draw_parameters(layout, scene)
self._draw_server_status(layout) self._draw_server_status(layout)
self._draw_cache_info(layout, context, seq_editor) self._draw_cache_info(layout, context, seq_editor)
self._draw_batch_controls(layout, context, seq_editor)
def _draw_parameters(self, layout, scene): def _draw_parameters(self, layout, scene):
"""Draw detection parameters.""" """Draw detection parameters."""
@@ -74,7 +82,6 @@ class SEQUENCER_PT_face_mask(Panel):
col = box.column(align=True) col = box.column(align=True)
col.prop(scene, "facemask_conf_threshold") col.prop(scene, "facemask_conf_threshold")
col.prop(scene, "facemask_iou_threshold") col.prop(scene, "facemask_iou_threshold")
col.prop(scene, "facemask_mask_scale")
def _draw_server_status(self, layout): def _draw_server_status(self, layout):
"""Draw server status and GPU info.""" """Draw server status and GPU info."""
@@ -183,6 +190,76 @@ class SEQUENCER_PT_face_mask(Panel):
icon='CANCEL', icon='CANCEL',
) )
def _draw_batch_progress(self, layout, wm, batch, generator, bake_generator):
"""Draw batch bake progress."""
box = layout.box()
if batch._mode == "mask_only":
box.label(text="Batch Generating Cache...", icon='RENDER_ANIMATION')
else:
box.label(text="Batch Baking...", icon='RENDER_ANIMATION')
# Overall progress
total = max(wm.batch_total, 1)
# Show n-1/total while current strip is in progress, n/total when moving to next
done_count = max(wm.batch_current - 1, 0)
overall_factor = done_count / total
box.progress(
factor=overall_factor,
text=f"{wm.batch_current} / {wm.batch_total}",
)
if wm.batch_current_name:
box.label(text=f"Strip: {wm.batch_current_name}")
# Inner progress (mask gen or bake)
if generator.is_running:
inner = wm.mask_progress / max(wm.mask_total, 1)
box.progress(
factor=inner,
text=f"Detecting: {wm.mask_progress} / {wm.mask_total}",
)
elif bake_generator.is_running:
inner = wm.bake_progress / max(wm.bake_total, 1)
box.progress(
factor=inner,
text=f"Baking: {wm.bake_progress} / {wm.bake_total}",
)
box.operator(
"sequencer.cancel_batch_bake",
text="Cancel Batch",
icon='CANCEL',
)
def _draw_batch_controls(self, layout, context, seq_editor):
"""Draw batch bake button when multiple MOVIE/IMAGE strips are selected."""
if not seq_editor:
return
selected_movies = [s for s in seq_editor.strips if s.select and s.type in {"MOVIE", "IMAGE"}]
if not selected_movies:
return
count = len(selected_movies)
image_count = sum(1 for s in selected_movies if s.type == "IMAGE")
video_count = sum(1 for s in selected_movies if s.type == "MOVIE")
label = f"Batch ({count} selected, image: {image_count}, video: {video_count})"
box = layout.box()
box.label(text=label, icon='RENDER_ANIMATION')
box.operator(
"sequencer.batch_bake_selected",
text="Batch Bake Selected",
icon='RENDER_ANIMATION',
)
box.operator(
"sequencer.batch_regenerate_cache",
text="Batch Regenerate Cache",
icon='FILE_REFRESH',
)
box.operator(
"sequencer.batch_restore_original",
text="Batch Restore Original",
icon='LOOP_BACK',
)
def _draw_generation_controls(self, layout, context, strip): def _draw_generation_controls(self, layout, context, strip):
"""Draw mask generation controls.""" """Draw mask generation controls."""
box = layout.box() box = layout.box()
@@ -192,31 +269,39 @@ class SEQUENCER_PT_face_mask(Panel):
row = box.row() row = box.row()
row.label(text=f"Strip: {strip.name}") row.label(text=f"Strip: {strip.name}")
detections_path = get_detections_path_for_strip(strip.name) has_mask = check_detection_cache(strip.name)
has_mask = bpy.path.abspath(detections_path) and os.path.exists(
bpy.path.abspath(detections_path)
)
if has_mask: if has_mask:
row = box.row() row = box.row()
row.label(text="✓ Detection cache exists", icon='CHECKMARK') row.label(text="✓ Detection cache exists", icon='CHECKMARK')
# Generate button # Generate / Regenerate button
if not has_mask:
box.operator( box.operator(
"sequencer.generate_face_mask", "sequencer.generate_face_mask",
text="Generate Detection Cache" if not has_mask else "Regenerate Cache", text="Generate Detection Cache",
icon='FACE_MAPS', icon='FACE_MAPS',
) )
else:
op = box.operator(
"sequencer.generate_face_mask",
text="Regenerate Cache",
icon='FILE_REFRESH',
)
op.force = True
if strip.type == 'MOVIE':
box.operator(
"sequencer.augment_pose_mask",
text="Augment with Pose",
icon='MOD_ARMATURE',
)
def _draw_blur_controls(self, layout, context, strip): def _draw_blur_controls(self, layout, context, strip):
"""Draw blur application controls.""" """Draw blur application controls."""
box = layout.box() box = layout.box()
box.label(text="Blur Bake", icon='MATFLUID') box.label(text="Blur Bake", icon='MATFLUID')
detections_path = get_detections_path_for_strip(strip.name) has_mask = check_detection_cache(strip.name)
has_mask = bpy.path.abspath(detections_path) and os.path.exists(
bpy.path.abspath(detections_path)
)
if not has_mask: if not has_mask:
box.label(text="Generate detection cache first", icon='INFO') box.label(text="Generate detection cache first", icon='INFO')
@@ -225,26 +310,43 @@ class SEQUENCER_PT_face_mask(Panel):
# Bake parameters # Bake parameters
col = box.column(align=True) col = box.column(align=True)
col.prop(context.scene, "facemask_bake_blur_size") col.prop(context.scene, "facemask_bake_blur_size")
col.prop(context.scene, "facemask_bake_display_scale")
if strip.type == "MOVIE":
col.prop(context.scene, "facemask_bake_format") col.prop(context.scene, "facemask_bake_format")
# Source status box.separator()
source_mode = strip.get("facemask_source_mode", "original")
if source_mode == "baked":
box.label(text="Source: Baked", icon='CHECKMARK')
else:
box.label(text="Source: Original", icon='FILE_MOVIE')
# Bake and restore buttons baked_path = strip.get("facemask_baked_filepath", "")
has_baked = bool(baked_path and os.path.exists(bpy.path.abspath(baked_path)))
source_mode = strip.get("facemask_source_mode", "original")
if not has_baked:
# 初回: Bakeのみ
box.operator( box.operator(
"sequencer.bake_and_swap_blur_source", "sequencer.bake_and_swap_blur_source",
text="Bake & Swap Source", text="Bake",
icon='RENDER_STILL', icon='RENDER_STILL',
) )
box.operator( else:
# Bake済み: ソース切り替え + Re-bake
row = box.row(align=True)
if source_mode == "baked":
row.operator(
"sequencer.restore_original_source", "sequencer.restore_original_source",
text="Restore Original Source", text="Restore Original",
icon='LOOP_BACK', icon='LOOP_BACK',
) )
else:
row.operator(
"sequencer.swap_to_baked_blur",
text="Swap to Baked",
icon='PLAY',
)
row.operator(
"sequencer.force_rebake_blur",
text="Re-bake",
icon='FILE_REFRESH',
)
# Registration # Registration
+257 -98
View File
@@ -1,28 +1,36 @@
""" """
YOLOv8 Face Detector using PyTorch with ROCm support. YOLOv8 Head Detector using CrowdHuman-trained model with PyTorch ROCm support.
This module provides high-performance face detection using Directly detects human heads (frontal, profile, rear) using the Owen718
YOLOv8-face model with AMD GPU (ROCm) acceleration. CrowdHuman YOLOv8 model, which was trained on dense crowd scenes.
""" """
import os import os
from typing import List, Tuple, Optional from typing import List, Tuple, Optional
from pathlib import Path
import numpy as np import numpy as np
class YOLOFaceDetector: def _download_model(dest_path: str):
""" """モデルが存在しない場合に手動ダウンロード手順を表示して例外を送出する。"""
YOLOv8 face detector with PyTorch ROCm support. gdrive_id = "1qlBmiEU4GBV13fxPhLZqjhjBbREvs8-m"
raise RuntimeError(
f"モデルファイルが見つかりません: {dest_path}\n"
"以下の手順でダウンロードしてください:\n"
f" 1. https://drive.google.com/file/d/{gdrive_id} を開く\n"
f" 2. ダウンロードしたファイルを {dest_path} に配置する"
)
Features:
- ROCm GPU acceleration for AMD GPUs class YOLOHeadDetector:
- High accuracy face detection """
- Automatic NMS for overlapping detections Head detector using CrowdHuman-trained YOLOv8 model with PyTorch ROCm support.
Directly detects heads (class 0: head) without pose estimation,
enabling robust detection of rear-facing, side-facing, and partially
visible people in dense crowd scenes.
""" """
# Default model path relative to this file DEFAULT_MODEL = os.path.join("models", "crowdhuman_yolov8_head.pt")
DEFAULT_MODEL = "yolov8n-face-lindevs.pt"
def __init__( def __init__(
self, self,
@@ -31,15 +39,6 @@ class YOLOFaceDetector:
iou_threshold: float = 0.45, iou_threshold: float = 0.45,
input_size: Tuple[int, int] = (640, 640), input_size: Tuple[int, int] = (640, 640),
): ):
"""
Initialize the YOLO face detector.
Args:
model_path: Path to PyTorch model file. If None, uses default model.
conf_threshold: Confidence threshold for detections
iou_threshold: IoU threshold for NMS
input_size: Model input size (width, height)
"""
self.conf_threshold = conf_threshold self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold self.iou_threshold = iou_threshold
self.input_size = input_size self.input_size = input_size
@@ -49,23 +48,20 @@ class YOLOFaceDetector:
@property @property
def model(self): def model(self):
"""Lazy-load YOLO model.""" """Lazy-load YOLO head detection model."""
if self._model is None: if self._model is None:
from ultralytics import YOLO from ultralytics import YOLO
import torch import torch
# Determine model path if self._model_path is not None:
if self._model_path is None: if not os.path.exists(self._model_path):
# Assuming models are in ../models relative to server/detector.py raise FileNotFoundError(f"Model not found: {self._model_path}")
models_dir = Path(__file__).parent.parent / "models"
model_path = str(models_dir / self.DEFAULT_MODEL)
else:
model_path = self._model_path model_path = self._model_path
else:
model_path = self.DEFAULT_MODEL
if not os.path.exists(model_path): if not os.path.exists(model_path):
raise FileNotFoundError(f"Model not found: {model_path}") _download_model(model_path)
# Detect device (ROCm GPU or CPU)
if torch.cuda.is_available(): if torch.cuda.is_available():
self._device = 'cuda' self._device = 'cuda'
device_name = torch.cuda.get_device_name(0) device_name = torch.cuda.get_device_name(0)
@@ -74,25 +70,32 @@ class YOLOFaceDetector:
self._device = 'cpu' self._device = 'cpu'
print("[FaceMask] Using CPU for inference (ROCm GPU not available)") print("[FaceMask] Using CPU for inference (ROCm GPU not available)")
# Load model (let Ultralytics handle device management)
try: try:
self._model = YOLO(model_path) self._model = YOLO(model_path)
# Don't call .to() - let predict() handle device assignment print(f"[FaceMask] Head detection model loaded: {model_path}")
print(f"[FaceMask] Model loaded, will use device: {self._device}") print(f"[FaceMask] Device: {self._device}")
except Exception as e: except Exception as e:
print(f"[FaceMask] Error loading model: {e}") print(f"[FaceMask] Error loading model: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
raise raise
print(f"[FaceMask] YOLO model loaded: {model_path}")
print(f"[FaceMask] Device: {self._device}")
return self._model return self._model
def _results_to_detections(self, result) -> List[Tuple[int, int, int, int, float]]:
"""Convert a single YOLO result to (x, y, w, h, conf) tuples."""
if result.boxes is None:
return []
detections = []
for box in result.boxes:
conf = float(box.conf[0].cpu().numpy())
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
detections.append((int(x1), int(y1), int(x2 - x1), int(y2 - y1), conf))
return detections
def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]: def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]:
""" """
Detect faces in a frame. Detect heads in a frame.
Args: Args:
frame: BGR image as numpy array (H, W, C) frame: BGR image as numpy array (H, W, C)
@@ -100,7 +103,6 @@ class YOLOFaceDetector:
Returns: Returns:
List of detections as (x, y, width, height, confidence) List of detections as (x, y, width, height, confidence)
""" """
# Run inference
import torch import torch
print(f"[FaceMask] Inference device: {self._device}, CUDA available: {torch.cuda.is_available()}") print(f"[FaceMask] Inference device: {self._device}, CUDA available: {torch.cuda.is_available()}")
try: try:
@@ -116,7 +118,6 @@ class YOLOFaceDetector:
print(f"[FaceMask] ERROR during inference: {e}") print(f"[FaceMask] ERROR during inference: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
# Fallback to CPU
print("[FaceMask] Falling back to CPU inference...") print("[FaceMask] Falling back to CPU inference...")
self._device = 'cpu' self._device = 'cpu'
results = self.model.predict( results = self.model.predict(
@@ -128,28 +129,13 @@ class YOLOFaceDetector:
device='cpu', device='cpu',
) )
# Extract detections if results:
detections = [] return self._results_to_detections(results[0])
if len(results) > 0 and results[0].boxes is not None: return []
boxes = results[0].boxes
for box in boxes:
# Get coordinates in xyxy format
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
conf = float(box.conf[0].cpu().numpy())
# Convert to x, y, width, height
x = int(x1)
y = int(y1)
w = int(x2 - x1)
h = int(y2 - y1)
detections.append((x, y, w, h, conf))
return detections
def detect_batch(self, frames: List[np.ndarray]) -> List[List[Tuple[int, int, int, int, float]]]: def detect_batch(self, frames: List[np.ndarray]) -> List[List[Tuple[int, int, int, int, float]]]:
""" """
Detect faces in multiple frames at once (batch processing). Detect heads in multiple frames at once (batch processing).
Args: Args:
frames: List of BGR images as numpy arrays (H, W, C) frames: List of BGR images as numpy arrays (H, W, C)
@@ -161,7 +147,6 @@ class YOLOFaceDetector:
if not frames: if not frames:
return [] return []
# Run batch inference
try: try:
results = self.model.predict( results = self.model.predict(
frames, frames,
@@ -175,7 +160,6 @@ class YOLOFaceDetector:
print(f"[FaceMask] ERROR during batch inference: {e}") print(f"[FaceMask] ERROR during batch inference: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
# Fallback to CPU
print("[FaceMask] Falling back to CPU inference...") print("[FaceMask] Falling back to CPU inference...")
self._device = 'cpu' self._device = 'cpu'
results = self.model.predict( results = self.model.predict(
@@ -187,28 +171,7 @@ class YOLOFaceDetector:
device='cpu', device='cpu',
) )
# Extract detections for each frame return [self._results_to_detections(r) for r in results]
all_detections = []
for result in results:
detections = []
if result.boxes is not None:
boxes = result.boxes
for box in boxes:
# Get coordinates in xyxy format
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
conf = float(box.conf[0].cpu().numpy())
# Convert to x, y, width, height
x = int(x1)
y = int(y1)
w = int(x2 - x1)
h = int(y2 - y1)
detections.append((x, y, w, h, conf))
all_detections.append(detections)
return all_detections
def generate_mask( def generate_mask(
self, self,
@@ -218,11 +181,11 @@ class YOLOFaceDetector:
feather_radius: int = 20, feather_radius: int = 20,
) -> np.ndarray: ) -> np.ndarray:
""" """
Generate a mask image from face detections. Generate a mask image from head detections.
Args: Args:
frame_shape: Shape of the original frame (height, width, channels) frame_shape: Shape of the original frame (height, width, channels)
detections: List of face detections (x, y, w, h, conf) detections: List of head detections (x, y, w, h, conf)
mask_scale: Scale factor for mask region mask_scale: Scale factor for mask region
feather_radius: Radius for edge feathering feather_radius: Radius for edge feathering
@@ -235,25 +198,19 @@ class YOLOFaceDetector:
mask = np.zeros((height, width), dtype=np.uint8) mask = np.zeros((height, width), dtype=np.uint8)
for (x, y, w, h, conf) in detections: for (x, y, w, h, conf) in detections:
# Scale the bounding box
center_x = x + w // 2 center_x = x + w // 2
center_y = y + h // 2 center_y = y + h // 2
scaled_w = int(w * mask_scale) scaled_w = int(w * mask_scale)
scaled_h = int(h * mask_scale) scaled_h = int(h * mask_scale)
# Draw ellipse for natural face shape
cv2.ellipse( cv2.ellipse(
mask, mask,
(center_x, center_y), (center_x, center_y),
(scaled_w // 2, scaled_h // 2), (scaled_w // 2, scaled_h // 2),
0, # angle 0, 0, 360,
0, 360, # arc 255, -1,
255, # color (white)
-1, # filled
) )
# Apply Gaussian blur for feathering
if feather_radius > 0 and len(detections) > 0: if feather_radius > 0 and len(detections) > 0:
kernel_size = feather_radius * 2 + 1 kernel_size = feather_radius * 2 + 1
mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0) mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
@@ -262,12 +219,214 @@ class YOLOFaceDetector:
# Singleton instance # Singleton instance
_detector: Optional[YOLOFaceDetector] = None _detector: Optional[YOLOHeadDetector] = None
def get_detector(**kwargs) -> YOLOFaceDetector: def get_detector(**kwargs) -> YOLOHeadDetector:
"""Get or create the global YOLO detector instance.""" """Get or create the global YOLO head detector instance."""
global _detector global _detector
if _detector is None: if _detector is None:
_detector = YOLOFaceDetector(**kwargs) _detector = YOLOHeadDetector(**kwargs)
return _detector return _detector
# ---------------------------------------------------------------------------
# Pose-based head detector (YOLOv8 pose estimation)
# ---------------------------------------------------------------------------
# COCO pose keypoint indices
_HEAD_KP = [0, 1, 2, 3, 4] # nose, left_eye, right_eye, left_ear, right_ear
_SHOULDER_KP = [5, 6] # left_shoulder, right_shoulder
_KP_CONF_THRESH = 0.3
def _head_bbox_from_pose(
kp_xy: np.ndarray,
kp_conf: np.ndarray,
person_x1: float,
person_y1: float,
person_x2: float,
person_y2: float,
) -> Tuple[int, int, int, int]:
"""
Estimate head bounding box (x, y, w, h) from COCO pose keypoints.
Strategy:
1. Use head keypoints (0-4: nose, eyes, ears) if visible.
2. Fall back to shoulder keypoints (5-6) to infer head position.
3. Last resort: use top of the person bounding box.
"""
person_w = max(person_x2 - person_x1, 1.0)
# --- Step 1: head keypoints ---
visible_head = [
(float(kp_xy[i][0]), float(kp_xy[i][1]))
for i in _HEAD_KP
if float(kp_conf[i]) > _KP_CONF_THRESH
]
if visible_head:
xs = [p[0] for p in visible_head]
ys = [p[1] for p in visible_head]
kp_x1, kp_y1 = min(xs), min(ys)
kp_x2, kp_y2 = max(xs), max(ys)
span = max(kp_x2 - kp_x1, kp_y2 - kp_y1, 1.0)
cx = (kp_x1 + kp_x2) / 2.0
cy = (kp_y1 + kp_y2) / 2.0
r = max(span * 0.5, person_w * 0.10)
x1 = int(cx - r)
y1 = int(cy - r)
x2 = int(cx + r)
y2 = int(cy + r)
return x1, y1, x2 - x1, y2 - y1
# --- Step 2: shoulder keypoints ---
visible_shoulder = [
(float(kp_xy[i][0]), float(kp_xy[i][1]))
for i in _SHOULDER_KP
if float(kp_conf[i]) > _KP_CONF_THRESH
]
if visible_shoulder:
cx = sum(p[0] for p in visible_shoulder) / len(visible_shoulder)
cy_sh = sum(p[1] for p in visible_shoulder) / len(visible_shoulder)
if len(visible_shoulder) == 2:
sh_width = abs(visible_shoulder[1][0] - visible_shoulder[0][0])
else:
sh_width = person_w * 0.5
r = max(sh_width * 0.3, person_w * 0.12)
cy = cy_sh - r * 1.3
x1 = int(cx - r)
y1 = int(cy - r)
x2 = int(cx + r)
y2 = int(cy + r)
return x1, y1, x2 - x1, y2 - y1
# --- Step 3: person bbox top ---
r = max(person_w * 0.15, 20.0)
cx = (person_x1 + person_x2) / 2.0
x1 = int(cx - r)
y1 = int(person_y1)
x2 = int(cx + r)
y2 = int(person_y1 + r * 2.0)
return x1, y1, x2 - x1, y2 - y1
class YOLOPoseHeadDetector:
"""
Head detector using YOLOv8 pose estimation with PyTorch ROCm support.
Extracts head bounding boxes from COCO pose keypoints (nose, eyes, ears).
yolov8l-pose.pt is auto-downloaded by Ultralytics on first use.
"""
DEFAULT_MODEL = os.path.join("models", "yolov8l-pose.pt")
def __init__(
self,
model_path: Optional[str] = None,
conf_threshold: float = 0.25,
iou_threshold: float = 0.45,
input_size: Tuple[int, int] = (640, 640),
):
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.input_size = input_size
self._model = None
self._model_path = model_path
self._device = None
@property
def model(self):
"""Lazy-load YOLO pose model."""
if self._model is None:
from ultralytics import YOLO
import torch
model_path = self._model_path if self._model_path is not None else self.DEFAULT_MODEL
if torch.cuda.is_available():
self._device = 'cuda'
device_name = torch.cuda.get_device_name(0)
print(f"[FaceMask] Using ROCm GPU for pose inference: {device_name}")
else:
self._device = 'cpu'
print("[FaceMask] Using CPU for pose inference (ROCm GPU not available)")
try:
self._model = YOLO(model_path)
print(f"[FaceMask] Pose model loaded: {model_path}")
print(f"[FaceMask] Device: {self._device}")
except Exception as e:
print(f"[FaceMask] Error loading pose model: {e}")
import traceback
traceback.print_exc()
raise
return self._model
def _results_to_detections(self, result) -> List[Tuple[int, int, int, int, float]]:
"""Convert a single YOLO pose result to (x, y, w, h, conf) tuples."""
detections = []
if result.boxes is None or result.keypoints is None:
return detections
boxes = result.boxes
keypoints = result.keypoints
for i, box in enumerate(boxes):
conf = float(box.conf[0].cpu().numpy())
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
kp_data = keypoints.data[i].cpu().numpy() # shape (17, 3): x, y, conf
kp_xy = kp_data[:, :2]
kp_conf = kp_data[:, 2]
hx, hy, hw, hh = _head_bbox_from_pose(
kp_xy, kp_conf,
float(x1), float(y1), float(x2), float(y2),
)
detections.append((hx, hy, hw, hh, conf))
return detections
def detect_batch(self, frames: List[np.ndarray]) -> List[List[Tuple[int, int, int, int, float]]]:
"""Detect heads in multiple frames at once (batch processing)."""
if not frames:
return []
try:
results = self.model.predict(
frames,
conf=self.conf_threshold,
iou=self.iou_threshold,
imgsz=self.input_size[0],
verbose=False,
device=self._device,
)
except Exception as e:
print(f"[FaceMask] ERROR during pose batch inference: {e}")
import traceback
traceback.print_exc()
print("[FaceMask] Falling back to CPU inference...")
self._device = 'cpu'
results = self.model.predict(
frames,
conf=self.conf_threshold,
iou=self.iou_threshold,
imgsz=self.input_size[0],
verbose=False,
device='cpu',
)
return [self._results_to_detections(r) for r in results]
# Pose detector singleton
_pose_detector: Optional[YOLOPoseHeadDetector] = None
def get_pose_detector(**kwargs) -> YOLOPoseHeadDetector:
"""Get or create the global YOLO pose head detector instance."""
global _pose_detector
if _pose_detector is None:
_pose_detector = YOLOPoseHeadDetector(**kwargs)
return _pose_detector
+684 -63
View File
@@ -31,28 +31,62 @@ def fix_library_path():
# Fix library path BEFORE any other imports # Fix library path BEFORE any other imports
fix_library_path() fix_library_path()
import queue import queue # noqa: E402
import threading import threading # noqa: E402
import uuid import uuid # noqa: E402
import traceback import traceback # noqa: E402
import subprocess import subprocess # noqa: E402
from typing import Dict, Optional, List from typing import Dict, Optional, List # noqa: E402
from pathlib import Path from pathlib import Path # noqa: E402
from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi import FastAPI, HTTPException, BackgroundTasks # noqa: E402
from pydantic import BaseModel from pydantic import BaseModel # noqa: E402
import uvicorn import uvicorn # noqa: E402
import cv2 import cv2 # noqa: E402
import numpy as np import numpy as np # noqa: E402
import msgpack import msgpack # noqa: E402
# Add project root to path for imports if needed # Add project root to path for imports if needed
sys.path.append(str(Path(__file__).parent.parent)) sys.path.append(str(Path(__file__).parent.parent))
from server.detector import get_detector from server.detector import get_detector, get_pose_detector # noqa: E402
app = FastAPI(title="Face Mask Inference Server") app = FastAPI(title="Face Mask Inference Server")
def _get_r_frame_rate(video_path: str) -> tuple:
"""ffprobe でコンテナ宣言の r_frame_rate を取得する。
Returns:
(fps_float, fps_str): fps_str は "120/1" のような分数文字列。
取得失敗時は (0.0, "")。
"""
try:
result = subprocess.run(
[
"ffprobe", "-v", "error",
"-select_streams", "v:0",
"-show_entries", "stream=r_frame_rate",
"-of", "default=noprint_wrappers=1:nokey=1",
video_path,
],
capture_output=True,
text=True,
timeout=10,
)
if result.returncode == 0:
rate_str = result.stdout.strip()
if "/" in rate_str:
num, den = rate_str.split("/")
fps_float = float(num) / float(den)
else:
fps_float = float(rate_str)
rate_str = str(fps_float)
return fps_float, rate_str
except Exception:
pass
return 0.0, ""
# GPU status cache # GPU status cache
_gpu_status_cache = None _gpu_status_cache = None
@@ -83,7 +117,10 @@ class GenerateRequest(BaseModel):
end_frame: int end_frame: int
conf_threshold: float = 0.5 conf_threshold: float = 0.5
iou_threshold: float = 0.45 iou_threshold: float = 0.45
mask_scale: float = 1.5
class VideoInfoRequest(BaseModel):
video_path: str
class BakeRequest(BaseModel): class BakeRequest(BaseModel):
@@ -91,9 +128,35 @@ class BakeRequest(BaseModel):
detections_path: str detections_path: str
output_path: str output_path: str
blur_size: int = 50 blur_size: int = 50
display_scale: float = 1.0
format: str = "mp4" 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 AugmentPoseRequest(BaseModel):
detections_path: str
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: class _FFmpegPipeWriter:
"""Write BGR frames to ffmpeg stdin.""" """Write BGR frames to ffmpeg stdin."""
@@ -138,8 +201,33 @@ def _build_ffmpeg_vaapi_writer(
fps: float, fps: float,
width: int, width: int,
height: int, height: int,
out_fps_str: str = "",
) -> _FFmpegPipeWriter: ) -> _FFmpegPipeWriter:
"""Create ffmpeg h264_vaapi writer with QP=24 (balanced quality/speed).""" """Create ffmpeg h264_vaapi writer with QP=24 (balanced quality/speed).
fps: ソース動画の avg_frame_raterawパイプの入力レート)
out_fps_str: 出力コンテナに宣言する r_frame_rate"120/1" 等)。
ソースと異なる場合は fps フィルタでフレームを補完する。
"""
# ソースの avg_fps と出力の r_fps が有意に異なる場合のみ fps フィルタを挿入
needs_fps_filter = bool(out_fps_str)
if needs_fps_filter:
try:
if "/" in out_fps_str:
num, den = out_fps_str.split("/")
out_fps_float = float(num) / float(den)
else:
out_fps_float = float(out_fps_str)
needs_fps_filter = abs(out_fps_float - fps) > 0.01
except ValueError:
needs_fps_filter = False
if needs_fps_filter:
vf = f"format=nv12,fps={out_fps_str},hwupload"
print(f"[FaceMask] fps filter: {fps:.3f} -> {out_fps_str}")
else:
vf = "format=nv12,hwupload"
cmd = [ cmd = [
"ffmpeg", "ffmpeg",
"-hide_banner", "-hide_banner",
@@ -160,7 +248,7 @@ def _build_ffmpeg_vaapi_writer(
"-", "-",
"-an", "-an",
"-vf", "-vf",
"format=nv12,hwupload", vf,
"-c:v", "-c:v",
"h264_vaapi", "h264_vaapi",
"-qp", "-qp",
@@ -176,13 +264,14 @@ def _build_video_writer(
fps: float, fps: float,
width: int, width: int,
height: int, height: int,
out_fps_str: str = "",
) -> object: ) -> object:
"""Create writer with VAAPI preference and OpenCV fallback.""" """Create writer with VAAPI preference and OpenCV fallback."""
format_key = fmt.lower() format_key = fmt.lower()
if format_key in {"mp4", "mov"}: if format_key in {"mp4", "mov"}:
try: try:
writer = _build_ffmpeg_vaapi_writer(output_path, fps, width, height) writer = _build_ffmpeg_vaapi_writer(output_path, fps, width, height, out_fps_str)
print("[FaceMask] Using output encoder: ffmpeg h264_vaapi (-qp 24)") print("[FaceMask] Using output encoder: ffmpeg h264_vaapi (-qp 24)")
return writer return writer
except Exception as e: except Exception as e:
@@ -238,6 +327,390 @@ def _scale_bbox(
return [x1, y1, out_w, out_h] 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 augment_pose_task(task_id: str, req: AugmentPoseRequest):
"""Background task: run pose estimation and merge results into existing cache."""
cap = None
try:
tasks[task_id].status = TaskStatus.PROCESSING
cancel_event = cancel_events.get(task_id)
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)
existing_frames: List[List[List[float]]] = payload.get("frames", [])
video_path = payload.get("video_path")
start_frame = int(payload.get("start_frame", 0))
total = len(existing_frames)
if not video_path:
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = "Cache does not contain video_path (image caches not supported)"
return
if not os.path.exists(video_path):
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = f"Video not found: {video_path}"
return
if total == 0:
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = "Cache has no frames"
return
tasks[task_id].total = total
detector = get_pose_detector(
conf_threshold=req.conf_threshold,
iou_threshold=req.iou_threshold,
)
_ = detector.model
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = "Failed to open video"
return
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
if start_frame > 0:
seek_ok = cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
if not seek_ok:
for _ in range(start_frame):
ret, _ = cap.read()
if not ret:
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = f"Failed to seek to start frame: {start_frame}"
return
frame_buffer: List[np.ndarray] = []
buffer_indices: List[int] = [] # existing_frames インデックス対応
current_count = 0
batch_size = 5
def process_pose_batch():
nonlocal current_count
if not frame_buffer:
return
batch_detections = detector.detect_batch(frame_buffer)
for idx, detections in zip(buffer_indices, batch_detections):
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 and bh > 0:
existing_frames[idx].append([bx, by, bw, bh, float(conf)])
current_count += 1
tasks[task_id].progress = current_count
frame_buffer.clear()
buffer_indices.clear()
for i 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
ret, frame = cap.read()
if not ret:
break
frame_buffer.append(frame)
buffer_indices.append(i)
if len(frame_buffer) >= batch_size:
process_pose_batch()
if frame_buffer:
process_pose_batch()
if tasks[task_id].status == TaskStatus.PROCESSING:
payload["frames"] = existing_frames
with open(req.detections_path, "wb") as f:
f.write(msgpack.packb(payload, use_bin_type=True))
tasks[task_id].status = TaskStatus.COMPLETED
tasks[task_id].result_path = req.detections_path
tasks[task_id].message = "Pose augmentation completed"
print(f"[FaceMask] Pose augmentation completed: {req.detections_path}")
except Exception as e:
tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = str(e)
traceback.print_exc()
finally:
if cap:
cap.release()
if task_id in cancel_events:
del cancel_events[task_id]
def process_video_task(task_id: str, req: GenerateRequest): def process_video_task(task_id: str, req: GenerateRequest):
"""Background task to detect faces and save bbox cache as msgpack.""" """Background task to detect faces and save bbox cache as msgpack."""
cap = None cap = None
@@ -305,20 +778,15 @@ def process_video_task(task_id: str, req: GenerateRequest):
for detections in batch_detections: for detections in batch_detections:
packed_detections: List[List[float]] = [] packed_detections: List[List[float]] = []
for x, y, w, h, conf in detections: for x, y, w, h, conf in detections:
scaled = _scale_bbox( # bboxをそのまま保存(表示スケールはBake時に適用)
int(x), bx, by, bw, bh = int(x), int(y), int(w), int(h)
int(y), bx = max(0, bx)
int(w), by = max(0, by)
int(h), bw = min(width - bx, bw)
float(req.mask_scale), bh = min(height - by, bh)
width, if bw <= 0 or bh <= 0:
height,
)
if scaled is None:
continue continue
packed_detections.append( packed_detections.append([bx, by, bw, bh, float(conf)])
[scaled[0], scaled[1], scaled[2], scaled[3], float(conf)]
)
frame_detections.append(packed_detections) frame_detections.append(packed_detections)
current_count += 1 current_count += 1
tasks[task_id].progress = current_count tasks[task_id].progress = current_count
@@ -356,7 +824,7 @@ def process_video_task(task_id: str, req: GenerateRequest):
"width": width, "width": width,
"height": height, "height": height,
"fps": fps, "fps": fps,
"mask_scale": float(req.mask_scale), "mask_scale": 1.0,
"frames": frame_detections, "frames": frame_detections,
} }
with open(output_msgpack_path, "wb") as f: with open(output_msgpack_path, "wb") as f:
@@ -402,6 +870,9 @@ def process_bake_task(task_id: str, req: BakeRequest):
tasks[task_id].message = "Invalid detections format: 'frames' is missing" tasks[task_id].message = "Invalid detections format: 'frames' is missing"
return return
# 検出キャッシュの開始フレーム(ソース動画のフレームインデックス)
det_start_frame = int(payload.get("start_frame", 0))
# Get video info # Get video info
temp_cap = cv2.VideoCapture(req.video_path) temp_cap = cv2.VideoCapture(req.video_path)
if not temp_cap.isOpened(): if not temp_cap.isOpened():
@@ -415,12 +886,22 @@ def process_bake_task(task_id: str, req: BakeRequest):
src_frames = int(temp_cap.get(cv2.CAP_PROP_FRAME_COUNT)) src_frames = int(temp_cap.get(cv2.CAP_PROP_FRAME_COUNT))
temp_cap.release() temp_cap.release()
# ffprobe で r_frame_rate を取得し、出力コンテナの宣言 FPS をソースに合わせる。
# 例: 120fps タイムベースで記録された 60fps 動画は r_frame_rate=120/1 だが
# cv2 は avg_frame_rate=60fps を返すため、Bake 後に Blender がFPSを別値で認識してしまう。
r_fps_float, r_fps_str = _get_r_frame_rate(req.video_path)
if r_fps_float > 0:
print(f"[FaceMask] r_frame_rate={r_fps_str}, avg_fps={src_fps:.3f}")
else:
r_fps_str = ""
if src_width <= 0 or src_height <= 0: if src_width <= 0 or src_height <= 0:
tasks[task_id].status = TaskStatus.FAILED tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = "Invalid source video dimensions" tasks[task_id].message = "Invalid source video dimensions"
return return
total = min(src_frames, len(frames_detections)) if src_frames > 0 else len(frames_detections) # ソース動画の全フレームを出力(スワップ後もトリム設定が正しく機能するよう)
total = src_frames if src_frames > 0 else (det_start_frame + len(frames_detections))
if total <= 0: if total <= 0:
tasks[task_id].status = TaskStatus.FAILED tasks[task_id].status = TaskStatus.FAILED
tasks[task_id].message = "Source/detections frame count is zero" tasks[task_id].message = "Source/detections frame count is zero"
@@ -435,9 +916,9 @@ def process_bake_task(task_id: str, req: BakeRequest):
blur_size = max(1, int(req.blur_size)) blur_size = max(1, int(req.blur_size))
if blur_size % 2 == 0: if blur_size % 2 == 0:
blur_size += 1 blur_size += 1
feather_radius = max(3, min(25, blur_size // 3)) display_scale = max(0.1, float(req.display_scale))
feather_kernel = feather_radius * 2 + 1 # blur_margin は境界問題回避のための計算用余白のみ(表示には使わない)
blur_margin = max(1, (blur_size // 2) + feather_radius) blur_margin = blur_size // 2
# Queues # Queues
queue_size = 8 queue_size = 8
@@ -451,52 +932,74 @@ def process_bake_task(task_id: str, req: BakeRequest):
def _reader_worker(): def _reader_worker():
"""Read frames from video.""" """Read frames from video."""
import time as _time
cap = cv2.VideoCapture(req.video_path) cap = cv2.VideoCapture(req.video_path)
if not cap.isOpened(): if not cap.isOpened():
error_holder["error"] = "Failed to open video in reader" error_holder["error"] = "Failed to open video in reader"
return return
t_read_total = 0.0
frame_count = 0
try: try:
for idx in range(total): for idx in range(total):
if cancel_event and cancel_event.is_set(): if cancel_event and cancel_event.is_set():
break break
t0 = _time.perf_counter()
ok, frame = cap.read() ok, frame = cap.read()
t_read_total += _time.perf_counter() - t0
if not ok: if not ok:
break break
read_queue.put((idx, frame)) read_queue.put((idx, frame))
frame_count += 1
except Exception as e: except Exception as e:
error_holder["error"] = f"Reader error: {e}" error_holder["error"] = f"Reader error: {e}"
finally: finally:
cap.release() cap.release()
read_queue.put(None) # Sentinel read_queue.put(None) # Sentinel
if frame_count > 0:
print(
f"[Perf/Reader] FINAL frame={frame_count}"
f" read_avg={t_read_total/frame_count*1000:.1f}ms"
f" throughput≈{frame_count/max(t_read_total,1e-9):.1f}fps"
)
def _processor_worker(): def _processor_worker():
"""Process frames with ROI blur.""" """Process frames with ROI blur."""
import time as _time
t_wait_total = 0.0
t_blur_total = 0.0
t_blend_total = 0.0
frame_count = 0
REPORT_INTERVAL = 50
try: try:
while True: while True:
if cancel_event and cancel_event.is_set(): if cancel_event and cancel_event.is_set():
process_queue.put(None) process_queue.put(None)
break break
t0 = _time.perf_counter()
item = read_queue.get() item = read_queue.get()
t_wait_total += _time.perf_counter() - t0
if item is None: if item is None:
process_queue.put(None) process_queue.put(None)
break break
idx, frame = item idx, frame = item
frame_boxes = frames_detections[idx] if idx < len(frames_detections) else [] det_idx = idx - det_start_frame
frame_boxes = frames_detections[det_idx] if 0 <= det_idx < len(frames_detections) else []
if not frame_boxes: if not frame_boxes:
process_queue.put((idx, frame)) process_queue.put((idx, frame))
frame_count += 1
continue continue
# ROI processing (same as original) # 各人物ごとに個別ROIで処理(全員まとめると離れた人物間が巨大ROIになるため)
min_x, min_y = src_width, src_height
max_x, max_y = 0, 0
valid_boxes = [] valid_boxes = []
for box in frame_boxes: for box in frame_boxes:
if not isinstance(box, list) or len(box) < 4: if not isinstance(box, list) or len(box) < 4:
continue continue
@@ -504,64 +1007,112 @@ def process_bake_task(task_id: str, req: BakeRequest):
if w <= 0 or h <= 0: if w <= 0 or h <= 0:
continue continue
valid_boxes.append((x, y, w, h)) valid_boxes.append((x, y, w, h))
min_x = min(min_x, x)
min_y = min(min_y, y)
max_x = max(max_x, x + w)
max_y = max(max_y, y + h)
if not valid_boxes: if not valid_boxes:
process_queue.put((idx, frame)) process_queue.put((idx, frame))
frame_count += 1
continue continue
roi_x1 = max(0, min_x - blur_margin) for x, y, w, h in valid_boxes:
roi_y1 = max(0, min_y - blur_margin) # display_scale で表示サイズを決定
roi_x2 = min(src_width, max_x + blur_margin) cx = x + w / 2
roi_y2 = min(src_height, max_y + blur_margin) 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は表示サイズ + blur_margin(計算用余白、境界問題回避のみ)
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_width = roi_x2 - roi_x1
roi_height = roi_y2 - roi_y1 roi_height = roi_y2 - roi_y1
if roi_width <= 0 or roi_height <= 0: if roi_width <= 0 or roi_height <= 0:
process_queue.put((idx, frame))
continue continue
# ブラーはROI全体で計算(余白があるので端の精度が保証される)
roi_src = frame[roi_y1:roi_y2, roi_x1:roi_x2]
# ダウンサンプル→blur→アップサンプル(同等のぼかしを1/4の計算量で実現)
t1 = _time.perf_counter()
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)
t_blur_total += _time.perf_counter() - t1
# 合成マスクはdisplay_scaleサイズの楕円のみ(featheringなし)
roi_mask = np.zeros((roi_height, roi_width), dtype=np.uint8) roi_mask = np.zeros((roi_height, roi_width), dtype=np.uint8)
for x, y, w, h in valid_boxes: center = (int(cx) - roi_x1, int(cy) - roi_y1)
center = (x + w // 2 - roi_x1, y + h // 2 - roi_y1) axes = (max(1, dw // 2), max(1, dh // 2))
axes = (max(1, w // 2), max(1, h // 2))
cv2.ellipse(roi_mask, center, axes, 0, 0, 360, 255, -1) cv2.ellipse(roi_mask, center, axes, 0, 0, 360, 255, -1)
roi_mask = cv2.GaussianBlur(roi_mask, (feather_kernel, feather_kernel), 0) # バイナリマスクなのでcopyToで高速合成(float32変換不要)
roi_src = frame[roi_y1:roi_y2, roi_x1:roi_x2] t2 = _time.perf_counter()
roi_blurred = cv2.GaussianBlur(roi_src, (blur_size, blur_size), 0) result = roi_src.copy()
cv2.copyTo(roi_blurred, roi_mask, result)
frame[roi_y1:roi_y2, roi_x1:roi_x2] = result
t_blend_total += _time.perf_counter() - t2
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
)
frame[roi_y1:roi_y2, roi_x1:roi_x2] = np.clip(roi_composed, 0, 255).astype(np.uint8)
process_queue.put((idx, frame)) process_queue.put((idx, frame))
frame_count += 1
if frame_count % REPORT_INTERVAL == 0:
n = max(frame_count, 1)
fps_proc = frame_count / max(t_wait_total + t_blur_total + t_blend_total, 1e-9)
print(
f"[Perf/Processor] frame={frame_count}"
f" wait={t_wait_total/n*1000:.1f}ms"
f" blur={t_blur_total/n*1000:.1f}ms"
f" blend={t_blend_total/n*1000:.1f}ms"
f" ROI={roi_width}x{roi_height}"
f" throughput≈{fps_proc:.1f}fps"
)
except Exception as e: except Exception as e:
error_holder["error"] = f"Processor error: {e}" error_holder["error"] = f"Processor error: {e}"
process_queue.put(None) process_queue.put(None)
finally:
if frame_count > 0:
n = max(frame_count, 1)
print(
f"[Perf/Processor] FINAL frame={frame_count}"
f" wait_avg={t_wait_total/n*1000:.1f}ms"
f" blur_avg={t_blur_total/n*1000:.1f}ms"
f" blend_avg={t_blend_total/n*1000:.1f}ms"
)
def _writer_worker(): def _writer_worker():
"""Write frames to output.""" """Write frames to output."""
import time as _time
t_wait_total = 0.0
t_write_total = 0.0
frame_count = 0
writer = None writer = None
try: try:
writer = _build_video_writer(req.output_path, req.format, src_fps, src_width, src_height) writer = _build_video_writer(req.output_path, req.format, src_fps, src_width, src_height, r_fps_str)
while True: while True:
if cancel_event and cancel_event.is_set(): if cancel_event and cancel_event.is_set():
break break
t0 = _time.perf_counter()
item = process_queue.get() item = process_queue.get()
t_wait_total += _time.perf_counter() - t0
if item is None: if item is None:
break break
idx, frame = item idx, frame = item
t1 = _time.perf_counter()
writer.write(frame) writer.write(frame)
t_write_total += _time.perf_counter() - t1
frame_count += 1
with progress_lock: with progress_lock:
current_progress[0] = idx + 1 current_progress[0] = idx + 1
@@ -575,6 +1126,13 @@ def process_bake_task(task_id: str, req: BakeRequest):
writer.release() writer.release()
except Exception as e: except Exception as e:
print(f"[FaceMask] Writer release error: {e}") print(f"[FaceMask] Writer release error: {e}")
if frame_count > 0:
n = max(frame_count, 1)
print(
f"[Perf/Writer] FINAL frame={frame_count}"
f" wait_avg={t_wait_total/n*1000:.1f}ms"
f" write_avg={t_write_total/n*1000:.1f}ms"
)
print( print(
f"[FaceMask] Starting blur bake: {req.video_path} + " f"[FaceMask] Starting blur bake: {req.video_path} + "
@@ -768,6 +1326,39 @@ def get_status():
"rocm_version": gpu_info["rocm_version"] "rocm_version": gpu_info["rocm_version"]
} }
@app.post("/video_info")
def get_video_info(req: VideoInfoRequest):
if not os.path.exists(req.video_path):
raise HTTPException(status_code=404, detail=f"Video not found: {req.video_path}")
cap = cv2.VideoCapture(req.video_path)
if not cap.isOpened():
raise HTTPException(status_code=400, detail="Failed to open video")
try:
avg_fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
finally:
cap.release()
# Blender は r_frame_rate でタイムライン配置を計算するため、
# cv2 の avg_frame_rate ではなく r_frame_rate を fps として返す。
# 例: 120fps タイムベース記録の 60fps 動画で r_frame_rate=120 を返すことで
# compute_strip_frame_range の fps_ratio が Blender の解釈と一致する。
r_fps_float, _ = _get_r_frame_rate(req.video_path)
fps = r_fps_float if r_fps_float > 0 else avg_fps
return {
"video_path": req.video_path,
"fps": fps,
"width": width,
"height": height,
"frame_count": frame_count,
}
@app.post("/generate", response_model=Task) @app.post("/generate", response_model=Task)
def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTasks): def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTasks):
task_id = str(uuid.uuid4()) task_id = str(uuid.uuid4())
@@ -789,6 +1380,36 @@ def bake_blur_endpoint(req: BakeRequest, background_tasks: BackgroundTasks):
background_tasks.add_task(process_bake_task, task_id, req) background_tasks.add_task(process_bake_task, task_id, req)
return task 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("/augment_pose", response_model=Task)
def augment_pose_endpoint(req: AugmentPoseRequest, 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(augment_pose_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) @app.get("/tasks/{task_id}", response_model=Task)
def get_task(task_id: str): def get_task(task_id: str):
if task_id not in tasks: if task_id not in tasks: