diff --git a/__init__.py b/__init__.py index 4b2d33b..583987a 100644 --- a/__init__.py +++ b/__init__.py @@ -15,21 +15,58 @@ bl_info = { def register(): """Register all extension components.""" + import bpy + from bpy.props import FloatProperty + from . import operators 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() panels.register() def unregister(): """Unregister all extension components.""" + import bpy from . import operators from . import panels - + panels.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__": register() diff --git a/core/utils.py b/core/utils.py new file mode 100644 index 0000000..701e5cd --- /dev/null +++ b/core/utils.py @@ -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" diff --git a/operators/__init__.py b/operators/__init__.py index bc3c68d..4b2005f 100644 --- a/operators/__init__.py +++ b/operators/__init__.py @@ -2,13 +2,16 @@ from . import generate_mask from . import apply_blur +from . import clear_cache def register(): generate_mask.register() apply_blur.register() + clear_cache.register() def unregister(): + clear_cache.unregister() apply_blur.unregister() generate_mask.unregister() diff --git a/operators/clear_cache.py b/operators/clear_cache.py new file mode 100644 index 0000000..ed00b83 --- /dev/null +++ b/operators/clear_cache.py @@ -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) diff --git a/operators/generate_mask.py b/operators/generate_mask.py index 8a926fb..ec868fc 100644 --- a/operators/generate_mask.py +++ b/operators/generate_mask.py @@ -7,7 +7,7 @@ from video strips in the Video Sequence Editor. import os import bpy -from bpy.props import FloatProperty, IntProperty +from bpy.props import IntProperty from bpy.types import Operator 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): """Generate face mask image sequence from video strip.""" - + bl_idname = "sequencer.generate_face_mask" bl_label = "Generate Face Mask" bl_description = "Detect faces and generate mask image sequence" 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 def poll(cls, context): """Check if operator can run.""" @@ -126,6 +101,11 @@ class SEQUENCER_OT_generate_face_mask(Operator): 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 generator.start( video_path=video_path, @@ -133,9 +113,9 @@ class SEQUENCER_OT_generate_face_mask(Operator): start_frame=0, # Frame indices in video end_frame=end_frame - start_frame, fps=fps, - conf_threshold=self.conf_threshold, - iou_threshold=self.iou_threshold, - mask_scale=self.mask_scale, + conf_threshold=conf_threshold, + iou_threshold=iou_threshold, + mask_scale=mask_scale, on_complete=on_complete, on_progress=on_progress, ) diff --git a/panels/vse_panel.py b/panels/vse_panel.py index c2e1d6f..678354b 100644 --- a/panels/vse_panel.py +++ b/panels/vse_panel.py @@ -9,6 +9,7 @@ import bpy from bpy.types import Panel 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): @@ -22,21 +23,29 @@ class SEQUENCER_PT_face_mask(Panel): def draw(self, context): layout = self.layout + scene = context.scene wm = context.window_manager seq_editor = context.scene.sequence_editor # Note: Blender 5.0 uses 'strips' instead of 'sequences' - + 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 if generator.is_running: self._draw_progress(layout, wm, generator) return - + # Show controls if strip selected if seq_editor and seq_editor.active_strip: strip = seq_editor.active_strip - + if strip.type in {'MOVIE', 'IMAGE'}: self._draw_generation_controls(layout, context, strip) self._draw_blur_controls(layout, context, strip) @@ -45,18 +54,96 @@ class SEQUENCER_PT_face_mask(Panel): else: 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): """Draw progress bar during generation.""" box = layout.box() box.label(text="Generating Masks...", icon='RENDER_ANIMATION') - + # Progress bar progress = wm.mask_progress / max(wm.mask_total, 1) box.progress( factor=progress, text=f"Frame {wm.mask_progress} / {wm.mask_total}", ) - + # Cancel button box.operator( "sequencer.cancel_mask_generation",