gear_machineandClaude
5777eeac20
- Enhanced error handling with specific exception types in detector and interpolation nodes - Reduced verbose console output in KeyPoint Denoiser components for cleaner operation - Improved coordinate validation and array bounds checking in detector nodes - Translated comments to English in params node for consistency - Maintained core functionality while improving user experience 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
341 lines
13 KiB
Python
341 lines
13 KiB
Python
"""
|
|
KeyPoint Denoiser Node for ComfyUI
|
|
ComfyUI node implementation of Kalman filter-based KeyPoint denoiser
|
|
"""
|
|
|
|
import copy
|
|
import traceback
|
|
from typing import List, Dict, Any, Tuple, Optional
|
|
|
|
# Support both relative and absolute imports for ComfyUI compatibility
|
|
try:
|
|
from .keypoint_denoiser import (
|
|
denoise_pose_keypoints_kalman,
|
|
DenoiserConfig,
|
|
get_config_by_name,
|
|
create_custom_config
|
|
)
|
|
except ImportError:
|
|
try:
|
|
from keypoint_denoiser import (
|
|
denoise_pose_keypoints_kalman,
|
|
DenoiserConfig,
|
|
get_config_by_name,
|
|
create_custom_config
|
|
)
|
|
except ImportError:
|
|
# Fallback for standalone testing
|
|
import sys
|
|
import os
|
|
current_dir = os.path.dirname(__file__)
|
|
keypoint_denoiser_dir = os.path.join(current_dir, 'keypoint_denoiser')
|
|
sys.path.insert(0, keypoint_denoiser_dir)
|
|
|
|
from core_algorithm import denoise_pose_keypoints_kalman
|
|
from config import DenoiserConfig, get_config_by_name, create_custom_config
|
|
|
|
class ProportionChangerKeypointDenoiser:
|
|
"""
|
|
KeyPoint denoiser node
|
|
|
|
Corrects KeyPoint jumps and detection errors in high-speed movements
|
|
such as dance videos using Kalman filter-based time-series noise removal
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pose_keypoint": ("POSE_KEYPOINT",),
|
|
"processing_mode": (
|
|
["balanced", "precision", "performance", "dance_optimized"],
|
|
{"default": "balanced"}
|
|
),
|
|
"denoising_strength": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.1}
|
|
),
|
|
},
|
|
"optional": {
|
|
"confidence_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.35, "min": 0.1, "max": 0.8, "step": 0.05}
|
|
),
|
|
"gate_threshold_multiplier": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.1}
|
|
),
|
|
"enable_bone_constraints": (
|
|
"BOOLEAN",
|
|
{"default": True}
|
|
),
|
|
"enable_rts_smoother": (
|
|
"BOOLEAN",
|
|
{"default": True}
|
|
),
|
|
"enable_verbose_logging": (
|
|
"BOOLEAN",
|
|
{"default": False}
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("POSE_KEYPOINT",)
|
|
RETURN_NAMES = ("denoised_pose_keypoint",)
|
|
FUNCTION = "denoise_keypoints"
|
|
CATEGORY = "ProportionChanger"
|
|
|
|
def denoise_keypoints(self,
|
|
pose_keypoint: List[Dict],
|
|
processing_mode: str = "balanced",
|
|
denoising_strength: float = 1.0,
|
|
confidence_threshold: float = 0.35,
|
|
gate_threshold_multiplier: float = 1.0,
|
|
enable_bone_constraints: bool = True,
|
|
enable_rts_smoother: bool = True,
|
|
enable_verbose_logging: bool = False) -> Tuple[List[Dict]]:
|
|
"""
|
|
Main processing of KeyPoint denoiser
|
|
|
|
Args:
|
|
pose_keypoint: Batch data in POSE_KEYPOINT format
|
|
processing_mode: Processing mode (balanced/precision/performance/dance_optimized)
|
|
denoising_strength: Denoising strength (0.0-3.0)
|
|
confidence_threshold: Confidence threshold
|
|
gate_threshold_multiplier: Gate threshold adjustment multiplier
|
|
enable_bone_constraints: Enable structure constraints
|
|
enable_rts_smoother: Enable RTS smoother
|
|
enable_verbose_logging: Enable verbose logging
|
|
|
|
Returns:
|
|
Tuple[List[Dict]]: Denoised POSE_KEYPOINT batch
|
|
"""
|
|
|
|
try:
|
|
# Input validation
|
|
if not pose_keypoint or len(pose_keypoint) == 0:
|
|
if enable_verbose_logging:
|
|
pass
|
|
return (pose_keypoint,)
|
|
|
|
if len(pose_keypoint) < 3:
|
|
if enable_verbose_logging:
|
|
pass
|
|
return (pose_keypoint,)
|
|
|
|
# Create configuration
|
|
base_config = get_config_by_name(processing_mode)
|
|
|
|
# Adjust configuration with custom parameters
|
|
custom_config = create_custom_config(
|
|
# ベース設定から開始
|
|
qv_base=base_config.qv_base * (denoising_strength ** 0.5),
|
|
gate_threshold=base_config.gate_threshold * gate_threshold_multiplier,
|
|
conf_min=confidence_threshold,
|
|
enable_bone_constraints=enable_bone_constraints,
|
|
use_rts_smoother=enable_rts_smoother,
|
|
enable_dance_protection=base_config.enable_dance_protection,
|
|
verbose_logging=enable_verbose_logging,
|
|
|
|
# デノイジング強度による調整
|
|
r_base=base_config.r_base / (denoising_strength + 0.5),
|
|
max_iterations=min(base_config.max_iterations, max(1, int(denoising_strength * 2))),
|
|
light_window_size=min(7, max(1, int(base_config.light_window_size * denoising_strength))),
|
|
)
|
|
|
|
if enable_verbose_logging:
|
|
pass
|
|
|
|
# Execute main denoising process
|
|
denoised_result = denoise_pose_keypoints_kalman(
|
|
pose_keypoint_batch=pose_keypoint,
|
|
config=custom_config
|
|
)
|
|
|
|
if enable_verbose_logging:
|
|
pass
|
|
|
|
return (denoised_result,)
|
|
|
|
except Exception as e:
|
|
error_msg = f"KeyPoint Denoiser error: {str(e)}"
|
|
pass
|
|
|
|
if enable_verbose_logging:
|
|
pass
|
|
traceback.print_exc()
|
|
print("-" * 60)
|
|
|
|
# Return original data on error
|
|
return (pose_keypoint,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""ノードの変更検出（常に実行するためにランダム値を返す）"""
|
|
import random
|
|
return random.random()
|
|
|
|
class ProportionChangerKeypointDenoiserAdvanced:
|
|
"""
|
|
KeyPointデノイザー上級者向けノード
|
|
|
|
より詳細なパラメータ制御が可能な版
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pose_keypoint": ("POSE_KEYPOINT",),
|
|
},
|
|
"optional": {
|
|
# カルマンフィルタパラメータ
|
|
"kalman_process_noise": (
|
|
"FLOAT",
|
|
{"default": 0.003, "min": 0.001, "max": 0.01, "step": 0.001}
|
|
),
|
|
"kalman_measurement_noise": (
|
|
"FLOAT",
|
|
{"default": 0.004, "min": 0.001, "max": 0.02, "step": 0.001}
|
|
),
|
|
"gate_threshold": (
|
|
"FLOAT",
|
|
{"default": 9.21, "min": 5.0, "max": 20.0, "step": 0.5}
|
|
),
|
|
"gate_recovery_threshold": (
|
|
"FLOAT",
|
|
{"default": 12.0, "min": 8.0, "max": 25.0, "step": 0.5}
|
|
),
|
|
|
|
# 信頼度・検出パラメータ
|
|
"confidence_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.35, "min": 0.1, "max": 0.8, "step": 0.05}
|
|
),
|
|
"orientation_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.7, "min": 0.5, "max": 0.95, "step": 0.05}
|
|
),
|
|
|
|
# ダンス動作保護
|
|
"enable_spin_protection": ("BOOLEAN", {"default": True}),
|
|
"enable_jump_protection": ("BOOLEAN", {"default": True}),
|
|
"enable_coordination_protection": ("BOOLEAN", {"default": True}),
|
|
|
|
# 構造投影パラメータ
|
|
"enable_bone_constraints": ("BOOLEAN", {"default": True}),
|
|
"bone_constraint_iterations": ("INT", {"default": 2, "min": 1, "max": 5}),
|
|
"projection_interval": ("INT", {"default": 3, "min": 1, "max": 10}),
|
|
|
|
# スムージングパラメータ
|
|
"enable_rts_smoother": ("BOOLEAN", {"default": True}),
|
|
"smoothing_window_size": ("INT", {"default": 3, "min": 1, "max": 7}),
|
|
|
|
# システム設定
|
|
"enable_verbose_logging": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("POSE_KEYPOINT",)
|
|
RETURN_NAMES = ("denoised_pose_keypoint",)
|
|
FUNCTION = "denoise_keypoints_advanced"
|
|
CATEGORY = "ProportionChanger"
|
|
|
|
def denoise_keypoints_advanced(self,
|
|
pose_keypoint: List[Dict],
|
|
kalman_process_noise: float = 0.003,
|
|
kalman_measurement_noise: float = 0.004,
|
|
gate_threshold: float = 9.21,
|
|
gate_recovery_threshold: float = 12.0,
|
|
confidence_threshold: float = 0.35,
|
|
orientation_threshold: float = 0.7,
|
|
enable_spin_protection: bool = True,
|
|
enable_jump_protection: bool = True,
|
|
enable_coordination_protection: bool = True,
|
|
enable_bone_constraints: bool = True,
|
|
bone_constraint_iterations: int = 2,
|
|
projection_interval: int = 3,
|
|
enable_rts_smoother: bool = True,
|
|
smoothing_window_size: int = 3,
|
|
enable_verbose_logging: bool = False) -> Tuple[List[Dict]]:
|
|
"""
|
|
上級者向けKeyPointデノイザー処理
|
|
"""
|
|
|
|
try:
|
|
# Input validation
|
|
if not pose_keypoint or len(pose_keypoint) < 3:
|
|
if enable_verbose_logging:
|
|
pass
|
|
return (pose_keypoint,)
|
|
|
|
# Create custom detailed configuration
|
|
advanced_config = create_custom_config(
|
|
# カルマンフィルタ
|
|
qv_base=kalman_process_noise**2,
|
|
r_base=kalman_measurement_noise**2,
|
|
gate_threshold=gate_threshold,
|
|
gate_threshold_recovery=gate_recovery_threshold,
|
|
|
|
# 検出パラメータ
|
|
conf_min=confidence_threshold,
|
|
orientation_threshold=orientation_threshold,
|
|
|
|
# ダンス保護
|
|
enable_dance_protection=(enable_spin_protection or enable_jump_protection or enable_coordination_protection),
|
|
spin_protection_factor=2.0 if enable_spin_protection else 1.0,
|
|
jump_protection_factor=1.8 if enable_jump_protection else 1.0,
|
|
coordination_protection_factor=1.5 if enable_coordination_protection else 1.0,
|
|
|
|
# 構造投影
|
|
enable_bone_constraints=enable_bone_constraints,
|
|
max_iterations=bone_constraint_iterations,
|
|
projection_interval=projection_interval,
|
|
|
|
# スムージング
|
|
use_rts_smoother=enable_rts_smoother,
|
|
light_window_size=smoothing_window_size,
|
|
|
|
# システム
|
|
verbose_logging=enable_verbose_logging
|
|
)
|
|
|
|
if enable_verbose_logging:
|
|
pass
|
|
|
|
# Execute main processing
|
|
denoised_result = denoise_pose_keypoints_kalman(
|
|
pose_keypoint_batch=pose_keypoint,
|
|
config=advanced_config
|
|
)
|
|
|
|
if enable_verbose_logging:
|
|
pass
|
|
|
|
return (denoised_result,)
|
|
|
|
except Exception as e:
|
|
error_msg = f"Advanced KeyPoint Denoiser error: {str(e)}"
|
|
pass
|
|
|
|
if enable_verbose_logging:
|
|
traceback.print_exc()
|
|
|
|
return (pose_keypoint,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
import random
|
|
return random.random()
|
|
|
|
|
|
# ノード登録用の辞書
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ProportionChangerKeypointDenoiser": ProportionChangerKeypointDenoiser,
|
|
"ProportionChangerKeypointDenoiserAdvanced": ProportionChangerKeypointDenoiserAdvanced,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ProportionChangerKeypointDenoiser": "KeyPoint Denoiser",
|
|
"ProportionChangerKeypointDenoiserAdvanced": "KeyPoint Denoiser (Advanced)",
|
|
} |