feat: 中間データのmsgpack移行
This commit is contained in:
+8
-22
@@ -12,6 +12,7 @@ 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"
|
||||
@@ -28,19 +29,9 @@ FORMAT_EXT = {
|
||||
}
|
||||
|
||||
|
||||
def _find_mask_strip(seq_editor, strip_name: str):
|
||||
return seq_editor.strips.get(f"{strip_name}_mask")
|
||||
|
||||
|
||||
def _resolve_mask_path(mask_strip) -> str:
|
||||
if mask_strip.type == "MOVIE":
|
||||
return bpy.path.abspath(mask_strip.filepath)
|
||||
return ""
|
||||
|
||||
|
||||
def _output_path(video_strip, mask_path: str, fmt: str) -> str:
|
||||
def _output_path(video_strip, detections_path: str, fmt: str) -> str:
|
||||
ext = FORMAT_EXT.get(fmt, "mp4")
|
||||
out_dir = os.path.dirname(mask_path)
|
||||
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}")
|
||||
|
||||
@@ -83,22 +74,17 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
|
||||
scene = context.scene
|
||||
video_strip = seq_editor.active_strip
|
||||
|
||||
mask_strip = _find_mask_strip(seq_editor, video_strip.name)
|
||||
if not mask_strip:
|
||||
self.report({"ERROR"}, f"Mask strip not found: {video_strip.name}_mask")
|
||||
return {"CANCELLED"}
|
||||
|
||||
video_path = bpy.path.abspath(video_strip.filepath)
|
||||
mask_path = _resolve_mask_path(mask_strip)
|
||||
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 mask_path or not os.path.exists(mask_path):
|
||||
self.report({"ERROR"}, f"Mask video not found: {mask_path}")
|
||||
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
|
||||
output_path = _output_path(video_strip, mask_path, bake_format)
|
||||
output_path = _output_path(video_strip, detections_path, bake_format)
|
||||
blur_size = int(scene.facemask_bake_blur_size)
|
||||
|
||||
# Reuse baked cache when parameters match and file still exists.
|
||||
@@ -164,7 +150,7 @@ class SEQUENCER_OT_bake_and_swap_blur_source(Operator):
|
||||
try:
|
||||
bake_generator.start(
|
||||
video_path=video_path,
|
||||
mask_path=mask_path,
|
||||
detections_path=detections_path,
|
||||
output_path=output_path,
|
||||
blur_size=blur_size,
|
||||
fmt=bake_format.lower(),
|
||||
|
||||
+29
-116
@@ -55,8 +55,7 @@ class SEQUENCER_OT_generate_face_mask(Operator):
|
||||
# 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 masks from {output_dir}")
|
||||
self._add_mask_strip(context, strip.name, output_dir)
|
||||
self.report({'INFO'}, f"Using cached detections from {output_dir}")
|
||||
return {'FINISHED'}
|
||||
|
||||
# Get frame range
|
||||
@@ -71,19 +70,25 @@ class SEQUENCER_OT_generate_face_mask(Operator):
|
||||
self.report({'WARNING'}, "Mask generation already in progress")
|
||||
return {'CANCELLED'}
|
||||
|
||||
# Store strip name for callback
|
||||
strip_name = strip.name
|
||||
|
||||
def on_complete(status, data):
|
||||
"""Called when mask generation completes."""
|
||||
wm = context.window_manager
|
||||
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":
|
||||
# Add mask strip to sequence editor
|
||||
self._add_mask_strip(context, strip_name, data)
|
||||
print(f"[FaceMask] Mask generation completed: {data}")
|
||||
elif status == "error":
|
||||
print(f"[FaceMask] Error: {data}")
|
||||
elif status == "cancelled":
|
||||
print("[FaceMask] Generation cancelled")
|
||||
|
||||
for area in context.screen.areas:
|
||||
if area.type == 'SEQUENCE_EDITOR':
|
||||
area.tag_redraw()
|
||||
|
||||
def on_progress(current, total):
|
||||
"""Called on progress updates."""
|
||||
@@ -143,119 +148,27 @@ class SEQUENCER_OT_generate_face_mask(Operator):
|
||||
if not os.path.exists(cache_dir):
|
||||
return False
|
||||
|
||||
# Check for MP4 video (new format)
|
||||
mask_video = os.path.join(cache_dir, "mask.mp4")
|
||||
if os.path.exists(mask_video):
|
||||
# Prefer frame-count verification when cv2 is available, but do not
|
||||
# hard-fail on Blender Python environments without cv2.
|
||||
try:
|
||||
import cv2
|
||||
detections_path = os.path.join(cache_dir, "detections.msgpack")
|
||||
if not os.path.exists(detections_path):
|
||||
return False
|
||||
|
||||
cap = cv2.VideoCapture(mask_video)
|
||||
if cap.isOpened():
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.release()
|
||||
# Accept cache if at least 90% of frames exist
|
||||
return frame_count >= expected_frames * 0.9
|
||||
cap.release()
|
||||
# Quick sanity check: non-empty file
|
||||
try:
|
||||
if os.path.getsize(detections_path) <= 0:
|
||||
return False
|
||||
except Exception:
|
||||
# Fallback: treat existing MP4 cache as valid when cv2 is unavailable.
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
# Fallback: check for PNG sequence (backward compatibility)
|
||||
mask_files = [f for f in os.listdir(cache_dir)
|
||||
if f.startswith("mask_") and f.endswith(".png")]
|
||||
# Optional frame count verification if msgpack is available
|
||||
try:
|
||||
import msgpack
|
||||
|
||||
# Accept cache if at least 90% of frames exist
|
||||
return len(mask_files) >= expected_frames * 0.9
|
||||
|
||||
def _add_mask_strip(self, context, source_strip_name: str, mask_path: str):
|
||||
"""Add mask video as a new strip.
|
||||
|
||||
Args:
|
||||
context: Blender context
|
||||
source_strip_name: Name of the source video strip
|
||||
mask_path: Path to mask video file or directory (for backward compatibility)
|
||||
"""
|
||||
scene = context.scene
|
||||
seq_editor = scene.sequence_editor
|
||||
|
||||
if not seq_editor:
|
||||
return
|
||||
|
||||
# Find source strip (Blender 5.0 uses 'strips' instead of 'sequences')
|
||||
source_strip = seq_editor.strips.get(source_strip_name)
|
||||
if not source_strip:
|
||||
return
|
||||
|
||||
# Check if mask_path is a video file or directory (backward compatibility)
|
||||
if os.path.isfile(mask_path):
|
||||
# New format: single MP4 file
|
||||
mask_video = mask_path
|
||||
else:
|
||||
# Old format: directory with PNG sequence (backward compatibility)
|
||||
mask_video = os.path.join(mask_path, "mask.mp4")
|
||||
if not os.path.exists(mask_video):
|
||||
# Fallback to PNG sequence
|
||||
mask_files = sorted([
|
||||
f for f in os.listdir(mask_path)
|
||||
if f.startswith("mask_") and f.endswith(".png")
|
||||
])
|
||||
if not mask_files:
|
||||
return
|
||||
first_mask = os.path.join(mask_path, mask_files[0])
|
||||
self._add_mask_strip_png_sequence(context, source_strip_name, mask_path, mask_files, first_mask)
|
||||
return
|
||||
|
||||
# Find an empty channel
|
||||
used_channels = {s.channel for s in seq_editor.strips}
|
||||
new_channel = source_strip.channel + 1
|
||||
while new_channel in used_channels:
|
||||
new_channel += 1
|
||||
|
||||
# Add movie strip (Blender 5.0 API)
|
||||
mask_strip = seq_editor.strips.new_movie(
|
||||
name=f"{source_strip_name}_mask",
|
||||
filepath=mask_video,
|
||||
channel=new_channel,
|
||||
frame_start=source_strip.frame_final_start,
|
||||
)
|
||||
|
||||
# Set blend mode for mask
|
||||
mask_strip.blend_type = 'ALPHA_OVER'
|
||||
mask_strip.blend_alpha = 0.5
|
||||
|
||||
def _add_mask_strip_png_sequence(self, context, source_strip_name, mask_dir, mask_files, first_mask):
|
||||
"""Backward compatibility: Add PNG sequence as mask strip."""
|
||||
scene = context.scene
|
||||
seq_editor = scene.sequence_editor
|
||||
source_strip = seq_editor.strips.get(source_strip_name)
|
||||
|
||||
if not source_strip:
|
||||
return
|
||||
|
||||
# Find an empty channel
|
||||
used_channels = {s.channel for s in seq_editor.strips}
|
||||
new_channel = source_strip.channel + 1
|
||||
while new_channel in used_channels:
|
||||
new_channel += 1
|
||||
|
||||
# Add image sequence (Blender 5.0 API)
|
||||
mask_strip = seq_editor.strips.new_image(
|
||||
name=f"{source_strip_name}_mask",
|
||||
filepath=first_mask,
|
||||
channel=new_channel,
|
||||
frame_start=source_strip.frame_final_start,
|
||||
)
|
||||
|
||||
# Add remaining frames
|
||||
for mask_file in mask_files[1:]:
|
||||
mask_strip.elements.append(mask_file)
|
||||
|
||||
# Set blend mode for mask
|
||||
mask_strip.blend_type = 'ALPHA_OVER'
|
||||
mask_strip.blend_alpha = 0.5
|
||||
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):
|
||||
|
||||
Reference in New Issue
Block a user