Compare commits
13
Commits
0d63b2ef6d
..
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
eb028ed278 | ||
|
|
de99aef9ad | ||
|
|
dc41327cea | ||
|
|
be65abc6b0 | ||
|
|
32e4fbceb2 | ||
|
|
0fdff5423e | ||
|
|
d67265aa39 | ||
|
|
a3de61d5ce | ||
|
|
da9de60697 | ||
|
|
9ce6ec99d3 | ||
|
|
08f20fa6fe | ||
|
|
920695696b | ||
|
|
914667edbf |
+10
-10
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
+179
-6
@@ -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.
|
||||||
@@ -133,7 +307,7 @@ class AsyncMaskGenerator:
|
|||||||
client = get_client()
|
client = get_client()
|
||||||
|
|
||||||
# Start task on server
|
# Start task on server
|
||||||
print(f"[FaceMask] Requesting generation on server...")
|
print("[FaceMask] Requesting generation on server...")
|
||||||
task_id = client.generate_mask(
|
task_id = client.generate_mask(
|
||||||
video_path=video_path,
|
video_path=video_path,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
@@ -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}")
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -5,13 +5,10 @@ Creates and manages compositing node trees that apply blur
|
|||||||
only to masked regions of a video strip.
|
only to masked regions of a video strip.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Optional, Tuple
|
|
||||||
|
|
||||||
|
|
||||||
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.
|
||||||
|
|
||||||
@@ -110,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.
|
||||||
|
|
||||||
@@ -125,8 +122,6 @@ def setup_strip_compositor_modifier(
|
|||||||
Returns:
|
Returns:
|
||||||
The created modifier
|
The created modifier
|
||||||
"""
|
"""
|
||||||
import bpy
|
|
||||||
|
|
||||||
# Add compositor modifier
|
# Add compositor modifier
|
||||||
modifier = strip.modifiers.new(
|
modifier = strip.modifiers.new(
|
||||||
name="FaceMaskBlur",
|
name="FaceMaskBlur",
|
||||||
@@ -153,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.
|
||||||
|
|
||||||
|
|||||||
+120
-3
@@ -14,7 +14,7 @@ import threading
|
|||||||
import time
|
import time
|
||||||
import urllib.error
|
import urllib.error
|
||||||
import urllib.request
|
import urllib.request
|
||||||
from typing import Any, Dict, Optional, Tuple
|
from typing import Any, Dict, Optional
|
||||||
|
|
||||||
|
|
||||||
class InferenceClient:
|
class InferenceClient:
|
||||||
@@ -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
-2
@@ -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.
|
||||||
@@ -93,8 +102,6 @@ def get_cache_info(strip_name: Optional[str] = None) -> Tuple[str, int, int]:
|
|||||||
Returns:
|
Returns:
|
||||||
Tuple of (cache_path, total_size_bytes, file_count)
|
Tuple of (cache_path, total_size_bytes, file_count)
|
||||||
"""
|
"""
|
||||||
import bpy
|
|
||||||
|
|
||||||
if strip_name:
|
if strip_name:
|
||||||
cache_path = get_cache_dir_for_strip(strip_name)
|
cache_path = get_cache_dir_for_strip(strip_name)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -73,9 +73,11 @@
|
|||||||
pip install --quiet -r "$PWD/requirements.txt"
|
pip install --quiet -r "$PWD/requirements.txt"
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# opencv-pythonが入っていた場合はheadlessに統一
|
# OpenCVは壊れやすいので、import失敗時のみheadlessを強制再導入
|
||||||
pip uninstall -y opencv-python opencv 2>/dev/null || true
|
if ! python -c "import cv2" >/dev/null 2>&1; then
|
||||||
pip install --quiet --upgrade opencv-python-headless
|
echo "[Setup] Repairing OpenCV (opencv-python-headless)..."
|
||||||
|
pip install --quiet --force-reinstall --no-cache-dir opencv-python-headless
|
||||||
|
fi
|
||||||
|
|
||||||
# Pythonパスにカレントディレクトリを追加
|
# Pythonパスにカレントディレクトリを追加
|
||||||
export PYTHONPATH="$PWD:$PYTHONPATH"
|
export PYTHONPATH="$PWD:$PYTHONPATH"
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
+208
-92
@@ -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,16 +51,162 @@ 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":
|
||||||
_reload_movie_strip(strip)
|
strip.directory = path
|
||||||
|
else:
|
||||||
|
strip.filepath = path
|
||||||
|
_reload_movie_strip(strip)
|
||||||
|
|
||||||
|
|
||||||
|
def _start_bake_impl(operator, context, force: bool = False, strip=None, on_complete_extra=None):
|
||||||
|
"""Bakeの共通実装。force=True でキャッシュを無視して再Bakeする。
|
||||||
|
|
||||||
|
strip: 処理対象のstrip。None の場合は active_strip を使用。
|
||||||
|
on_complete_extra: 非同期Bake完了時に追加で呼ばれるコールバック (status, data)。
|
||||||
|
キャッシュヒット即時完了の場合は呼ばれない。
|
||||||
|
MOVIE / IMAGE 両対応。
|
||||||
|
"""
|
||||||
|
seq_editor = context.scene.sequence_editor
|
||||||
|
scene = context.scene
|
||||||
|
video_strip = strip if strip is not None else seq_editor.active_strip
|
||||||
|
is_image = video_strip.type == "IMAGE"
|
||||||
|
|
||||||
|
detections_path = get_detections_path_for_strip(video_strip.name)
|
||||||
|
if not os.path.exists(detections_path):
|
||||||
|
operator.report({"ERROR"}, f"Detection cache not found: {detections_path}")
|
||||||
|
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
|
||||||
|
output_path = _output_path(video_strip, detections_path, bake_format)
|
||||||
|
original_source = video_path
|
||||||
|
|
||||||
|
if not force:
|
||||||
|
# パラメータが一致するキャッシュがあればswapのみ
|
||||||
|
cached_baked_path = video_strip.get(KEY_BAKED)
|
||||||
|
cached_blur_size = video_strip.get(KEY_BLUR_SIZE)
|
||||||
|
cached_display_scale = video_strip.get(KEY_DISPLAY_SCALE)
|
||||||
|
try:
|
||||||
|
cached_blur_size_int = int(cached_blur_size)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
cached_blur_size_int = None
|
||||||
|
try:
|
||||||
|
cached_display_scale_f = float(cached_display_scale)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
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_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":
|
||||||
|
video_strip[KEY_MODE] = "baked"
|
||||||
|
_set_strip_source(video_strip, cached_baked_path)
|
||||||
|
operator.report({"INFO"}, "Using cached baked blur")
|
||||||
|
return {"FINISHED"}
|
||||||
|
|
||||||
|
bake_generator = get_bake_generator()
|
||||||
|
wm = context.window_manager
|
||||||
|
|
||||||
|
def on_complete(status, data):
|
||||||
|
strip = context.scene.sequence_editor.strips.get(video_strip.name)
|
||||||
|
if not strip:
|
||||||
|
print(f"[FaceMask] Bake complete but strip no longer exists: {video_strip.name}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if status == "done":
|
||||||
|
result = data or (output_dir if is_image else output_path)
|
||||||
|
current_mode = strip.get(KEY_MODE, "original")
|
||||||
|
if not strip.get(KEY_ORIGINAL) or current_mode != "baked":
|
||||||
|
strip[KEY_ORIGINAL] = original_source
|
||||||
|
strip[KEY_BAKED] = result
|
||||||
|
strip[KEY_MODE] = "baked"
|
||||||
|
strip[KEY_BLUR_SIZE] = blur_size
|
||||||
|
strip[KEY_DISPLAY_SCALE] = display_scale
|
||||||
|
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":
|
||||||
|
print(f"[FaceMask] Bake failed: {data}")
|
||||||
|
elif status == "cancelled":
|
||||||
|
print("[FaceMask] Bake cancelled")
|
||||||
|
|
||||||
|
for area in context.screen.areas:
|
||||||
|
if area.type == "SEQUENCE_EDITOR":
|
||||||
|
area.tag_redraw()
|
||||||
|
|
||||||
|
if on_complete_extra:
|
||||||
|
on_complete_extra(status, data)
|
||||||
|
|
||||||
|
def on_progress(current, total):
|
||||||
|
wm.bake_progress = current
|
||||||
|
wm.bake_total = max(total, 1)
|
||||||
|
for area in context.screen.areas:
|
||||||
|
if area.type == "SEQUENCE_EDITOR":
|
||||||
|
area.tag_redraw()
|
||||||
|
|
||||||
|
wm.bake_progress = 0
|
||||||
|
wm.bake_total = 1
|
||||||
|
|
||||||
|
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(
|
||||||
|
video_path=video_path,
|
||||||
|
detections_path=detections_path,
|
||||||
|
output_path=output_path,
|
||||||
|
blur_size=blur_size,
|
||||||
|
display_scale=display_scale,
|
||||||
|
fmt=bake_format.lower(),
|
||||||
|
on_complete=on_complete,
|
||||||
|
on_progress=on_progress,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
operator.report({"ERROR"}, f"Failed to start bake: {e}")
|
||||||
|
return {"CANCELLED"}
|
||||||
|
|
||||||
|
operator.report({"INFO"}, "Started blur bake in background")
|
||||||
|
return {"FINISHED"}
|
||||||
|
|
||||||
|
|
||||||
class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
|
class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
|
||||||
"""Bake masked blur and replace active strip source with baked video."""
|
"""Bake masked blur (reuse cache if parameters match)."""
|
||||||
|
|
||||||
bl_idname = "sequencer.bake_and_swap_blur_source"
|
bl_idname = "sequencer.bake_and_swap_blur_source"
|
||||||
bl_label = "Bake & Swap Source"
|
bl_label = "Bake"
|
||||||
bl_description = "Bake masked blur to video and swap active strip source"
|
bl_description = "Bake masked blur to video and swap active strip source"
|
||||||
bl_options = {"REGISTER", "UNDO"}
|
bl_options = {"REGISTER", "UNDO"}
|
||||||
|
|
||||||
@@ -61,107 +214,66 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
|
|||||||
def poll(cls, context):
|
def poll(cls, context):
|
||||||
if not context.scene.sequence_editor:
|
if not context.scene.sequence_editor:
|
||||||
return False
|
return False
|
||||||
# Prevent overlapping heavy tasks
|
|
||||||
if get_mask_generator().is_running:
|
if get_mask_generator().is_running:
|
||||||
return False
|
return False
|
||||||
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
|
||||||
return bool(strip and strip.type == "MOVIE")
|
return bool(strip and strip.type in {"MOVIE", "IMAGE"})
|
||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
seq_editor = context.scene.sequence_editor
|
return _start_bake_impl(self, context, force=False)
|
||||||
scene = context.scene
|
|
||||||
video_strip = seq_editor.active_strip
|
|
||||||
|
|
||||||
video_path = bpy.path.abspath(video_strip.filepath)
|
|
||||||
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):
|
|
||||||
self.report({"ERROR"}, f"Detection cache not found: {detections_path}")
|
|
||||||
return {"CANCELLED"}
|
|
||||||
|
|
||||||
bake_format = scene.facemask_bake_format
|
class SEQUENCER_OT_force_rebake_blur(Operator):
|
||||||
output_path = _output_path(video_strip, detections_path, bake_format)
|
"""Force re-bake, ignoring any existing cached result."""
|
||||||
blur_size = int(scene.facemask_bake_blur_size)
|
|
||||||
|
|
||||||
# Reuse baked cache when parameters match and file still exists.
|
bl_idname = "sequencer.force_rebake_blur"
|
||||||
cached_baked_path = video_strip.get(KEY_BAKED)
|
bl_label = "Re-bake"
|
||||||
cached_format = video_strip.get(KEY_FORMAT)
|
bl_description = "Discard cached bake and re-bake from scratch"
|
||||||
cached_blur_size = video_strip.get(KEY_BLUR_SIZE)
|
bl_options = {"REGISTER", "UNDO"}
|
||||||
try:
|
|
||||||
cached_blur_size_int = int(cached_blur_size)
|
|
||||||
except (TypeError, ValueError):
|
|
||||||
cached_blur_size_int = None
|
|
||||||
if (
|
|
||||||
cached_baked_path
|
|
||||||
and os.path.exists(cached_baked_path)
|
|
||||||
and cached_format == bake_format
|
|
||||||
and cached_blur_size_int == blur_size
|
|
||||||
):
|
|
||||||
if video_strip.get(KEY_MODE) != "baked":
|
|
||||||
video_strip[KEY_MODE] = "baked"
|
|
||||||
_set_strip_source(video_strip, cached_baked_path)
|
|
||||||
self.report({"INFO"}, "Using cached baked blur")
|
|
||||||
return {"FINISHED"}
|
|
||||||
|
|
||||||
bake_generator = get_bake_generator()
|
@classmethod
|
||||||
wm = context.window_manager
|
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 on_complete(status, data):
|
def execute(self, context):
|
||||||
strip = context.scene.sequence_editor.strips.get(video_strip.name)
|
return _start_bake_impl(self, context, force=True)
|
||||||
if not strip:
|
|
||||||
print(f"[FaceMask] Bake complete but strip no longer exists: {video_strip.name}")
|
|
||||||
return
|
|
||||||
|
|
||||||
if status == "done":
|
|
||||||
result_path = data or output_path
|
|
||||||
original_path = strip.get(KEY_ORIGINAL)
|
|
||||||
current_mode = strip.get(KEY_MODE, "original")
|
|
||||||
if not original_path or current_mode != "baked":
|
|
||||||
strip[KEY_ORIGINAL] = video_path
|
|
||||||
strip[KEY_BAKED] = result_path
|
|
||||||
strip[KEY_MODE] = "baked"
|
|
||||||
strip[KEY_FORMAT] = bake_format
|
|
||||||
strip[KEY_BLUR_SIZE] = blur_size
|
|
||||||
_set_strip_source(strip, result_path)
|
|
||||||
print(f"[FaceMask] Bake completed and source swapped: {result_path}")
|
|
||||||
elif status == "error":
|
|
||||||
print(f"[FaceMask] Bake failed: {data}")
|
|
||||||
elif status == "cancelled":
|
|
||||||
print("[FaceMask] Bake cancelled")
|
|
||||||
|
|
||||||
for area in context.screen.areas:
|
class SEQUENCER_OT_swap_to_baked_blur(Operator):
|
||||||
if area.type == "SEQUENCE_EDITOR":
|
"""Swap active strip source to already-baked video (no re-bake)."""
|
||||||
area.tag_redraw()
|
|
||||||
|
|
||||||
def on_progress(current, total):
|
bl_idname = "sequencer.swap_to_baked_blur"
|
||||||
wm.bake_progress = current
|
bl_label = "Swap to Baked"
|
||||||
wm.bake_total = max(total, 1)
|
bl_description = "Switch active strip source to the baked video without re-baking"
|
||||||
for area in context.screen.areas:
|
bl_options = {"REGISTER", "UNDO"}
|
||||||
if area.type == "SEQUENCE_EDITOR":
|
|
||||||
area.tag_redraw()
|
|
||||||
|
|
||||||
wm.bake_progress = 0
|
@classmethod
|
||||||
wm.bake_total = 1
|
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))
|
||||||
|
|
||||||
try:
|
def execute(self, context):
|
||||||
bake_generator.start(
|
strip = context.scene.sequence_editor.active_strip
|
||||||
video_path=video_path,
|
baked_path = strip.get(KEY_BAKED)
|
||||||
detections_path=detections_path,
|
_set_strip_source(strip, baked_path)
|
||||||
output_path=output_path,
|
strip[KEY_MODE] = "baked"
|
||||||
blur_size=blur_size,
|
self.report({"INFO"}, "Swapped to baked source")
|
||||||
fmt=bake_format.lower(),
|
|
||||||
on_complete=on_complete,
|
|
||||||
on_progress=on_progress,
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
self.report({"ERROR"}, f"Failed to start bake: {e}")
|
|
||||||
return {"CANCELLED"}
|
|
||||||
|
|
||||||
self.report({"INFO"}, "Started blur bake in background")
|
|
||||||
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,
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -29,7 +29,6 @@ class SEQUENCER_OT_clear_mask_cache(Operator):
|
|||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
total_size = 0
|
total_size = 0
|
||||||
cleared_count = 0
|
|
||||||
|
|
||||||
if self.all_strips:
|
if self.all_strips:
|
||||||
# Clear all cache directories
|
# Clear all cache directories
|
||||||
@@ -48,7 +47,6 @@ class SEQUENCER_OT_clear_mask_cache(Operator):
|
|||||||
# Delete cache directory
|
# Delete cache directory
|
||||||
try:
|
try:
|
||||||
shutil.rmtree(cache_root)
|
shutil.rmtree(cache_root)
|
||||||
cleared_count = len(os.listdir(cache_root)) if os.path.exists(cache_root) else 0
|
|
||||||
self.report({'INFO'}, f"Cleared all cache ({self._format_size(total_size)})")
|
self.report({'INFO'}, f"Cleared all cache ({self._format_size(total_size)})")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.report({'ERROR'}, f"Failed to clear cache: {e}")
|
self.report({'ERROR'}, f"Failed to clear cache: {e}")
|
||||||
|
|||||||
+201
-111
@@ -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):
|
||||||
@@ -21,57 +109,43 @@ class SEQUENCER_OT_generate_face_mask(Operator):
|
|||||||
bl_label = "Generate Face Mask"
|
bl_label = "Generate Face Mask"
|
||||||
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":
|
||||||
@@ -89,105 +163,122 @@ class SEQUENCER_OT_generate_face_mask(Operator):
|
|||||||
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()
|
||||||
|
|
||||||
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
|
|
||||||
wm = context.window_manager
|
|
||||||
wm.mask_progress = 0
|
|
||||||
wm.mask_total = end_frame - start_frame + 1
|
|
||||||
|
|
||||||
# Get parameters from scene properties
|
|
||||||
conf_threshold = scene.facemask_conf_threshold
|
|
||||||
iou_threshold = scene.facemask_iou_threshold
|
|
||||||
mask_scale = scene.facemask_mask_scale
|
|
||||||
|
|
||||||
# Start generation
|
try:
|
||||||
generator.start(
|
start_mask_gen_for_strip(context, strip, on_complete, on_progress)
|
||||||
video_path=video_path,
|
except RuntimeError as e:
|
||||||
output_dir=output_dir,
|
self.report({'WARNING'}, str(e))
|
||||||
start_frame=0, # Frame indices in video
|
return {'CANCELLED'}
|
||||||
end_frame=end_frame - start_frame,
|
except Exception as e:
|
||||||
fps=fps,
|
self.report({'ERROR'}, f"Failed to start mask generation: {e}")
|
||||||
conf_threshold=conf_threshold,
|
return {'CANCELLED'}
|
||||||
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."""
|
||||||
|
|
||||||
bl_idname = "sequencer.cancel_mask_generation"
|
bl_idname = "sequencer.cancel_mask_generation"
|
||||||
bl_label = "Cancel Mask Generation"
|
bl_label = "Cancel Mask Generation"
|
||||||
bl_description = "Cancel the current mask generation process"
|
bl_description = "Cancel the current mask generation process"
|
||||||
bl_options = {'REGISTER'}
|
bl_options = {'REGISTER'}
|
||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
generator = get_generator()
|
generator = get_generator()
|
||||||
|
|
||||||
if generator.is_running:
|
if generator.is_running:
|
||||||
generator.cancel()
|
generator.cancel()
|
||||||
self.report({'INFO'}, "Mask generation cancelled")
|
self.report({'INFO'}, "Mask generation cancelled")
|
||||||
else:
|
else:
|
||||||
self.report({'WARNING'}, "No mask generation in progress")
|
self.report({'WARNING'}, "No mask generation in progress")
|
||||||
|
|
||||||
|
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'}
|
return {'FINISHED'}
|
||||||
|
|
||||||
|
|
||||||
@@ -195,22 +286,21 @@ class SEQUENCER_OT_cancel_mask_generation(Operator):
|
|||||||
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,
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
def register():
|
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
|
||||||
|
|
||||||
for cls in reversed(classes):
|
for cls in reversed(classes):
|
||||||
bpy.utils.unregister_class(cls)
|
bpy.utils.unregister_class(cls)
|
||||||
|
|||||||
+138
-36
@@ -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."""
|
||||||
@@ -182,7 +189,77 @@ class SEQUENCER_PT_face_mask(Panel):
|
|||||||
text="Cancel",
|
text="Cancel",
|
||||||
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,32 +269,40 @@ 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
|
||||||
op = box.operator(
|
if not has_mask:
|
||||||
"sequencer.generate_face_mask",
|
box.operator(
|
||||||
text="Generate Detection Cache" if not has_mask else "Regenerate Cache",
|
"sequencer.generate_face_mask",
|
||||||
icon='FACE_MAPS',
|
text="Generate Detection Cache",
|
||||||
)
|
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')
|
||||||
return
|
return
|
||||||
@@ -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_format")
|
col.prop(context.scene, "facemask_bake_display_scale")
|
||||||
|
if strip.type == "MOVIE":
|
||||||
|
col.prop(context.scene, "facemask_bake_format")
|
||||||
|
|
||||||
# Source status
|
box.separator()
|
||||||
|
|
||||||
|
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")
|
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
|
if not has_baked:
|
||||||
box.operator(
|
# 初回: Bakeのみ
|
||||||
"sequencer.bake_and_swap_blur_source",
|
box.operator(
|
||||||
text="Bake & Swap Source",
|
"sequencer.bake_and_swap_blur_source",
|
||||||
icon='RENDER_STILL',
|
text="Bake",
|
||||||
)
|
icon='RENDER_STILL',
|
||||||
box.operator(
|
)
|
||||||
"sequencer.restore_original_source",
|
else:
|
||||||
text="Restore Original Source",
|
# Bake済み: ソース切り替え + Re-bake
|
||||||
icon='LOOP_BACK',
|
row = box.row(align=True)
|
||||||
)
|
if source_mode == "baked":
|
||||||
|
row.operator(
|
||||||
|
"sequencer.restore_original_source",
|
||||||
|
text="Restore Original",
|
||||||
|
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
|
||||||
|
|||||||
+258
-99
@@ -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):
|
||||||
|
_download_model(model_path)
|
||||||
|
|
||||||
if not os.path.exists(model_path):
|
|
||||||
raise FileNotFoundError(f"Model not found: {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
|
||||||
|
|||||||
+919
-119
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user