""" Bake-and-swap blur operators for VSE. This module bakes masked blur into a regular video file using the inference server, then swaps the active strip's source filepath to the baked result. """ import os import bpy from bpy.props import IntProperty from bpy.types import Operator from ..core.async_bake_generator import get_bake_generator from ..core.async_generator import get_generator as get_mask_generator from ..core.utils import get_detections_path_for_strip KEY_ORIGINAL = "facemask_original_filepath" KEY_BAKED = "facemask_baked_filepath" KEY_MODE = "facemask_source_mode" KEY_FORMAT = "facemask_bake_format" KEY_BLUR_SIZE = "facemask_bake_blur_size" KEY_DISPLAY_SCALE = "facemask_bake_display_scale" FORMAT_EXT = { "MP4": "mp4", "AVI": "avi", "MOV": "mov", } def _output_path(video_strip, detections_path: str, fmt: str) -> str: ext = FORMAT_EXT.get(fmt, "mp4") out_dir = os.path.dirname(detections_path) safe_name = video_strip.name.replace("/", "_").replace("\\", "_") 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): if hasattr(strip, "reload"): try: strip.reload() except Exception: pass def _set_strip_source(strip, path: str): if strip.type == "IMAGE": 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): """Bake masked blur (reuse cache if parameters match).""" bl_idname = "sequencer.bake_and_swap_blur_source" bl_label = "Bake" bl_description = "Bake masked blur to video and swap active strip source" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context): if not context.scene.sequence_editor: return False if get_mask_generator().is_running: return False if get_bake_generator().is_running: return False strip = context.scene.sequence_editor.active_strip return bool(strip and strip.type in {"MOVIE", "IMAGE"}) def execute(self, context): return _start_bake_impl(self, context, force=False) class SEQUENCER_OT_force_rebake_blur(Operator): """Force re-bake, ignoring any existing cached result.""" bl_idname = "sequencer.force_rebake_blur" bl_label = "Re-bake" bl_description = "Discard cached bake and re-bake from scratch" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context): if not context.scene.sequence_editor: return False if get_mask_generator().is_running: return False if get_bake_generator().is_running: return False strip = context.scene.sequence_editor.active_strip return bool(strip and strip.type in {"MOVIE", "IMAGE"}) def execute(self, context): return _start_bake_impl(self, context, force=True) class SEQUENCER_OT_swap_to_baked_blur(Operator): """Swap active strip source to already-baked video (no re-bake).""" bl_idname = "sequencer.swap_to_baked_blur" bl_label = "Swap to Baked" bl_description = "Switch active strip source to the baked video without re-baking" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context): if not context.scene.sequence_editor: return False if get_bake_generator().is_running: return False strip = context.scene.sequence_editor.active_strip if not strip or strip.type not in {"MOVIE", "IMAGE"}: return False baked_path = strip.get(KEY_BAKED) return bool(baked_path and os.path.exists(baked_path)) def execute(self, context): strip = context.scene.sequence_editor.active_strip baked_path = strip.get(KEY_BAKED) _set_strip_source(strip, baked_path) strip[KEY_MODE] = "baked" self.report({"INFO"}, "Swapped to baked source") return {"FINISHED"} class SEQUENCER_OT_restore_original_source(Operator): """Restore active strip source filepath to original video.""" bl_idname = "sequencer.restore_original_source" bl_label = "Restore Original" bl_description = "Restore active strip to original source filepath" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context): if not context.scene.sequence_editor: return False if get_bake_generator().is_running: return False strip = context.scene.sequence_editor.active_strip if not strip or strip.type not in {"MOVIE", "IMAGE"}: return False if strip.get(KEY_MODE, "original") == "original": return False return bool(strip.get(KEY_ORIGINAL)) def execute(self, context): strip = context.scene.sequence_editor.active_strip original_path = strip.get(KEY_ORIGINAL) if not original_path: self.report({"ERROR"}, "Original source path is not stored") return {"CANCELLED"} if not os.path.exists(original_path): self.report({"ERROR"}, f"Original source not found: {original_path}") return {"CANCELLED"} _set_strip_source(strip, original_path) strip[KEY_MODE] = "original" self.report({"INFO"}, "Restored original source") return {"FINISHED"} class SEQUENCER_OT_apply_mask_blur(Operator): """Compatibility alias: run bake-and-swap blur workflow.""" bl_idname = "sequencer.apply_mask_blur" bl_label = "Apply Mask Blur" bl_description = "Compatibility alias for Bake" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context): return SEQUENCER_OT_bake_and_swap_blur_source.poll(context) def execute(self, context): return bpy.ops.sequencer.bake_and_swap_blur_source("EXEC_DEFAULT") class SEQUENCER_OT_cancel_bake_blur(Operator): """Cancel ongoing blur bake.""" bl_idname = "sequencer.cancel_bake_blur" bl_label = "Cancel Blur Bake" bl_description = "Cancel current blur bake process" bl_options = {"REGISTER"} def execute(self, context): bake_generator = get_bake_generator() if bake_generator.is_running: bake_generator.cancel() self.report({"INFO"}, "Blur bake cancelled") else: self.report({"WARNING"}, "No blur bake in progress") return {"FINISHED"} classes = [ 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_cancel_bake_blur, SEQUENCER_OT_apply_mask_blur, ] def register(): for cls in classes: bpy.utils.register_class(cls) bpy.types.WindowManager.bake_progress = IntProperty(default=0) bpy.types.WindowManager.bake_total = IntProperty(default=0) def unregister(): del bpy.types.WindowManager.bake_progress del bpy.types.WindowManager.bake_total for cls in reversed(classes): bpy.utils.unregister_class(cls)