姿勢推定と頭部検知を合成する手段を追加

operators/generate_mask.py

@@ -205,10 +205,88 @@ class SEQUENCER_OT_cancel_mask_generation(Operator):
         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'}
+
+
 # Registration
 classes = [
     SEQUENCER_OT_generate_face_mask,
     SEQUENCER_OT_cancel_mask_generation,
+    SEQUENCER_OT_augment_pose_mask,
 ]