This commit is contained in:
Keisuke Hirata 2026-02-12 22:03:02 +09:00
parent f2665a49dd
commit eeb8400727
6 changed files with 382 additions and 38 deletions

View File

@ -15,21 +15,58 @@ bl_info = {
def register(): def register():
"""Register all extension components.""" """Register all extension components."""
import bpy
from bpy.props import FloatProperty
from . import operators from . import operators
from . import panels from . import panels
# Register scene properties for face detection parameters
bpy.types.Scene.facemask_conf_threshold = FloatProperty(
name="Confidence",
description="YOLO confidence threshold (higher = fewer false positives)",
default=0.5,
min=0.1,
max=1.0,
step=0.01,
)
bpy.types.Scene.facemask_iou_threshold = FloatProperty(
name="IOU Threshold",
description="Non-maximum suppression IOU threshold",
default=0.45,
min=0.1,
max=1.0,
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,
)
operators.register() operators.register()
panels.register() panels.register()
def unregister(): def unregister():
"""Unregister all extension components.""" """Unregister all extension components."""
import bpy
from . import operators from . import operators
from . import panels from . import panels
panels.unregister() panels.unregister()
operators.unregister() operators.unregister()
# Unregister scene properties
del bpy.types.Scene.facemask_conf_threshold
del bpy.types.Scene.facemask_iou_threshold
del bpy.types.Scene.facemask_mask_scale
if __name__ == "__main__": if __name__ == "__main__":
register() register()

111
core/utils.py Normal file
View File

@ -0,0 +1,111 @@
"""
Utility functions for Face Mask extension.
Provides helper functions for server status, cache info, etc.
"""
import os
import urllib.request
import urllib.error
import json
import tempfile
from typing import Dict, Tuple, Optional
def get_server_status() -> Dict:
"""
Get server status and GPU information.
Returns:
dict: {
'running': bool,
'gpu_available': bool,
'gpu_device': str or None,
'gpu_count': int,
'rocm_version': str or None,
}
"""
result = {
'running': False,
'gpu_available': False,
'gpu_device': None,
'gpu_count': 0,
'rocm_version': None,
}
try:
with urllib.request.urlopen("http://127.0.0.1:8181/status", timeout=1) as response:
data = json.loads(response.read().decode('utf-8'))
result['running'] = data.get('status') == 'running'
result['gpu_available'] = data.get('gpu_available', False)
result['gpu_device'] = data.get('gpu_device')
result['gpu_count'] = data.get('gpu_count', 0)
result['rocm_version'] = data.get('rocm_version')
except (urllib.error.URLError, ConnectionRefusedError, TimeoutError):
result['running'] = False
return result
def get_cache_info(strip_name: Optional[str] = None) -> Tuple[str, int, int]:
"""
Get cache directory information.
Args:
strip_name: If provided, get info for specific strip. Otherwise, get info for all cache.
Returns:
Tuple of (cache_path, total_size_bytes, file_count)
"""
import bpy
blend_file = bpy.data.filepath
if strip_name:
# Get cache for specific strip
if blend_file:
project_dir = os.path.dirname(blend_file)
cache_path = os.path.join(project_dir, ".mask_cache", strip_name)
else:
cache_path = os.path.join(tempfile.gettempdir(), "blender_mask_cache", strip_name)
else:
# Get cache root
if blend_file:
project_dir = os.path.dirname(blend_file)
cache_path = os.path.join(project_dir, ".mask_cache")
else:
cache_path = os.path.join(tempfile.gettempdir(), "blender_mask_cache")
# Calculate size and count
total_size = 0
file_count = 0
if os.path.exists(cache_path):
for root, dirs, files in os.walk(cache_path):
for file in files:
if file.endswith('.png'): # Only count mask images
file_path = os.path.join(root, file)
try:
total_size += os.path.getsize(file_path)
file_count += 1
except OSError:
pass
return cache_path, total_size, file_count
def format_size(size_bytes: int) -> str:
"""
Format bytes to human-readable size.
Args:
size_bytes: Size in bytes
Returns:
Formatted string (e.g., "1.5 MB")
"""
for unit in ['B', 'KB', 'MB', 'GB']:
if size_bytes < 1024.0:
return f"{size_bytes:.1f} {unit}"
size_bytes /= 1024.0
return f"{size_bytes:.1f} TB"

View File

@ -2,13 +2,16 @@
from . import generate_mask from . import generate_mask
from . import apply_blur from . import apply_blur
from . import clear_cache
def register(): def register():
generate_mask.register() generate_mask.register()
apply_blur.register() apply_blur.register()
clear_cache.register()
def unregister(): def unregister():
clear_cache.unregister()
apply_blur.unregister() apply_blur.unregister()
generate_mask.unregister() generate_mask.unregister()

126
operators/clear_cache.py Normal file
View File

@ -0,0 +1,126 @@
"""
Clear Cache Operator.
Provides operators to clear mask cache directories.
"""
import os
import shutil
import bpy
from bpy.types import Operator
from bpy.props import BoolProperty
class SEQUENCER_OT_clear_mask_cache(Operator):
"""Clear mask cache directories."""
bl_idname = "sequencer.clear_mask_cache"
bl_label = "Clear Mask Cache"
bl_description = "Delete cached mask images"
bl_options = {'REGISTER', 'UNDO'}
all_strips: BoolProperty(
name="All Strips",
description="Clear cache for all strips (otherwise only current strip)",
default=False,
)
def execute(self, context):
import tempfile
blend_file = bpy.data.filepath
total_size = 0
cleared_count = 0
if self.all_strips:
# Clear all cache directories
if blend_file:
# Project cache
project_dir = os.path.dirname(blend_file)
cache_root = os.path.join(project_dir, ".mask_cache")
else:
# Temp cache
cache_root = os.path.join(tempfile.gettempdir(), "blender_mask_cache")
if os.path.exists(cache_root):
# Calculate size before deletion
for root, dirs, files in os.walk(cache_root):
for file in files:
file_path = os.path.join(root, file)
try:
total_size += os.path.getsize(file_path)
except OSError:
pass
# Delete cache directory
try:
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)})")
except Exception as e:
self.report({'ERROR'}, f"Failed to clear cache: {e}")
return {'CANCELLED'}
else:
self.report({'INFO'}, "No cache to clear")
return {'FINISHED'}
else:
# Clear cache for active strip only
seq_editor = context.scene.sequence_editor
if not seq_editor or not seq_editor.active_strip:
self.report({'WARNING'}, "No strip selected")
return {'CANCELLED'}
strip = seq_editor.active_strip
if blend_file:
project_dir = os.path.dirname(blend_file)
cache_dir = os.path.join(project_dir, ".mask_cache", strip.name)
else:
cache_dir = os.path.join(tempfile.gettempdir(), "blender_mask_cache", strip.name)
if os.path.exists(cache_dir):
# Calculate size
for root, dirs, files in os.walk(cache_dir):
for file in files:
file_path = os.path.join(root, file)
try:
total_size += os.path.getsize(file_path)
except OSError:
pass
# Delete
try:
shutil.rmtree(cache_dir)
self.report({'INFO'}, f"Cleared cache for {strip.name} ({self._format_size(total_size)})")
except Exception as e:
self.report({'ERROR'}, f"Failed to clear cache: {e}")
return {'CANCELLED'}
else:
self.report({'INFO'}, f"No cache for {strip.name}")
return {'FINISHED'}
return {'FINISHED'}
def _format_size(self, size_bytes):
"""Format bytes to human-readable size."""
for unit in ['B', 'KB', 'MB', 'GB']:
if size_bytes < 1024.0:
return f"{size_bytes:.1f} {unit}"
size_bytes /= 1024.0
return f"{size_bytes:.1f} TB"
# Registration
classes = [
SEQUENCER_OT_clear_mask_cache,
]
def register():
for cls in classes:
bpy.utils.register_class(cls)
def unregister():
for cls in reversed(classes):
bpy.utils.unregister_class(cls)

View File

@ -7,7 +7,7 @@ from video strips in the Video Sequence Editor.
import os import os
import bpy import bpy
from bpy.props import FloatProperty, IntProperty from bpy.props import IntProperty
from bpy.types import Operator from bpy.types import Operator
from ..core.async_generator import get_generator from ..core.async_generator import get_generator
@ -15,37 +15,12 @@ from ..core.async_generator import get_generator
class SEQUENCER_OT_generate_face_mask(Operator): class SEQUENCER_OT_generate_face_mask(Operator):
"""Generate face mask image sequence from video strip.""" """Generate face mask image sequence from video strip."""
bl_idname = "sequencer.generate_face_mask" bl_idname = "sequencer.generate_face_mask"
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'}
# YOLO Detection parameters
conf_threshold: FloatProperty(
name="Confidence",
description="YOLO confidence threshold (higher = fewer false positives)",
default=0.25,
min=0.1,
max=1.0,
)
iou_threshold: FloatProperty(
name="IOU Threshold",
description="Non-maximum suppression IOU threshold",
default=0.45,
min=0.1,
max=1.0,
)
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,
)
@classmethod @classmethod
def poll(cls, context): def poll(cls, context):
"""Check if operator can run.""" """Check if operator can run."""
@ -126,6 +101,11 @@ class SEQUENCER_OT_generate_face_mask(Operator):
wm.mask_progress = 0 wm.mask_progress = 0
wm.mask_total = end_frame - start_frame + 1 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 # Start generation
generator.start( generator.start(
video_path=video_path, video_path=video_path,
@ -133,9 +113,9 @@ class SEQUENCER_OT_generate_face_mask(Operator):
start_frame=0, # Frame indices in video start_frame=0, # Frame indices in video
end_frame=end_frame - start_frame, end_frame=end_frame - start_frame,
fps=fps, fps=fps,
conf_threshold=self.conf_threshold, conf_threshold=conf_threshold,
iou_threshold=self.iou_threshold, iou_threshold=iou_threshold,
mask_scale=self.mask_scale, mask_scale=mask_scale,
on_complete=on_complete, on_complete=on_complete,
on_progress=on_progress, on_progress=on_progress,
) )

View File

@ -9,6 +9,7 @@ import bpy
from bpy.types import Panel from bpy.types import Panel
from ..core.async_generator import get_generator from ..core.async_generator import get_generator
from ..core.utils import get_server_status, get_cache_info, format_size
class SEQUENCER_PT_face_mask(Panel): class SEQUENCER_PT_face_mask(Panel):
@ -22,21 +23,29 @@ class SEQUENCER_PT_face_mask(Panel):
def draw(self, context): def draw(self, context):
layout = self.layout layout = self.layout
scene = context.scene
wm = context.window_manager wm = context.window_manager
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'
generator = get_generator() generator = get_generator()
# Always show parameters and status
self._draw_parameters(layout, scene)
self._draw_server_status(layout)
self._draw_cache_info(layout, seq_editor)
layout.separator()
# Show progress if generating # Show progress if generating
if generator.is_running: if generator.is_running:
self._draw_progress(layout, wm, generator) self._draw_progress(layout, wm, generator)
return return
# Show controls if strip selected # Show controls if strip selected
if seq_editor and seq_editor.active_strip: if seq_editor and seq_editor.active_strip:
strip = seq_editor.active_strip strip = seq_editor.active_strip
if strip.type in {'MOVIE', 'IMAGE'}: if strip.type in {'MOVIE', 'IMAGE'}:
self._draw_generation_controls(layout, context, strip) self._draw_generation_controls(layout, context, strip)
self._draw_blur_controls(layout, context, strip) self._draw_blur_controls(layout, context, strip)
@ -45,18 +54,96 @@ class SEQUENCER_PT_face_mask(Panel):
else: else:
layout.label(text="No strip selected") layout.label(text="No strip selected")
def _draw_parameters(self, layout, scene):
"""Draw detection parameters."""
box = layout.box()
box.label(text="Parameters", icon='PREFERENCES')
col = box.column(align=True)
col.prop(scene, "facemask_conf_threshold")
col.prop(scene, "facemask_iou_threshold")
col.prop(scene, "facemask_mask_scale")
def _draw_server_status(self, layout):
"""Draw server status and GPU info."""
box = layout.box()
box.label(text="Server Status", icon='SYSTEM')
status = get_server_status()
# Server status
row = box.row()
if status['running']:
row.label(text="Server:", icon='CHECKMARK')
row.label(text="Running")
else:
row.label(text="Server:", icon='ERROR')
row.label(text="Stopped")
# GPU status
if status['running']:
row = box.row()
if status['gpu_available']:
row.label(text="GPU:", icon='CHECKMARK')
gpu_name = status['gpu_device'] or "Available"
# Truncate long GPU names
if len(gpu_name) > 25:
gpu_name = gpu_name[:22] + "..."
row.label(text=gpu_name)
else:
row.label(text="GPU:", icon='ERROR')
row.label(text="Not Available")
def _draw_cache_info(self, layout, seq_editor):
"""Draw cache information and clear button."""
box = layout.box()
box.label(text="Cache", icon='FILE_CACHE')
# Get cache info
if seq_editor and seq_editor.active_strip:
strip_name = seq_editor.active_strip.name
cache_path, total_size, file_count = get_cache_info(strip_name)
else:
cache_path, total_size, file_count = get_cache_info()
# Cache info
row = box.row()
row.label(text="Size:")
row.label(text=format_size(total_size))
row = box.row()
row.label(text="Files:")
row.label(text=str(file_count))
# Clear cache buttons
row = box.row(align=True)
if seq_editor and seq_editor.active_strip:
op = row.operator(
"sequencer.clear_mask_cache",
text="Clear Strip Cache",
icon='TRASH',
)
op.all_strips = False
op = row.operator(
"sequencer.clear_mask_cache",
text="Clear All",
icon='TRASH',
)
op.all_strips = True
def _draw_progress(self, layout, wm, generator): def _draw_progress(self, layout, wm, generator):
"""Draw progress bar during generation.""" """Draw progress bar during generation."""
box = layout.box() box = layout.box()
box.label(text="Generating Masks...", icon='RENDER_ANIMATION') box.label(text="Generating Masks...", icon='RENDER_ANIMATION')
# Progress bar # Progress bar
progress = wm.mask_progress / max(wm.mask_total, 1) progress = wm.mask_progress / max(wm.mask_total, 1)
box.progress( box.progress(
factor=progress, factor=progress,
text=f"Frame {wm.mask_progress} / {wm.mask_total}", text=f"Frame {wm.mask_progress} / {wm.mask_total}",
) )
# Cancel button # Cancel button
box.operator( box.operator(
"sequencer.cancel_mask_generation", "sequencer.cancel_mask_generation",