From 6e4f14ab5e5f439648472d11db6d7ad165f3335b Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:50:56 -0400 Subject: [PATCH 01/20] coords unused imports --- node_wrappers/coordinates/conversion_nodes.py | 2 -- node_wrappers/coordinates/position_tracking.py | 4 +--- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/node_wrappers/coordinates/conversion_nodes.py b/node_wrappers/coordinates/conversion_nodes.py index 99f2cf4..863830a 100644 --- a/node_wrappers/coordinates/conversion_nodes.py +++ b/node_wrappers/coordinates/conversion_nodes.py @@ -4,9 +4,7 @@ Node wrappers for coordinate conversion in ComfyUI. These nodes provide user-friendly interfaces for the coordinate system functionality. """ -import torch import logging -from typing import Union, List, Tuple from ...src.coordinates import CoordinateSystem diff --git a/node_wrappers/coordinates/position_tracking.py b/node_wrappers/coordinates/position_tracking.py index 0c1df32..6d8f9d0 100644 --- a/node_wrappers/coordinates/position_tracking.py +++ b/node_wrappers/coordinates/position_tracking.py @@ -1,11 +1,9 @@ import logging -from typing import Optional, List, Any, Tuple, Union +from typing import List, Tuple, Union import math import torch import numpy as np import cv2 -import collections -import typing # Add typing for Dict hint #TODO: consider moving base nodes to src From 8653cc3916455b96cc52a95d9fdb4546bf106902 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:51:04 -0400 Subject: [PATCH 02/20] common unused imports --- node_wrappers/mediapipe_vision/common/base_detector_node.py | 3 +-- .../mediapipe_vision/common/base_visualization_nodes.py | 3 +-- node_wrappers/mediapipe_vision/common/model_loader.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/node_wrappers/mediapipe_vision/common/base_detector_node.py b/node_wrappers/mediapipe_vision/common/base_detector_node.py index 662c681..80b2097 100644 --- a/node_wrappers/mediapipe_vision/common/base_detector_node.py +++ b/node_wrappers/mediapipe_vision/common/base_detector_node.py @@ -1,8 +1,7 @@ """Base node wrapper for MediaPipe Vision detectors.""" -import torch import logging -from typing import Any, Dict, List, Tuple, Optional, Type, Callable, Union +from typing import Dict logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py index f258cd0..7470124 100644 --- a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py +++ b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py @@ -1,7 +1,6 @@ import logging -from typing import List, Dict, Tuple, Any, Optional, Type +from typing import List, Any, Optional import cv2 -import numpy as np from ....src.utils.image import convert_to_cv2, convert_to_tensor diff --git a/node_wrappers/mediapipe_vision/common/model_loader.py b/node_wrappers/mediapipe_vision/common/model_loader.py index 76ae1b6..36f3ec0 100644 --- a/node_wrappers/mediapipe_vision/common/model_loader.py +++ b/node_wrappers/mediapipe_vision/common/model_loader.py @@ -6,7 +6,7 @@ import logging import traceback # For better error logging # Add import for MEDIAPIPE_MODELS dictionary and the get_model_description function -from ....src.mediapipe_vision.model_loader import MEDIAPIPE_MODELS, get_model_description +from ....src.mediapipe_vision.model_loader import get_model_description logger = logging.getLogger(__name__) From 997e4b89868efb8cf85c5c139aa19464e07ab03b Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:51:15 -0400 Subject: [PATCH 03/20] face detection unused imports --- .../face/detection/face_detection_visualization.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py index 69b20f9..0422c24 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py @@ -1,6 +1,4 @@ # visualize_face_detection.py -import torch -import numpy as np import cv2 from ...common.base_visualization_nodes import BaseDetectionVisualizationNode From 4f042b4825f5b541bf533acad3bd70ece44cdc8c Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:51:27 -0400 Subject: [PATCH 04/20] face landmark unused imports --- node_wrappers/mediapipe_vision/face/landmark/blendshape.py | 2 +- node_wrappers/mediapipe_vision/face/landmark/face_masking.py | 4 +--- node_wrappers/mediapipe_vision/face/landmark/head_pose.py | 2 +- node_wrappers/mediapipe_vision/face/landmark/texture_warp.py | 2 -- 4 files changed, 3 insertions(+), 7 deletions(-) diff --git a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py index 0a1c1fd..39e6a08 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py +++ b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py @@ -1,6 +1,6 @@ # node_wrappers/face/blendshape.py import logging -from typing import Optional, List, Any, Dict, Tuple +from typing import Optional from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alias from .....src.mediapipe_vision.utils.blendshape_utils import get_blendshape_categories # Utility to get all category names diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py index 54dd39f..106f2e8 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py @@ -1,9 +1,7 @@ # node_wrappers/face/masking.py import torch -import numpy as np -import cv2 import logging -from typing import List, Dict, Tuple +from typing import List from .....src.mediapipe_vision.types import LandmarkPoint from .....src.utils.image import create_mask_from_points diff --git a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py index a0990c5..22b09e9 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py +++ b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py @@ -2,7 +2,7 @@ import logging import numpy as np from scipy.spatial.transform import Rotation as R -from typing import Optional, List, Any, Tuple, Dict +from typing import Optional, Dict from .....src.mediapipe_vision.types import TRANSFORM_MATRIX_LIST from .....src.utils.math import scale_value diff --git a/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py b/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py index abd7875..c6f1190 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py +++ b/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py @@ -1,7 +1,5 @@ # node_wrappers/face/texture_warp.py import logging -import torch -import numpy as np logger = logging.getLogger(__name__) From 274ef6287505bb644bc8fce245cab19cae9bedbc Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:51:37 -0400 Subject: [PATCH 05/20] hand gesture unused imports --- .../mediapipe_vision/hand/gesture/gesture_visualiztion.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py index c93a659..43202e5 100644 --- a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py +++ b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py @@ -1,6 +1,4 @@ # visualize_gestures.py -import torch -import numpy as np import cv2 from ...common.base_visualization_nodes import BaseDetectionVisualizationNode From cb67962394f0e2270e33df43cdf21798bfaa4d33 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:51:50 -0400 Subject: [PATCH 06/20] hand landmark unused imports --- .../mediapipe_vision/hand/landmark/hand_landmark_masking.py | 1 - .../hand/landmark/hand_landmark_visualization.py | 1 - node_wrappers/mediapipe_vision/hand/landmark/position.py | 2 -- 3 files changed, 4 deletions(-) diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py index b5279e5..e0e77e1 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py @@ -1,6 +1,5 @@ # node_wrappers/hand/masking.py import torch -import numpy as np import logging from typing import List diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py index 3f7f73c..d806896 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py @@ -2,7 +2,6 @@ import mediapipe as mp from typing import Optional, List, Any -from .....src.mediapipe_vision.types import HandLandmarksResult # Import specific result type if needed for extraction from ...common.base_visualization_nodes import BaseLandmarkVisualizationNode import logging diff --git a/node_wrappers/mediapipe_vision/hand/landmark/position.py b/node_wrappers/mediapipe_vision/hand/landmark/position.py index ab87983..9b53858 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/position.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/position.py @@ -1,11 +1,9 @@ # node_wrappers/hand/position.py import logging -from typing import List, Tuple, Any, Optional # Import the base class from .....src.mediapipe_vision.common.position import LandmarkPositionBase # Import the specific landmark type for clarity and potential future use -from .....src.mediapipe_vision.types import HAND_LANDMARKS # Import the landmark definitions from .....src.mediapipe_vision.landmark_definitions import HAND_LANDMARK_TOOLTIP From 482fa1ed670c54358b670ad0142c4c17c5312d94 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:01 -0400 Subject: [PATCH 07/20] holistic unused imports --- node_wrappers/mediapipe_vision/holistic/holistic_landmark.py | 3 --- .../holistic/holistic_landmark_visualization.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py index 818573e..7ce863b 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py @@ -1,7 +1,4 @@ import torch -import cv2 -import mediapipe as mp -import numpy as np from ....src.mediapipe_vision.holistic_landmark.detector import HolisticLandmarkDetector from ..common.base_detector_node import BaseMediaPipeDetectorNode import logging diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py index 7ea52c5..2070b66 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py @@ -2,7 +2,7 @@ import logging import mediapipe as mp import numpy as np import cv2 -from typing import List, Dict, Tuple, Any, Optional +from typing import List, Tuple, Optional import torch logger = logging.getLogger(__name__) From fd22bc213665b87aa56cec4d96463107dfa44cbf Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:11 -0400 Subject: [PATCH 08/20] segmentation unused imports --- .../image_segmentation/image_segmentation_masking.py | 1 - 1 file changed, 1 deletion(-) diff --git a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py index e643a66..e4ee1e2 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py @@ -1,5 +1,4 @@ import torch -import numpy as np import logging From ee4ec8a70b9a3722c49c580d72170aabe924d488 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:19 -0400 Subject: [PATCH 09/20] object detection unused imports --- .../object_detection/object_detection_visualization.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py index ba065cb..86f7a07 100644 --- a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py @@ -1,9 +1,6 @@ # visualize_object_detection.py -import torch -import numpy as np import cv2 -from ....src.utils.image import convert_to_cv2, convert_to_tensor from ..common.base_visualization_nodes import BaseDetectionVisualizationNode import logging From d8f7f24c3ca70f8f94dac24205c4d608b36df0e8 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:25 -0400 Subject: [PATCH 10/20] pose unused imports --- node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py | 4 +--- .../mediapipe_vision/pose/pose_landmark_visualizer.py | 1 - node_wrappers/mediapipe_vision/pose/position.py | 2 -- 3 files changed, 1 insertion(+), 6 deletions(-) diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py index 6af580c..edd2660 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py @@ -1,9 +1,7 @@ # node_wrappers/pose/masking.py import torch -import numpy as np -import cv2 import logging -from typing import List, Dict, Tuple +from typing import List from ....src.mediapipe_vision.types import PoseLandmarksResult from ....src.utils.image import create_mask_from_points diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py index 4b81e70..964ed0f 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py @@ -2,7 +2,6 @@ import mediapipe as mp from typing import Optional, List, Any -from ....src.mediapipe_vision.types import PoseLandmarksResult # Import specific result type if needed for extraction from ..common.base_visualization_nodes import BaseLandmarkVisualizationNode import logging diff --git a/node_wrappers/mediapipe_vision/pose/position.py b/node_wrappers/mediapipe_vision/pose/position.py index db68049..042517b 100644 --- a/node_wrappers/mediapipe_vision/pose/position.py +++ b/node_wrappers/mediapipe_vision/pose/position.py @@ -1,10 +1,8 @@ import logging -from typing import List, Tuple, Any, Optional # Import the base class from ....src.mediapipe_vision.common.position import LandmarkPositionBase # Import the specific landmark type for clarity and potential future use -from ....src.mediapipe_vision.types import POSE_LANDMARKS # Import the landmark definitions from ....src.mediapipe_vision.landmark_definitions import POSE_LANDMARK_TOOLTIP From 0640934a4526e309675bf4ad5e866e5f593b26d7 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:37 -0400 Subject: [PATCH 11/20] rtn unused imports --- node_wrappers/realtimenodes/controls/mask_controls.py | 1 - node_wrappers/realtimenodes/utils/utility_controls.py | 4 ---- 2 files changed, 5 deletions(-) diff --git a/node_wrappers/realtimenodes/controls/mask_controls.py b/node_wrappers/realtimenodes/controls/mask_controls.py index b7f92a1..df44e3c 100644 --- a/node_wrappers/realtimenodes/controls/mask_controls.py +++ b/node_wrappers/realtimenodes/controls/mask_controls.py @@ -1,6 +1,5 @@ import torch import numpy as np -import cv2 from ....src.realtimenodes.control_base import ControlNodeBase from ....src.realtimenodes.mask_controls import MaskControlMixin diff --git a/node_wrappers/realtimenodes/utils/utility_controls.py b/node_wrappers/realtimenodes/utils/utility_controls.py index 92d08e5..be203a9 100644 --- a/node_wrappers/realtimenodes/utils/utility_controls.py +++ b/node_wrappers/realtimenodes/utils/utility_controls.py @@ -1,11 +1,7 @@ -import os -import sys import time import numpy as np import cv2 import torch -import random -from ....src.utils.general import AlwaysEqualProxy # # Add package root to Python path # <-- REMOVED # sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from ....src.realtimenodes.control_base import ControlNodeBase From f54059810c9cee2790977b1c2a897105b0f7d7a8 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:52:55 -0400 Subject: [PATCH 12/20] src coords unused imports --- src/coordinates/coordinate_delta.py | 6 +----- src/coordinates/coordinate_system.py | 3 +-- src/coordinates/drawing_engine.py | 2 +- 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/src/coordinates/coordinate_delta.py b/src/coordinates/coordinate_delta.py index 25cfb39..f0768e5 100644 --- a/src/coordinates/coordinate_delta.py +++ b/src/coordinates/coordinate_delta.py @@ -1,16 +1,12 @@ import logging -from typing import Optional, List, Any, Tuple, Union +from typing import Optional import math -import torch -import numpy as np -import cv2 import collections import typing # Add typing for Dict hint #TODO: consider moving base nodes to src # Import from the new consolidated utilities -from ...src.utils.math import scale_value logger = logging.getLogger(__name__) diff --git a/src/coordinates/coordinate_system.py b/src/coordinates/coordinate_system.py index a00b240..6374939 100644 --- a/src/coordinates/coordinate_system.py +++ b/src/coordinates/coordinate_system.py @@ -6,8 +6,7 @@ spaces (normalized, pixel) with support for both single values and batches. """ import torch -import numpy as np -from typing import Union, List, Tuple, Optional, Dict, Any +from typing import Union, List, Tuple, Optional import logging logger = logging.getLogger(__name__) diff --git a/src/coordinates/drawing_engine.py b/src/coordinates/drawing_engine.py index c00689e..bd9fdad 100644 --- a/src/coordinates/drawing_engine.py +++ b/src/coordinates/drawing_engine.py @@ -1,7 +1,7 @@ import torch import numpy as np import cv2 -from typing import Union, List, Tuple, Dict, Optional, Any +from typing import Union, List, Tuple import logging from .coordinate_system import CoordinateSystem From e9e6824dd49963d05893cece69788ef4e05ad74d Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:53:04 -0400 Subject: [PATCH 13/20] src utils unused imports --- src/utils/image.py | 4 ++-- src/utils/math.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/utils/image.py b/src/utils/image.py index 5caa2fb..0c4b815 100644 --- a/src/utils/image.py +++ b/src/utils/image.py @@ -8,7 +8,7 @@ particularly between ComfyUI tensor format (BHWC) and OpenCV format. import torch import numpy as np import cv2 -from typing import List, Union, Tuple +from typing import List, Tuple def convert_to_cv2(tensor: torch.Tensor) -> list: """Converts a ComfyUI IMAGE tensor (BHWC, float32, 0-1) to a list of cv2 images (BGR, uint8).""" @@ -61,7 +61,7 @@ def create_mask_from_points(height: int, width: int, points: List[Tuple[int, int np_points = np.array(points, dtype=np.int32) hull = cv2.convexHull(np_points) cv2.fillPoly(mask, [hull], 1) - except Exception as e: + except Exception: return torch.from_numpy(mask).float().unsqueeze(0).to(device) mask_tensor = torch.from_numpy(mask).float().unsqueeze(0).to(device) diff --git a/src/utils/math.py b/src/utils/math.py index 4f985d9..debec70 100644 --- a/src/utils/math.py +++ b/src/utils/math.py @@ -8,7 +8,7 @@ Contains functions for mathematical operations including: """ import math -from typing import Optional, Tuple, Any +from typing import Optional, Tuple def scale_value(value: Optional[float], input_min: float, input_max: float, From c71d5f9be2e2724cf9c7c2a7f17fdf564d957e81 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:53:34 -0400 Subject: [PATCH 14/20] src detectors unused imports --- src/mediapipe_vision/face_detection/detector.py | 2 +- src/mediapipe_vision/face_landmark/detector.py | 2 +- src/mediapipe_vision/face_stylization/detector.py | 1 - src/mediapipe_vision/hand_landmark/detector.py | 2 +- src/mediapipe_vision/holistic_landmark/detector.py | 2 +- src/mediapipe_vision/image_embedding/detector.py | 2 +- src/mediapipe_vision/interactive_segmentation/detector.py | 4 ---- 7 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/mediapipe_vision/face_detection/detector.py b/src/mediapipe_vision/face_detection/detector.py index 0e40e6e..44f73f4 100644 --- a/src/mediapipe_vision/face_detection/detector.py +++ b/src/mediapipe_vision/face_detection/detector.py @@ -9,7 +9,7 @@ import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List, Tuple +from typing import List # Import new types from ..types import FaceDetectionResult, BoundingBox, FaceKeypoint from ...utils.timing import TimestampProvider # Import from utils diff --git a/src/mediapipe_vision/face_landmark/detector.py b/src/mediapipe_vision/face_landmark/detector.py index 73b448d..ca9d165 100644 --- a/src/mediapipe_vision/face_landmark/detector.py +++ b/src/mediapipe_vision/face_landmark/detector.py @@ -6,7 +6,7 @@ from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions from typing import List, Tuple, Optional, Dict # Import new types -from ..types import LandmarkPoint, FaceLandmarksResult +from ..types import LandmarkPoint from ...utils.timing import TimestampProvider # Import TimestampProvider class FaceLandmarkDetector: diff --git a/src/mediapipe_vision/face_stylization/detector.py b/src/mediapipe_vision/face_stylization/detector.py index 706a0de..cb23d71 100644 --- a/src/mediapipe_vision/face_stylization/detector.py +++ b/src/mediapipe_vision/face_stylization/detector.py @@ -3,7 +3,6 @@ import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision -from typing import List, Optional from ...utils.timing import TimestampProvider class FaceStylizer: diff --git a/src/mediapipe_vision/hand_landmark/detector.py b/src/mediapipe_vision/hand_landmark/detector.py index 693764d..613b3c2 100644 --- a/src/mediapipe_vision/hand_landmark/detector.py +++ b/src/mediapipe_vision/hand_landmark/detector.py @@ -4,7 +4,7 @@ import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List, Tuple, Optional +from typing import List # Import new types from ..types import LandmarkPoint, HandLandmarksResult from ...utils.timing import TimestampProvider # Import TimestampProvider diff --git a/src/mediapipe_vision/holistic_landmark/detector.py b/src/mediapipe_vision/holistic_landmark/detector.py index 33b157c..5d5bbfb 100644 --- a/src/mediapipe_vision/holistic_landmark/detector.py +++ b/src/mediapipe_vision/holistic_landmark/detector.py @@ -8,7 +8,7 @@ import mediapipe as mp import numpy as np import torch import cv2 -from typing import List, Tuple, Optional, Dict, Any +from typing import List, Tuple, Optional, Any # Import needed types from ..types import LandmarkPoint, HolisticLandmarksResult diff --git a/src/mediapipe_vision/image_embedding/detector.py b/src/mediapipe_vision/image_embedding/detector.py index 41aad69..eb7152c 100644 --- a/src/mediapipe_vision/image_embedding/detector.py +++ b/src/mediapipe_vision/image_embedding/detector.py @@ -4,7 +4,7 @@ import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List, Optional +from typing import List # Import new types from ..types import ImageEmbedderResult # Import TimestampProvider from the new location diff --git a/src/mediapipe_vision/interactive_segmentation/detector.py b/src/mediapipe_vision/interactive_segmentation/detector.py index c9aafd4..1ce0f1a 100644 --- a/src/mediapipe_vision/interactive_segmentation/detector.py +++ b/src/mediapipe_vision/interactive_segmentation/detector.py @@ -3,10 +3,6 @@ import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python -from mediapipe.tasks.python import vision -from mediapipe.tasks.python.vision.interactive_segmenter import ( - InteractiveSegmenter, -) from mediapipe.tasks.python.vision import interactive_segmenter from mediapipe.tasks.python.components.containers import keypoint From 8e5974c10e57a8bab4dc57f1d2fa4505c572574a Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:53:43 -0400 Subject: [PATCH 15/20] src utils unused imports --- src/mediapipe_vision/utils/blendshape_utils.py | 1 - src/mediapipe_vision/utils/delta_utils.py | 1 - src/mediapipe_vision/utils/model_utils.py | 1 - 3 files changed, 3 deletions(-) diff --git a/src/mediapipe_vision/utils/blendshape_utils.py b/src/mediapipe_vision/utils/blendshape_utils.py index 5de35cc..e3b564e 100644 --- a/src/mediapipe_vision/utils/blendshape_utils.py +++ b/src/mediapipe_vision/utils/blendshape_utils.py @@ -1,5 +1,4 @@ # src/utils/misc_utils.py -import mediapipe as mp import logging from typing import List diff --git a/src/mediapipe_vision/utils/delta_utils.py b/src/mediapipe_vision/utils/delta_utils.py index 6642527..eea8f36 100644 --- a/src/mediapipe_vision/utils/delta_utils.py +++ b/src/mediapipe_vision/utils/delta_utils.py @@ -8,7 +8,6 @@ import math from typing import Optional -from ...utils.math import scale_value from ..types import LandmarkPoint # Assuming LandmarkPoint is in src/types.py #TODO: refactor to common after unified coordinate system def calculate_euclidean_delta(p1: Optional[LandmarkPoint], p2: Optional[LandmarkPoint]) -> Optional[float]: diff --git a/src/mediapipe_vision/utils/model_utils.py b/src/mediapipe_vision/utils/model_utils.py index 499d75f..c29bc46 100644 --- a/src/mediapipe_vision/utils/model_utils.py +++ b/src/mediapipe_vision/utils/model_utils.py @@ -2,7 +2,6 @@ import os import logging -from typing import Dict, Any, Optional # Configure logging logging.basicConfig(level=logging.INFO) From dcbbe7cc3fd9b1f82bc26344c713d31e70ee8176 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:54:06 -0400 Subject: [PATCH 16/20] src common postition unused imports --- src/mediapipe_vision/common/position.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/mediapipe_vision/common/position.py b/src/mediapipe_vision/common/position.py index 2498b47..f8f831e 100644 --- a/src/mediapipe_vision/common/position.py +++ b/src/mediapipe_vision/common/position.py @@ -1,23 +1,14 @@ import logging -from typing import Optional, List, Any, Tuple, Union -import math -import torch -import numpy as np -import cv2 -import collections -import typing # Add typing for Dict hint +from typing import Optional, List, Tuple #TODO: consider moving base nodes to src # Assuming types are correctly imported relative to this file's location from ..types import ( - FACE_LANDMARKS, HAND_LANDMARKS, POSE_LANDMARKS, - LandmarkPoint, FaceLandmarksResult, HandLandmarksResult, PoseLandmarksResult + LandmarkPoint ) # Import from the new consolidated utilities -from ...utils.math import scale_value -from ...utils.general import AlwaysEqualProxy logger = logging.getLogger(__name__) From a0818dd78fd32220675a05aa6231ff487ba3d647 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 17:54:29 -0400 Subject: [PATCH 17/20] src types, misc nodes unused imports --- node_wrappers/utils/misc_nodes.py | 3 --- src/mediapipe_vision/types.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/node_wrappers/utils/misc_nodes.py b/node_wrappers/utils/misc_nodes.py index d7ffcde..21f7dc3 100644 --- a/node_wrappers/utils/misc_nodes.py +++ b/node_wrappers/utils/misc_nodes.py @@ -1,5 +1,4 @@ import torch -import comfy.utils import numpy as np import base64 import re @@ -7,8 +6,6 @@ import math from io import BytesIO from PIL import Image, ImageFont, ImageDraw, ImageOps import nodes -import random -from torchvision import transforms from ...src.utils.general import AlwaysEqualProxy MAX_RESOLUTION = nodes.MAX_RESOLUTION # Get the same max resolution as core nodes diff --git a/src/mediapipe_vision/types.py b/src/mediapipe_vision/types.py index fa304ba..f3d2147 100644 --- a/src/mediapipe_vision/types.py +++ b/src/mediapipe_vision/types.py @@ -1,4 +1,4 @@ -from typing import List, Optional, NamedTuple, Tuple +from typing import List, Optional, NamedTuple import numpy as np # Added numpy for matrix type # Define a structure for a single landmark point From ea32681c1b9cbadf662fb7bed92278ec6979a8ee Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 18:00:55 -0400 Subject: [PATCH 18/20] sort imports --- __init__.py | 2 +- node_wrappers/coordinates/__init__.py | 6 +----- node_wrappers/coordinates/drawing_nodes.py | 8 ++++---- node_wrappers/coordinates/position_tracking.py | 12 +++++++----- .../common/base_visualization_nodes.py | 3 ++- .../mediapipe_vision/common/model_loader.py | 9 ++++++--- .../face/detection/face_detection.py | 10 ++++++---- .../face/detection/face_detection_mask.py | 5 +++-- .../face/detection/face_detection_visualization.py | 4 +++- .../mediapipe_vision/face/landmark/blendshape.py | 11 +++++++---- .../face/landmark/face_landmark_detection.py | 7 +++++-- .../face_landmark_detection_visualization.py | 5 +++-- .../mediapipe_vision/face/landmark/face_masking.py | 5 +++-- .../mediapipe_vision/face/landmark/head_pose.py | 7 +++++-- .../mediapipe_vision/face/landmark/position.py | 1 + .../face/stylization/face_stylization.py | 8 +++++--- .../hand/gesture/gesture_visualiztion.py | 3 ++- .../mediapipe_vision/hand/gesture/guesture.py | 8 +++++--- .../mediapipe_vision/hand/landmark/hand_landmark.py | 10 ++++++---- .../hand/landmark/hand_landmark_masking.py | 6 ++++-- .../hand/landmark/hand_landmark_visualization.py | 5 +++-- .../mediapipe_vision/hand/landmark/position.py | 1 + node_wrappers/mediapipe_vision/holistic/__init__.py | 3 +-- .../mediapipe_vision/holistic/holistic_landmark.py | 4 +++- .../holistic/holistic_landmark_visualization.py | 5 +++-- .../image_embedding/image_embeddings.py | 8 +++++--- .../mediapipe_vision/image_embedding/utils.py | 9 ++++++--- .../image_segmentation/image_segmentation.py | 12 +++++++----- .../image_segmentation/image_segmentation_masking.py | 2 +- .../interactive_segmentation.py | 12 ++++++++---- .../object_detection/object_detection.py | 10 ++++++---- .../object_detection_visualization.py | 3 ++- node_wrappers/mediapipe_vision/pose/pose_landmark.py | 10 ++++++---- .../mediapipe_vision/pose/pose_landmark_masking.py | 4 +++- .../pose/pose_landmark_visualizer.py | 5 +++-- node_wrappers/mediapipe_vision/pose/position.py | 3 ++- .../realtimenodes/controls/mask_controls.py | 4 +++- .../realtimenodes/controls/motion_controls.py | 9 ++++++--- .../realtimenodes/controls/sequence_controls.py | 1 + .../controls/state_management_controls.py | 4 +++- .../realtimenodes/controls/value_controls.py | 1 + node_wrappers/realtimenodes/logic/logic.py | 1 + .../realtimenodes/media/media_pipe_nodes.py | 7 ++++--- .../realtimenodes/utils/utility_controls.py | 5 ++++- node_wrappers/utils/misc_nodes.py | 12 ++++++++---- src/coordinates/coordinate_delta.py | 8 ++++---- src/coordinates/coordinate_system.py | 5 +++-- src/coordinates/drawing_engine.py | 9 +++++---- src/mediapipe_vision/__init__.py | 2 +- src/mediapipe_vision/common/position.py | 7 ++----- src/mediapipe_vision/face_detection/detector.py | 10 +++++++--- src/mediapipe_vision/face_landmark/detector.py | 8 ++++++-- src/mediapipe_vision/face_stylization/detector.py | 2 ++ src/mediapipe_vision/gesture_recognition/detector.py | 10 +++++++--- src/mediapipe_vision/hand_landmark/detector.py | 10 +++++++--- src/mediapipe_vision/holistic_landmark/__init__.py | 2 +- src/mediapipe_vision/holistic_landmark/detector.py | 11 +++++++---- src/mediapipe_vision/image_embedding/detector.py | 10 +++++++--- src/mediapipe_vision/image_segmentation/segmenter.py | 6 ++++-- .../interactive_segmentation/detector.py | 8 ++++---- src/mediapipe_vision/model_loader.py | 9 +++++---- src/mediapipe_vision/object_detection/detector.py | 10 +++++++--- src/mediapipe_vision/pose_landmark/detector.py | 8 ++++++-- src/mediapipe_vision/types.py | 6 ++++-- src/mediapipe_vision/utils/__init__.py | 4 ++-- src/mediapipe_vision/utils/delta_utils.py | 4 +++- src/mediapipe_vision/utils/download_utils.py | 7 ++++--- src/mediapipe_vision/utils/model_utils.py | 2 +- src/realtimenodes/control_base.py | 3 ++- src/realtimenodes/mask_controls.py | 1 + src/realtimenodes/patterns/movement_patterns.py | 4 +++- src/utils/__init__.py | 6 +++++- src/utils/image.py | 8 +++++--- src/utils/math.py | 1 + src/utils/similar_image_filter.py | 6 ++++-- src/utils/timing.py | 1 + 76 files changed, 291 insertions(+), 167 deletions(-) diff --git a/__init__.py b/__init__.py index 2bbc2a8..a9dd5cb 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ -import os import importlib import inspect +import os import re import traceback diff --git a/node_wrappers/coordinates/__init__.py b/node_wrappers/coordinates/__init__.py index 3683012..d832f58 100644 --- a/node_wrappers/coordinates/__init__.py +++ b/node_wrappers/coordinates/__init__.py @@ -3,8 +3,4 @@ Coordinate system node wrappers for ComfyUI. This package provides node wrappers for the coordinate system functionality, allowing users to convert between coordinate spaces and create coordinate objects. -""" - -from .conversion_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file +""" \ No newline at end of file diff --git a/node_wrappers/coordinates/drawing_nodes.py b/node_wrappers/coordinates/drawing_nodes.py index ada1bda..da726a8 100644 --- a/node_wrappers/coordinates/drawing_nodes.py +++ b/node_wrappers/coordinates/drawing_nodes.py @@ -1,9 +1,9 @@ -import torch import logging -from typing import Union, List +from typing import List, Union -from ...src.coordinates import CoordinateSystem -from ...src.coordinates import DrawingEngine +import torch + +from ...src.coordinates import CoordinateSystem, DrawingEngine logger = logging.getLogger(__name__) diff --git a/node_wrappers/coordinates/position_tracking.py b/node_wrappers/coordinates/position_tracking.py index 6d8f9d0..ecaf422 100644 --- a/node_wrappers/coordinates/position_tracking.py +++ b/node_wrappers/coordinates/position_tracking.py @@ -1,15 +1,17 @@ import logging -from typing import List, Tuple, Union import math -import torch -import numpy as np +from typing import List, Tuple, Union + import cv2 +import numpy as np +import torch + +from ...src.coordinates.coordinate_delta import BaseCoordinateDelta #TODO: consider moving base nodes to src - # Import from the new consolidated utilities from ...src.utils.math import scale_value -from ...src.coordinates.coordinate_delta import BaseCoordinateDelta + logger = logging.getLogger(__name__) # --- Position Delta Control Nodes (Int/Float) (Handling Lists) --- diff --git a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py index 7470124..fae1a64 100644 --- a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py +++ b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py @@ -1,5 +1,6 @@ import logging -from typing import List, Any, Optional +from typing import Any, List, Optional + import cv2 from ....src.utils.image import convert_to_cv2, convert_to_tensor diff --git a/node_wrappers/mediapipe_vision/common/model_loader.py b/node_wrappers/mediapipe_vision/common/model_loader.py index 36f3ec0..701ed13 100644 --- a/node_wrappers/mediapipe_vision/common/model_loader.py +++ b/node_wrappers/mediapipe_vision/common/model_loader.py @@ -1,12 +1,15 @@ """Node wrapper for MediaPipe Vision model loader.""" # Import necessary components from the source directory -from ....src.mediapipe_vision.model_loader import MediaPipeModelLoader, get_available_models import logging -import traceback # For better error logging +import traceback # For better error logging # Add import for MEDIAPIPE_MODELS dictionary and the get_model_description function -from ....src.mediapipe_vision.model_loader import get_model_description +from ....src.mediapipe_vision.model_loader import ( + MediaPipeModelLoader, + get_available_models, + get_model_description, +) logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection.py b/node_wrappers/mediapipe_vision/face/detection/face_detection.py index 5ce91ff..e3e15cf 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection.py @@ -1,9 +1,11 @@ -import torch -from .....src.mediapipe_vision.face_detection.detector import FaceDetector -from ...common.model_loader import MediaPipeModelLoaderBaseNode -from ...common.base_detector_node import BaseMediaPipeDetectorNode import logging +import torch + +from .....src.mediapipe_vision.face_detection.detector import FaceDetector +from ...common.base_detector_node import BaseMediaPipeDetectorNode +from ...common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection" # Inherit from Base Loader diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py b/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py index b319d26..c0538f2 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py @@ -1,8 +1,9 @@ # face_detection_mask.py -import torch -import numpy as np import logging + +import numpy as np import scipy.ndimage +import torch logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection" diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py index 0422c24..70fee8a 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py @@ -1,8 +1,10 @@ # visualize_face_detection.py +import logging + import cv2 + from ...common.base_visualization_nodes import BaseDetectionVisualizationNode -import logging logger = logging.getLogger(__name__) class VisualizeFaceDetections(BaseDetectionVisualizationNode): diff --git a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py index 39e6a08..68cb85f 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py +++ b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py @@ -2,11 +2,14 @@ import logging from typing import Optional -from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alias -from .....src.mediapipe_vision.utils.blendshape_utils import get_blendshape_categories # Utility to get all category names - # Import base trigger/control logic if we adapt base classes later, or implement directly for now -from .....src.mediapipe_vision.common.position import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons +from .....src.mediapipe_vision.common.position import ( + FLOAT_EQUALITY_TOLERANCE, # For trigger node comparisons +) +from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alias +from .....src.mediapipe_vision.utils.blendshape_utils import ( + get_blendshape_categories, # Utility to get all category names +) from .....src.utils.math import scale_value logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py index dba20f2..700324d 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py @@ -1,11 +1,14 @@ """Node wrapper for face detection.""" +import logging + import torch + from .....src.mediapipe_vision.face_landmark.detector import FaceLandmarkDetector +from ...common.base_detector_node import BaseMediaPipeDetectorNode + # Import Base Classes from ...common.model_loader import MediaPipeModelLoaderBaseNode -from ...common.base_detector_node import BaseMediaPipeDetectorNode -import logging logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py index 6297f52..cd82436 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py @@ -1,10 +1,11 @@ # visualize_face.py +import logging +from typing import Any, List, Optional + import mediapipe as mp -from typing import Optional, List, Any from ...common.base_visualization_nodes import BaseLandmarkVisualizationNode -import logging logger = logging.getLogger(__name__) # Define standard face connections diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py index 106f2e8..b7c2620 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py @@ -1,10 +1,11 @@ # node_wrappers/face/masking.py -import torch import logging from typing import List +import torch + from .....src.mediapipe_vision.types import LandmarkPoint -from .....src.utils.image import create_mask_from_points +from .....src.utils.image import create_mask_from_points logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Masking" diff --git a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py index 22b09e9..ab3eba9 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py +++ b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py @@ -1,12 +1,15 @@ # node_wrappers/face/head_pose.py import logging +from typing import Dict, Optional + import numpy as np from scipy.spatial.transform import Rotation as R -from typing import Optional, Dict +from .....src.mediapipe_vision.common.position import ( + FLOAT_EQUALITY_TOLERANCE, # For trigger node comparisons +) from .....src.mediapipe_vision.types import TRANSFORM_MATRIX_LIST from .....src.utils.math import scale_value -from .....src.mediapipe_vision.common.position import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/face/landmark/position.py b/node_wrappers/mediapipe_vision/face/landmark/position.py index 01f58d6..97e15ba 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/position.py +++ b/node_wrappers/mediapipe_vision/face/landmark/position.py @@ -1,4 +1,5 @@ import logging + from .....src.mediapipe_vision.common.position import LandmarkPositionBase from .....src.mediapipe_vision.landmark_definitions import FACE_LANDMARK_TOOLTIP diff --git a/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py b/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py index 2f5bfbd..4398456 100644 --- a/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py +++ b/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py @@ -1,10 +1,12 @@ -import torch import logging -# Import Base Loader and Detector -from ...common.model_loader import MediaPipeModelLoaderBaseNode +import torch + from .....src.mediapipe_vision.face_stylization.detector import FaceStylizer +# Import Base Loader and Detector +from ...common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceStylization" # --- Model Loader --- diff --git a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py index 43202e5..b07498d 100644 --- a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py +++ b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py @@ -1,9 +1,10 @@ # visualize_gestures.py +import logging + import cv2 from ...common.base_visualization_nodes import BaseDetectionVisualizationNode -import logging logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/GestureVisualization" diff --git a/node_wrappers/mediapipe_vision/hand/gesture/guesture.py b/node_wrappers/mediapipe_vision/hand/gesture/guesture.py index af1f4d8..97eccd0 100644 --- a/node_wrappers/mediapipe_vision/hand/gesture/guesture.py +++ b/node_wrappers/mediapipe_vision/hand/gesture/guesture.py @@ -1,10 +1,12 @@ -import torch import logging +import torch + +from .....src.mediapipe_vision.gesture_recognition.detector import GestureRecognizer +from ...common.base_detector_node import BaseMediaPipeDetectorNode + # Import Base Loader and Detector from ...common.model_loader import MediaPipeModelLoaderBaseNode -from ...common.base_detector_node import BaseMediaPipeDetectorNode -from .....src.mediapipe_vision.gesture_recognition.detector import GestureRecognizer logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/GestureRecognition" diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py index c0e4bc2..e0552b9 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py @@ -1,9 +1,11 @@ -import torch -from .....src.mediapipe_vision.hand_landmark.detector import HandLandmarkDetector -from ...common.model_loader import MediaPipeModelLoaderBaseNode -from ...common.base_detector_node import BaseMediaPipeDetectorNode import logging +import torch + +from .....src.mediapipe_vision.hand_landmark.detector import HandLandmarkDetector +from ...common.base_detector_node import BaseMediaPipeDetectorNode +from ...common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark" diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py index e0e77e1..cba76ec 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py @@ -1,9 +1,11 @@ # node_wrappers/hand/masking.py -import torch import logging -from typing import List +from typing import List + +import torch from .....src.mediapipe_vision.types import HandLandmarksResult + # Updated import path from .....src.utils.image import create_mask_from_points diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py index d806896..ede6d69 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py @@ -1,10 +1,11 @@ # visualize_hand.py +import logging +from typing import Any, List, Optional + import mediapipe as mp -from typing import Optional, List, Any from ...common.base_visualization_nodes import BaseLandmarkVisualizationNode -import logging logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Visualization" # Define standard hand connections diff --git a/node_wrappers/mediapipe_vision/hand/landmark/position.py b/node_wrappers/mediapipe_vision/hand/landmark/position.py index 9b53858..6bee4b9 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/position.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/position.py @@ -3,6 +3,7 @@ import logging # Import the base class from .....src.mediapipe_vision.common.position import LandmarkPositionBase + # Import the specific landmark type for clarity and potential future use # Import the landmark definitions from .....src.mediapipe_vision.landmark_definitions import HAND_LANDMARK_TOOLTIP diff --git a/node_wrappers/mediapipe_vision/holistic/__init__.py b/node_wrappers/mediapipe_vision/holistic/__init__.py index 99fe05f..82eb3f5 100644 --- a/node_wrappers/mediapipe_vision/holistic/__init__.py +++ b/node_wrappers/mediapipe_vision/holistic/__init__.py @@ -6,5 +6,4 @@ to this is released, this module will be updated to use the new api. """ # Import modules to ensure they're registered -from . import holistic_landmark -from . import holistic_landmark_visualization \ No newline at end of file +from . import holistic_landmark, holistic_landmark_visualization diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py index 7ce863b..731fe86 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py @@ -1,7 +1,9 @@ +import logging + import torch + from ....src.mediapipe_vision.holistic_landmark.detector import HolisticLandmarkDetector from ..common.base_detector_node import BaseMediaPipeDetectorNode -import logging logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Holistic" diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py index 2070b66..757b4e0 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py @@ -1,8 +1,9 @@ import logging +from typing import List, Optional, Tuple + +import cv2 import mediapipe as mp import numpy as np -import cv2 -from typing import List, Tuple, Optional import torch logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py b/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py index d68616c..2e675bc 100644 --- a/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py +++ b/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py @@ -1,10 +1,12 @@ -import torch import logging +import torch + +from ....src.mediapipe_vision.image_embedding.detector import ImageEmbedder +from ..common.base_detector_node import BaseMediaPipeDetectorNode + # Import Base Loader and Detector from ..common.model_loader import MediaPipeModelLoaderBaseNode -from ..common.base_detector_node import BaseMediaPipeDetectorNode -from ....src.mediapipe_vision.image_embedding.detector import ImageEmbedder logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ImageEmbedding" diff --git a/node_wrappers/mediapipe_vision/image_embedding/utils.py b/node_wrappers/mediapipe_vision/image_embedding/utils.py index e8c95b7..5e4e08b 100644 --- a/node_wrappers/mediapipe_vision/image_embedding/utils.py +++ b/node_wrappers/mediapipe_vision/image_embedding/utils.py @@ -1,9 +1,12 @@ -import torch -import torch.nn.functional as F import logging from typing import List, Optional -from ....src.mediapipe_vision.types import ImageEmbedderResult # Assuming this is the relevant type +import torch +import torch.nn.functional as F + +from ....src.mediapipe_vision.types import ( + ImageEmbedderResult, # Assuming this is the relevant type +) logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py index 4623ede..c5e8090 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py @@ -1,16 +1,18 @@ """Node wrapper for image segmentation""" -import torch -import numpy as np -import cv2 # For visualization import logging +import cv2 # For visualization +import numpy as np +import torch + # Imports from this project from ....src.mediapipe_vision.image_segmentation.segmenter import ImageSegmenter -# Import Base Classes -from ..common.model_loader import MediaPipeModelLoaderBaseNode from ..common.base_detector_node import BaseMediaPipeDetectorNode +# Import Base Classes +from ..common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ImageSegmentation" diff --git a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py index e4ee1e2..dfa9dc7 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py @@ -1,6 +1,6 @@ -import torch import logging +import torch logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ImageSegmentation/Masking" diff --git a/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py b/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py index ddb5400..93938c3 100644 --- a/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py +++ b/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py @@ -1,11 +1,15 @@ -import torch -import logging import json +import logging + +import torch + +from ....src.mediapipe_vision.interactive_segmentation.detector import ( + InteractiveSegmenterProcessor, +) +from ....src.mediapipe_vision. types import PointOfInterest # Import Base Loader and Detector from ..common.model_loader import MediaPipeModelLoaderBaseNode -from ....src.mediapipe_vision.interactive_segmentation.detector import InteractiveSegmenterProcessor -from ....src.mediapipe_vision. types import PointOfInterest logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/InteractiveSegmentation" diff --git a/node_wrappers/mediapipe_vision/object_detection/object_detection.py b/node_wrappers/mediapipe_vision/object_detection/object_detection.py index 3c19131..ea25a92 100644 --- a/node_wrappers/mediapipe_vision/object_detection/object_detection.py +++ b/node_wrappers/mediapipe_vision/object_detection/object_detection.py @@ -1,9 +1,11 @@ -import torch -from ....src.mediapipe_vision.object_detection.detector import ObjectDetector -from ..common.model_loader import MediaPipeModelLoaderBaseNode -from ..common.base_detector_node import BaseMediaPipeDetectorNode import logging +import torch + +from ....src.mediapipe_vision.object_detection.detector import ObjectDetector +from ..common.base_detector_node import BaseMediaPipeDetectorNode +from ..common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ObjectDetection" diff --git a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py index 86f7a07..0b50833 100644 --- a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py @@ -1,9 +1,10 @@ # visualize_object_detection.py +import logging + import cv2 from ..common.base_visualization_nodes import BaseDetectionVisualizationNode -import logging logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ObjectDetection/Visualization" diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark.py b/node_wrappers/mediapipe_vision/pose/pose_landmark.py index a45d708..40397f6 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark.py @@ -1,9 +1,11 @@ -import torch -from ....src.mediapipe_vision.pose_landmark.detector import PoseLandmarkDetector -from ..common.model_loader import MediaPipeModelLoaderBaseNode -from ..common.base_detector_node import BaseMediaPipeDetectorNode import logging +import torch + +from ....src.mediapipe_vision.pose_landmark.detector import PoseLandmarkDetector +from ..common.base_detector_node import BaseMediaPipeDetectorNode +from ..common.model_loader import MediaPipeModelLoaderBaseNode + logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Pose" diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py index edd2660..f5a308f 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py @@ -1,10 +1,12 @@ # node_wrappers/pose/masking.py -import torch import logging from typing import List +import torch + from ....src.mediapipe_vision.types import PoseLandmarksResult from ....src.utils.image import create_mask_from_points + # Removed GetImageDimensionsNode import logger = logging.getLogger(__name__) diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py index 964ed0f..630a7e8 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py @@ -1,10 +1,11 @@ # visualize_pose.py +import logging +from typing import Any, List, Optional + import mediapipe as mp -from typing import Optional, List, Any from ..common.base_visualization_nodes import BaseLandmarkVisualizationNode -import logging logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Pose/Visualization" # Define standard pose connections diff --git a/node_wrappers/mediapipe_vision/pose/position.py b/node_wrappers/mediapipe_vision/pose/position.py index 042517b..589539d 100644 --- a/node_wrappers/mediapipe_vision/pose/position.py +++ b/node_wrappers/mediapipe_vision/pose/position.py @@ -1,7 +1,8 @@ import logging # Import the base class -from ....src.mediapipe_vision.common.position import LandmarkPositionBase +from ....src.mediapipe_vision.common.position import LandmarkPositionBase + # Import the specific landmark type for clarity and potential future use # Import the landmark definitions from ....src.mediapipe_vision.landmark_definitions import POSE_LANDMARK_TOOLTIP diff --git a/node_wrappers/realtimenodes/controls/mask_controls.py b/node_wrappers/realtimenodes/controls/mask_controls.py index df44e3c..c35ef38 100644 --- a/node_wrappers/realtimenodes/controls/mask_controls.py +++ b/node_wrappers/realtimenodes/controls/mask_controls.py @@ -1,8 +1,10 @@ -import torch import numpy as np +import torch + from ....src.realtimenodes.control_base import ControlNodeBase from ....src.realtimenodes.mask_controls import MaskControlMixin + class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin): """Node that maintains a mask that repulses or attracts to input masks""" @classmethod diff --git a/node_wrappers/realtimenodes/controls/motion_controls.py b/node_wrappers/realtimenodes/controls/motion_controls.py index f3d62a0..b340329 100644 --- a/node_wrappers/realtimenodes/controls/motion_controls.py +++ b/node_wrappers/realtimenodes/controls/motion_controls.py @@ -1,9 +1,12 @@ -from ....src.realtimenodes.control_base import ControlNodeBase +import time from enum import Enum + +import cv2 import numpy as np import torch -import cv2 -import time + +from ....src.realtimenodes.control_base import ControlNodeBase + class ROIAction(Enum): # Behavioral actions diff --git a/node_wrappers/realtimenodes/controls/sequence_controls.py b/node_wrappers/realtimenodes/controls/sequence_controls.py index 37a0cae..64f3c95 100644 --- a/node_wrappers/realtimenodes/controls/sequence_controls.py +++ b/node_wrappers/realtimenodes/controls/sequence_controls.py @@ -1,5 +1,6 @@ from ....src.realtimenodes.control_base import ControlNodeBase + class SequenceControlBase(ControlNodeBase): """Base class for sequence-based controls""" diff --git a/node_wrappers/realtimenodes/controls/state_management_controls.py b/node_wrappers/realtimenodes/controls/state_management_controls.py index 09c0d82..bad46d6 100644 --- a/node_wrappers/realtimenodes/controls/state_management_controls.py +++ b/node_wrappers/realtimenodes/controls/state_management_controls.py @@ -1,7 +1,9 @@ -from ....src.realtimenodes.control_base import ControlNodeBase import copy + +from ....src.realtimenodes.control_base import ControlNodeBase from ....src.utils.general import AlwaysEqualProxy + class StateResetNode(ControlNodeBase): """Node that resets all control node states when triggered""" diff --git a/node_wrappers/realtimenodes/controls/value_controls.py b/node_wrappers/realtimenodes/controls/value_controls.py index 99f6aea..90950c7 100644 --- a/node_wrappers/realtimenodes/controls/value_controls.py +++ b/node_wrappers/realtimenodes/controls/value_controls.py @@ -1,6 +1,7 @@ from ....src.realtimenodes.control_base import ControlNodeBase from ....src.realtimenodes.patterns.movement_patterns import MOVEMENT_PATTERNS + class ValueControlBase(ControlNodeBase): """Base class for float and integer control nodes""" diff --git a/node_wrappers/realtimenodes/logic/logic.py b/node_wrappers/realtimenodes/logic/logic.py index 7fcb23e..2f56b42 100644 --- a/node_wrappers/realtimenodes/logic/logic.py +++ b/node_wrappers/realtimenodes/logic/logic.py @@ -1,6 +1,7 @@ from ....src.realtimenodes.control_base import ControlNodeBase from ....src.utils.general import AlwaysEqualProxy + class LazyCondition(ControlNodeBase): DESCRIPTION = "Uses lazy evaluation to truly skip execution of unused paths. Maintains state of the last value to circumvent feedback loops." diff --git a/node_wrappers/realtimenodes/media/media_pipe_nodes.py b/node_wrappers/realtimenodes/media/media_pipe_nodes.py index fc8c489..a860d75 100644 --- a/node_wrappers/realtimenodes/media/media_pipe_nodes.py +++ b/node_wrappers/realtimenodes/media/media_pipe_nodes.py @@ -1,7 +1,8 @@ -import torch -import numpy as np -import mediapipe as mp import cv2 +import mediapipe as mp +import numpy as np +import torch + class HandTrackingNode: @classmethod diff --git a/node_wrappers/realtimenodes/utils/utility_controls.py b/node_wrappers/realtimenodes/utils/utility_controls.py index be203a9..6bab29a 100644 --- a/node_wrappers/realtimenodes/utils/utility_controls.py +++ b/node_wrappers/realtimenodes/utils/utility_controls.py @@ -1,12 +1,15 @@ import time -import numpy as np + import cv2 +import numpy as np import torch + # # Add package root to Python path # <-- REMOVED # sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from ....src.realtimenodes.control_base import ControlNodeBase from ....src.utils.similar_image_filter import SimilarImageFilter + class FPSMonitor(ControlNodeBase): """Generates an FPS overlay as an image and mask""" diff --git a/node_wrappers/utils/misc_nodes.py b/node_wrappers/utils/misc_nodes.py index 21f7dc3..c9a60a3 100644 --- a/node_wrappers/utils/misc_nodes.py +++ b/node_wrappers/utils/misc_nodes.py @@ -1,12 +1,16 @@ -import torch -import numpy as np import base64 -import re import math +import re from io import BytesIO -from PIL import Image, ImageFont, ImageDraw, ImageOps + +import numpy as np +import torch +from PIL import Image, ImageDraw, ImageFont, ImageOps + import nodes + from ...src.utils.general import AlwaysEqualProxy + MAX_RESOLUTION = nodes.MAX_RESOLUTION # Get the same max resolution as core nodes class DTypeConverter: diff --git a/src/coordinates/coordinate_delta.py b/src/coordinates/coordinate_delta.py index f0768e5..c306819 100644 --- a/src/coordinates/coordinate_delta.py +++ b/src/coordinates/coordinate_delta.py @@ -1,8 +1,8 @@ -import logging -from typing import Optional -import math import collections -import typing # Add typing for Dict hint +import logging +import math +import typing # Add typing for Dict hint +from typing import Optional #TODO: consider moving base nodes to src diff --git a/src/coordinates/coordinate_system.py b/src/coordinates/coordinate_system.py index 6374939..a21120b 100644 --- a/src/coordinates/coordinate_system.py +++ b/src/coordinates/coordinate_system.py @@ -5,9 +5,10 @@ This module provides a centralized system for handling coordinates across differ spaces (normalized, pixel) with support for both single values and batches. """ -import torch -from typing import Union, List, Tuple, Optional import logging +from typing import List, Optional, Tuple, Union + +import torch logger = logging.getLogger(__name__) diff --git a/src/coordinates/drawing_engine.py b/src/coordinates/drawing_engine.py index bd9fdad..39ed0d4 100644 --- a/src/coordinates/drawing_engine.py +++ b/src/coordinates/drawing_engine.py @@ -1,8 +1,9 @@ -import torch -import numpy as np -import cv2 -from typing import Union, List, Tuple import logging +from typing import List, Tuple, Union + +import cv2 +import numpy as np +import torch from .coordinate_system import CoordinateSystem diff --git a/src/mediapipe_vision/__init__.py b/src/mediapipe_vision/__init__.py index d96faea..556ba3f 100644 --- a/src/mediapipe_vision/__init__.py +++ b/src/mediapipe_vision/__init__.py @@ -1,4 +1,4 @@ """Source code for MediaPipe Vision tasks implementation.""" # Import utility modules -from . import utils \ No newline at end of file +from . import utils diff --git a/src/mediapipe_vision/common/position.py b/src/mediapipe_vision/common/position.py index f8f831e..323811c 100644 --- a/src/mediapipe_vision/common/position.py +++ b/src/mediapipe_vision/common/position.py @@ -1,12 +1,9 @@ import logging -from typing import Optional, List, Tuple +from typing import List, Optional, Tuple #TODO: consider moving base nodes to src - # Assuming types are correctly imported relative to this file's location -from ..types import ( - LandmarkPoint -) +from ..types import LandmarkPoint # Import from the new consolidated utilities diff --git a/src/mediapipe_vision/face_detection/detector.py b/src/mediapipe_vision/face_detection/detector.py index 44f73f4..75272ce 100644 --- a/src/mediapipe_vision/face_detection/detector.py +++ b/src/mediapipe_vision/face_detection/detector.py @@ -3,16 +3,20 @@ This module contains the implementation of MediaPipe Face Detection functionality. """ +from typing import List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List + +from ...utils.timing import TimestampProvider # Import from utils + # Import new types -from ..types import FaceDetectionResult, BoundingBox, FaceKeypoint -from ...utils.timing import TimestampProvider # Import from utils +from ..types import BoundingBox, FaceDetectionResult, FaceKeypoint + class FaceDetector: """Detects faces in an image using MediaPipe FaceDetector.""" diff --git a/src/mediapipe_vision/face_landmark/detector.py b/src/mediapipe_vision/face_landmark/detector.py index ca9d165..f04601b 100644 --- a/src/mediapipe_vision/face_landmark/detector.py +++ b/src/mediapipe_vision/face_landmark/detector.py @@ -1,13 +1,17 @@ +from typing import Dict, List, Optional, Tuple + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List, Tuple, Optional, Dict + +from ...utils.timing import TimestampProvider # Import TimestampProvider + # Import new types from ..types import LandmarkPoint -from ...utils.timing import TimestampProvider # Import TimestampProvider + class FaceLandmarkDetector: """Detects face landmarks and blendshapes using MediaPipe FaceLandmarker.""" diff --git a/src/mediapipe_vision/face_stylization/detector.py b/src/mediapipe_vision/face_stylization/detector.py index cb23d71..03fdfaa 100644 --- a/src/mediapipe_vision/face_stylization/detector.py +++ b/src/mediapipe_vision/face_stylization/detector.py @@ -3,8 +3,10 @@ import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision + from ...utils.timing import TimestampProvider + class FaceStylizer: """Applies stylization to faces using MediaPipe FaceStylizer.""" diff --git a/src/mediapipe_vision/gesture_recognition/detector.py b/src/mediapipe_vision/gesture_recognition/detector.py index 565274a..3257720 100644 --- a/src/mediapipe_vision/gesture_recognition/detector.py +++ b/src/mediapipe_vision/gesture_recognition/detector.py @@ -1,13 +1,17 @@ +from typing import List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List + +from ...utils.timing import TimestampProvider # Import TimestampProvider + # Import new types -from ..types import GestureRecognitionResult, GestureCategory -from ...utils.timing import TimestampProvider # Import TimestampProvider +from ..types import GestureCategory, GestureRecognitionResult + class GestureRecognizer: """Recognizes hand gestures using MediaPipe GestureRecognizer.""" diff --git a/src/mediapipe_vision/hand_landmark/detector.py b/src/mediapipe_vision/hand_landmark/detector.py index 613b3c2..d577dec 100644 --- a/src/mediapipe_vision/hand_landmark/detector.py +++ b/src/mediapipe_vision/hand_landmark/detector.py @@ -1,13 +1,17 @@ +from typing import List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List + +from ...utils.timing import TimestampProvider # Import TimestampProvider + # Import new types -from ..types import LandmarkPoint, HandLandmarksResult -from ...utils.timing import TimestampProvider # Import TimestampProvider +from ..types import HandLandmarksResult, LandmarkPoint + class HandLandmarkDetector: """Detects hand landmarks using MediaPipe HandLandmarker.""" diff --git a/src/mediapipe_vision/holistic_landmark/__init__.py b/src/mediapipe_vision/holistic_landmark/__init__.py index 12faee9..d306730 100644 --- a/src/mediapipe_vision/holistic_landmark/__init__.py +++ b/src/mediapipe_vision/holistic_landmark/__init__.py @@ -1,3 +1,3 @@ """Holistic landmark detection module using MediaPipe Legacy API.""" -from .detector import HolisticLandmarkDetector \ No newline at end of file +from .detector import HolisticLandmarkDetector diff --git a/src/mediapipe_vision/holistic_landmark/detector.py b/src/mediapipe_vision/holistic_landmark/detector.py index 5d5bbfb..8f7176f 100644 --- a/src/mediapipe_vision/holistic_landmark/detector.py +++ b/src/mediapipe_vision/holistic_landmark/detector.py @@ -4,15 +4,18 @@ This implementation specifically handles the MediaPipe Holistic solution which is currently only available through the legacy API, not the Task API. """ +from typing import Any, List, Optional, Tuple + +import cv2 import mediapipe as mp import numpy as np import torch -import cv2 -from typing import List, Tuple, Optional, Any + +from ...utils.timing import TimestampProvider # Import needed types -from ..types import LandmarkPoint, HolisticLandmarksResult -from ...utils.timing import TimestampProvider +from ..types import HolisticLandmarksResult, LandmarkPoint + class HolisticLandmarkDetector: """Detects holistic landmarks combining face, pose and hand tracking. diff --git a/src/mediapipe_vision/image_embedding/detector.py b/src/mediapipe_vision/image_embedding/detector.py index eb7152c..721e34c 100644 --- a/src/mediapipe_vision/image_embedding/detector.py +++ b/src/mediapipe_vision/image_embedding/detector.py @@ -1,15 +1,19 @@ +from typing import List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List -# Import new types -from ..types import ImageEmbedderResult + # Import TimestampProvider from the new location from ...utils.timing import TimestampProvider +# Import new types +from ..types import ImageEmbedderResult + + class ImageEmbedder: """Generates image embeddings using MediaPipe ImageEmbedder.""" diff --git a/src/mediapipe_vision/image_segmentation/segmenter.py b/src/mediapipe_vision/image_segmentation/segmenter.py index f78b91c..8b8e318 100644 --- a/src/mediapipe_vision/image_segmentation/segmenter.py +++ b/src/mediapipe_vision/image_segmentation/segmenter.py @@ -3,14 +3,16 @@ This module contains the implementation of MediaPipe Image Segmentation functionality. """ +import logging +import platform # Needed for OS check + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.vision import ImageSegmenter -import platform # Needed for OS check -import logging + from ...utils.timing import TimestampProvider logger = logging.getLogger(__name__) diff --git a/src/mediapipe_vision/interactive_segmentation/detector.py b/src/mediapipe_vision/interactive_segmentation/detector.py index 1ce0f1a..cab2ec8 100644 --- a/src/mediapipe_vision/interactive_segmentation/detector.py +++ b/src/mediapipe_vision/interactive_segmentation/detector.py @@ -1,16 +1,16 @@ import logging +from typing import List, Optional + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python -from mediapipe.tasks.python.vision import interactive_segmenter from mediapipe.tasks.python.components.containers import keypoint +from mediapipe.tasks.python.vision import interactive_segmenter -from typing import List, Optional -from ..types import PointOfInterest from ...utils.timing import TimestampProvider +from ..types import PointOfInterest -import logging logger = logging.getLogger(__name__) NormalizedKeypoint = keypoint.NormalizedKeypoint diff --git a/src/mediapipe_vision/model_loader.py b/src/mediapipe_vision/model_loader.py index de8b42d..c1cb9cb 100644 --- a/src/mediapipe_vision/model_loader.py +++ b/src/mediapipe_vision/model_loader.py @@ -3,12 +3,13 @@ Handles model definition, download, and path retrieval. """ -import os import logging -import urllib.request -import folder_paths # ComfyUI path helper -from typing import Dict, Any, Optional, List +import os import traceback +import urllib.request +from typing import Any, Dict, List, Optional + +import folder_paths # ComfyUI path helper # Configure logging # logging.basicConfig(level=logging.INFO) # Avoid reconfiguring root logger diff --git a/src/mediapipe_vision/object_detection/detector.py b/src/mediapipe_vision/object_detection/detector.py index a4ab82b..b822ffc 100644 --- a/src/mediapipe_vision/object_detection/detector.py +++ b/src/mediapipe_vision/object_detection/detector.py @@ -1,14 +1,18 @@ +from typing import List + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List -# Import new types -from ..types import ObjectDetectionResult, ObjectDetectionCategory, BoundingBox + from ...utils.timing import TimestampProvider +# Import new types +from ..types import BoundingBox, ObjectDetectionCategory, ObjectDetectionResult + + class ObjectDetector: """Detects objects using MediaPipe ObjectDetector.""" diff --git a/src/mediapipe_vision/pose_landmark/detector.py b/src/mediapipe_vision/pose_landmark/detector.py index b4e3c5e..1eb5c53 100644 --- a/src/mediapipe_vision/pose_landmark/detector.py +++ b/src/mediapipe_vision/pose_landmark/detector.py @@ -1,13 +1,17 @@ +from typing import List, Optional, Tuple + import mediapipe as mp import numpy as np import torch from mediapipe.tasks import python from mediapipe.tasks.python import vision from mediapipe.tasks.python.core.base_options import BaseOptions -from typing import List, Tuple, Optional + +from ...utils.timing import TimestampProvider + # Import new types from ..types import LandmarkPoint, PoseLandmarksResult -from ...utils.timing import TimestampProvider + class PoseLandmarkDetector: """Detects pose landmarks in an image using MediaPipe PoseLandmarker.""" diff --git a/src/mediapipe_vision/types.py b/src/mediapipe_vision/types.py index f3d2147..171a629 100644 --- a/src/mediapipe_vision/types.py +++ b/src/mediapipe_vision/types.py @@ -1,5 +1,7 @@ -from typing import List, Optional, NamedTuple -import numpy as np # Added numpy for matrix type +from typing import List, NamedTuple, Optional + +import numpy as np # Added numpy for matrix type + # Define a structure for a single landmark point class LandmarkPoint(NamedTuple): diff --git a/src/mediapipe_vision/utils/__init__.py b/src/mediapipe_vision/utils/__init__.py index 5505374..82af41d 100644 --- a/src/mediapipe_vision/utils/__init__.py +++ b/src/mediapipe_vision/utils/__init__.py @@ -3,12 +3,12 @@ # DEPRECATED: These imports are maintained only for backward compatibility. # For new code, please import directly from src.utils -from ...utils.math import scale_value from ...utils.image import convert_to_cv2, convert_to_tensor +from ...utils.math import scale_value +from .delta_utils import calculate_euclidean_delta # Keep only the functions specific to mediapipe_vision from .download_utils import download_model -from .delta_utils import calculate_euclidean_delta __all__ = [ "convert_to_cv2", diff --git a/src/mediapipe_vision/utils/delta_utils.py b/src/mediapipe_vision/utils/delta_utils.py index eea8f36..c3d03b0 100644 --- a/src/mediapipe_vision/utils/delta_utils.py +++ b/src/mediapipe_vision/utils/delta_utils.py @@ -8,7 +8,9 @@ import math from typing import Optional -from ..types import LandmarkPoint # Assuming LandmarkPoint is in src/types.py +from ..types import LandmarkPoint # Assuming LandmarkPoint is in src/types.py + + #TODO: refactor to common after unified coordinate system def calculate_euclidean_delta(p1: Optional[LandmarkPoint], p2: Optional[LandmarkPoint]) -> Optional[float]: """Calculates the Euclidean distance between two LandmarkPoint objects.""" diff --git a/src/mediapipe_vision/utils/download_utils.py b/src/mediapipe_vision/utils/download_utils.py index 4a87889..d65c19f 100644 --- a/src/mediapipe_vision/utils/download_utils.py +++ b/src/mediapipe_vision/utils/download_utils.py @@ -1,10 +1,11 @@ """Utilities for downloading MediaPipe models.""" -import os -import requests import logging +import os +from typing import Any, Dict, Optional, Tuple + +import requests from tqdm import tqdm -from typing import Dict, Optional, Any, Tuple from .model_utils import MEDIAPIPE_MODELS, get_model_path diff --git a/src/mediapipe_vision/utils/model_utils.py b/src/mediapipe_vision/utils/model_utils.py index c29bc46..dd6fa4a 100644 --- a/src/mediapipe_vision/utils/model_utils.py +++ b/src/mediapipe_vision/utils/model_utils.py @@ -1,7 +1,7 @@ """Utilities for managing MediaPipe Vision models.""" -import os import logging +import os # Configure logging logging.basicConfig(level=logging.INFO) diff --git a/src/realtimenodes/control_base.py b/src/realtimenodes/control_base.py index 9af10e3..9e29c82 100644 --- a/src/realtimenodes/control_base.py +++ b/src/realtimenodes/control_base.py @@ -1,5 +1,6 @@ -from abc import ABC, abstractmethod import time +from abc import ABC, abstractmethod + class Pattern(ABC): """Base class for all patterns (movement, timing, etc)""" diff --git a/src/realtimenodes/mask_controls.py b/src/realtimenodes/mask_controls.py index e7bb791..e03c579 100644 --- a/src/realtimenodes/mask_controls.py +++ b/src/realtimenodes/mask_controls.py @@ -1,5 +1,6 @@ import numpy as np + class MaskControlMixin: """Mixin providing common mask functionality""" def create_circle_mask(self, height, width, center_y, center_x, size): diff --git a/src/realtimenodes/patterns/movement_patterns.py b/src/realtimenodes/patterns/movement_patterns.py index 96e9abf..0444053 100644 --- a/src/realtimenodes/patterns/movement_patterns.py +++ b/src/realtimenodes/patterns/movement_patterns.py @@ -1,7 +1,9 @@ -from ..control_base import Pattern import math import random +from ..control_base import Pattern + + class StaticPattern(Pattern): def get_name(self) -> str: return "static" diff --git a/src/utils/__init__.py b/src/utils/__init__.py index ed501ed..4f93b20 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -9,7 +9,11 @@ from .general import AlwaysEqualProxy from .image import convert_to_cv2, convert_to_tensor, create_mask_from_points # Math utilities -from .math import scale_value, calculate_euclidean_distance, calculate_euclidean_distance_2d +from .math import ( + calculate_euclidean_distance, + calculate_euclidean_distance_2d, + scale_value, +) # Timing utilities from .timing import TimestampProvider diff --git a/src/utils/image.py b/src/utils/image.py index 0c4b815..90d7f24 100644 --- a/src/utils/image.py +++ b/src/utils/image.py @@ -5,11 +5,13 @@ Contains functions for converting between different image formats, particularly between ComfyUI tensor format (BHWC) and OpenCV format. """ -import torch -import numpy as np -import cv2 from typing import List, Tuple +import cv2 +import numpy as np +import torch + + def convert_to_cv2(tensor: torch.Tensor) -> list: """Converts a ComfyUI IMAGE tensor (BHWC, float32, 0-1) to a list of cv2 images (BGR, uint8).""" if tensor.ndim != 4 or tensor.shape[3] != 3: diff --git a/src/utils/math.py b/src/utils/math.py index debec70..60bffb2 100644 --- a/src/utils/math.py +++ b/src/utils/math.py @@ -10,6 +10,7 @@ Contains functions for mathematical operations including: import math from typing import Optional, Tuple + def scale_value(value: Optional[float], input_min: float, input_max: float, output_min: float, output_max: float, diff --git a/src/utils/similar_image_filter.py b/src/utils/similar_image_filter.py index 66a4481..2d9bb25 100644 --- a/src/utils/similar_image_filter.py +++ b/src/utils/similar_image_filter.py @@ -1,6 +1,8 @@ -import torch -from typing import Optional import random +from typing import Optional + +import torch + class SimilarImageFilter: def __init__(self, threshold: float = 0.98, max_skip_frame: float = 10) -> None: diff --git a/src/utils/timing.py b/src/utils/timing.py index 3c6bcf9..2460e69 100644 --- a/src/utils/timing.py +++ b/src/utils/timing.py @@ -4,6 +4,7 @@ Timing and timestamp utilities for RealTimeNodes. import time + class TimestampProvider: """ Provides real-time, monotonic timestamps in milliseconds. From c0068cff2ac0bc6a3b96ed54ea30da0cf01b2036 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 18:12:11 -0400 Subject: [PATCH 19/20] ruff format --- __init__.py | 58 +- node_wrappers/__init__.py | 2 +- node_wrappers/coordinates/__init__.py | 2 +- node_wrappers/coordinates/conversion_nodes.py | 61 +- node_wrappers/coordinates/drawing_nodes.py | 278 ++-- .../coordinates/position_tracking.py | 254 ++-- .../common/base_detector_node.py | 83 +- .../common/base_visualization_nodes.py | 209 +-- .../mediapipe_vision/common/model_loader.py | 60 +- .../mediapipe_vision/face/__init__.py | 1 - .../face/detection/face_detection.py | 51 +- .../face/detection/face_detection_mask.py | 170 ++- .../detection/face_detection_visualization.py | 118 +- .../face/landmark/blendshape.py | 175 ++- .../face/landmark/face_landmark_detection.py | 117 +- .../face_landmark_detection_visualization.py | 14 +- .../face/landmark/face_masking.py | 99 +- .../face/landmark/head_pose.py | 205 ++- .../face/landmark/position.py | 16 +- .../face/landmark/texture_warp.py | 63 +- .../face/stylization/face_stylization.py | 57 +- .../mediapipe_vision/hand/__init__.py | 1 - .../hand/gesture/gesture_visualiztion.py | 114 +- .../mediapipe_vision/hand/gesture/guesture.py | 98 +- .../hand/landmark/hand_landmark.py | 93 +- .../hand/landmark/hand_landmark_masking.py | 71 +- .../landmark/hand_landmark_visualization.py | 13 +- .../hand/landmark/position.py | 17 +- .../holistic/holistic_landmark.py | 106 +- .../holistic_landmark_visualization.py | 177 +-- .../image_embedding/__init__.py | 1 - .../image_embedding/image_embeddings.py | 62 +- .../mediapipe_vision/image_embedding/utils.py | 64 +- .../image_segmentation/__init__.py | 1 - .../image_segmentation/image_segmentation.py | 204 +-- .../image_segmentation_masking.py | 48 +- .../interactive_segmentation/__init__.py | 1 - .../interactive_segmentation.py | 110 +- .../object_detection/__init__.py | 1 - .../object_detection/object_detection.py | 59 +- .../object_detection_visualization.py | 87 +- .../mediapipe_vision/pose/__init__.py | 1 - .../mediapipe_vision/pose/pose_landmark.py | 106 +- .../pose/pose_landmark_masking.py | 86 +- .../pose/pose_landmark_visualizer.py | 15 +- .../mediapipe_vision/pose/position.py | 17 +- node_wrappers/realtimenodes/__init__.py | 2 +- .../realtimenodes/controls/__init__.py | 2 +- .../realtimenodes/controls/mask_controls.py | 233 ++-- .../realtimenodes/controls/motion_controls.py | 216 ++-- .../controls/sequence_controls.py | 104 +- .../controls/state_management_controls.py | 106 +- .../realtimenodes/controls/value_controls.py | 240 ++-- node_wrappers/realtimenodes/logic/logic.py | 105 +- node_wrappers/realtimenodes/media/__init__.py | 2 +- .../realtimenodes/media/media_pipe_nodes.py | 127 +- node_wrappers/realtimenodes/utils/__init__.py | 2 +- .../realtimenodes/utils/utility_controls.py | 185 ++- node_wrappers/utils/misc_nodes.py | 1117 +++++++++-------- pyproject.toml | 3 + src/__init__.py | 2 +- src/coordinates/__init__.py | 2 +- src/coordinates/coordinate_delta.py | 27 +- src/coordinates/coordinate_system.py | 207 +-- src/coordinates/drawing_engine.py | 244 ++-- src/mediapipe_vision/common/position.py | 74 +- .../face_detection/__init__.py | 2 +- .../face_detection/detector.py | 73 +- .../face_landmark/__init__.py | 2 +- .../face_landmark/detector.py | 125 +- .../face_stylization/__init__.py | 2 +- .../face_stylization/detector.py | 26 +- .../gesture_recognition/__init__.py | 2 +- .../gesture_recognition/detector.py | 103 +- .../hand_landmark/__init__.py | 2 +- .../hand_landmark/detector.py | 121 +- .../holistic_landmark/detector.py | 147 +-- .../image_classification/__init__.py | 2 +- .../image_embedding/__init__.py | 2 +- .../image_embedding/detector.py | 60 +- .../image_segmentation/__init__.py | 2 +- .../image_segmentation/segmenter.py | 99 +- .../interactive_segmentation/__init__.py | 1 - .../interactive_segmentation/detector.py | 43 +- src/mediapipe_vision/landmark_definitions.py | 9 +- src/mediapipe_vision/model_loader.py | 194 +-- .../object_detection/__init__.py | 2 +- .../object_detection/detector.py | 77 +- .../pose_landmark/__init__.py | 2 +- .../pose_landmark/detector.py | 155 ++- src/mediapipe_vision/types.py | 102 +- src/mediapipe_vision/utils/__init__.py | 10 +- .../utils/blendshape_utils.py | 71 +- src/mediapipe_vision/utils/delta_utils.py | 11 +- src/mediapipe_vision/utils/download_utils.py | 62 +- src/mediapipe_vision/utils/model_utils.py | 59 +- src/realtimenodes/__init__.py | 2 +- src/realtimenodes/control_base.py | 40 +- src/realtimenodes/mask_controls.py | 12 +- .../patterns/movement_patterns.py | 46 +- src/utils/__init__.py | 2 +- src/utils/general.py | 6 +- src/utils/image.py | 26 +- src/utils/math.py | 39 +- src/utils/similar_image_filter.py | 4 +- src/utils/timing.py | 3 +- 106 files changed, 4914 insertions(+), 3480 deletions(-) diff --git a/__init__.py b/__init__.py index a9dd5cb..b783a0d 100644 --- a/__init__.py +++ b/__init__.py @@ -11,14 +11,15 @@ NODE_DISPLAY_NAME_MAPPINGS = {} print("\033[93m" + "[ComfyUI_RealtimeNodes] Loading nodes...") + def load_nodes_from_directory(package_path): """Dynamically loads nodes from Python files in a directory.""" - loaded_nodes = {} + loaded_nodes = {} loaded_display_names = {} - + # Calculate absolute path to the package directory - package_abs_path = os.path.join(os.path.dirname(__file__), *package_path.split('.')) - + package_abs_path = os.path.join(os.path.dirname(__file__), *package_path.split(".")) + if not os.path.isdir(package_abs_path): print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Warning: Directory not found: {package_abs_path}") return loaded_nodes, loaded_display_names @@ -26,59 +27,63 @@ def load_nodes_from_directory(package_path): print(f"\033[92m" + f"[ComfyUI_RealtimeNodes] Searching for nodes in: {package_abs_path}") for filename in os.listdir(package_abs_path): - if filename.endswith(".py") and not filename.startswith('__'): + if filename.endswith(".py") and not filename.startswith("__"): module_name = filename[:-3] - full_module_path = f".{package_path}.{module_name}" # Relative import path - + full_module_path = f".{package_path}.{module_name}" # Relative import path + try: module = importlib.import_module(full_module_path, package=__name__) - + # Find classes defined in the module for name, obj in inspect.getmembers(module): if inspect.isclass(obj): # Check if it's a ComfyUI node (heuristic: has INPUT_TYPES) - if hasattr(obj, 'INPUT_TYPES') and callable(obj.INPUT_TYPES): - # Exclude base classes/mixins if they are explicitly named or lack a CATEGORY - if not name.endswith("Base") and not name.endswith("Mixin") and hasattr(obj, 'CATEGORY'): + if hasattr(obj, "INPUT_TYPES") and callable(obj.INPUT_TYPES): + # Exclude base classes/mixins if they are explicitly named or lack a CATEGORY + if not name.endswith("Base") and not name.endswith("Mixin") and hasattr(obj, "CATEGORY"): loaded_nodes[name] = obj - + # Generate display name (similar to original logic) - display_name = ' '.join(word.capitalize() for word in re.findall(r'[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+', name)) + display_name = " ".join( + word.capitalize() for word in re.findall(r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+", name) + ) suffix = " 🕒🅡🅣🅝" if not display_name.endswith(suffix): display_name += suffix loaded_display_names[name] = display_name print(f"\033[94m" + f"[ComfyUI_RealtimeNodes] - Loaded node: {name} -> {display_name}") - + except ImportError as e: print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error importing module {full_module_path}: {e}") traceback.print_exc() except Exception as e: print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error processing module {full_module_path}: {e}") traceback.print_exc() - + return loaded_nodes, loaded_display_names + # --- Main Loading Logic --- # Get base path for the node_wrappers directory base_dir_path = os.path.join(os.path.dirname(__file__), WRAPPER_DIR) + # Recursively scan all directories and subdirectories def scan_for_modules(base_path, base_package): """Recursively scan directories for module files""" try: for item in os.listdir(base_path): item_path = os.path.join(base_path, item) - + # Skip __pycache__ and files/dirs starting with __ - if item.startswith('__'): + if item.startswith("__"): continue - + if os.path.isdir(item_path): # Handle subdirectory subdir_package = f"{base_package}.{item}" - + # First load .py files in this directory try: nodes, display_names = load_nodes_from_directory(subdir_package) @@ -87,7 +92,7 @@ def scan_for_modules(base_path, base_package): except Exception as e: print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error processing directory {subdir_package}: {e}") traceback.print_exc() - + # Then recursively process subdirectories subdir_path = os.path.join(base_path, item) scan_for_modules(subdir_path, subdir_package) @@ -95,6 +100,7 @@ def scan_for_modules(base_path, base_package): print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error scanning directory {base_path}: {e}") traceback.print_exc() + # Start recursive scanning from the base directory print(f"\033[92m" + f"[ComfyUI_RealtimeNodes] Scanning node_wrappers directory for modules...") scan_for_modules(base_dir_path, WRAPPER_DIR) @@ -104,19 +110,19 @@ suffix = " 🕒🅡🅣🅝" for node_name in NODE_CLASS_MAPPINGS.keys(): # Convert camelCase or snake_case to Title Case if node_name not in NODE_DISPLAY_NAME_MAPPINGS: - display_name = ' '.join(word.capitalize() for word in re.findall(r'[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+', node_name)) + display_name = " ".join(word.capitalize() for word in re.findall(r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+", node_name)) else: display_name = NODE_DISPLAY_NAME_MAPPINGS[node_name] - + # Add the suffix if it's not already present if not display_name.endswith(suffix): display_name += suffix - + # Assign the final display name to the mappings NODE_DISPLAY_NAME_MAPPINGS[node_name] = display_name # --- Original Web Directory and Export --- -WEB_DIRECTORY = "./web/js" # Adjusted path if needed +WEB_DIRECTORY = "./web/js" # Adjusted path if needed -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] -print("\033[92m" + f"[ComfyUI_RealtimeNodes] Loaded {len(NODE_CLASS_MAPPINGS)} nodes.") \ No newline at end of file +print("\033[92m" + f"[ComfyUI_RealtimeNodes] Loaded {len(NODE_CLASS_MAPPINGS)} nodes.") diff --git a/node_wrappers/__init__.py b/node_wrappers/__init__.py index da22290..6e7633b 100644 --- a/node_wrappers/__init__.py +++ b/node_wrappers/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/node_wrappers/coordinates/__init__.py b/node_wrappers/coordinates/__init__.py index d832f58..88d2322 100644 --- a/node_wrappers/coordinates/__init__.py +++ b/node_wrappers/coordinates/__init__.py @@ -3,4 +3,4 @@ Coordinate system node wrappers for ComfyUI. This package provides node wrappers for the coordinate system functionality, allowing users to convert between coordinate spaces and create coordinate objects. -""" \ No newline at end of file +""" diff --git a/node_wrappers/coordinates/conversion_nodes.py b/node_wrappers/coordinates/conversion_nodes.py index 863830a..fd7c373 100644 --- a/node_wrappers/coordinates/conversion_nodes.py +++ b/node_wrappers/coordinates/conversion_nodes.py @@ -10,10 +10,12 @@ from ...src.coordinates import CoordinateSystem logger = logging.getLogger(__name__) + class CoordinateConverterNode: """ Fast coordinate conversion between different coordinate spaces. """ + CATEGORY = "Realtime Nodes/Coordinates" FUNCTION = "convert_coordinates" RETURN_TYPES = ("FLOAT", "FLOAT") @@ -40,20 +42,20 @@ class CoordinateConverterNode: try: # Get dimensions from image dimensions = CoordinateSystem.get_dimensions_from_tensor(image_for_dimensions) - + # Map string inputs to CoordinateSystem constants space_map = { "pixel": CoordinateSystem.PIXEL, "normalized": CoordinateSystem.NORMALIZED, } - + from_space_const = space_map[from_space] to_space_const = space_map[to_space] - + # Convert coordinates x_out = CoordinateSystem.convert(x, dimensions[0], from_space_const, to_space_const) y_out = CoordinateSystem.convert(y, dimensions[1], from_space_const, to_space_const) - + return (x_out, y_out) except Exception as e: logger.error(f"Error in coordinate conversion: {e}") @@ -64,6 +66,7 @@ class Point2DNode: """ Creates a 2D point with coordinate space awareness. """ + CATEGORY = "Realtime Nodes/Coordinates" FUNCTION = "create_point" RETURN_TYPES = ("POINT",) @@ -78,26 +81,29 @@ class Point2DNode: "space": (["normalized", "pixel"],), }, "optional": { - "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image for dimensions (required for pixel space)"}), - } + "image_for_dimensions": ( + "IMAGE", + {"tooltip": "Reference image for dimensions (required for pixel space)"}, + ), + }, } def create_point(self, x, y, space, image_for_dimensions=None): """Create a Point object.""" from ...src.coordinates import Point - + space_map = { "pixel": CoordinateSystem.PIXEL, "normalized": CoordinateSystem.NORMALIZED, } - + space_const = space_map[space] - + # Validate dimensions if using pixel space if space == "pixel" and image_for_dimensions is None: logger.warning("Pixel space selected but no image provided for dimensions. Using normalized space.") space_const = CoordinateSystem.NORMALIZED - + return (Point(x, y, None, space_const),) @@ -105,6 +111,7 @@ class PointListNode: """ Creates a list of points from coordinate lists. """ + CATEGORY = "Realtime Nodes/Coordinates" FUNCTION = "create_point_list" RETURN_TYPES = ("POINT_LIST",) @@ -114,32 +121,44 @@ class PointListNode: def INPUT_TYPES(cls): return { "required": { - "x_coords": ("FLOAT", {"default": [0.25, 0.5, 0.75], "forceInput": True, "tooltip": "List of X coordinates"}), - "y_coords": ("FLOAT", {"default": [0.25, 0.5, 0.75], "forceInput": True, "tooltip": "List of Y coordinates"}), + "x_coords": ( + "FLOAT", + {"default": [0.25, 0.5, 0.75], "forceInput": True, "tooltip": "List of X coordinates"}, + ), + "y_coords": ( + "FLOAT", + {"default": [0.25, 0.5, 0.75], "forceInput": True, "tooltip": "List of Y coordinates"}, + ), "space": (["normalized", "pixel"],), }, "optional": { - "z_coords": ("FLOAT", {"default": None, "forceInput": True, "tooltip": "List of Z coordinates (optional)"}), - "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image for dimensions (required for pixel space)"}), - } + "z_coords": ( + "FLOAT", + {"default": None, "forceInput": True, "tooltip": "List of Z coordinates (optional)"}, + ), + "image_for_dimensions": ( + "IMAGE", + {"tooltip": "Reference image for dimensions (required for pixel space)"}, + ), + }, } def create_point_list(self, x_coords, y_coords, space, z_coords=None, image_for_dimensions=None): """Create a PointList object.""" from ...src.coordinates import PointList - + space_map = { "pixel": CoordinateSystem.PIXEL, "normalized": CoordinateSystem.NORMALIZED, } - + space_const = space_map[space] - + # Validate dimensions if using pixel space if space == "pixel" and image_for_dimensions is None: logger.warning("Pixel space selected but no image provided for dimensions. Using normalized space.") space_const = CoordinateSystem.NORMALIZED - + # Ensure inputs are lists if not isinstance(x_coords, list): x_coords = [x_coords] @@ -147,7 +166,7 @@ class PointListNode: y_coords = [y_coords] if z_coords is not None and not isinstance(z_coords, list): z_coords = [z_coords] - + return (PointList.from_coordinates(x_coords, y_coords, z_coords, space_const),) @@ -163,4 +182,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CoordinateConverter": "Coordinate Converter", "Point2D": "Create 2D Point", "PointList": "Create Point List", -} \ No newline at end of file +} diff --git a/node_wrappers/coordinates/drawing_nodes.py b/node_wrappers/coordinates/drawing_nodes.py index da726a8..f6da394 100644 --- a/node_wrappers/coordinates/drawing_nodes.py +++ b/node_wrappers/coordinates/drawing_nodes.py @@ -7,11 +7,13 @@ from ...src.coordinates import CoordinateSystem, DrawingEngine logger = logging.getLogger(__name__) + class RTDrawPointsNode: """ High-performance point drawing node optimized for real-time applications. Draws points (circles) on images with minimal overhead. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" FUNCTION = "draw_points" RETURN_TYPES = ("IMAGE",) @@ -22,40 +24,64 @@ class RTDrawPointsNode: return { "required": { "image": ("IMAGE",), - "x": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "X coordinate(s) (0.0-1.0 normalized by default)"}), - "y": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Y coordinate(s) (0.0-1.0 normalized by default)"}), + "x": ( + "FLOAT", + {"default": 0.5, "forceInput": True, "tooltip": "X coordinate(s) (0.0-1.0 normalized by default)"}, + ), + "y": ( + "FLOAT", + {"default": 0.5, "forceInput": True, "tooltip": "Y coordinate(s) (0.0-1.0 normalized by default)"}, + ), "radius": ("INT", {"default": 5, "min": 1, "max": 100, "tooltip": "Radius of points in pixels"}), "color_hex": ("STRING", {"default": "#FF0000", "tooltip": "Color in #RRGGBB format"}), }, "optional": { - "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), - "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast", "tooltip": "How to map coords to batch images"}) - } + "is_normalized": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}, + ), + "batch_mapping": ( + ["broadcast", "one-to-one", "all-on-first"], + {"default": "broadcast", "tooltip": "How to map coords to batch images"}, + ), + }, } - def draw_points(self, image: torch.Tensor, x: Union[float, List[float]], y: Union[float, List[float]], - radius: int, color_hex: str, is_normalized: bool = True, batch_mapping: str = "broadcast"): + def draw_points( + self, + image: torch.Tensor, + x: Union[float, List[float]], + y: Union[float, List[float]], + radius: int, + color_hex: str, + is_normalized: bool = True, + batch_mapping: str = "broadcast", + ): """Draw points efficiently using the DrawingEngine.""" # Use the coordinate system to ensure proper handling space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL dimensions = CoordinateSystem.get_dimensions_from_tensor(image) - + drawing_engine = DrawingEngine() - return (drawing_engine.draw_points( - image=image, - points_x=x, - points_y=y, - radius=radius, - color=color_hex, - is_normalized=is_normalized, - batch_mapping=batch_mapping - ),) + return ( + drawing_engine.draw_points( + image=image, + points_x=x, + points_y=y, + radius=radius, + color=color_hex, + is_normalized=is_normalized, + batch_mapping=batch_mapping, + ), + ) + class RTDrawLinesNode: """ High-performance line drawing node optimized for real-time applications. Draws lines on images with minimal overhead. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" FUNCTION = "draw_lines" RETURN_TYPES = ("IMAGE",) @@ -66,59 +92,110 @@ class RTDrawLinesNode: return { "required": { "image": ("IMAGE",), - "x1": ("FLOAT", {"default": 0.25, "forceInput": True, "tooltip": "Start X coordinate(s) (0.0-1.0 normalized by default)"}), - "y1": ("FLOAT", {"default": 0.25, "forceInput": True, "tooltip": "Start Y coordinate(s) (0.0-1.0 normalized by default)"}), - "x2": ("FLOAT", {"default": 0.75, "forceInput": True, "tooltip": "End X coordinate(s) (0.0-1.0 normalized by default)"}), - "y2": ("FLOAT", {"default": 0.75, "forceInput": True, "tooltip": "End Y coordinate(s) (0.0-1.0 normalized by default)"}), + "x1": ( + "FLOAT", + { + "default": 0.25, + "forceInput": True, + "tooltip": "Start X coordinate(s) (0.0-1.0 normalized by default)", + }, + ), + "y1": ( + "FLOAT", + { + "default": 0.25, + "forceInput": True, + "tooltip": "Start Y coordinate(s) (0.0-1.0 normalized by default)", + }, + ), + "x2": ( + "FLOAT", + { + "default": 0.75, + "forceInput": True, + "tooltip": "End X coordinate(s) (0.0-1.0 normalized by default)", + }, + ), + "y2": ( + "FLOAT", + { + "default": 0.75, + "forceInput": True, + "tooltip": "End Y coordinate(s) (0.0-1.0 normalized by default)", + }, + ), "thickness": ("INT", {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness in pixels"}), "color_hex": ("STRING", {"default": "#00FF00", "tooltip": "Color in #RRGGBB format"}), }, "optional": { - "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), + "is_normalized": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}, + ), "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}), "draw_endpoints": ("BOOLEAN", {"default": False, "tooltip": "Draw circles at line endpoints"}), - "point_radius": ("INT", {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of endpoint circles in pixels"}), + "point_radius": ( + "INT", + {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of endpoint circles in pixels"}, + ), "draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on line"}), "label_text": ("STRING", {"default": "Line", "tooltip": "Text to display as label"}), "label_position": (["Midpoint", "Start", "End"], {"default": "Midpoint"}), - "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}) - } + "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}), + }, } - def draw_lines(self, image: torch.Tensor, - x1: Union[float, List[float]], y1: Union[float, List[float]], - x2: Union[float, List[float]], y2: Union[float, List[float]], - thickness: int, color_hex: str, - is_normalized: bool = True, batch_mapping: str = "broadcast", - draw_endpoints: bool = False, point_radius: int = 3, - draw_label: bool = False, label_text: str = "Line", - label_position: str = "Midpoint", font_scale: float = 0.5): + def draw_lines( + self, + image: torch.Tensor, + x1: Union[float, List[float]], + y1: Union[float, List[float]], + x2: Union[float, List[float]], + y2: Union[float, List[float]], + thickness: int, + color_hex: str, + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_endpoints: bool = False, + point_radius: int = 3, + draw_label: bool = False, + label_text: str = "Line", + label_position: str = "Midpoint", + font_scale: float = 0.5, + ): """Draw lines efficiently using the DrawingEngine.""" # Use the coordinate system to ensure proper handling space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL dimensions = CoordinateSystem.get_dimensions_from_tensor(image) - + drawing_engine = DrawingEngine() - return (drawing_engine.draw_lines( - image=image, - x1=x1, y1=y1, x2=x2, y2=y2, - thickness=thickness, - color=color_hex, - is_normalized=is_normalized, - batch_mapping=batch_mapping, - draw_endpoints=draw_endpoints, - point_radius=point_radius, - draw_label=draw_label, - label_text=label_text, - label_position=label_position, - font_scale=font_scale - ),) + return ( + drawing_engine.draw_lines( + image=image, + x1=x1, + y1=y1, + x2=x2, + y2=y2, + thickness=thickness, + color=color_hex, + is_normalized=is_normalized, + batch_mapping=batch_mapping, + draw_endpoints=draw_endpoints, + point_radius=point_radius, + draw_label=draw_label, + label_text=label_text, + label_position=label_position, + font_scale=font_scale, + ), + ) + class RTDrawPolygonNode: """ High-performance polygon drawing node optimized for real-time applications. Draws polygons on images with minimal overhead. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing" FUNCTION = "draw_polygon" RETURN_TYPES = ("IMAGE",) @@ -129,58 +206,91 @@ class RTDrawPolygonNode: return { "required": { "image": ("IMAGE",), - "points_x": ("FLOAT", {"default": [0.25, 0.75, 0.5], "forceInput": True, "tooltip": "List of X coordinates for polygon vertices"}), - "points_y": ("FLOAT", {"default": [0.25, 0.25, 0.75], "forceInput": True, "tooltip": "List of Y coordinates for polygon vertices"}), + "points_x": ( + "FLOAT", + { + "default": [0.25, 0.75, 0.5], + "forceInput": True, + "tooltip": "List of X coordinates for polygon vertices", + }, + ), + "points_y": ( + "FLOAT", + { + "default": [0.25, 0.25, 0.75], + "forceInput": True, + "tooltip": "List of Y coordinates for polygon vertices", + }, + ), "color_hex": ("STRING", {"default": "#0000FF", "tooltip": "Color in #RRGGBB format"}), - "thickness": ("INT", {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness for polygon outline"}), + "thickness": ( + "INT", + {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness for polygon outline"}, + ), }, "optional": { "fill": ("BOOLEAN", {"default": False, "tooltip": "Whether to fill the polygon"}), - "is_normalized": ("BOOLEAN", {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}), + "is_normalized": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"}, + ), "batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}), "draw_vertices": ("BOOLEAN", {"default": False, "tooltip": "Draw points at polygon vertices"}), "vertex_radius": ("INT", {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of vertex points"}), "draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on polygon"}), "label_text": ("STRING", {"default": "Polygon", "tooltip": "Text to display as label"}), - "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}) - } + "font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}), + }, } - def draw_polygon(self, image: torch.Tensor, - points_x: List[float], points_y: List[float], - color_hex: str, thickness: int, - fill: bool = False, is_normalized: bool = True, - batch_mapping: str = "broadcast", draw_vertices: bool = False, - vertex_radius: int = 3, draw_label: bool = False, - label_text: str = "Polygon", font_scale: float = 0.5): + def draw_polygon( + self, + image: torch.Tensor, + points_x: List[float], + points_y: List[float], + color_hex: str, + thickness: int, + fill: bool = False, + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_vertices: bool = False, + vertex_radius: int = 3, + draw_label: bool = False, + label_text: str = "Polygon", + font_scale: float = 0.5, + ): """Draw polygon efficiently using the DrawingEngine.""" # Use the coordinate system to ensure proper handling space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL dimensions = CoordinateSystem.get_dimensions_from_tensor(image) - + drawing_engine = DrawingEngine() - return (drawing_engine.draw_polygon( - image=image, - points_x=points_x, - points_y=points_y, - color=color_hex, - thickness=thickness, - fill=fill, - is_normalized=is_normalized, - batch_mapping=batch_mapping, - draw_vertices=draw_vertices, - vertex_radius=vertex_radius, - draw_label=draw_label, - label_text=label_text, - label_position="Center", # Only center supported for polygons currently - font_scale=font_scale - ),) + return ( + drawing_engine.draw_polygon( + image=image, + points_x=points_x, + points_y=points_y, + color=color_hex, + thickness=thickness, + fill=fill, + is_normalized=is_normalized, + batch_mapping=batch_mapping, + draw_vertices=draw_vertices, + vertex_radius=vertex_radius, + draw_label=draw_label, + label_text=label_text, + label_position="Center", # Only center supported for polygons currently + font_scale=font_scale, + ), + ) + class RTCoordinateConverterNode: """ Fast coordinate conversion between normalized (0-1) and pixel space. Optimized for real-time applications with minimal overhead. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Coordinates" FUNCTION = "convert_coordinates" RETURN_TYPES = ("FLOAT", "FLOAT") @@ -197,20 +307,22 @@ class RTCoordinateConverterNode: } } - def convert_coordinates(self, x: Union[float, List[float]], y: Union[float, List[float]], - image_for_dimensions: torch.Tensor, mode: str): + def convert_coordinates( + self, x: Union[float, List[float]], y: Union[float, List[float]], image_for_dimensions: torch.Tensor, mode: str + ): """Convert coordinates using the new coordinate system.""" dimensions = CoordinateSystem.get_dimensions_from_tensor(image_for_dimensions) - + if mode == "Pixel to Normalized": x_out = CoordinateSystem.normalize(x, dimensions[0], CoordinateSystem.PIXEL) y_out = CoordinateSystem.normalize(y, dimensions[1], CoordinateSystem.PIXEL) else: # "Normalized to Pixel" x_out = CoordinateSystem.denormalize(x, dimensions[0], CoordinateSystem.PIXEL) y_out = CoordinateSystem.denormalize(y, dimensions[1], CoordinateSystem.PIXEL) - + return (x_out, y_out) + # Register nodes NODE_CLASS_MAPPINGS = { "RTDrawPoints": RTDrawPointsNode, @@ -224,4 +336,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RTDrawLines": "Draw Lines (Fast)", "RTDrawPolygon": "Draw Polygon (Fast)", "RTCoordinateConverter": "Coordinate Converter (Fast)", -} \ No newline at end of file +} diff --git a/node_wrappers/coordinates/position_tracking.py b/node_wrappers/coordinates/position_tracking.py index ecaf422..e9bd273 100644 --- a/node_wrappers/coordinates/position_tracking.py +++ b/node_wrappers/coordinates/position_tracking.py @@ -8,24 +8,26 @@ import torch from ...src.coordinates.coordinate_delta import BaseCoordinateDelta -#TODO: consider moving base nodes to src +# TODO: consider moving base nodes to src # Import from the new consolidated utilities from ...src.utils.math import scale_value logger = logging.getLogger(__name__) + # --- Position Delta Control Nodes (Int/Float) (Handling Lists) --- class CoordinateDeltaControlNode(BaseCoordinateDelta): """Base class for position-based delta control nodes (Int/Float). - Handles single floats or lists of coordinates. - Outputs FLOAT type (actual value is float or list[float]). + Handles single floats or lists of coordinates. + Outputs FLOAT type (actual value is float or list[float]). """ + # Output semantic type is FLOAT (or INT for subclass) RETURN_TYPES = ("FLOAT",) - RETURN_NAMES = ("value",) # Singular name + RETURN_NAMES = ("value",) # Singular name IS_INT_OUTPUT = False DEFAULT_DELTA_MIN = 0.0 - DEFAULT_DELTA_MAX = 0.1 + DEFAULT_DELTA_MAX = 0.1 DEFAULT_OUTPUT_MIN = 0.0 DEFAULT_OUTPUT_MAX = 1.0 DEFAULT_CLAMP = True @@ -42,7 +44,15 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): "x": ("FLOAT", {"default": 0.0, "forceInput": True}), "y": ("FLOAT", {"default": 0.0, "forceInput": True}), "z": ("FLOAT", {"default": 0.0, "forceInput": True}), - "window_size": ("INT", { "default": 1, "min": 1, "max": MAX_POSITION_HISTORY - 1, "tooltip": "Number of frames back for delta"}), + "window_size": ( + "INT", + { + "default": 1, + "min": 1, + "max": MAX_POSITION_HISTORY - 1, + "tooltip": "Number of frames back for delta", + }, + ), "delta_min": ("FLOAT", {"default": cls.DEFAULT_DELTA_MIN, "step": 0.01}), "delta_max": ("FLOAT", {"default": cls.DEFAULT_DELTA_MAX, "step": 0.01}), output_min_key: (output_type_str, {"default": cls.DEFAULT_OUTPUT_MIN}), @@ -55,13 +65,13 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): output_type_str = "INT" if self.IS_INT_OUTPUT else "FLOAT" output_min_key = f"output_min_{output_type_str.lower()}" output_max_key = f"output_max_{output_type_str.lower()}" - x_in, y_in, z_in = kwargs.get('x'), kwargs.get('y'), kwargs.get('z') - window_size = kwargs.get('window_size', 1) - delta_min = kwargs.get('delta_min') - delta_max = kwargs.get('delta_max') + x_in, y_in, z_in = kwargs.get("x"), kwargs.get("y"), kwargs.get("z") + window_size = kwargs.get("window_size", 1) + delta_min = kwargs.get("delta_min") + delta_max = kwargs.get("delta_max") output_min = kwargs.get(output_min_key) output_max = kwargs.get(output_max_key) - clamp = kwargs.get('clamp') + clamp = kwargs.get("clamp") fallback_value = int(output_min) if self.IS_INT_OUTPUT else float(output_min) if None in [x_in, y_in, z_in, delta_min, delta_max, output_min, output_max, clamp]: @@ -81,11 +91,11 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): # logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}") for k in keys_to_remove: del self._histories[k] - + if current_batch_size == 0: - logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.") - return ([],) - + logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.") + return ([],) + # Process List Input results = [] for i in range(current_batch_size): @@ -94,17 +104,17 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.") results.append(fallback_value) continue - + raw_delta = self._get_position_delta(float(xi), float(yi), float(zi), window_size, item_index=i) scaled_value = scale_value(raw_delta, delta_min, delta_max, output_min, output_max, clamp) - + if self.IS_INT_OUTPUT: output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min) else: output_val = scaled_value if scaled_value is not None else float(output_min) results.append(output_val) return (results,) - + else: # Single Float Input: Remove all histories except for index 0 keys_to_remove = [k for k in self._histories if k != 0] @@ -112,43 +122,47 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): # logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}") for k in keys_to_remove: del self._histories[k] - + # Process Single Float Input if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]): - logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.") - return (fallback_value,) - + logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.") + return (fallback_value,) + raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0) scaled_value = scale_value(raw_delta, delta_min, delta_max, output_min, output_max, clamp) - + if self.IS_INT_OUTPUT: output_val = int(round(scaled_value)) if scaled_value is not None else int(output_min) else: output_val = scaled_value if scaled_value is not None else float(output_min) return (output_val,) + class CoordinateDeltaIntControlNode(CoordinateDeltaControlNode): - RETURN_TYPES = ("INT",) # Semantic type is INT + RETURN_TYPES = ("INT",) # Semantic type is INT RETURN_NAMES = ("value",) IS_INT_OUTPUT = True DEFAULT_OUTPUT_MIN = 0 DEFAULT_OUTPUT_MAX = 100 + class CoordinateDeltaFloatControlNode(CoordinateDeltaControlNode): - RETURN_TYPES = ("FLOAT",) # Semantic type is FLOAT + RETURN_TYPES = ("FLOAT",) # Semantic type is FLOAT RETURN_NAMES = ("value",) IS_INT_OUTPUT = False DEFAULT_OUTPUT_MIN = 0.0 DEFAULT_OUTPUT_MAX = 1.0 + # --- Position Delta Trigger Node (Handling Lists) --- class CoordinateDeltaTriggerNode(BaseCoordinateDelta): """Outputs a BOOLEAN trigger if position delta crosses a threshold. - Handles single floats or lists of coordinates. - Outputs BOOLEAN type (actual value is boolean or list[boolean]). + Handles single floats or lists of coordinates. + Outputs BOOLEAN type (actual value is boolean or list[boolean]). """ - RETURN_TYPES = ("BOOLEAN",) # Semantic type - RETURN_NAMES = ("trigger",) # Singular name + + RETURN_TYPES = ("BOOLEAN",) # Semantic type + RETURN_NAMES = ("trigger",) # Singular name DEFAULT_THRESHOLD = 0.05 CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Delta" @@ -167,15 +181,15 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): } def execute(self, **kwargs): - x_in, y_in, z_in = kwargs.get('x'), kwargs.get('y'), kwargs.get('z') - window_size = kwargs.get('window_size', 1) - threshold = kwargs.get('threshold') - condition = kwargs.get('condition') + x_in, y_in, z_in = kwargs.get("x"), kwargs.get("y"), kwargs.get("z") + window_size = kwargs.get("window_size", 1) + threshold = kwargs.get("threshold") + condition = kwargs.get("condition") fallback_value = False if None in [x_in, y_in, z_in, threshold, condition]: - logger.error(f"{self.__class__.__name__}: Missing required arguments.") - return (fallback_value,) + logger.error(f"{self.__class__.__name__}: Missing required arguments.") + return (fallback_value,) # --- Input Type Handling & History Cleanup --- is_list = isinstance(x_in, list) @@ -190,11 +204,11 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): # logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}") for k in keys_to_remove: del self._histories[k] - + if current_batch_size == 0: - logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.") - return ([],) - + logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.") + return ([],) + # Process List Input results = [] for i in range(current_batch_size): @@ -203,39 +217,48 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.") results.append(fallback_value) continue - + raw_delta = self._get_position_delta(float(xi), float(yi), float(zi), window_size, item_index=i) triggered = fallback_value if raw_delta is not None: - if condition == "Above": triggered = raw_delta > threshold - elif condition == "Below": triggered = raw_delta < threshold - elif condition == "Equals": triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE - elif condition == "Not Equals": triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE + if condition == "Above": + triggered = raw_delta > threshold + elif condition == "Below": + triggered = raw_delta < threshold + elif condition == "Equals": + triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE + elif condition == "Not Equals": + triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE results.append(triggered) return (results,) - + else: - # Single Float Input: Remove all histories except for index 0 + # Single Float Input: Remove all histories except for index 0 keys_to_remove = [k for k in self._histories if k != 0] if keys_to_remove: # logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}") for k in keys_to_remove: del self._histories[k] - + # Process Single Float Input if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]): - logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.") - return (fallback_value,) - + logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.") + return (fallback_value,) + raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0) triggered = fallback_value if raw_delta is not None: - if condition == "Above": triggered = raw_delta > threshold - elif condition == "Below": triggered = raw_delta < threshold - elif condition == "Equals": triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE - elif condition == "Not Equals": triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE + if condition == "Above": + triggered = raw_delta > threshold + elif condition == "Below": + triggered = raw_delta < threshold + elif condition == "Equals": + triggered = abs(raw_delta - threshold) < FLOAT_EQUALITY_TOLERANCE + elif condition == "Not Equals": + triggered = abs(raw_delta - threshold) >= FLOAT_EQUALITY_TOLERANCE return (triggered,) + # --- Landmark Proximity Node (Handles Float or List Inputs) --- class CoordinateProximityNode: """ @@ -243,11 +266,12 @@ class CoordinateProximityNode: Handles single coordinate sets or lists thereof. Outputs FLOAT type (actual value is float or list[float]). """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Utils" FUNCTION = "calculate_distances" # Output semantic type is FLOAT RETURN_TYPES = ("FLOAT",) - RETURN_NAMES = ("distance",) # Singular name + RETURN_NAMES = ("distance",) # Singular name @classmethod def INPUT_TYPES(cls): @@ -263,60 +287,67 @@ class CoordinateProximityNode: } } - def calculate_distances(self, x1: Union[float, List[float]], y1: Union[float, List[float]], z1: Union[float, List[float]], - x2: Union[float, List[float]], y2: Union[float, List[float]], z2: Union[float, List[float]]) -> Union[float, List[float]]: # Return type hint updated - + def calculate_distances( + self, + x1: Union[float, List[float]], + y1: Union[float, List[float]], + z1: Union[float, List[float]], + x2: Union[float, List[float]], + y2: Union[float, List[float]], + z2: Union[float, List[float]], + ) -> Union[float, List[float]]: # Return type hint updated # --- Input Type Handling --- inputs = [x1, y1, z1, x2, y2, z2] - is_list_input = isinstance(x1, list) + is_list_input = isinstance(x1, list) fallback_value = 0.0 if any(isinstance(inp, list) != is_list_input for inp in inputs): - raise ValueError("All coordinate inputs (x1..z2) must be of the same type (all floats or all lists).") + raise ValueError("All coordinate inputs (x1..z2) must be of the same type (all floats or all lists).") if is_list_input: # Process List Input list_len = len(x1) if any(len(inp) != list_len for inp in inputs if isinstance(inp, list)): - raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.") + raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.") if list_len == 0: logger.warning("Input coordinate lists are empty. Returning empty distance list.") - return ([],) # Return empty list for FLOAT type + return ([],) # Return empty list for FLOAT type num_calcs = list_len x1_list, y1_list, z1_list = x1, y1, z1 x2_list, y2_list, z2_list = x2, y2, z2 - + distances = [] for i in range(num_calcs): try: _x1, _y1, _z1 = x1_list[i], y1_list[i], z1_list[i] _x2, _y2, _z2 = x2_list[i], y2_list[i], z2_list[i] - + if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _z1, _x2, _y2, _z2]): - logger.warning(f"Invalid coordinate type found at index {i}. Using fallback distance {fallback_value}.") - distances.append(fallback_value) - continue - - dist = math.sqrt((_x2 - _x1)**2 + (_y2 - _y1)**2 + (_z2 - _z1)**2) + logger.warning(f"Invalid coordinate type found at index {i}. Using fallback distance {fallback_value}.") + distances.append(fallback_value) + continue + + dist = math.sqrt((_x2 - _x1) ** 2 + (_y2 - _y1) ** 2 + (_z2 - _z1) ** 2) distances.append(dist) except Exception as e: logger.error(f"Error calculating distance at index {i}: {e}", exc_info=True) - distances.append(fallback_value) - return (distances,) # Return list of results - + distances.append(fallback_value) + return (distances,) # Return list of results + else: # Process Single Float Input _x1, _y1, _z1 = x1, y1, z1 _x2, _y2, _z2 = x2, y2, z2 if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _z1, _x2, _y2, _z2]): - logger.error(f"Invalid coordinate types for single float input.") - return (fallback_value,) # Return single fallback + logger.error(f"Invalid coordinate types for single float input.") + return (fallback_value,) # Return single fallback try: - dist = math.sqrt((_x2 - _x1)**2 + (_y2 - _y1)**2 + (_z2 - _z1)**2) - return (dist,) # Return single float value + dist = math.sqrt((_x2 - _x1) ** 2 + (_y2 - _y1) ** 2 + (_z2 - _z1) ** 2) + return (dist,) # Return single float value except Exception as e: logger.error(f"Error calculating distance for single input: {e}", exc_info=True) - return (fallback_value,) # Return single fallback + return (fallback_value,) # Return single fallback + class MaskFromCoordinate: """ @@ -324,25 +355,50 @@ class MaskFromCoordinate: Handles single (x, y) floats or lists (x_list, y_list) as input. Outputs a batch of masks (N, 1, H, W). """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Masking" - FUNCTION = "create_masks" # Keep plural name for clarity + FUNCTION = "create_masks" # Keep plural name for clarity RETURN_TYPES = ("MASK",) @classmethod def INPUT_TYPES(cls): return { "required": { - "x": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Normalized X coordinate (0.0-1.0) or List thereof"}), - "y": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Normalized Y coordinate (0.0-1.0) or List thereof"}), + "x": ( + "FLOAT", + { + "default": 0.5, + "forceInput": True, + "tooltip": "Normalized X coordinate (0.0-1.0) or List thereof", + }, + ), + "y": ( + "FLOAT", + { + "default": 0.5, + "forceInput": True, + "tooltip": "Normalized Y coordinate (0.0-1.0) or List thereof", + }, + ), "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the output mask size"}), - "radius": ("INT", {"default": 25, "min": 1, "max": 1024, "tooltip": "Radius of the mask shape in pixels"}), + "radius": ( + "INT", + {"default": 25, "min": 1, "max": 1024, "tooltip": "Radius of the mask shape in pixels"}, + ), "shape": (["Circle", "Square"], {"default": "Circle"}), "invert": ("BOOLEAN", {"default": False}), } } - def create_masks(self, x: Union[float, List[float]], y: Union[float, List[float]], image_for_dimensions: torch.Tensor, radius: int, shape: str, invert: bool) -> Tuple[torch.Tensor]: - + def create_masks( + self, + x: Union[float, List[float]], + y: Union[float, List[float]], + image_for_dimensions: torch.Tensor, + radius: int, + shape: str, + invert: bool, + ) -> Tuple[torch.Tensor]: # --- Input Type Handling --- is_x_list = isinstance(x, list) is_y_list = isinstance(y, list) @@ -356,14 +412,14 @@ class MaskFromCoordinate: y_list = y if len(x_list) != len(y_list): raise ValueError(f"x_list (len {len(x_list)}) and y_list (len {len(y_list)}) must have the same length.") - if not x_list: # Handle empty list - logger.warning("Input coordinate lists are empty. Returning empty mask batch.") - # Determine shape for empty tensor - if image_for_dimensions.dim() != 4: - logger.error("Input image must be BHWC format to determine empty mask size.") - return (torch.zeros((0, 1, 64, 64), dtype=torch.float32),) - _, height, width, _ = image_for_dimensions.shape - return (torch.zeros((0, 1, height, width), dtype=torch.float32, device=image_for_dimensions.device),) + if not x_list: # Handle empty list + logger.warning("Input coordinate lists are empty. Returning empty mask batch.") + # Determine shape for empty tensor + if image_for_dimensions.dim() != 4: + logger.error("Input image must be BHWC format to determine empty mask size.") + return (torch.zeros((0, 1, 64, 64), dtype=torch.float32),) + _, height, width, _ = image_for_dimensions.shape + return (torch.zeros((0, 1, height, width), dtype=torch.float32, device=image_for_dimensions.device),) num_coords = len(x_list) else: # Input is float - wrap in lists for unified processing @@ -371,7 +427,7 @@ class MaskFromCoordinate: y_list = [y] num_coords = 1 - # --- Get Dimensions --- + # --- Get Dimensions --- if image_for_dimensions.dim() != 4: logger.error("Input image must be BHWC format.") # Return empty batch matching input size @@ -379,7 +435,7 @@ class MaskFromCoordinate: _, height, width, _ = image_for_dimensions.shape device = image_for_dimensions.device - # --- Generate Masks --- + # --- Generate Masks --- output_masks = [] for i in range(num_coords): xi = x_list[i] @@ -387,10 +443,12 @@ class MaskFromCoordinate: # Validate individual coordinates if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)): - logger.warning(f"Invalid coordinate type at index {i} (x={type(xi)}, y={type(yi)}). Skipping mask generation for this index.") - mask_tensor = torch.zeros((1, height, width), dtype=torch.float32, device=device) - output_masks.append(mask_tensor) - continue + logger.warning( + f"Invalid coordinate type at index {i} (x={type(xi)}, y={type(yi)}). Skipping mask generation for this index." + ) + mask_tensor = torch.zeros((1, height, width), dtype=torch.float32, device=device) + output_masks.append(mask_tensor) + continue # Calculate pixel coordinates px = max(0, min(width - 1, int(xi * width))) @@ -407,7 +465,7 @@ class MaskFromCoordinate: bottom_right = (min(width - 1, px + half_size), min(height - 1, py + half_size)) cv2.rectangle(mask_np, top_left, bottom_right, 1, -1) else: - cv2.circle(mask_np, center, radius, 1, -1) # Default circle + cv2.circle(mask_np, center, radius, 1, -1) # Default circle if invert: mask_np = 1 - mask_np diff --git a/node_wrappers/mediapipe_vision/common/base_detector_node.py b/node_wrappers/mediapipe_vision/common/base_detector_node.py index 80b2097..3a48344 100644 --- a/node_wrappers/mediapipe_vision/common/base_detector_node.py +++ b/node_wrappers/mediapipe_vision/common/base_detector_node.py @@ -5,29 +5,30 @@ from typing import Dict logger = logging.getLogger(__name__) + class BaseMediaPipeDetectorNode: """Base class for all MediaPipe detector nodes in ComfyUI. - + This class captures common patterns across detector nodes like: - Detector instance management - Model validation - Common input parameters - Cleanup in __del__ - + Parameter Naming Standards for Subclasses: ----------------------------------------- To maintain consistency across detector types, use these standardized parameter names: - - - min_confidence: Main confidence threshold for detection (replaces various min_detection_confidence, + + - min_confidence: Main confidence threshold for detection (replaces various min_detection_confidence, score_threshold, etc.) - min_tracking_confidence: Confidence threshold for tracking objects across video frames - min_presence_confidence: Confidence threshold for presence of an object in a frame - max_results: Maximum number of detections to return (replaces num_poses, num_hands, etc.) - - When implementing specific detector nodes, map these standardized parameter names to + + When implementing specific detector nodes, map these standardized parameter names to the appropriate underlying MediaPipe API parameters in the detect() method. """ - + # To be overridden by subclasses DETECTOR_CLASS = None # Subclasses must set to their specific detector class MODEL_INFO_TYPE = "MEDIAPIPE_MODEL_INFO" # Override with specific model info type @@ -37,12 +38,12 @@ class BaseMediaPipeDetectorNode: RETURN_NAMES = tuple() # Must be set by subclass FUNCTION = "detect" # Default function name, can be overridden DESCRIPTION = "" # Override with specific node description - + def __init__(self): """Initialize with empty detector and model path.""" self._detector = None self._model_path = None - + @classmethod def INPUT_TYPES(cls): """Define common input types for detector nodes.""" @@ -50,84 +51,90 @@ class BaseMediaPipeDetectorNode: "required": { "image": ("IMAGE", {"tooltip": "Input image for detection"}), "model_info": (cls.MODEL_INFO_TYPE, {"tooltip": f"Model loaded from the appropriate loader node"}), - "running_mode": (["image", "video"], { - "default": "video", - "tooltip": "Processing mode: 'image' for single images, 'video' for sequences/streams" - }), - "delegate": (["cpu", "gpu"], { - "default": "cpu", - "tooltip": "Processing device - GPU is faster but CPU is more compatible with all systems" - }), + "running_mode": ( + ["image", "video"], + { + "default": "video", + "tooltip": "Processing mode: 'image' for single images, 'video' for sequences/streams", + }, + ), + "delegate": ( + ["cpu", "gpu"], + { + "default": "cpu", + "tooltip": "Processing device - GPU is faster but CPU is more compatible with all systems", + }, + ), } } # Subclasses should extend this with additional parameters - + def validate_model_info(self, model_info: Dict) -> str: """Validate model_info and return model path.""" if not isinstance(model_info, dict): raise ValueError(f"Invalid model_info provided. Expected dictionary, got {type(model_info)}") - - task_type = model_info.get('task_type') + + task_type = model_info.get("task_type") if task_type != self.EXPECTED_TASK_TYPE: raise ValueError(f"Invalid model_info provided. Expected task_type '{self.EXPECTED_TASK_TYPE}', got '{task_type}'") - - model_path = model_info.get('model_path') + + model_path = model_info.get("model_path") if not model_path: raise ValueError("Model path not found or invalid in model_info") - + return model_path - + def initialize_or_update_detector(self, model_path: str): """Initialize detector or update if model path changed.""" if self.DETECTOR_CLASS is None: raise NotImplementedError("Subclasses must set DETECTOR_CLASS to their specific detector class") - + if self._detector is None or self._model_path != model_path: # Close existing detector if needed - if self._detector and hasattr(self._detector, 'close'): + if self._detector and hasattr(self._detector, "close"): try: logger.info(f"Closing existing detector for {self._model_path}") self._detector.close() except Exception as e: logger.warning(f"Error closing detector: {e}") - + # Create new detector logger.info(f"Creating new {self.DETECTOR_CLASS.__name__} instance for {model_path}") self._detector = self.DETECTOR_CLASS(model_path) self._model_path = model_path - elif not hasattr(self._detector, '_detector_instance'): + elif not hasattr(self._detector, "_detector_instance"): # Re-initialize if internal instance is missing logger.warning(f"Re-initializing detector wrapper due to missing internal instance") self._detector = self.DETECTOR_CLASS(model_path) self._model_path = model_path - + return self._detector - + def detect(self, **kwargs): """ Main detection function, should be implemented by subclasses. This is a placeholder that subclasses can use as a template. """ # Basic implementation that subclasses can override - + # 1. Extract and validate model_info - model_info = kwargs.get('model_info') + model_info = kwargs.get("model_info") model_path = self.validate_model_info(model_info) - + # 2. Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # 3. Call detector's detect method with appropriate parameters # This part should be implemented by subclasses - + raise NotImplementedError("Subclasses must implement or override the detect method") - + def __del__(self): """Clean up resources by closing detector.""" - if hasattr(self, '_detector') and self._detector and hasattr(self._detector, 'close'): + if hasattr(self, "_detector") and self._detector and hasattr(self._detector, "close"): try: logger.info(f"Closing detector instance in __del__ for {self._model_path}") self._detector.close() except Exception as e: logger.warning(f"Error closing detector in __del__: {e}") - self._detector = None \ No newline at end of file + self._detector = None diff --git a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py index fae1a64..688d5bb 100644 --- a/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py +++ b/node_wrappers/mediapipe_vision/common/base_visualization_nodes.py @@ -12,24 +12,26 @@ DEFAULT_FONT = cv2.FONT_HERSHEY_SIMPLEX DEFAULT_FONT_SCALE = 0.5 DEFAULT_FONT_THICKNESS = 1 + class BaseLandmarkVisualizationNode: """Base class for landmark visualization nodes.""" + CATEGORY = "Realtime Nodes/MediaPipe Vision/Visualization" RETURN_TYPES = ("IMAGE",) FUNCTION = "draw_visualizations" - - # --- To be overridden by subclasses --- - LANDMARKS_TYPE = "ANY" # The specific ComfyUI landmark type string (e.g., "FACE_LANDMARKS") - CONNECTIONS = None # MediaPipe connections list (e.g., mp.solutions.face_mesh.FACEMESH_TESSELATION) + + # --- To be overridden by subclasses --- + LANDMARKS_TYPE = "ANY" # The specific ComfyUI landmark type string (e.g., "FACE_LANDMARKS") + CONNECTIONS = None # MediaPipe connections list (e.g., mp.solutions.face_mesh.FACEMESH_TESSELATION) DEFAULT_DRAW_POINTS = True DEFAULT_POINT_COLOR = "#FFFFFF" DEFAULT_POINT_RADIUS = 3 DEFAULT_DRAW_CONNECTIONS = True DEFAULT_CONNECTION_COLOR = "#00FF00" DEFAULT_CONNECTION_THICKNESS = 2 - SUPPORTS_VISIBILITY = False # Does this landmark type provide visibility scores? - # --- End Override Section --- - + SUPPORTS_VISIBILITY = False # Does this landmark type provide visibility scores? + # --- End Override Section --- + @classmethod def INPUT_TYPES(cls): """Generate input types based on class properties.""" @@ -45,29 +47,35 @@ class BaseLandmarkVisualizationNode: "draw_connections": ("BOOLEAN", {"default": cls.DEFAULT_DRAW_CONNECTIONS}), "connection_color": ("STRING", {"default": cls.DEFAULT_CONNECTION_COLOR}), "connection_thickness": ("INT", {"default": cls.DEFAULT_CONNECTION_THICKNESS, "min": 1, "max": 10}), - } + }, } - + if cls.SUPPORTS_VISIBILITY: - inputs["optional"]["min_visibility"] = ("FLOAT", { - "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, - "tooltip": "Minimum visibility score to draw a point/connection endpoint" - }) - + inputs["optional"]["min_visibility"] = ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Minimum visibility score to draw a point/connection endpoint", + }, + ) + return inputs - + def draw_visualizations(self, landmarks, image, **kwargs): """Main execution function: handles batching, image conversion, and calls drawing logic.""" - + # Extract relevant optional arguments using class defaults - draw_points = kwargs.get('draw_points', self.DEFAULT_DRAW_POINTS) - point_color = kwargs.get('point_color', self.DEFAULT_POINT_COLOR) - point_radius = kwargs.get('point_radius', self.DEFAULT_POINT_RADIUS) - draw_connections = kwargs.get('draw_connections', self.DEFAULT_DRAW_CONNECTIONS) - connection_color = kwargs.get('connection_color', self.DEFAULT_CONNECTION_COLOR) - connection_thickness = kwargs.get('connection_thickness', self.DEFAULT_CONNECTION_THICKNESS) - min_visibility = kwargs.get('min_visibility', 0.5) # Default if not supported/provided - + draw_points = kwargs.get("draw_points", self.DEFAULT_DRAW_POINTS) + point_color = kwargs.get("point_color", self.DEFAULT_POINT_COLOR) + point_radius = kwargs.get("point_radius", self.DEFAULT_POINT_RADIUS) + draw_connections = kwargs.get("draw_connections", self.DEFAULT_DRAW_CONNECTIONS) + connection_color = kwargs.get("connection_color", self.DEFAULT_CONNECTION_COLOR) + connection_thickness = kwargs.get("connection_thickness", self.DEFAULT_CONNECTION_THICKNESS) + min_visibility = kwargs.get("min_visibility", 0.5) # Default if not supported/provided + if not landmarks: logger.warning(f"{self.__class__.__name__}: No landmarks provided.") return (image,) @@ -77,7 +85,9 @@ class BaseLandmarkVisualizationNode: batch_size = len(cv2_images) if len(landmarks) != batch_size: - logger.error(f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Landmark sets: {len(landmarks)}. Cannot proceed reliably.") + logger.error( + f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Landmark sets: {len(landmarks)}. Cannot proceed reliably." + ) return (image,) # Convert hex colors to BGR tuples @@ -89,55 +99,73 @@ class BaseLandmarkVisualizationNode: landmarks_for_image = landmarks[i] # List of landmark results for this image item if not landmarks_for_image: - processed_cv2_images.append(cv2_image) # Append unmodified image + processed_cv2_images.append(cv2_image) # Append unmodified image continue # Iterate through all detections (faces, hands, poses) in this image item for single_detection_result in landmarks_for_image: - self.draw_single_detection(cv2_image, width, height, single_detection_result, - draw_points, point_color_bgr, point_radius, - draw_connections, connection_color_bgr, connection_thickness, - min_visibility) - + self.draw_single_detection( + cv2_image, + width, + height, + single_detection_result, + draw_points, + point_color_bgr, + point_radius, + draw_connections, + connection_color_bgr, + connection_thickness, + min_visibility, + ) + processed_cv2_images.append(cv2_image) output_image = convert_to_tensor(processed_cv2_images) return (output_image,) - + def _hex_to_bgr(self, hex_color): """Convert hex color string to BGR tuple.""" - hex_color = hex_color.lstrip('#') - r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + hex_color = hex_color.lstrip("#") + r, g, b = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) return (b, g, r) # OpenCV uses BGR - - def draw_single_detection(self, cv2_image, width: int, height: int, - detection_result: Any, - draw_points: bool, point_color, point_radius: int, - draw_connections: bool, connection_color, connection_thickness: int, - min_visibility: float): + + def draw_single_detection( + self, + cv2_image, + width: int, + height: int, + detection_result: Any, + draw_points: bool, + point_color, + point_radius: int, + draw_connections: bool, + connection_color, + connection_thickness: int, + min_visibility: float, + ): """Draws landmarks and connections for a single detected object (face, hand, pose).""" - + # Extract the actual list of landmarks (this might vary slightly based on MediaPipe result structure) landmark_list = self.get_landmark_list(detection_result) if not landmark_list: - return # Skip if no landmarks found for this detection - + return # Skip if no landmarks found for this detection + points = {} # Process landmarks: calculate coordinates and optionally draw points for lm in landmark_list: # Check visibility if supported is_visible = True if self.SUPPORTS_VISIBILITY: - vis = lm.visibility if hasattr(lm, 'visibility') and lm.visibility is not None else 1.0 - is_visible = vis >= min_visibility - + vis = lm.visibility if hasattr(lm, "visibility") and lm.visibility is not None else 1.0 + is_visible = vis >= min_visibility + if is_visible: cx, cy = int(lm.x * width), int(lm.y * height) - points[lm.index] = (cx, cy) # Store denormalized point coordinates - + points[lm.index] = (cx, cy) # Store denormalized point coordinates + if draw_points: cv2.circle(cv2_image, (cx, cy), point_radius, point_color, -1) - + # Draw connections if enabled and applicable if draw_connections and self.CONNECTIONS and points: for connection in self.CONNECTIONS: @@ -149,37 +177,38 @@ class BaseLandmarkVisualizationNode: def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]: """Abstract method to extract the list of landmark objects from a single detection result. - Needs to be implemented by subclasses based on the specific MediaPipe result structure. - e.g., for hands/pose: return detection_result.landmarks - for face: return detection_result (it's already the list) + Needs to be implemented by subclasses based on the specific MediaPipe result structure. + e.g., for hands/pose: return detection_result.landmarks + for face: return detection_result (it's already the list) """ raise NotImplementedError("Subclasses must implement get_landmark_list") class BaseDetectionVisualizationNode: """Base class for visualization nodes that need to draw bounding boxes and labels. - + This can be used for object detection, face detection, and similar tasks that require drawing boxes and text labels with a consistent font rendering approach. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/Visualization" RETURN_TYPES = ("IMAGE",) FUNCTION = "draw_visualizations" - + # Default font settings FONT = DEFAULT_FONT FONT_SCALE = DEFAULT_FONT_SCALE FONT_THICKNESS = DEFAULT_FONT_THICKNESS - + @classmethod def get_text_dimensions(cls, text, font_scale=None, thickness=None): """Gets text dimensions using cv2. - + Args: text: The text to measure font_scale: The font scale (defaults to cls.FONT_SCALE) thickness: The font thickness (defaults to cls.FONT_THICKNESS) - + Returns: Tuple of (width, height) in pixels """ @@ -187,27 +216,33 @@ class BaseDetectionVisualizationNode: font_scale = cls.FONT_SCALE if thickness is None: thickness = cls.FONT_THICKNESS - + try: - (text_width, text_height), baseline = cv2.getTextSize( - text, cls.FONT, font_scale, thickness - ) + (text_width, text_height), baseline = cv2.getTextSize(text, cls.FONT, font_scale, thickness) return text_width, text_height + baseline except Exception as e: logger.error(f"Error measuring text dimensions: {e}") return 100, 20 # Default fallback - + def _hex_to_bgr(self, hex_color): """Convert hex color string to BGR tuple.""" - hex_color = hex_color.lstrip('#') - r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + hex_color = hex_color.lstrip("#") + r, g, b = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) return (b, g, r) # OpenCV uses BGR - - def draw_text_with_background(self, image, text, position, - text_color="#FFFFFF", bg_color="#000000", - padding=(4, 2), font_scale=None, thickness=None): + + def draw_text_with_background( + self, + image, + text, + position, + text_color="#FFFFFF", + bg_color="#000000", + padding=(4, 2), + font_scale=None, + thickness=None, + ): """Draws text with a background rectangle. - + Args: image: cv2 image to draw on text: Text to draw @@ -222,46 +257,38 @@ class BaseDetectionVisualizationNode: font_scale = self.FONT_SCALE if thickness is None: thickness = self.FONT_THICKNESS - + # Convert hex colors to BGR tuples text_color_bgr = self._hex_to_bgr(text_color) bg_color_bgr = self._hex_to_bgr(bg_color) - + x, y = position text_width, text_height = self.get_text_dimensions(text, font_scale, thickness) - + # Draw background rectangle x_pad, y_pad = padding cv2.rectangle( image, (x, y), - (x + text_width + x_pad*2, y + text_height + y_pad*2), + (x + text_width + x_pad * 2, y + text_height + y_pad * 2), bg_color_bgr, - -1 # Filled rectangle + -1, # Filled rectangle ) - + # Draw text - cv2.putText( - image, - text, - (x + x_pad, y + text_height), - self.FONT, - font_scale, - text_color_bgr, - thickness - ) - + cv2.putText(image, text, (x + x_pad, y + text_height), self.FONT, font_scale, text_color_bgr, thickness) + return text_width, text_height - + def process_batch(self, detections, image, drawing_function, **kwargs): """Process a batch of images and detections. - + Args: detections: List of detection results (one per image) image: Batch of images as tensor drawing_function: Function to call for each image **kwargs: Additional arguments for the drawing function - + Returns: Tuple with processed image tensor """ @@ -274,7 +301,9 @@ class BaseDetectionVisualizationNode: batch_size = len(cv2_images) if len(detections) != batch_size: - logger.error(f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Detection sets: {len(detections)}. Cannot proceed.") + logger.error( + f"{self.__class__.__name__}: Batch size mismatch. Images: {batch_size}, Detection sets: {len(detections)}. Cannot proceed." + ) return (image,) for i, cv2_image in enumerate(cv2_images): @@ -290,4 +319,4 @@ class BaseDetectionVisualizationNode: processed_cv2_images.append(cv2_image) output_image = convert_to_tensor(processed_cv2_images) - return (output_image,) \ No newline at end of file + return (output_image,) diff --git a/node_wrappers/mediapipe_vision/common/model_loader.py b/node_wrappers/mediapipe_vision/common/model_loader.py index 701ed13..b19014b 100644 --- a/node_wrappers/mediapipe_vision/common/model_loader.py +++ b/node_wrappers/mediapipe_vision/common/model_loader.py @@ -13,36 +13,37 @@ from ....src.mediapipe_vision.model_loader import ( logger = logging.getLogger(__name__) + # Define a base class for MediaPipe model loading nodes class MediaPipeModelLoaderBaseNode: """Base ComfyUI node for loading MediaPipe Vision models. - + Inheriting classes MUST define the class variable TASK_TYPE. This node handles selecting the model variant for a specific task, triggers the download/loading via the src.model_loader, and returns a dictionary containing model information needed by other nodes. """ - + # Define common category for all MediaPipe loaders CATEGORY = "Realtime Nodes/MediaPipe Vision/Loaders" # Inheriting classes MUST set this - TASK_TYPE = None # e.g., "face_detector", "hand_landmarker" + TASK_TYPE = None # e.g., "face_detector", "hand_landmarker" @classmethod def INPUT_TYPES(cls): """Define input types for the node.""" if cls.TASK_TYPE is None: - raise NotImplementedError("Subclasses of MediaPipeModelLoaderBaseNode must define TASK_TYPE") + raise NotImplementedError("Subclasses of MediaPipeModelLoaderBaseNode must define TASK_TYPE") task_variants = [] - default_variant = "default" # Sensible default + default_variant = "default" # Sensible default try: task_variants = get_available_models(cls.TASK_TYPE) if not task_variants: - logger.warning(f"No model variants found for task '{cls.TASK_TYPE}'. Defaulting to ['default'].") - task_variants = ["default"] # Fallback variants + logger.warning(f"No model variants found for task '{cls.TASK_TYPE}'. Defaulting to ['default'].") + task_variants = ["default"] # Fallback variants default_variant = task_variants[0] if task_variants else default_variant except Exception as e: @@ -53,7 +54,7 @@ class MediaPipeModelLoaderBaseNode: # Create tooltip with model descriptions tooltip = f"Select model variant for {cls.TASK_TYPE}." - + # Add model descriptions if available tooltip += "\n\nModel Variants:" for variant in task_variants: @@ -63,54 +64,55 @@ class MediaPipeModelLoaderBaseNode: return { "required": { # Dropdown for selecting the model variant for the specific task - "model_variant": (task_variants, { - "default": default_variant, - "tooltip": tooltip - }), + "model_variant": (task_variants, {"default": default_variant, "tooltip": tooltip}), } } - + # Return a dictionary with model info, not the loader object itself - RETURN_TYPES = ("MEDIAPIPE_MODEL_INFO",) # Default, can be overridden by subclasses for specificity + RETURN_TYPES = ("MEDIAPIPE_MODEL_INFO",) # Default, can be overridden by subclasses for specificity RETURN_NAMES = ("model_info",) FUNCTION = "load_and_get_info" - def load_and_get_info(self, model_variant: str): # Removed task_type argument + def load_and_get_info(self, model_variant: str): # Removed task_type argument """ Loads the specified MediaPipe model using the src loader and returns a dictionary with model information (path, task, variant). Uses the TASK_TYPE defined by the subclass. """ if self.TASK_TYPE is None: - raise NotImplementedError("TASK_TYPE not defined in the subclass.") + raise NotImplementedError("TASK_TYPE not defined in the subclass.") - task_type = self.TASK_TYPE # Use the class variable + task_type = self.TASK_TYPE # Use the class variable logger.info(f"Attempting to load MediaPipe model: Task='{task_type}', Variant='{model_variant}'") - + # Use a temporary loader instance just to trigger the load/download logic # The actual detector nodes will instantiate their own detectors using the path loader = MediaPipeModelLoader() # Pass the class's task_type here success = loader.load_model(task_type, model_variant) - + if not success: error_msg = f"Failed to load MediaPipe model for Task='{task_type}', Variant='{model_variant}'" logger.error(error_msg) # It's often better to raise an error to stop the workflow raise RuntimeError(error_msg) - + # Get the essential model information - model_info = loader.get_model_info() # Should return {'model_path': ..., 'task_type': ..., 'model_variant': ...} - + model_info = loader.get_model_info() # Should return {'model_path': ..., 'task_type': ..., 'model_variant': ...} + # Ensure the info dict is populated correctly - get_model_info should ideally return all needed keys - if 'model_path' not in model_info or not model_info['model_path']: - raise RuntimeError(f"Failed to get valid model path from loader for {task_type}/{model_variant}") - if 'task_type' not in model_info: model_info['task_type'] = task_type - if 'model_variant' not in model_info: model_info['model_variant'] = model_variant - - logger.info(f"Successfully loaded model info: Path='{model_info.get('model_path')}', Task='{model_info.get('task_type')}', Variant='{model_info.get('model_variant')}'") - + if "model_path" not in model_info or not model_info["model_path"]: + raise RuntimeError(f"Failed to get valid model path from loader for {task_type}/{model_variant}") + if "task_type" not in model_info: + model_info["task_type"] = task_type + if "model_variant" not in model_info: + model_info["model_variant"] = model_variant + + logger.info( + f"Successfully loaded model info: Path='{model_info.get('model_path')}', Task='{model_info.get('task_type')}', Variant='{model_info.get('model_variant')}'" + ) + # Return the dictionary containing model path and other relevant info # The specific RETURN_TYPES are handled by the subclass definition return (model_info,) diff --git a/node_wrappers/mediapipe_vision/face/__init__.py b/node_wrappers/mediapipe_vision/face/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/face/__init__.py +++ b/node_wrappers/mediapipe_vision/face/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection.py b/node_wrappers/mediapipe_vision/face/detection/face_detection.py index e3e15cf..e3228b7 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection.py @@ -8,20 +8,24 @@ from ...common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection" + + # Inherit from Base Loader class MediaPipeFaceDetectorModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Face Detector models.""" + TASK_TYPE = "face_detector" - RETURN_TYPES = ("FACE_DETECTOR_MODEL_INFO",) # Specific return type + RETURN_TYPES = ("FACE_DETECTOR_MODEL_INFO",) # Specific return type RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Face Detector model for finding faces in images. Required before using the Face Detector node." # INPUT_TYPES and FUNCTION inherited + # Detection/Landmarking Nodes class MediaPipeFaceDetectorNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Face Detection.""" - + # Define class variables required by the base class DETECTOR_CLASS = FaceDetector MODEL_INFO_TYPE = "FACE_DETECTOR_MODEL_INFO" @@ -31,40 +35,45 @@ class MediaPipeFaceDetectorNode(BaseMediaPipeDetectorNode): FUNCTION = "detect" CATEGORY = _category DESCRIPTION = "Detects faces in images using MediaPipe's face detection technology. Provides coordinates and confidence scores for all faces found in the image." - + @classmethod def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add face detection specific parameters - inputs["required"].update({ - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, - "tooltip": "Minimum confidence threshold for face detection - lower values find more faces but may include false positives"}), - }) - + inputs["required"].update( + { + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Minimum confidence threshold for face detection - lower values find more faces but may include false positives", + }, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, min_confidence: float, - running_mode: str, delegate: str): + def detect(self, image: torch.Tensor, model_info: dict, min_confidence: float, running_mode: str, delegate: str): """Performs face detection with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Perform detection with all parameters - batch_results = detector.detect( - image, - min_detection_confidence=min_confidence, - running_mode=running_mode, - delegate=delegate - ) - + batch_results = detector.detect(image, min_detection_confidence=min_confidence, running_mode=running_mode, delegate=delegate) + return (batch_results,) + # Define mappings for ComfyUI NODE_CLASS_MAPPINGS = { "MediaPipeFaceDetectorModelLoader": MediaPipeFaceDetectorModelLoaderNode, diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py b/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py index c0538f2..9828c11 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection_mask.py @@ -8,67 +8,133 @@ import torch logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection" + class FaceDetectionToMask: """Converts face detection bounding boxes to masks that can be used in ComfyUI.""" + CATEGORY = _category DESCRIPTION = "Creates masks from face detections that can be used for selective processing, inpainting, or compositing. Generate masks based on face bounding boxes or facial keypoints." - + @classmethod def INPUT_TYPES(cls): return { "required": { - "face_detections": ("FACE_DETECTIONS", {"tooltip": "Face detection results from the Face Detector node"}), - "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1, - "tooltip": "Width of the output mask in pixels"}), - "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1, - "tooltip": "Height of the output mask in pixels"}), - "mask_type": (["Bounding Box", "Keypoints", "Both"], {"default": "Bounding Box", - "tooltip": "Type of mask to create - Bounding Box (rectangle around face), Keypoints (dots at facial points), or Both"}), - "face_selection": (["All Faces", "First Face", "Largest Face", "Highest Confidence Face"], - {"default": "All Faces", - "tooltip": "Which faces to include in the mask - all detected faces, only the first face, largest face, or most confident detection"}), - "feather_amount": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1, - "tooltip": "How much to soften the edges of the mask in pixels - higher values create smoother transitions"}), - "keypoint_radius": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1, - "tooltip": "Size of each keypoint in the mask when using Keypoints or Both mask types"}), - "invert": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, inverts the mask (white becomes black, black becomes white)"}), - "score_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence score for including face detections - higher values filter out less confident detections"}) + "face_detections": ( + "FACE_DETECTIONS", + {"tooltip": "Face detection results from the Face Detector node"}, + ), + "width": ( + "INT", + {"default": 512, "min": 1, "max": 8192, "step": 1, "tooltip": "Width of the output mask in pixels"}, + ), + "height": ( + "INT", + { + "default": 512, + "min": 1, + "max": 8192, + "step": 1, + "tooltip": "Height of the output mask in pixels", + }, + ), + "mask_type": ( + ["Bounding Box", "Keypoints", "Both"], + { + "default": "Bounding Box", + "tooltip": "Type of mask to create - Bounding Box (rectangle around face), Keypoints (dots at facial points), or Both", + }, + ), + "face_selection": ( + ["All Faces", "First Face", "Largest Face", "Highest Confidence Face"], + { + "default": "All Faces", + "tooltip": "Which faces to include in the mask - all detected faces, only the first face, largest face, or most confident detection", + }, + ), + "feather_amount": ( + "INT", + { + "default": 0, + "min": 0, + "max": 100, + "step": 1, + "tooltip": "How much to soften the edges of the mask in pixels - higher values create smoother transitions", + }, + ), + "keypoint_radius": ( + "INT", + { + "default": 5, + "min": 1, + "max": 50, + "step": 1, + "tooltip": "Size of each keypoint in the mask when using Keypoints or Both mask types", + }, + ), + "invert": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, inverts the mask (white becomes black, black becomes white)", + }, + ), + "score_threshold": ( + "FLOAT", + { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence score for including face detections - higher values filter out less confident detections", + }, + ), }, "optional": { - "reference_image": ("IMAGE", {"tooltip": "Optional image to match dimensions - mask will be created with the same size as this image"}) - } + "reference_image": ( + "IMAGE", + {"tooltip": "Optional image to match dimensions - mask will be created with the same size as this image"}, + ) + }, } RETURN_TYPES = ("MASK",) FUNCTION = "create_mask" - - def create_mask(self, face_detections, width, height, mask_type, face_selection, - feather_amount, keypoint_radius, invert, score_threshold, reference_image=None): - + + def create_mask( + self, + face_detections, + width, + height, + mask_type, + face_selection, + feather_amount, + keypoint_radius, + invert, + score_threshold, + reference_image=None, + ): # If reference image provided, use its dimensions if reference_image is not None: height, width = reference_image.shape[1], reference_image.shape[2] - + batch_size = len(face_detections) masks = [] - + for i in range(batch_size): # Create empty mask mask = torch.zeros((height, width), dtype=torch.float32, device="cpu") detections_for_image = face_detections[i] - + if not detections_for_image: masks.append(mask) continue - + # Filter detections by score threshold filtered_detections = [] for detection in detections_for_image: if detection.score is None or detection.score >= score_threshold: filtered_detections.append(detection) - + # Select which faces to include based on face_selection parameter selected_detections = [] if face_selection == "All Faces": @@ -77,8 +143,7 @@ class FaceDetectionToMask: selected_detections = [filtered_detections[0]] elif face_selection == "Largest Face" and filtered_detections: # Find the detection with the largest area - largest_detection = max(filtered_detections, - key=lambda d: d.bounding_box.width * d.bounding_box.height) + largest_detection = max(filtered_detections, key=lambda d: d.bounding_box.width * d.bounding_box.height) selected_detections = [largest_detection] elif face_selection == "Highest Confidence Face" and filtered_detections: # Find the detection with the highest confidence score @@ -90,67 +155,64 @@ class FaceDetectionToMask: else: # If no detections have scores, fall back to first detection selected_detections = [filtered_detections[0]] - + for detection in selected_detections: # Get bounding box bbox = detection.bounding_box x1, y1 = max(0, bbox.origin_x), max(0, bbox.origin_y) - x2, y2 = min(width-1, bbox.origin_x + bbox.width), min(height-1, bbox.origin_y + bbox.height) - + x2, y2 = min(width - 1, bbox.origin_x + bbox.width), min(height - 1, bbox.origin_y + bbox.height) + if x2 <= x1 or y2 <= y1: continue # Skip invalid boxes - + # Fill bounding box region if requested if mask_type in ["Bounding Box", "Both"]: - mask[y1:y2+1, x1:x2+1] = 1.0 - + mask[y1 : y2 + 1, x1 : x2 + 1] = 1.0 + # Add keypoints if requested if mask_type in ["Keypoints", "Both"] and detection.keypoints: for kp in detection.keypoints: # Keypoints are in normalized coordinates, convert to pixel coordinates cx, cy = int(kp.x * width), int(kp.y * height) - + # Draw circle for each keypoint for y in range(max(0, cy - keypoint_radius), min(height, cy + keypoint_radius + 1)): for x in range(max(0, cx - keypoint_radius), min(width, cx + keypoint_radius + 1)): # Check if point is within radius - if (x - cx) ** 2 + (y - cy) ** 2 <= keypoint_radius ** 2: + if (x - cx) ** 2 + (y - cy) ** 2 <= keypoint_radius**2: mask[y, x] = 1.0 - + # Feather the edges if requested if feather_amount > 0: # Convert mask to numpy for scipy operations mask_np = mask.numpy() - + # Calculate distance transform if invert: distance = scipy.ndimage.distance_transform_edt(mask_np == 0) else: distance = scipy.ndimage.distance_transform_edt(mask_np > 0) - + # Create feathered edge feather_mask = np.clip(distance / feather_amount, 0, 1) - + if invert: feather_mask = 1.0 - feather_mask - + # Convert back to tensor mask = torch.from_numpy(feather_mask).float() elif invert: # Simple inversion without feathering mask = 1.0 - mask - + masks.append(mask) - + # Stack all masks into a batch mask_batch = torch.stack(masks) return (mask_batch,) -# Node registration -NODE_CLASS_MAPPINGS = { - "FaceDetectionToMask": FaceDetectionToMask -} -NODE_DISPLAY_NAME_MAPPINGS = { - "FaceDetectionToMask": "Face Detection to Mask (MediaPipe)" -} \ No newline at end of file +# Node registration +NODE_CLASS_MAPPINGS = {"FaceDetectionToMask": FaceDetectionToMask} + +NODE_DISPLAY_NAME_MAPPINGS = {"FaceDetectionToMask": "Face Detection to Mask (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py index 70fee8a..d9dd9fd 100644 --- a/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/face/detection/face_detection_visualization.py @@ -7,8 +7,10 @@ from ...common.base_visualization_nodes import BaseDetectionVisualizationNode logger = logging.getLogger(__name__) + class VisualizeFaceDetections(BaseDetectionVisualizationNode): """Draws face detection bounding boxes and keypoints onto images.""" + CATEGORY = "Realtime Nodes/MediaPipe Vision/Face/FaceDetection/Visualization" DESCRIPTION = "Visualizes face detections by drawing boxes and keypoints onto images. Allows customization of colors and sizes to highlight detected faces." @@ -16,29 +18,57 @@ class VisualizeFaceDetections(BaseDetectionVisualizationNode): def INPUT_TYPES(cls): return { "required": { - "face_detections": ("FACE_DETECTIONS", {"tooltip": "Face detection results from the Face Detector node"}), + "face_detections": ( + "FACE_DETECTIONS", + {"tooltip": "Face detection results from the Face Detector node"}, + ), "image": ("IMAGE", {"tooltip": "The image on which to draw face detection visualizations"}), }, "optional": { - "box_color": ("STRING", {"default": "#00FF00", "tooltip": "Hex color code for the face bounding boxes (default: green)"}), - "box_thickness": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Thickness of the face bounding box lines"}), - "keypoint_color": ("STRING", {"default": "#FF0000", "tooltip": "Hex color code for facial keypoints (default: red)"}), - "keypoint_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of the dots representing facial keypoints"}), - "label_color": ("STRING", {"default": "#FFFFFF", "tooltip": "Hex color code for the text labels (default: white)"}), - "label_bg_color": ("STRING", {"default": "#00FF00", "tooltip": "Hex color code for the label background (default: green)"}), - } + "box_color": ( + "STRING", + {"default": "#00FF00", "tooltip": "Hex color code for the face bounding boxes (default: green)"}, + ), + "box_thickness": ( + "INT", + {"default": 2, "min": 1, "max": 10, "tooltip": "Thickness of the face bounding box lines"}, + ), + "keypoint_color": ( + "STRING", + {"default": "#FF0000", "tooltip": "Hex color code for facial keypoints (default: red)"}, + ), + "keypoint_radius": ( + "INT", + {"default": 2, "min": 1, "max": 10, "tooltip": "Size of the dots representing facial keypoints"}, + ), + "label_color": ( + "STRING", + {"default": "#FFFFFF", "tooltip": "Hex color code for the text labels (default: white)"}, + ), + "label_bg_color": ( + "STRING", + {"default": "#00FF00", "tooltip": "Hex color code for the label background (default: green)"}, + ), + }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "draw_visualizations" - def draw_visualizations(self, face_detections, image, - box_color="#00FF00", box_thickness=2, - keypoint_color="#FF0000", keypoint_radius=2, - label_color="#FFFFFF", label_bg_color="#00FF00"): + def draw_visualizations( + self, + face_detections, + image, + box_color="#00FF00", + box_thickness=2, + keypoint_color="#FF0000", + keypoint_radius=2, + label_color="#FFFFFF", + label_bg_color="#00FF00", + ): """Main function that processes the batch of images and detections.""" return self.process_batch( - face_detections, + face_detections, image, self.draw_faces, box_color=box_color, @@ -46,26 +76,35 @@ class VisualizeFaceDetections(BaseDetectionVisualizationNode): keypoint_color=keypoint_color, keypoint_radius=keypoint_radius, label_color=label_color, - label_bg_color=label_bg_color + label_bg_color=label_bg_color, ) - - def draw_faces(self, image, width, height, detections, - box_color="#00FF00", box_thickness=2, - keypoint_color="#FF0000", keypoint_radius=2, - label_color="#FFFFFF", label_bg_color="#00FF00"): + + def draw_faces( + self, + image, + width, + height, + detections, + box_color="#00FF00", + box_thickness=2, + keypoint_color="#FF0000", + keypoint_radius=2, + label_color="#FFFFFF", + label_bg_color="#00FF00", + ): """Callback function to draw faces on a single image.""" # Convert hex colors to BGR tuples for cv2 box_color_bgr = self._hex_to_bgr(box_color) keypoint_color_bgr = self._hex_to_bgr(keypoint_color) - + for detection in detections: bbox = detection.bounding_box x1, y1 = bbox.origin_x, bbox.origin_y x2, y2 = bbox.origin_x + bbox.width, bbox.origin_y + bbox.height - + # Draw bounding box cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), box_color_bgr, box_thickness) - + # Draw Keypoints if detection.keypoints: for kp in detection.keypoints: @@ -74,34 +113,25 @@ class VisualizeFaceDetections(BaseDetectionVisualizationNode): # Prepare label text (confidence score) label_text = f"Face: {detection.score:.2f}" if detection.score is not None else "Face" - + # Draw label with background try: - self.draw_text_with_background( - image, - label_text, - (int(x1), int(y1)), - text_color=label_color, - bg_color=label_bg_color - ) + self.draw_text_with_background(image, label_text, (int(x1), int(y1)), text_color=label_color, bg_color=label_bg_color) except Exception as e: logger.error(f"Error drawing text label: {e}") # Fallback using basic cv2 text cv2.putText( - image, - label_text, - (int(x1) + 2, int(y1) + 15), - cv2.FONT_HERSHEY_SIMPLEX, - 0.5, - self._hex_to_bgr(label_color), - 1 + image, + label_text, + (int(x1) + 2, int(y1) + 15), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + self._hex_to_bgr(label_color), + 1, ) -# Node registration -NODE_CLASS_MAPPINGS = { - "VisualizeFaceDetections": VisualizeFaceDetections -} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizeFaceDetections": "Visualize Face Detections (MediaPipe)" -} \ No newline at end of file +# Node registration +NODE_CLASS_MAPPINGS = {"VisualizeFaceDetections": VisualizeFaceDetections} + +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizeFaceDetections": "Visualize Face Detections (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py index 68cb85f..7b4c5fd 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/blendshape.py +++ b/node_wrappers/mediapipe_vision/face/landmark/blendshape.py @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) ALL_BLENDSHAPE_CATEGORIES = get_blendshape_categories() if not ALL_BLENDSHAPE_CATEGORIES: logger.warning("Could not retrieve blendshape categories. Blendshape nodes might not function correctly.") - ALL_BLENDSHAPE_CATEGORIES = ["_neutral"] # Provide a fallback + ALL_BLENDSHAPE_CATEGORIES = ["_neutral"] # Provide a fallback _category = "Realtime Nodes/MediaPipe Vision/Face/Blendshape" @@ -28,8 +28,8 @@ _category = "Realtime Nodes/MediaPipe Vision/Face/Blendshape" def get_blendshape_score(blendshapes_result: BLENDSHAPES_LIST, category_name: str) -> Optional[float]: """Extracts the score for a specific blendshape category from the result.""" if not blendshapes_result or not blendshapes_result[0] or not blendshapes_result[0][0]: - #logger.debug(f"No blendshapes found in result.") - return None # No blendshape data + # logger.debug(f"No blendshapes found in result.") + return None # No blendshape data # BLENDSHAPES structure: List[List[Dict[str, float]]] # We typically care about the first face in the first batch item @@ -44,12 +44,15 @@ def get_blendshape_score(blendshapes_result: BLENDSHAPES_LIST, category_name: st elif isinstance(first_face_blendshapes, list): # Try to handle the list format with objects that have category_name attribute try: - target_blendshape = next((bs for bs in first_face_blendshapes if hasattr(bs, 'category_name') and bs.category_name == category_name), None) + target_blendshape = next( + (bs for bs in first_face_blendshapes if hasattr(bs, "category_name") and bs.category_name == category_name), + None, + ) if target_blendshape: return target_blendshape.score except Exception as e: logger.error(f"Error processing blendshape list: {e}") - + logger.warning(f"Blendshape category '{category_name}' not found in results. Type: {type(first_face_blendshapes)}") return None @@ -57,6 +60,7 @@ def get_blendshape_score(blendshapes_result: BLENDSHAPES_LIST, category_name: st # --- Control Nodes --- class BlendshapeControlFloatNode: """Outputs a scaled FLOAT based on the score of a selected blendshape.""" + CATEGORY = _category RETURN_TYPES = ("FLOAT",) FUNCTION = "execute" @@ -66,20 +70,57 @@ class BlendshapeControlFloatNode: def INPUT_TYPES(cls): return { "required": { - "blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}), - "blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "mouthSmileLeft", - "tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)"}), + "blendshapes": ( + "BLENDSHAPES_LIST", + {"tooltip": "Facial expression data from the Face Landmarker node"}, + ), + "blendshape_name": ( + ALL_BLENDSHAPE_CATEGORIES, + { + "default": "mouthSmileLeft", + "tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)", + }, + ), # Blendshape scores are typically 0.0 to 1.0, but allow flexibility - "score_min": ("FLOAT", {"default": 0.0, "step": 0.01, - "tooltip": "The minimum facial expression intensity to consider (usually 0.0)"}), - "score_max": ("FLOAT", {"default": 1.0, "step": 0.01, - "tooltip": "The maximum facial expression intensity to consider (usually 1.0)"}), - "output_min_float": ("FLOAT", {"default": 0.0, "step": 0.01, - "tooltip": "The minimum output value when expression is at or below score_min"}), - "output_max_float": ("FLOAT", {"default": 1.0, "step": 0.01, - "tooltip": "The maximum output value when expression is at or above score_max"}), - "clamp": ("BOOLEAN", {"default": True, - "tooltip": "When enabled, restricts output values to stay within min and max range"}), + "score_min": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "The minimum facial expression intensity to consider (usually 0.0)", + }, + ), + "score_max": ( + "FLOAT", + { + "default": 1.0, + "step": 0.01, + "tooltip": "The maximum facial expression intensity to consider (usually 1.0)", + }, + ), + "output_min_float": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "The minimum output value when expression is at or below score_min", + }, + ), + "output_max_float": ( + "FLOAT", + { + "default": 1.0, + "step": 0.01, + "tooltip": "The maximum output value when expression is at or above score_max", + }, + ), + "clamp": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, restricts output values to stay within min and max range", + }, + ), } } @@ -92,8 +133,10 @@ class BlendshapeControlFloatNode: output_val = scaled_value if scaled_value is not None else float(output_min_float) return (output_val,) + class BlendshapeControlIntNode: """Outputs a scaled INT based on the score of a selected blendshape.""" + CATEGORY = _category RETURN_TYPES = ("INT",) FUNCTION = "execute" @@ -103,19 +146,54 @@ class BlendshapeControlIntNode: def INPUT_TYPES(cls): return { "required": { - "blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}), - "blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "mouthSmileLeft", - "tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)"}), - "score_min": ("FLOAT", {"default": 0.0, "step": 0.01, - "tooltip": "The minimum facial expression intensity to consider (usually 0.0)"}), - "score_max": ("FLOAT", {"default": 1.0, "step": 0.01, - "tooltip": "The maximum facial expression intensity to consider (usually 1.0)"}), - "output_min_int": ("INT", {"default": 0, - "tooltip": "The minimum integer output value when expression is at or below score_min"}), - "output_max_int": ("INT", {"default": 100, - "tooltip": "The maximum integer output value when expression is at or above score_max"}), - "clamp": ("BOOLEAN", {"default": True, - "tooltip": "When enabled, restricts output values to stay within min and max range"}), + "blendshapes": ( + "BLENDSHAPES_LIST", + {"tooltip": "Facial expression data from the Face Landmarker node"}, + ), + "blendshape_name": ( + ALL_BLENDSHAPE_CATEGORIES, + { + "default": "mouthSmileLeft", + "tooltip": "The specific facial expression to monitor (e.g., smile, frown, raised eyebrows)", + }, + ), + "score_min": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "The minimum facial expression intensity to consider (usually 0.0)", + }, + ), + "score_max": ( + "FLOAT", + { + "default": 1.0, + "step": 0.01, + "tooltip": "The maximum facial expression intensity to consider (usually 1.0)", + }, + ), + "output_min_int": ( + "INT", + { + "default": 0, + "tooltip": "The minimum integer output value when expression is at or below score_min", + }, + ), + "output_max_int": ( + "INT", + { + "default": 100, + "tooltip": "The maximum integer output value when expression is at or above score_max", + }, + ), + "clamp": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, restricts output values to stay within min and max range", + }, + ), } } @@ -132,6 +210,7 @@ class BlendshapeControlIntNode: # --- Trigger Node --- class BlendshapeTriggerNode: """Outputs a BOOLEAN trigger if a blendshape score crosses a threshold.""" + CATEGORY = _category RETURN_TYPES = ("BOOLEAN",) FUNCTION = "execute" @@ -141,13 +220,31 @@ class BlendshapeTriggerNode: def INPUT_TYPES(cls): return { "required": { - "blendshapes": ("BLENDSHAPES_LIST", {"tooltip": "Facial expression data from the Face Landmarker node"}), - "blendshape_name": (ALL_BLENDSHAPE_CATEGORIES, {"default": "jawOpen", - "tooltip": "The specific facial expression to monitor (e.g., smile, frown, open mouth)"}), - "threshold": ("FLOAT", {"default": 0.5, "step": 0.01, - "tooltip": "The expression intensity threshold value that triggers the condition"}), - "condition": (["Above", "Below", "Equals", "Not Equals"], - {"tooltip": "The condition that determines when to trigger - Above (expression > threshold), Below (expression < threshold), etc."}), + "blendshapes": ( + "BLENDSHAPES_LIST", + {"tooltip": "Facial expression data from the Face Landmarker node"}, + ), + "blendshape_name": ( + ALL_BLENDSHAPE_CATEGORIES, + { + "default": "jawOpen", + "tooltip": "The specific facial expression to monitor (e.g., smile, frown, open mouth)", + }, + ), + "threshold": ( + "FLOAT", + { + "default": 0.5, + "step": 0.01, + "tooltip": "The expression intensity threshold value that triggers the condition", + }, + ), + "condition": ( + ["Above", "Below", "Equals", "Not Equals"], + { + "tooltip": "The condition that determines when to trigger - Above (expression > threshold), Below (expression < threshold), etc." + }, + ), } } @@ -179,4 +276,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "BlendshapeControlFloat": "Blendshape Control (Float)", "BlendshapeControlInt": "Blendshape Control (Int)", "BlendshapeTrigger": "Blendshape Trigger", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py index 700324d..8fc9dde 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection.py @@ -13,11 +13,14 @@ from ...common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark" + + # Inherit from Base Loader class MediaPipeFaceLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Face Landmarker models.""" + TASK_TYPE = "face_landmarker" - RETURN_TYPES = ("FACE_LANDMARKER_MODEL_INFO",) # Specific return type + RETURN_TYPES = ("FACE_LANDMARKER_MODEL_INFO",) # Specific return type RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Face Landmarker model which detects facial landmarks, expressions and face orientation. Required before using the Face Landmarker node." @@ -41,37 +44,90 @@ class MediaPipeFaceLandmarkerNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add face-specific parameters - inputs["required"].update({ - "max_results": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1, - "tooltip": "Maximum number of faces to analyze - higher values detect landmarks on more faces but use more processing power"}), - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for detecting faces - lower finds more faces but may include false positives"}), - "min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence that a face is present - lower values may detect faces that are less clear"}), - "min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence for tracking a face between frames - lower values maintain tracking longer but may drift"}), - "output_blendshapes": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, outputs facial expression data (smile, frown, etc.) useful for animation or emotion analysis"}), - "output_transform_matrix": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, outputs 3D face orientation data (head pose) useful for AR/VR applications"}), - }) - + inputs["required"].update( + { + "max_results": ( + "INT", + { + "default": 1, + "min": 1, + "max": 10, + "step": 1, + "tooltip": "Maximum number of faces to analyze - higher values detect landmarks on more faces but use more processing power", + }, + ), + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for detecting faces - lower finds more faces but may include false positives", + }, + ), + "min_presence_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence that a face is present - lower values may detect faces that are less clear", + }, + ), + "min_tracking_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence for tracking a face between frames - lower values maintain tracking longer but may drift", + }, + ), + "output_blendshapes": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, outputs facial expression data (smile, frown, etc.) useful for animation or emotion analysis", + }, + ), + "output_transform_matrix": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, outputs 3D face orientation data (head pose) useful for AR/VR applications", + }, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, max_results: int, - min_confidence: float, min_presence_confidence: float, - min_tracking_confidence: float, output_blendshapes: bool, - output_transform_matrix: bool, running_mode: str, delegate: str): + def detect( + self, + image: torch.Tensor, + model_info: dict, + max_results: int, + min_confidence: float, + min_presence_confidence: float, + min_tracking_confidence: float, + output_blendshapes: bool, + output_transform_matrix: bool, + running_mode: str, + delegate: str, + ): """Performs face landmark detection with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Perform detection with all parameters face_landmarks_batch, blendshapes_batch, matrices_batch = detector.detect( image, @@ -82,16 +138,17 @@ class MediaPipeFaceLandmarkerNode(BaseMediaPipeDetectorNode): output_blendshapes=output_blendshapes, output_transform_matrix=output_transform_matrix, running_mode=running_mode, - delegate=delegate + delegate=delegate, ) - + # Return the results with optional empty lists for disabled outputs return ( - face_landmarks_batch, - blendshapes_batch if output_blendshapes else [], - matrices_batch if output_transform_matrix else [] + face_landmarks_batch, + blendshapes_batch if output_blendshapes else [], + matrices_batch if output_transform_matrix else [], ) + # Define mappings for ComfyUI NODE_CLASS_MAPPINGS = { "MediaPipeFaceLandmarkerModelLoader": MediaPipeFaceLandmarkerModelLoaderNode, @@ -101,4 +158,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeFaceLandmarkerModelLoader": "Load Face Landmarker Model (MediaPipe)", "MediaPipeFaceLandmarker": "Face Landmarker (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py index cd82436..0c37d41 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_landmark_detection_visualization.py @@ -11,6 +11,7 @@ logger = logging.getLogger(__name__) # Define standard face connections FACE_CONNECTIONS = mp.solutions.face_mesh.FACEMESH_TESSELATION + class VisualizeFaceLandmarks(BaseLandmarkVisualizationNode): LANDMARKS_TYPE = "FACE_LANDMARKS" CONNECTIONS = FACE_CONNECTIONS @@ -22,15 +23,12 @@ class VisualizeFaceLandmarks(BaseLandmarkVisualizationNode): SUPPORTS_VISIBILITY = False CATEGORY = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Visualization" DESCRIPTION = "Visualizes facial landmarks by drawing a mesh of points and connections on the face. Creates a detailed facial wireframe that shows the precise contours and features of detected faces." - + def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]: - return detection_result + return detection_result + # Node registration -NODE_CLASS_MAPPINGS = { - "VisualizeFaceLandmarks": VisualizeFaceLandmarks -} +NODE_CLASS_MAPPINGS = {"VisualizeFaceLandmarks": VisualizeFaceLandmarks} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizeFaceLandmarks": "Visualize Face Landmarks (MediaPipe)" -} \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizeFaceLandmarks": "Visualize Face Landmarks (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py index b7c2620..067d39e 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/face_masking.py +++ b/node_wrappers/mediapipe_vision/face/landmark/face_masking.py @@ -9,9 +9,46 @@ from .....src.utils.image import create_mask_from_points logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Masking" -# --- Face Part Definitions --- -FACE_PART_INDICES = { # NOTE: These indices are EXAMPLES and NEED VERIFICATION! - "FACE_OVAL": [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109], +# --- Face Part Definitions --- +FACE_PART_INDICES = { # NOTE: These indices are EXAMPLES and NEED VERIFICATION! + "FACE_OVAL": [ + 10, + 338, + 297, + 332, + 284, + 251, + 389, + 356, + 454, + 323, + 361, + 288, + 397, + 365, + 379, + 378, + 400, + 377, + 152, + 148, + 176, + 149, + 150, + 136, + 172, + 58, + 132, + 93, + 234, + 127, + 162, + 21, + 54, + 103, + 67, + 109, + ], "LIPS_OUTER": [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 62, 185, 40, 39, 37, 0, 267, 269, 270, 409, 292], "LIPS_INNER": [78, 95, 88, 178, 87, 14, 317, 402, 318, 324, 308, 76, 191, 81, 82, 13, 312, 311, 310, 415, 306], "LEFT_EYE": [33, 7, 163, 144, 145, 153, 154, 155, 133, 246, 161, 160, 159, 158, 157, 173], @@ -22,9 +59,11 @@ FACE_PART_INDICES = { # NOTE: These indices are EXAMPLES and NEED VERIFICATION! "NOSE_TIP_LOWER": [5, 4, 1, 19, 94, 2], } -# --- Node Definition --- + +# --- Node Definition --- class MaskFromFaceLandmarksNode: """Creates a mask for a specific facial part from Face Landmarks.""" + CATEGORY = _category DESCRIPTION = "Creates precise masks for specific facial features like eyes, lips, nose, or the entire face. Perfect for targeted editing, inpainting, or creating effects on specific facial regions." @@ -33,10 +72,21 @@ class MaskFromFaceLandmarksNode: part_names = list(FACE_PART_INDICES.keys()) return { "required": { - "face_landmarks": ("FACE_LANDMARKS", {"tooltip": "Facial landmark data from the Face Landmarker node containing the precise locations of facial features"}), - "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}), - "part_name": (part_names, {"default": "FACE_OVAL", - "tooltip": "Choose which facial feature to create a mask for - FACE_OVAL (whole face), LIPS_OUTER (mouth), LEFT_EYE, etc."}), + "face_landmarks": ( + "FACE_LANDMARKS", + {"tooltip": "Facial landmark data from the Face Landmarker node containing the precise locations of facial features"}, + ), + "image_for_dimensions": ( + "IMAGE", + {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}, + ), + "part_name": ( + part_names, + { + "default": "FACE_OVAL", + "tooltip": "Choose which facial feature to create a mask for - FACE_OVAL (whole face), LIPS_OUTER (mouth), LEFT_EYE, etc.", + }, + ), } } @@ -46,7 +96,7 @@ class MaskFromFaceLandmarksNode: def generate_mask(self, face_landmarks: List[List[List[LandmarkPoint]]], image_for_dimensions: torch.Tensor, part_name: str): if part_name not in FACE_PART_INDICES: raise ValueError(f"Unknown face part name: '{part_name}'. Available: {list(FACE_PART_INDICES.keys())}") - + # Get dimensions from image if image_for_dimensions.dim() != 4: raise ValueError("image_for_dimensions must be BHWC format.") @@ -57,22 +107,24 @@ class MaskFromFaceLandmarksNode: # Validate batch sizes match if len(face_landmarks) != input_batch_size: - logger.warning(f"Batch size mismatch between face_landmarks ({len(face_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.") + logger.warning( + f"Batch size mismatch between face_landmarks ({len(face_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size." + ) # Adjust effective batch size based on landmarks if mismatched, though ideally they match batch_size = len(face_landmarks) else: - batch_size = input_batch_size - + batch_size = input_batch_size + required_indices = FACE_PART_INDICES[part_name] output_masks = [] for i in range(batch_size): - landmarks_for_image = face_landmarks[i] # List[List[LandmarkPoint]] (Faces in image) + landmarks_for_image = face_landmarks[i] # List[List[LandmarkPoint]] (Faces in image) combined_points_for_image = [] # Combine points from all detected faces for the selected part - for face_result_landmarks in landmarks_for_image: # face_result_landmarks is List[LandmarkPoint] - points_dict = {lm.index: lm for lm in face_result_landmarks} # Iterate directly over the list + for face_result_landmarks in landmarks_for_image: # face_result_landmarks is List[LandmarkPoint] + points_dict = {lm.index: lm for lm in face_result_landmarks} # Iterate directly over the list part_points = [] for idx in required_indices: if idx in points_dict: @@ -82,27 +134,28 @@ class MaskFromFaceLandmarksNode: part_points.append((px, py)) # Removed warning here to avoid spam if landmarks are just missing combined_points_for_image.extend(part_points) - + # Create mask from all collected points for this image mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device) output_masks.append(mask_tensor) if not output_masks: - logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") - output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) + logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") + output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) else: - # Stack along batch dimension, ensure channel dim is present - output_batch = torch.stack(output_masks, dim=0) # Shape becomes (B, 1, H, W) - + # Stack along batch dimension, ensure channel dim is present + output_batch = torch.stack(output_masks, dim=0) # Shape becomes (B, 1, H, W) + return (output_batch,) + # Remove GetImageDimensionsNode class definition -# --- Mappings --- +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MaskFromFaceLandmarks": MaskFromFaceLandmarksNode, } NODE_DISPLAY_NAME_MAPPINGS = { "MaskFromFaceLandmarks": "Mask From Face Landmarks (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py index ab3eba9..c8efb44 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/head_pose.py +++ b/node_wrappers/mediapipe_vision/face/landmark/head_pose.py @@ -15,6 +15,7 @@ logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/HeadPose" + # --- Helper Functions --- def get_first_transform_matrix(matrix_result: TRANSFORM_MATRIX_LIST) -> Optional[np.ndarray]: """Extracts the first 4x4 transformation matrix from the result.""" @@ -27,6 +28,7 @@ def get_first_transform_matrix(matrix_result: TRANSFORM_MATRIX_LIST) -> Optional # FACE_TRANSFORM_MATRICES structure: List[List[np.ndarray]] return matrix_result[0][0] + def decompose_transform_matrix(matrix: np.ndarray) -> Optional[Dict[str, float]]: """Decomposes a 4x4 matrix into translation (x, y, z) and Euler angles (pitch, yaw, roll).""" if matrix is None or matrix.shape != (4, 4): @@ -39,16 +41,16 @@ def decompose_transform_matrix(matrix: np.ndarray) -> Optional[Dict[str, float]] # Convert rotation matrix to Euler angles (ZYX convention common for head pose) # Note: Scipy uses intrinsic rotations. ZYX order corresponds to yaw, pitch, roll. rotation = R.from_matrix(rotation_matrix) - euler_angles_rad = rotation.as_euler('zyx', degrees=False) + euler_angles_rad = rotation.as_euler("zyx", degrees=False) # Convert radians to degrees for easier interpretation euler_angles_deg = np.degrees(euler_angles_rad) - + return { "x": translation[0], "y": translation[1], "z": translation[2], - "yaw": euler_angles_deg[0], # Rotation around Y (vertical) axis - "pitch": euler_angles_deg[1], # Rotation around X (side-to-side) axis + "yaw": euler_angles_deg[0], # Rotation around Y (vertical) axis + "pitch": euler_angles_deg[1], # Rotation around X (side-to-side) axis "roll": euler_angles_deg[2], # Rotation around Z (front-to-back) axis } except Exception as e: @@ -59,8 +61,10 @@ def decompose_transform_matrix(matrix: np.ndarray) -> Optional[Dict[str, float]] # --- Control Node --- HEAD_POSE_COMPONENTS = ["x", "y", "z", "pitch", "yaw", "roll"] + class HeadPoseControlFloatNode: """Outputs a scaled FLOAT based on a selected component (translation/rotation) of the head pose.""" + CATEGORY = _category RETURN_TYPES = ("FLOAT",) FUNCTION = "execute" @@ -70,39 +74,80 @@ class HeadPoseControlFloatNode: def INPUT_TYPES(cls): # Example ranges - these might need tuning based on expected matrix values default_ranges = { - "x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (-1.0, 0.5), # Example translation ranges - "pitch": (-45.0, 45.0), "yaw": (-60.0, 60.0), "roll": (-45.0, 45.0) # Example angle ranges (degrees) + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-1.0, 0.5), # Example translation ranges + "pitch": (-45.0, 45.0), + "yaw": (-60.0, 60.0), + "roll": (-45.0, 45.0), # Example angle ranges (degrees) } default_component = "yaw" default_min, default_max = default_ranges[default_component] - + return { "required": { - "face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}), - "component": (HEAD_POSE_COMPONENTS, {"default": default_component, - "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}), - "value_min": ("FLOAT", {"default": default_min, "step": 0.1, - "tooltip": "The minimum value of the selected movement component to consider"}), - "value_max": ("FLOAT", {"default": default_max, "step": 0.1, - "tooltip": "The maximum value of the selected movement component to consider"}), - "output_min_float": ("FLOAT", {"default": 0.0, "step": 0.01, - "tooltip": "The minimum output value when head position is at or below value_min"}), - "output_max_float": ("FLOAT", {"default": 1.0, "step": 0.01, - "tooltip": "The maximum output value when head position is at or above value_max"}), - "clamp": ("BOOLEAN", {"default": True, - "tooltip": "When enabled, restricts output values to stay within min and max range"}), + "face_transform_matrices": ( + "TRANSFORM_MATRIX_LIST", + {"tooltip": "3D head orientation data from the Face Landmarker node"}, + ), + "component": ( + HEAD_POSE_COMPONENTS, + { + "default": default_component, + "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)", + }, + ), + "value_min": ( + "FLOAT", + { + "default": default_min, + "step": 0.1, + "tooltip": "The minimum value of the selected movement component to consider", + }, + ), + "value_max": ( + "FLOAT", + { + "default": default_max, + "step": 0.1, + "tooltip": "The maximum value of the selected movement component to consider", + }, + ), + "output_min_float": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "The minimum output value when head position is at or below value_min", + }, + ), + "output_max_float": ( + "FLOAT", + { + "default": 1.0, + "step": 0.01, + "tooltip": "The maximum output value when head position is at or above value_max", + }, + ), + "clamp": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, restricts output values to stay within min and max range", + }, + ), } } def execute(self, face_transform_matrices, component, value_min, value_max, output_min_float, output_max_float, clamp): matrix = get_first_transform_matrix(face_transform_matrices) decomposed = decompose_transform_matrix(matrix) - + raw_value = None if decomposed and component in decomposed: raw_value = decomposed[component] else: - logger.debug(f"Could not get head pose component '{component}'") + logger.debug(f"Could not get head pose component '{component}'") scaled_value = scale_value(raw_value, value_min, value_max, output_min_float, output_max_float, clamp) @@ -110,8 +155,10 @@ class HeadPoseControlFloatNode: output_val = scaled_value if scaled_value is not None else float(output_min_float) return (output_val,) + class HeadPoseControlIntNode: """Outputs a scaled INT based on a selected component (translation/rotation) of the head pose.""" + CATEGORY = _category RETURN_TYPES = ("INT",) FUNCTION = "execute" @@ -121,40 +168,79 @@ class HeadPoseControlIntNode: def INPUT_TYPES(cls): # Share the same default range logic as Float node default_ranges = { - "x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (-1.0, 0.5), - "pitch": (-45.0, 45.0), "yaw": (-60.0, 60.0), "roll": (-45.0, 45.0) + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-1.0, 0.5), + "pitch": (-45.0, 45.0), + "yaw": (-60.0, 60.0), + "roll": (-45.0, 45.0), } default_component = "yaw" default_min, default_max = default_ranges[default_component] - + return { "required": { - "face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}), - "component": (HEAD_POSE_COMPONENTS, {"default": default_component, - "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}), - "value_min": ("FLOAT", {"default": default_min, "step": 0.1, - "tooltip": "The minimum value of the selected movement component to consider"}), - "value_max": ("FLOAT", {"default": default_max, "step": 0.1, - "tooltip": "The maximum value of the selected movement component to consider"}), - "output_min_int": ("INT", {"default": 0, - "tooltip": "The minimum integer output value when head position is at or below value_min"}), - "output_max_int": ("INT", {"default": 100, - "tooltip": "The maximum integer output value when head position is at or above value_max"}), - "clamp": ("BOOLEAN", {"default": True, - "tooltip": "When enabled, restricts output values to stay within min and max range"}), + "face_transform_matrices": ( + "TRANSFORM_MATRIX_LIST", + {"tooltip": "3D head orientation data from the Face Landmarker node"}, + ), + "component": ( + HEAD_POSE_COMPONENTS, + { + "default": default_component, + "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)", + }, + ), + "value_min": ( + "FLOAT", + { + "default": default_min, + "step": 0.1, + "tooltip": "The minimum value of the selected movement component to consider", + }, + ), + "value_max": ( + "FLOAT", + { + "default": default_max, + "step": 0.1, + "tooltip": "The maximum value of the selected movement component to consider", + }, + ), + "output_min_int": ( + "INT", + { + "default": 0, + "tooltip": "The minimum integer output value when head position is at or below value_min", + }, + ), + "output_max_int": ( + "INT", + { + "default": 100, + "tooltip": "The maximum integer output value when head position is at or above value_max", + }, + ), + "clamp": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, restricts output values to stay within min and max range", + }, + ), } } def execute(self, face_transform_matrices, component, value_min, value_max, output_min_int, output_max_int, clamp): matrix = get_first_transform_matrix(face_transform_matrices) decomposed = decompose_transform_matrix(matrix) - + raw_value = None if decomposed and component in decomposed: raw_value = decomposed[component] else: logger.debug(f"Could not get head pose component '{component}'") - + scaled_value = scale_value(raw_value, value_min, value_max, output_min_int, output_max_int, clamp) # Default to min output value on error @@ -165,6 +251,7 @@ class HeadPoseControlIntNode: # --- Trigger Node --- class HeadPoseTriggerNode: """Outputs a BOOLEAN trigger if a head pose component crosses a threshold.""" + CATEGORY = _category RETURN_TYPES = ("BOOLEAN",) FUNCTION = "execute" @@ -174,25 +261,43 @@ class HeadPoseTriggerNode: def INPUT_TYPES(cls): return { "required": { - "face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node"}), - "component": (HEAD_POSE_COMPONENTS, {"default": "yaw", - "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)"}), - "threshold": ("FLOAT", {"default": 30.0, "step": 0.5, - "tooltip": "The angle or position value that triggers the condition (example: 30 degrees of head turn)"}), - "condition": (["Above", "Below", "Equals", "Not Equals"], - {"tooltip": "The condition that determines when to trigger - Above (movement > threshold), Below (movement < threshold), etc."}), + "face_transform_matrices": ( + "TRANSFORM_MATRIX_LIST", + {"tooltip": "3D head orientation data from the Face Landmarker node"}, + ), + "component": ( + HEAD_POSE_COMPONENTS, + { + "default": "yaw", + "tooltip": "Which aspect of head movement to track - yaw (left/right), pitch (up/down), roll (tilt), or position (x/y/z)", + }, + ), + "threshold": ( + "FLOAT", + { + "default": 30.0, + "step": 0.5, + "tooltip": "The angle or position value that triggers the condition (example: 30 degrees of head turn)", + }, + ), + "condition": ( + ["Above", "Below", "Equals", "Not Equals"], + { + "tooltip": "The condition that determines when to trigger - Above (movement > threshold), Below (movement < threshold), etc." + }, + ), } } def execute(self, face_transform_matrices, component, threshold, condition): matrix = get_first_transform_matrix(face_transform_matrices) decomposed = decompose_transform_matrix(matrix) - + raw_value = None if decomposed and component in decomposed: raw_value = decomposed[component] else: - logger.debug(f"Could not get head pose component '{component}' for trigger") + logger.debug(f"Could not get head pose component '{component}' for trigger") triggered = False if raw_value is not None: @@ -219,4 +324,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "HeadPoseControlFloat": "Head Pose Control (Float)", "HeadPoseControlInt": "Head Pose Control (Int)", "HeadPoseTrigger": "Head Pose Trigger", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/landmark/position.py b/node_wrappers/mediapipe_vision/face/landmark/position.py index 97e15ba..5e7e5f6 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/position.py +++ b/node_wrappers/mediapipe_vision/face/landmark/position.py @@ -3,13 +3,15 @@ import logging from .....src.mediapipe_vision.common.position import LandmarkPositionBase from .....src.mediapipe_vision.landmark_definitions import FACE_LANDMARK_TOOLTIP -logger = logging.getLogger(__name__) +logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceLandmark/Position" + class FaceLandmarkPositionNode(LandmarkPositionBase): """Extracts position lists (x, y, z, vis, pres) for a specific landmark index - from FACE_LANDMARKS across the batch.""" + from FACE_LANDMARKS across the batch.""" + CATEGORY = _category # Define the specific input type string this node expects LANDMARKS_TYPE = "FACE_LANDMARKS" @@ -20,12 +22,16 @@ class FaceLandmarkPositionNode(LandmarkPositionBase): # Define the specific input type for this node return { "required": { - "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), # Use the class variable + "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), # Use the class variable "landmark_index": ("INT", {"default": 0, "min": 0, "tooltip": FACE_LANDMARK_TOOLTIP}), - "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the face detection to use (0=first detected face)"}), + "result_index": ( + "INT", + {"default": 0, "min": 0, "tooltip": "Index of the face detection to use (0=first detected face)"}, + ), } } + # --- Node Mappings --- NODE_CLASS_MAPPINGS = { "FaceLandmarkPosition": FaceLandmarkPositionNode, @@ -33,4 +39,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "FaceLandmarkPosition": "Face Landmark Position Extractor (Batch)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py b/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py index c6f1190..ffb8626 100644 --- a/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py +++ b/node_wrappers/mediapipe_vision/face/landmark/texture_warp.py @@ -5,11 +5,13 @@ logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/Experimental" + class FaceTextureWarpNode: """(Experimental) Attempts to warp a texture based on the face transform matrix. - NOTE: This is a conceptual placeholder. Actual warping requires significant - integration with sampling or image processing libraries. + NOTE: This is a conceptual placeholder. Actual warping requires significant + integration with sampling or image processing libraries. """ + CATEGORY = _category RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" @@ -19,35 +21,66 @@ class FaceTextureWarpNode: def INPUT_TYPES(cls): return { "required": { - "target_image": ("IMAGE", {"tooltip": "The base image containing a face onto which the texture will be applied"}), - "texture_image": ("IMAGE", {"tooltip": "The texture image to be warped onto the face (could be another face, pattern, or design)"}), - "face_transform_matrices": ("TRANSFORM_MATRIX_LIST", {"tooltip": "3D head orientation data from the Face Landmarker node to determine warping angles"}), + "target_image": ( + "IMAGE", + {"tooltip": "The base image containing a face onto which the texture will be applied"}, + ), + "texture_image": ( + "IMAGE", + {"tooltip": "The texture image to be warped onto the face (could be another face, pattern, or design)"}, + ), + "face_transform_matrices": ( + "TRANSFORM_MATRIX_LIST", + {"tooltip": "3D head orientation data from the Face Landmarker node to determine warping angles"}, + ), # Potentially add mask input for blending? }, "optional": { - "offset_x": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Horizontal positioning adjustment for the texture overlay"}), - "offset_y": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Vertical positioning adjustment for the texture overlay"}), - "scale": ("FLOAT", {"default": 1.0, "step": 0.01, "tooltip": "Size adjustment for the texture overlay (1.0 = original size)"}), - } + "offset_x": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "Horizontal positioning adjustment for the texture overlay", + }, + ), + "offset_y": ( + "FLOAT", + { + "default": 0.0, + "step": 0.01, + "tooltip": "Vertical positioning adjustment for the texture overlay", + }, + ), + "scale": ( + "FLOAT", + { + "default": 1.0, + "step": 0.01, + "tooltip": "Size adjustment for the texture overlay (1.0 = original size)", + }, + ), + }, } def execute(self, target_image, texture_image, face_transform_matrices, offset_x=0.0, offset_y=0.0, scale=1.0): logger.warning(f"{self.__class__.__name__} is experimental and does not perform actual warping yet.") - #TODO: implement this - # --- Conceptual Logic --- + # TODO: implement this + # --- Conceptual Logic --- # 1. Get the primary face transform matrix # 2. Potentially combine with offset/scale parameters # 3. Define or load a canonical 3D face model mesh. # 4. Project the canonical mesh onto the 2D target_image plane using the matrix. # 5. Determine the UV coordinates on the texture_image corresponding to the projected vertices. - # 6. Perform perspective warping of the texture_image onto the target_image using the - # calculated transformation (e.g., using OpenCV's getPerspectiveTransform and warpPerspective, + # 6. Perform perspective warping of the texture_image onto the target_image using the + # calculated transformation (e.g., using OpenCV's getPerspectiveTransform and warpPerspective, # or more advanced 3D rendering techniques). # 7. Blend the warped texture onto the target image (potentially using a face mask). - + # For now, just return the original target image return (target_image,) + # --- Mappings --- NODE_CLASS_MAPPINGS = { "FaceTextureWarp": FaceTextureWarpNode, @@ -55,4 +88,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "FaceTextureWarp": "Face Texture Warp (Experimental)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py b/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py index 4398456..c5f648d 100644 --- a/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py +++ b/node_wrappers/mediapipe_vision/face/stylization/face_stylization.py @@ -9,15 +9,20 @@ from ...common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Face/FaceStylization" -# --- Model Loader --- + + +# --- Model Loader --- class MediaPipeFaceStylizerModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Face Stylizer models.""" - TASK_TYPE = "face_stylizer" # Need to add this task to MEDIAPIPE_MODELS + + TASK_TYPE = "face_stylizer" # Need to add this task to MEDIAPIPE_MODELS RETURN_TYPES = ("FACE_STYLIZER_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Face Stylizer model that can transform face images into artistic styles. Required before using the Face Stylizer node." -# --- Stylizer Node --- + + +# --- Stylizer Node --- class MediaPipeFaceStylizerNode: """ComfyUI node for MediaPipe Face Stylization.""" @@ -30,11 +35,14 @@ class MediaPipeFaceStylizerNode: return { "required": { "image": ("IMAGE", {"tooltip": "Input image containing faces to be artistically stylized"}), - "model_info": ("FACE_STYLIZER_MODEL_INFO", {"tooltip": "Face stylizer model loaded from the 'Load Face Stylizer Model' node"}), + "model_info": ( + "FACE_STYLIZER_MODEL_INFO", + {"tooltip": "Face stylizer model loaded from the 'Load Face Stylizer Model' node"}, + ), } } - RETURN_TYPES = ("IMAGE",) # Output is a stylized image + RETURN_TYPES = ("IMAGE",) # Output is a stylized image RETURN_NAMES = ("stylized_image",) FUNCTION = "stylize_face" CATEGORY = _category @@ -42,36 +50,41 @@ class MediaPipeFaceStylizerNode: def stylize_face(self, image: torch.Tensor, model_info: dict): """Applies face stylization.""" - - task_type = model_info.get('task_type') + + task_type = model_info.get("task_type") expected_task_type = MediaPipeFaceStylizerModelLoaderNode.TASK_TYPE if not isinstance(model_info, dict) or task_type != expected_task_type: - raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but got '{task_type}'.") - model_path = model_info.get('model_path') + raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but got '{task_type}'.") + model_path = model_info.get("model_path") if not model_path: - raise ValueError("Model path not found or invalid in model_info.") - + raise ValueError("Model path not found or invalid in model_info.") + # Manage detector instance if self._detector is None or self._model_path != model_path: - if self._detector and hasattr(self._detector, 'close'): - try: self._detector.close() - except Exception as e: logger.warning(f"Error closing FaceStylizer: {e}") + if self._detector and hasattr(self._detector, "close"): + try: + self._detector.close() + except Exception as e: + logger.warning(f"Error closing FaceStylizer: {e}") logger.info(f"Creating new FaceStylizer instance for {model_path}") self._detector = FaceStylizer(model_path) self._model_path = model_path - + # Call the detector's stylize method stylized_image_batch = self._detector.stylize(image) - + return (stylized_image_batch,) def __del__(self): - if hasattr(self, '_detector') and self._detector and hasattr(self._detector, 'close'): - try: self._detector.close() - except Exception as e: logger.warning(f"Error closing FaceStylizer in __del__: {e}") - self._detector = None + if hasattr(self, "_detector") and self._detector and hasattr(self._detector, "close"): + try: + self._detector.close() + except Exception as e: + logger.warning(f"Error closing FaceStylizer in __del__: {e}") + self._detector = None -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MediaPipeFaceStylizerModelLoader": MediaPipeFaceStylizerModelLoaderNode, "MediaPipeFaceStylizer": MediaPipeFaceStylizerNode, @@ -80,4 +93,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeFaceStylizerModelLoader": "Load Face Stylizer Model (MediaPipe)", "MediaPipeFaceStylizer": "Face Stylizer (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/hand/__init__.py b/node_wrappers/mediapipe_vision/hand/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/hand/__init__.py +++ b/node_wrappers/mediapipe_vision/hand/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py index b07498d..d3bb3c3 100644 --- a/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py +++ b/node_wrappers/mediapipe_vision/hand/gesture/gesture_visualiztion.py @@ -9,11 +9,13 @@ logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/GestureVisualization" + class VisualizeGestureRecognitions(BaseDetectionVisualizationNode): """Draws recognized gesture labels onto images near the detected hand.""" + CATEGORY = _category DESCRIPTION = "Visualizes recognized hand gestures by drawing text labels next to hands in the image. Displays gesture names like 'Thumbs Up', 'Victory', 'Open Palm' etc. near each detected hand." - + # Override font scale for gesture labels (larger than default) FONT_SCALE = 0.7 # Larger than default for gesture labels @@ -21,39 +23,61 @@ class VisualizeGestureRecognitions(BaseDetectionVisualizationNode): def INPUT_TYPES(cls): return { "required": { - "gesture_recognitions": ("GESTURE_RECOGNITIONS", {"tooltip": "Gesture recognition results from the Gesture Recognizer node"}), - "hand_landmarks": ("HAND_LANDMARKS", {"tooltip": "Hand landmark data from the Hand Landmarker node - needed to position the text labels"}), + "gesture_recognitions": ( + "GESTURE_RECOGNITIONS", + {"tooltip": "Gesture recognition results from the Gesture Recognizer node"}, + ), + "hand_landmarks": ( + "HAND_LANDMARKS", + {"tooltip": "Hand landmark data from the Hand Landmarker node - needed to position the text labels"}, + ), "image": ("IMAGE", {"tooltip": "The image on which to draw gesture labels"}), }, "optional": { - "label_color": ("STRING", {"default": "#FFFFFF", "tooltip": "Hex color code for the gesture label text (default: white)"}), - "label_bg_color": ("STRING", {"default": "#0000FF", "tooltip": "Hex color code for the label background (default: blue)"}), - "y_offset": ("INT", {"default": -30, "step": 5, "tooltip": "Vertical position adjustment for text labels relative to the wrist - negative moves up, positive moves down"}), - } + "label_color": ( + "STRING", + {"default": "#FFFFFF", "tooltip": "Hex color code for the gesture label text (default: white)"}, + ), + "label_bg_color": ( + "STRING", + {"default": "#0000FF", "tooltip": "Hex color code for the label background (default: blue)"}, + ), + "y_offset": ( + "INT", + { + "default": -30, + "step": 5, + "tooltip": "Vertical position adjustment for text labels relative to the wrist - negative moves up, positive moves down", + }, + ), + }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "draw_visualizations" - def draw_visualizations(self, gesture_recognitions, hand_landmarks, image, - label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30): + def draw_visualizations( + self, gesture_recognitions, hand_landmarks, image, label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30 + ): """Main validation and batch processing function.""" if not gesture_recognitions: logger.warning("VisualizeGestureRecognitions: No gestures provided.") return (image,) if not hand_landmarks: - logger.warning("VisualizeGestureRecognitions: Hand landmarks required for positioning text.") - return (image,) - + logger.warning("VisualizeGestureRecognitions: Hand landmarks required for positioning text.") + return (image,) + # Basic batch size check batch_size = len(image) if len(gesture_recognitions) != batch_size or len(hand_landmarks) != batch_size: - logger.error(f"VisualizeGestureRecognitions: Batch size mismatch. Images: {batch_size}, Gestures: {len(gesture_recognitions)}, Landmarks: {len(hand_landmarks)}. Cannot proceed.") + logger.error( + f"VisualizeGestureRecognitions: Batch size mismatch. Images: {batch_size}, Gestures: {len(gesture_recognitions)}, Landmarks: {len(hand_landmarks)}. Cannot proceed." + ) return (image,) - + # Pair gestures with landmarks for convenience paired_data = list(zip(gesture_recognitions, hand_landmarks)) - + # Use the base class process_batch method with our custom drawing function return self.process_batch( paired_data, # We're passing paired data instead of just detections @@ -61,51 +85,50 @@ class VisualizeGestureRecognitions(BaseDetectionVisualizationNode): self.draw_gestures, label_color=label_color, label_bg_color=label_bg_color, - y_offset=y_offset + y_offset=y_offset, ) - - def draw_gestures(self, image, width, height, paired_data, - label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30): + + def draw_gestures(self, image, width, height, paired_data, label_color="#FFFFFF", label_bg_color="#0000FF", y_offset=-30): """Draw gestures for a single image.""" gestures_for_image, landmarks_for_image = paired_data - + # Check if number of detected hands matches for gestures and landmarks if len(gestures_for_image) != len(landmarks_for_image): - logger.warning(f"VisualizeGestureRecognitions: Mismatch between detected gestures and landmarks. Skipping labels.") - return - + logger.warning(f"VisualizeGestureRecognitions: Mismatch between detected gestures and landmarks. Skipping labels.") + return + if not gestures_for_image: return # Iterate through each detected hand/gesture set for this image for hand_idx, gesture_result in enumerate(gestures_for_image): landmark_result = landmarks_for_image[hand_idx] - + # Get top gesture category top_gesture = "N/A" if gesture_result.gestures: top_gesture = gesture_result.gestures[0].display_name or gesture_result.gestures[0].category_name or "N/A" - + # Use wrist landmark (index 0) for positioning text if not landmark_result.landmarks or len(landmark_result.landmarks) == 0: logger.warning(f"VisualizeGestureRecognitions: No landmarks found for this hand to position text.") continue - - wrist_lm = landmark_result.landmarks[0] # Assuming index 0 is wrist + + wrist_lm = landmark_result.landmarks[0] # Assuming index 0 is wrist cx = int(wrist_lm.x * width) - cy = int(wrist_lm.y * height) + y_offset # Apply offset + cy = int(wrist_lm.y * height) + y_offset # Apply offset label_text = f"{top_gesture}" - + # Draw label with background try: # Get text size text_width, text_height = self.get_text_dimensions(label_text, self.FONT_SCALE) - + # Center the text around the wrist point text_x = cx - (text_width // 2) text_y = cy - (text_height // 2) - + # Draw the text with background self.draw_text_with_background( image, @@ -113,27 +136,24 @@ class VisualizeGestureRecognitions(BaseDetectionVisualizationNode): (text_x, text_y), text_color=label_color, bg_color=label_bg_color, - font_scale=self.FONT_SCALE + font_scale=self.FONT_SCALE, ) except Exception as e: logger.error(f"Error drawing text label for gesture: {e}") # Fallback: Draw text directly cv2.putText( - image, - label_text, - (cx, cy), - cv2.FONT_HERSHEY_SIMPLEX, - self.FONT_SCALE, - self._hex_to_bgr(label_color), - 1, - cv2.LINE_AA + image, + label_text, + (cx, cy), + cv2.FONT_HERSHEY_SIMPLEX, + self.FONT_SCALE, + self._hex_to_bgr(label_color), + 1, + cv2.LINE_AA, ) -# Node registration -NODE_CLASS_MAPPINGS = { - "VisualizeGestureRecognitions": VisualizeGestureRecognitions -} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizeGestureRecognitions": "Visualize Gestures (MediaPipe)" -} \ No newline at end of file +# Node registration +NODE_CLASS_MAPPINGS = {"VisualizeGestureRecognitions": VisualizeGestureRecognitions} + +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizeGestureRecognitions": "Visualize Gestures (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/hand/gesture/guesture.py b/node_wrappers/mediapipe_vision/hand/gesture/guesture.py index 97eccd0..4a529f0 100644 --- a/node_wrappers/mediapipe_vision/hand/gesture/guesture.py +++ b/node_wrappers/mediapipe_vision/hand/gesture/guesture.py @@ -10,15 +10,20 @@ from ...common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/GestureRecognition" -# --- Model Loader --- + + +# --- Model Loader --- class MediaPipeGestureRecognizerModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Gesture Recognizer models.""" - TASK_TYPE = "gesture_recognizer" # Need to add this task to MEDIAPIPE_MODELS in src/model_loader.py + + TASK_TYPE = "gesture_recognizer" # Need to add this task to MEDIAPIPE_MODELS in src/model_loader.py RETURN_TYPES = ("GESTURE_RECOGNIZER_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Gesture Recognizer model for identifying hand gestures in images, such as thumbs up, open palm, pointing, etc. Required before using the Gesture Recognizer node." -# --- Recognizer Node --- + + +# --- Recognizer Node --- class MediaPipeGestureRecognizerNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Gesture Recognition.""" @@ -36,32 +41,74 @@ class MediaPipeGestureRecognizerNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add gesture-specific parameters - inputs["required"].update({ - "max_results": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1, - "tooltip": "Maximum number of hands to analyze in the image - higher values can detect more hands but use more processing power"}), - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives"}), - "min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence threshold for tracking hands between frames - lower values maintain tracking longer but may be less accurate"}), - "min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear"}), - }) - + inputs["required"].update( + { + "max_results": ( + "INT", + { + "default": 2, + "min": 1, + "max": 4, + "step": 1, + "tooltip": "Maximum number of hands to analyze in the image - higher values can detect more hands but use more processing power", + }, + ), + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives", + }, + ), + "min_tracking_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence threshold for tracking hands between frames - lower values maintain tracking longer but may be less accurate", + }, + ), + "min_presence_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear", + }, + ), + } + ) + return inputs - def recognize_gestures(self, image: torch.Tensor, model_info: dict, max_results: int, - min_confidence: float, min_tracking_confidence: float, min_presence_confidence: float, - running_mode: str, delegate: str): + def recognize_gestures( + self, + image: torch.Tensor, + model_info: dict, + max_results: int, + min_confidence: float, + min_tracking_confidence: float, + min_presence_confidence: float, + running_mode: str, + delegate: str, + ): """Performs gesture recognition on the input image.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Call the detector's recognize method with standardized parameters gesture_results_batch = detector.recognize( image, @@ -70,12 +117,13 @@ class MediaPipeGestureRecognizerNode(BaseMediaPipeDetectorNode): min_tracking_confidence=min_tracking_confidence, min_presence_confidence=min_presence_confidence, running_mode=running_mode, - delegate=delegate + delegate=delegate, ) - + return (gesture_results_batch,) -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MediaPipeGestureRecognizerModelLoader": MediaPipeGestureRecognizerModelLoaderNode, "MediaPipeGestureRecognizer": MediaPipeGestureRecognizerNode, @@ -84,4 +132,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeGestureRecognizerModelLoader": "Load Gesture Recognizer Model (MediaPipe)", "MediaPipeGestureRecognizer": "Gesture Recognizer (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py index e0552b9..e308ba1 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark.py @@ -9,14 +9,18 @@ from ...common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark" + + class MediaPipeHandLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Hand Landmarker models.""" + TASK_TYPE = "hand_landmarker" RETURN_TYPES = ("HAND_LANDMARKER_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Hand Landmarker model for detecting hand positions and gestures in images. Required before using the Hand Landmarker node." - + + class MediaPipeHandLandmarkerNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Hand Landmark Detection.""" @@ -34,32 +38,74 @@ class MediaPipeHandLandmarkerNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add hand-specific parameters - inputs["required"].update({ - "max_results": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1, - "tooltip": "Maximum number of hands to detect and track in the image - higher values can detect more hands but use more processing power"}), - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives"}), - "min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear"}), - "min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence for tracking a hand between frames - lower values maintain tracking longer but may drift"}), - }) - + inputs["required"].update( + { + "max_results": ( + "INT", + { + "default": 2, + "min": 1, + "max": 4, + "step": 1, + "tooltip": "Maximum number of hands to detect and track in the image - higher values can detect more hands but use more processing power", + }, + ), + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for hand detection - lower values detect more hands but may include false positives", + }, + ), + "min_presence_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence that a hand is present - lower values may detect hands that are less clear", + }, + ), + "min_tracking_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence for tracking a hand between frames - lower values maintain tracking longer but may drift", + }, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, max_results: int, - min_confidence: float, min_presence_confidence: float, - min_tracking_confidence: float, running_mode: str, delegate: str): + def detect( + self, + image: torch.Tensor, + model_info: dict, + max_results: int, + min_confidence: float, + min_presence_confidence: float, + min_tracking_confidence: float, + running_mode: str, + delegate: str, + ): """Performs hand landmark detection with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Perform detection with all parameters hand_landmarks_batch = detector.detect( image, @@ -68,18 +114,19 @@ class MediaPipeHandLandmarkerNode(BaseMediaPipeDetectorNode): min_presence_confidence=min_presence_confidence, min_tracking_confidence=min_tracking_confidence, running_mode=running_mode, - delegate=delegate + delegate=delegate, ) - + # Extract handedness information handedness_batch = [] if hand_landmarks_batch: for image_hands in hand_landmarks_batch: handedness_per_image = [hand.handedness for hand in image_hands if hand.handedness is not None] handedness_batch.append(handedness_per_image) - + return (hand_landmarks_batch, handedness_batch) + NODE_CLASS_MAPPINGS = { "MediaPipeHandLandmarkerModelLoader": MediaPipeHandLandmarkerModelLoaderNode, "MediaPipeHandLandmarker": MediaPipeHandLandmarkerNode, @@ -88,4 +135,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeHandLandmarkerModelLoader": "Load Hand Landmarker Model (MediaPipe)", "MediaPipeHandLandmarker": "Hand Landmarker (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py index cba76ec..37f8e47 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_masking.py @@ -14,20 +14,22 @@ _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Masking" # --- Hand Part Definitions (Based on standard MediaPipe Hand Indices) --- # These indices need verification HAND_PART_INDICES = { - "PALM": [0, 1, 5, 9, 13, 17], # Example - "THUMB": [1, 2, 3, 4], # Example - "INDEX_FINGER": [5, 6, 7, 8], # Example - "MIDDLE_FINGER": [9, 10, 11, 12], # Example - "RING_FINGER": [13, 14, 15, 16], # Example - "PINKY_FINGER": [17, 18, 19, 20], # Example - "WRIST": [0], # Single point, convex hull won't work - "ALL_FINGERS": list(range(1, 21)), # Example excluding wrist - "FULL_HAND": list(range(21)), # All points + "PALM": [0, 1, 5, 9, 13, 17], # Example + "THUMB": [1, 2, 3, 4], # Example + "INDEX_FINGER": [5, 6, 7, 8], # Example + "MIDDLE_FINGER": [9, 10, 11, 12], # Example + "RING_FINGER": [13, 14, 15, 16], # Example + "PINKY_FINGER": [17, 18, 19, 20], # Example + "WRIST": [0], # Single point, convex hull won't work + "ALL_FINGERS": list(range(1, 21)), # Example excluding wrist + "FULL_HAND": list(range(21)), # All points } -# --- Node Definition --- + +# --- Node Definition --- class MaskFromHandLandmarksNode: """Creates a mask for a specific hand part from Hand Landmarks.""" + CATEGORY = _category DESCRIPTION = "Creates precise masks for specific hand parts like individual fingers, palm, or the entire hand. Perfect for targeted editing or creating effects on hands and fingers." @@ -36,10 +38,23 @@ class MaskFromHandLandmarksNode: part_names = list(HAND_PART_INDICES.keys()) return { "required": { - "hand_landmarks": ("HAND_LANDMARKS", {"tooltip": "Hand landmark data from the Hand Landmarker node containing the precise locations of hand and finger joints"}), - "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}), - "part_name": (part_names, {"default": "FULL_HAND", - "tooltip": "Choose which part of the hand to create a mask for - FULL_HAND (entire hand), THUMB, INDEX_FINGER, PALM, etc."}), + "hand_landmarks": ( + "HAND_LANDMARKS", + { + "tooltip": "Hand landmark data from the Hand Landmarker node containing the precise locations of hand and finger joints" + }, + ), + "image_for_dimensions": ( + "IMAGE", + {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}, + ), + "part_name": ( + part_names, + { + "default": "FULL_HAND", + "tooltip": "Choose which part of the hand to create a mask for - FULL_HAND (entire hand), THUMB, INDEX_FINGER, PALM, etc.", + }, + ), } } @@ -49,7 +64,7 @@ class MaskFromHandLandmarksNode: def generate_mask(self, hand_landmarks: List[List[HandLandmarksResult]], image_for_dimensions: torch.Tensor, part_name: str): if part_name not in HAND_PART_INDICES: raise ValueError(f"Unknown hand part name: '{part_name}'. Available: {list(HAND_PART_INDICES.keys())}") - + # Get dimensions from image if image_for_dimensions.dim() != 4: raise ValueError("image_for_dimensions must be BHWC format.") @@ -60,16 +75,18 @@ class MaskFromHandLandmarksNode: # Validate batch sizes match if len(hand_landmarks) != input_batch_size: - logger.warning(f"Batch size mismatch between hand_landmarks ({len(hand_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.") - batch_size = len(hand_landmarks) + logger.warning( + f"Batch size mismatch between hand_landmarks ({len(hand_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size." + ) + batch_size = len(hand_landmarks) else: - batch_size = input_batch_size + batch_size = input_batch_size required_indices = HAND_PART_INDICES[part_name] output_masks = [] for i in range(batch_size): - landmarks_for_image = hand_landmarks[i] # List[HandLandmarksResult] + landmarks_for_image = hand_landmarks[i] # List[HandLandmarksResult] combined_points_for_image = [] # Combine points from all detected hands for the selected part @@ -84,25 +101,25 @@ class MaskFromHandLandmarksNode: part_points.append((px, py)) # No warning for missing index combined_points_for_image.extend(part_points) - + # Create mask from all collected points for this image mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device) output_masks.append(mask_tensor) if not output_masks: - logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") - output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) + logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") + output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) else: - output_batch = torch.stack(output_masks, dim=0) - + output_batch = torch.stack(output_masks, dim=0) + return (output_batch,) -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MaskFromHandLandmarks": MaskFromHandLandmarksNode, - } NODE_DISPLAY_NAME_MAPPINGS = { "MaskFromHandLandmarks": "Mask From Hand Landmarks (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py index ede6d69..cdbd909 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/hand_landmark_visualization.py @@ -11,9 +11,10 @@ _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Visualization" # Define standard hand connections HAND_CONNECTIONS = mp.solutions.hands.HAND_CONNECTIONS + class VisualizeHandLandmarks(BaseLandmarkVisualizationNode): """Draws Hand landmarks and connections onto images.""" - + # Override base class properties LANDMARKS_TYPE = "HAND_LANDMARKS" CONNECTIONS = HAND_CONNECTIONS @@ -25,16 +26,12 @@ class VisualizeHandLandmarks(BaseLandmarkVisualizationNode): SUPPORTS_VISIBILITY = False CATEGORY = _category DESCRIPTION = "Visualizes hand landmarks by drawing points and connections on detected hands. Creates a wireframe overlay showing finger joints, palm, and hand structure." - + def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]: return detection_result.landmarks # Node registration -NODE_CLASS_MAPPINGS = { - "VisualizeHandLandmarks": VisualizeHandLandmarks -} +NODE_CLASS_MAPPINGS = {"VisualizeHandLandmarks": VisualizeHandLandmarks} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizeHandLandmarks": "Visualize Hand Landmarks (MediaPipe)" -} \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizeHandLandmarks": "Visualize Hand Landmarks (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/hand/landmark/position.py b/node_wrappers/mediapipe_vision/hand/landmark/position.py index 6bee4b9..d166c6b 100644 --- a/node_wrappers/mediapipe_vision/hand/landmark/position.py +++ b/node_wrappers/mediapipe_vision/hand/landmark/position.py @@ -11,9 +11,11 @@ from .....src.mediapipe_vision.landmark_definitions import HAND_LANDMARK_TOOLTIP logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Hand/HandLandmark/Position" + class HandLandmarkPositionNode(LandmarkPositionBase): """Extracts position lists (x, y, z, vis, pres) for a specific landmark index - from HAND_LANDMARKS across the batch.""" + from HAND_LANDMARKS across the batch.""" + CATEGORY = _category # Define the specific input type string this node expects LANDMARKS_TYPE = "HAND_LANDMARKS" @@ -26,11 +28,18 @@ class HandLandmarkPositionNode(LandmarkPositionBase): "required": { "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), "landmark_index": ("INT", {"default": 8, "min": 0, "max": 20, "tooltip": HAND_LANDMARK_TOOLTIP}), - "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the hand detection to use (0=first detected hand)"}), - "use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}), + "result_index": ( + "INT", + {"default": 0, "min": 0, "tooltip": "Index of the hand detection to use (0=first detected hand)"}, + ), + "use_world_coordinates": ( + "BOOLEAN", + {"default": False, "tooltip": "Use world coordinates if available"}, + ), } } + # --- Node Mappings --- NODE_CLASS_MAPPINGS = { "HandLandmarkPosition": HandLandmarkPositionNode, @@ -38,4 +47,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "HandLandmarkPosition": "Hand Landmark Position Extractor (Batch)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py index 731fe86..6f416a6 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark.py @@ -8,14 +8,15 @@ from ..common.base_detector_node import BaseMediaPipeDetectorNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Holistic" + class MediaPipeHolisticLandmarkerNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Holistic Landmark Detection. - + This node uses the legacy MediaPipe Holistic API that combines face, pose and hand tracking in a single detector. - + Note: This uses the legacy MediaPipe Python Solution API while we wait for an official - Holistic Task implementation, and thus it stands out among other implementations in this project. + Holistic Task implementation, and thus it stands out among other implementations in this project. It does not require a model file to be loaded separately. """ @@ -23,7 +24,10 @@ class MediaPipeHolisticLandmarkerNode(BaseMediaPipeDetectorNode): DETECTOR_CLASS = HolisticLandmarkDetector MODEL_INFO_TYPE = None # Not used - legacy solution doesn't need model file EXPECTED_TASK_TYPE = None # Not used - legacy solution doesn't need model info - RETURN_TYPES = ("HOLISTIC_LANDMARKS", "MASK_LIST",) + RETURN_TYPES = ( + "HOLISTIC_LANDMARKS", + "MASK_LIST", + ) RETURN_NAMES = ("landmarks", "segmentation_masks") FUNCTION = "detect_holistic" CATEGORY = _category @@ -40,36 +44,78 @@ class MediaPipeHolisticLandmarkerNode(BaseMediaPipeDetectorNode): return { "required": { "image": ("IMAGE", {"tooltip": "Input image for detection"}), - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for detection - lower values detect more poses but may include false positives"}), - "min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for tracking - lower values maintain tracking longer but may be less accurate"}), - "model_complexity": ("INT", {"default": 1, "min": 0, "max": 2, "step": 1, - "tooltip": "Model complexity (0=Lite, 1=Full, 2=Heavy) - higher values are more accurate but slower"}), - "enable_segmentation": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, generates masks that separate people from the background"}), - "refine_face_landmarks": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, provides more detailed face landmarks around eyes and lips"}), - "static_image_mode": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, treats each frame as independent - more accurate but slower and doesn't benefit from tracking"}), + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for detection - lower values detect more poses but may include false positives", + }, + ), + "min_tracking_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for tracking - lower values maintain tracking longer but may be less accurate", + }, + ), + "model_complexity": ( + "INT", + { + "default": 1, + "min": 0, + "max": 2, + "step": 1, + "tooltip": "Model complexity (0=Lite, 1=Full, 2=Heavy) - higher values are more accurate but slower", + }, + ), + "enable_segmentation": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, generates masks that separate people from the background", + }, + ), + "refine_face_landmarks": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, provides more detailed face landmarks around eyes and lips", + }, + ), + "static_image_mode": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, treats each frame as independent - more accurate but slower and doesn't benefit from tracking", + }, + ), } } - def detect_holistic(self, image: torch.Tensor, - min_confidence: float, - min_tracking_confidence: float, - model_complexity: int, - enable_segmentation: bool, - refine_face_landmarks: bool, - static_image_mode: bool): + def detect_holistic( + self, + image: torch.Tensor, + min_confidence: float, + min_tracking_confidence: float, + model_complexity: int, + enable_segmentation: bool, + refine_face_landmarks: bool, + static_image_mode: bool, + ): """Performs holistic landmark detection with the configured parameters. - + Unlike other detectors in this codebase, this doesn't need a model_info parameter as it uses the built-in legacy Holistic API. """ # Initialize detector (without model path) detector = self.initialize_or_update_detector(None) - + # Perform detection with all parameters holistic_landmarks_batch, segmentation_masks_batch = detector.detect( image, @@ -78,17 +124,17 @@ class MediaPipeHolisticLandmarkerNode(BaseMediaPipeDetectorNode): model_complexity=model_complexity, enable_segmentation=enable_segmentation, refine_face_landmarks=refine_face_landmarks, - static_image_mode=static_image_mode + static_image_mode=static_image_mode, ) - + return (holistic_landmarks_batch, segmentation_masks_batch or []) - + def initialize_or_update_detector(self, model_path=None): """Override to handle the case where model_path is not needed.""" if self._detector is None: self._detector = self.DETECTOR_CLASS() return self._detector - + def validate_model_info(self, model_info=None): """Override to handle the case where model_info is not needed.""" # No validation needed for legacy API @@ -102,4 +148,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeHolisticLandmarker": "**CAUTION** Holistic Landmarker (MediaPipe Legacy)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py index 757b4e0..06589b6 100644 --- a/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py +++ b/node_wrappers/mediapipe_vision/holistic/holistic_landmark_visualization.py @@ -14,19 +14,20 @@ POSE_CONNECTIONS = mp.solutions.pose.POSE_CONNECTIONS HAND_CONNECTIONS = mp.solutions.hands.HAND_CONNECTIONS FACE_CONNECTIONS = mp.solutions.face_mesh.FACEMESH_TESSELATION + class MediaPipeHolisticVisualizerNode: """ComfyUI node for visualizing holistic landmarks using direct drawing. - + Since holistic landmarks combine multiple landmark types (face, pose, hands), this doesn't use the BaseLandmarkVisualizationNode but implements a similar approach. """ - + CATEGORY = _category RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "visualize_holistic" DESCRIPTION = "Visualizes holistic landmarks by drawing face mesh, pose skeleton, and hand landmarks on the input image." - + @classmethod def INPUT_TYPES(cls): return { @@ -40,39 +41,61 @@ class MediaPipeHolisticVisualizerNode: "draw_face": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing face landmarks"}), "draw_pose": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing pose landmarks"}), "draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Toggle drawing hand landmarks"}), - "face_point_radius": ("INT", {"default": 1, "min": 1, "max": 5, "tooltip": "Size of face landmark points"}), - "pose_point_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of pose landmark points"}), - "hand_point_radius": ("INT", {"default": 2, "min": 1, "max": 10, "tooltip": "Size of hand landmark points"}), - "connection_thickness": ("INT", {"default": 1, "min": 1, "max": 5, "tooltip": "Thickness of connection lines"}), - "min_visibility": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, - "tooltip": "Minimum visibility score to draw a pose landmark"}) + "face_point_radius": ( + "INT", + {"default": 1, "min": 1, "max": 5, "tooltip": "Size of face landmark points"}, + ), + "pose_point_radius": ( + "INT", + {"default": 2, "min": 1, "max": 10, "tooltip": "Size of pose landmark points"}, + ), + "hand_point_radius": ( + "INT", + {"default": 2, "min": 1, "max": 10, "tooltip": "Size of hand landmark points"}, + ), + "connection_thickness": ( + "INT", + {"default": 1, "min": 1, "max": 5, "tooltip": "Thickness of connection lines"}, + ), + "min_visibility": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Minimum visibility score to draw a pose landmark", + }, + ), } } - + def _hex_to_rgb(self, hex_color): """Convert hex color string to RGB tuple for OpenCV (BGR format).""" - hex_color = hex_color.lstrip('#') + hex_color = hex_color.lstrip("#") # Convert to BGR for OpenCV - r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + r, g, b = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) return (b, g, r) # Return in BGR format for OpenCV - - def visualize_holistic(self, - image, - landmarks, - face_color, - pose_color, - left_hand_color, - right_hand_color, - draw_face, - draw_pose, - draw_hands, - face_point_radius, - pose_point_radius, - hand_point_radius, - connection_thickness, - min_visibility): + + def visualize_holistic( + self, + image, + landmarks, + face_color, + pose_color, + left_hand_color, + right_hand_color, + draw_face, + draw_pose, + draw_hands, + face_point_radius, + pose_point_radius, + hand_point_radius, + connection_thickness, + min_visibility, + ): """Visualize holistic landmarks using OpenCV for faster performance.""" - + if not landmarks: logger.warning("No holistic landmarks provided.") return (image,) @@ -90,107 +113,111 @@ class MediaPipeHolisticVisualizerNode: pose_color_bgr = self._hex_to_rgb(pose_color) left_hand_color_bgr = self._hex_to_rgb(left_hand_color) right_hand_color_bgr = self._hex_to_rgb(right_hand_color) - + for i in range(batch_size): # Clone the tensor to avoid modifying the original - img_tensor = image[i].clone() - + img_tensor = image[i].clone() + # Convert tensor to numpy array scaled to 0-255 range for OpenCV np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) - + # Convert to BGR for OpenCV processing cv_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR) - + if i < len(landmarks): holistic_result = landmarks[i] # One holistic result per image - + # Draw face landmarks if requested if draw_face and holistic_result.face_landmarks: self._draw_landmarks_cv2( cv_image, - holistic_result.face_landmarks, + holistic_result.face_landmarks, FACE_CONNECTIONS, - face_color_bgr, face_point_radius, connection_thickness, - min_visibility=None # Face doesn't use visibility + face_color_bgr, + face_point_radius, + connection_thickness, + min_visibility=None, # Face doesn't use visibility ) - + # Draw pose landmarks if requested if draw_pose and holistic_result.pose_landmarks: self._draw_landmarks_cv2( cv_image, - holistic_result.pose_landmarks, + holistic_result.pose_landmarks, POSE_CONNECTIONS, - pose_color_bgr, pose_point_radius, connection_thickness, - min_visibility=min_visibility + pose_color_bgr, + pose_point_radius, + connection_thickness, + min_visibility=min_visibility, ) - + # Draw left hand landmarks if requested if draw_hands and holistic_result.left_hand_landmarks: self._draw_landmarks_cv2( cv_image, - holistic_result.left_hand_landmarks, + holistic_result.left_hand_landmarks, HAND_CONNECTIONS, - left_hand_color_bgr, hand_point_radius, connection_thickness, - min_visibility=None # Hands don't use visibility + left_hand_color_bgr, + hand_point_radius, + connection_thickness, + min_visibility=None, # Hands don't use visibility ) - + # Draw right hand landmarks if requested if draw_hands and holistic_result.right_hand_landmarks: self._draw_landmarks_cv2( cv_image, - holistic_result.right_hand_landmarks, + holistic_result.right_hand_landmarks, HAND_CONNECTIONS, - right_hand_color_bgr, hand_point_radius, connection_thickness, - min_visibility=None # Hands don't use visibility + right_hand_color_bgr, + hand_point_radius, + connection_thickness, + min_visibility=None, # Hands don't use visibility ) - + # Convert back to RGB and tensor format cv_image_rgb = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) result_tensor = torch.from_numpy(cv_image_rgb.astype(np.float32) / 255.0) result_images.append(result_tensor) - + return (torch.stack(result_images),) - - def _draw_landmarks_cv2(self, - image: np.ndarray, - landmark_list: List, - connections: List, - color: Tuple[int, int, int], - point_radius: int, - connection_thickness: int, - min_visibility: Optional[float] = None): + + def _draw_landmarks_cv2( + self, + image: np.ndarray, + landmark_list: List, + connections: List, + color: Tuple[int, int, int], + point_radius: int, + connection_thickness: int, + min_visibility: Optional[float] = None, + ): """Helper method to draw landmarks with OpenCV for faster performance.""" - + height, width = image.shape[:2] points = {} - + # Process landmarks and draw points for lm in landmark_list: # Check visibility if required is_visible = True - if min_visibility is not None and hasattr(lm, 'visibility') and lm.visibility is not None: + if min_visibility is not None and hasattr(lm, "visibility") and lm.visibility is not None: is_visible = lm.visibility >= min_visibility - + if is_visible: # Convert normalized coordinates to pixel coordinates cx, cy = int(lm.x * width), int(lm.y * height) points[lm.index] = (cx, cy) - + # Draw point as a filled circle cv2.circle(image, (cx, cy), point_radius, color, -1) - + # Draw connections if connections and points: for connection in connections: start_idx, end_idx = connection if start_idx in points and end_idx in points: - cv2.line( - image, - points[start_idx], - points[end_idx], - color, - connection_thickness - ) + cv2.line(image, points[start_idx], points[end_idx], color, connection_thickness) # Define mappings for ComfyUI @@ -200,4 +227,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeHolisticVisualizer": "Holistic Visualizer (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/image_embedding/__init__.py b/node_wrappers/mediapipe_vision/image_embedding/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/image_embedding/__init__.py +++ b/node_wrappers/mediapipe_vision/image_embedding/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py b/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py index 2e675bc..0bdc5e4 100644 --- a/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py +++ b/node_wrappers/mediapipe_vision/image_embedding/image_embeddings.py @@ -11,15 +11,18 @@ from ..common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ImageEmbedding" -# --- Model Loader --- + +# --- Model Loader --- class MediaPipeImageEmbedderModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Image Embedder models.""" - TASK_TYPE = "image_embedder" # Need to add this task to MEDIAPIPE_MODELS + + TASK_TYPE = "image_embedder" # Need to add this task to MEDIAPIPE_MODELS RETURN_TYPES = ("IMAGE_EMBEDDER_MODEL_INFO",) RETURN_NAMES = ("model_info",) - CATEGORY = _category + CATEGORY = _category -# --- Embedder Node --- + +# --- Embedder Node --- class MediaPipeImageEmbedderNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Image Embedding.""" @@ -36,39 +39,44 @@ class MediaPipeImageEmbedderNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add embedder-specific parameters - inputs["required"].update({ - "l2_normalize": ("BOOLEAN", {"default": True, - "tooltip": "Apply L2 normalization to the embedding"}), - "quantize": ("BOOLEAN", {"default": False, - "tooltip": "Output quantized (uint8) embedding instead of float"}), - }) - + inputs["required"].update( + { + "l2_normalize": ("BOOLEAN", {"default": True, "tooltip": "Apply L2 normalization to the embedding"}), + "quantize": ( + "BOOLEAN", + {"default": False, "tooltip": "Output quantized (uint8) embedding instead of float"}, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, l2_normalize: bool, - quantize: bool, running_mode: str, delegate: str): + def detect( + self, + image: torch.Tensor, + model_info: dict, + l2_normalize: bool, + quantize: bool, + running_mode: str, + delegate: str, + ): """Generates image embeddings with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Generate embeddings with all parameters - embeddings_batch = detector.embed( - image, - l2_normalize=l2_normalize, - quantize=quantize, - running_mode=running_mode, - delegate=delegate - ) - + embeddings_batch = detector.embed(image, l2_normalize=l2_normalize, quantize=quantize, running_mode=running_mode, delegate=delegate) + return (embeddings_batch,) -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MediaPipeImageEmbedderModelLoader": MediaPipeImageEmbedderModelLoaderNode, "MediaPipeImageEmbedder": MediaPipeImageEmbedderNode, @@ -77,4 +85,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeImageEmbedderModelLoader": "Load Image Embedder Model (MediaPipe)", "MediaPipeImageEmbedder": "Image Embedder (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/image_embedding/utils.py b/node_wrappers/mediapipe_vision/image_embedding/utils.py index 5e4e08b..0373530 100644 --- a/node_wrappers/mediapipe_vision/image_embedding/utils.py +++ b/node_wrappers/mediapipe_vision/image_embedding/utils.py @@ -10,20 +10,30 @@ from ....src.mediapipe_vision.types import ( logger = logging.getLogger(__name__) + class ReshapeMediaPipeEmbeddingNode: """Reshapes MediaPipe embedding for experimental use as ComfyUI Conditioning. - WARNING: Semantically different from standard CLIP conditioning. - Pads/truncates the embedding vector and replicates it across 77 tokens. + WARNING: Semantically different from standard CLIP conditioning. + Pads/truncates the embedding vector and replicates it across 77 tokens. """ + CATEGORY = "Realtime Nodes/MediaPipe Vision/ImageEmbedding/_Experimental" @classmethod def INPUT_TYPES(cls): return { "required": { - "image_embeddings": ("IMAGE_EMBEDDINGS",), # List[List[ImageEmbedderResult]] - "target_dimension": ("INT", {"default": 768, "min": 64, "max": 4096, "step": 8, - "tooltip": "Target dimension D for the reshaped (1, 77, D) tensor (e.g., 768 for SD1.5, 1024/1280 for SDXL)"}), + "image_embeddings": ("IMAGE_EMBEDDINGS",), # List[List[ImageEmbedderResult]] + "target_dimension": ( + "INT", + { + "default": 768, + "min": 64, + "max": 4096, + "step": 8, + "tooltip": "Target dimension D for the reshaped (1, 77, D) tensor (e.g., 768 for SD1.5, 1024/1280 for SDXL)", + }, + ), } } @@ -33,11 +43,11 @@ class ReshapeMediaPipeEmbeddingNode: def reshape_embedding(self, image_embeddings: List[List[ImageEmbedderResult]], target_dimension: int): if not image_embeddings or not image_embeddings[0]: raise ValueError("Input image_embeddings are empty.") - + # --- Process First Embedding in Batch/Head --- # Assuming we use the first embedding result from the first image in the batch first_embedding_result = image_embeddings[0][0] - + if first_embedding_result.float_embedding: embedding_vector = first_embedding_result.float_embedding elif first_embedding_result.quantized_embedding: @@ -46,12 +56,12 @@ class ReshapeMediaPipeEmbeddingNode: raise NotImplementedError("Reshaping quantized embeddings is not directly supported. Use float embeddings.") else: raise ValueError("No valid embedding found in the input.") - + # Convert to tensor embed_tensor = torch.tensor(embedding_vector, dtype=torch.float32) orig_dim = embed_tensor.shape[0] - # --- Pad or Truncate --- + # --- Pad or Truncate --- if orig_dim < target_dimension: # Pad with zeros padding = torch.zeros(target_dimension - orig_dim, dtype=torch.float32) @@ -61,35 +71,37 @@ class ReshapeMediaPipeEmbeddingNode: final_vector = embed_tensor[:target_dimension] else: final_vector = embed_tensor - - # --- Reshape and Replicate for Conditioning --- + + # --- Reshape and Replicate for Conditioning --- # Reshape to (1, 1, D) cond_vec = final_vector.unsqueeze(0).unsqueeze(0) # Replicate to (1, 77, D) cond_batch = cond_vec.repeat(1, 77, 1) - # --- Create Pooled Output --- + # --- Create Pooled Output --- # Use the padded/truncated vector as the pooled output (1, D) pooled_output = final_vector.unsqueeze(0) - # --- Package as CONDITIONING --- + # --- Package as CONDITIONING --- # Standard ComfyUI conditioning format: List[List[Union[torch.Tensor, Dict[str, torch.Tensor]]]] conditioning = [[cond_batch, {"pooled_output": pooled_output}]] logger.warning("Reshaped MediaPipe embedding to CONDITIONING format. Use experimentally.") return (conditioning,) -# --- Comparison Node --- + +# --- Comparison Node --- class CompareMediaPipeEmbeddingsNode: """Calculates the cosine similarity between two MediaPipe image embeddings.""" + CATEGORY = "Realtime Nodes/MediaPipe Vision/ImageEmbedding" @classmethod def INPUT_TYPES(cls): return { "required": { - "image_embeddings_1": ("IMAGE_EMBEDDINGS",), - "image_embeddings_2": ("IMAGE_EMBEDDINGS",), + "image_embeddings_1": ("IMAGE_EMBEDDINGS",), + "image_embeddings_2": ("IMAGE_EMBEDDINGS",), } } @@ -97,9 +109,7 @@ class CompareMediaPipeEmbeddingsNode: RETURN_NAMES = ("cosine_similarity",) FUNCTION = "compare_embeddings" - def compare_embeddings(self, image_embeddings_1: List[List[ImageEmbedderResult]], - image_embeddings_2: List[List[ImageEmbedderResult]]): - + def compare_embeddings(self, image_embeddings_1: List[List[ImageEmbedderResult]], image_embeddings_2: List[List[ImageEmbedderResult]]): def _extract_vector(embeddings: List[List[ImageEmbedderResult]]) -> Optional[torch.Tensor]: """Helper to extract the first float embedding vector.""" if not embeddings or not embeddings[0]: @@ -116,8 +126,8 @@ class CompareMediaPipeEmbeddingsNode: if vec1 is None or vec2 is None: logger.error("Could not extract valid float embedding vectors for comparison.") - return (0.0,) # Return 0 similarity on error - + return (0.0,) # Return 0 similarity on error + # Ensure vectors have the same dimension (pad/truncate shorter one) dim1, dim2 = vec1.shape[0], vec2.shape[0] if dim1 != dim2: @@ -127,21 +137,21 @@ class CompareMediaPipeEmbeddingsNode: vec1 = F.pad(vec1, (0, max_dim - dim1)) elif dim2 < max_dim: vec2 = F.pad(vec2, (0, max_dim - dim2)) - + # Calculate Cosine Similarity # Normalize vectors before dot product for cosine similarity similarity = F.cosine_similarity(vec1.unsqueeze(0), vec2.unsqueeze(0), dim=1) - + return (similarity.item(),) -# --- Mappings --- +# --- Mappings --- NODE_CLASS_MAPPINGS = { "ReshapeMediaPipeEmbedding": ReshapeMediaPipeEmbeddingNode, - "CompareMediaPipeEmbeddings": CompareMediaPipeEmbeddingsNode, # Add new node + "CompareMediaPipeEmbeddings": CompareMediaPipeEmbeddingsNode, # Add new node } NODE_DISPLAY_NAME_MAPPINGS = { "ReshapeMediaPipeEmbedding": "Reshape MediaPipe Embedding (Experimental)", - "CompareMediaPipeEmbeddings": "Compare MediaPipe Embeddings", # Add new node name -} \ No newline at end of file + "CompareMediaPipeEmbeddings": "Compare MediaPipe Embeddings", # Add new node name +} diff --git a/node_wrappers/mediapipe_vision/image_segmentation/__init__.py b/node_wrappers/mediapipe_vision/image_segmentation/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/__init__.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py index c5e8090..fd451c8 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation.py @@ -28,34 +28,41 @@ MULTICLASS_NAMES = { # Uses MULTICLASS_NAMES indices CLASS_COLORS_BGR = { - MULTICLASS_NAMES["Background"]: (0, 0, 0), # Black - MULTICLASS_NAMES["Hair"]: (255, 0, 0), # Blue - MULTICLASS_NAMES["Body-skin"]: (0, 255, 255), # Yellow - MULTICLASS_NAMES["Face-skin"]: (255, 200, 100), # Light Blue - MULTICLASS_NAMES["Clothes"]: (0, 255, 0), # Green - MULTICLASS_NAMES["Accessories/Other"]: (255, 0, 255), # Magenta + MULTICLASS_NAMES["Background"]: (0, 0, 0), # Black + MULTICLASS_NAMES["Hair"]: (255, 0, 0), # Blue + MULTICLASS_NAMES["Body-skin"]: (0, 255, 255), # Yellow + MULTICLASS_NAMES["Face-skin"]: (255, 200, 100), # Light Blue + MULTICLASS_NAMES["Clothes"]: (0, 255, 0), # Green + MULTICLASS_NAMES["Accessories/Other"]: (255, 0, 255), # Magenta } -# --- End Constants --- +# --- End Constants --- + # Inherit from the base loader class MediaPipeImageSegmenterModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Image Segmenter models.""" + TASK_TYPE = "image_segmenter" RETURN_TYPES = ("IMAGE_SEGMENTER_MODEL_INFO",) RETURN_NAMES = ("model_info",) # INPUT_TYPES and FUNCTION inherited - CATEGORY = _category + CATEGORY = _category + class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Image Segmentation, adapted from Stream-Pack. - Manages the ImageSegmenter instance for potential reuse based on config. + Manages the ImageSegmenter instance for potential reuse based on config. """ - + # Define class variables required by the base class DETECTOR_CLASS = ImageSegmenter MODEL_INFO_TYPE = "IMAGE_SEGMENTER_MODEL_INFO" EXPECTED_TASK_TYPE = "image_segmenter" - RETURN_TYPES = ("MASK", "IMAGE", "MP_INT_MASK",) + RETURN_TYPES = ( + "MASK", + "IMAGE", + "MP_INT_MASK", + ) RETURN_NAMES = ("mask", "visualization", "multiclass_segments") FUNCTION = "detect" CATEGORY = _category @@ -64,37 +71,54 @@ class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add segmentation-specific parameters - inputs["required"].update({ - "output_confidence_masks": ("BOOLEAN", {"default": False, - "tooltip": "Output confidence mask (0-1) instead of category mask. Disables multiclass/visualization."}), - "threshold": ("FLOAT", { - "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Confidence threshold for confidence mask output mode. Pixels below threshold become 0." - }), - "generate_visualization": ("BOOLEAN", {"default": False, - "tooltip": "Generate a colored visualization image (only works in category mask mode)." - }), - }) - + inputs["required"].update( + { + "output_confidence_masks": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Output confidence mask (0-1) instead of category mask. Disables multiclass/visualization.", + }, + ), + "threshold": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Confidence threshold for confidence mask output mode. Pixels below threshold become 0.", + }, + ), + "generate_visualization": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Generate a colored visualization image (only works in category mask mode).", + }, + ), + } + ) + # Rename the delegate parameter for consistency with original implementation inputs["required"]["delegate_mode"] = inputs["required"].pop("delegate") - + return inputs - # --- Helper methods from Stream-Pack --- + # --- Helper methods from Stream-Pack --- def is_multiclass_model(self, model_info: dict): # Improved check based on variant or path - model_variant = model_info.get('model_variant', '').lower() - model_path = model_info.get('model_path', '').lower() + model_variant = model_info.get("model_variant", "").lower() + model_path = model_info.get("model_path", "").lower() # Add known multiclass identifiers - return 'multiclass' in model_variant or 'multiclass' in model_path + return "multiclass" in model_variant or "multiclass" in model_path def create_visualization(self, category_mask_np, is_multiclass=False): """Generates a colored visualization from a category mask.""" if category_mask_np is None: - return None + return None if not is_multiclass: # Simple binary visualization (white on black) vis_image = np.where(category_mask_np > 0, 255, 0).astype(np.uint8) @@ -104,38 +128,47 @@ class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode): height, width = category_mask_np.shape vis_image = np.zeros((height, width, 3), dtype=np.uint8) for class_name, class_id in MULTICLASS_NAMES.items(): - color_bgr = CLASS_COLORS_BGR.get(class_id, (128, 128, 128)) # Default grey - vis_image[category_mask_np == class_id] = color_bgr + color_bgr = CLASS_COLORS_BGR.get(class_id, (128, 128, 128)) # Default grey + vis_image[category_mask_np == class_id] = color_bgr return vis_image - # --- End Helper methods --- - def detect(self, image: torch.Tensor, model_info: dict, output_confidence_masks: bool, - threshold: float, generate_visualization: bool, running_mode: str, delegate_mode: str): + # --- End Helper methods --- + + def detect( + self, + image: torch.Tensor, + model_info: dict, + output_confidence_masks: bool, + threshold: float, + generate_visualization: bool, + running_mode: str, + delegate_mode: str, + ): """Performs image segmentation with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Determine which output type the detector needs request_confidence = output_confidence_masks # Need category mask if confidence isn't requested OR if visualization is needed - request_category = not output_confidence_masks or generate_visualization - + request_category = not output_confidence_masks or generate_visualization + if not request_confidence and not request_category: - # Should not happen with the logic above, but as a safeguard - raise ValueError("Internal logic error: No mask type determined for segmenter.") + # Should not happen with the logic above, but as a safeguard + raise ValueError("Internal logic error: No mask type determined for segmenter.") # Call segmenter's segment method # Returns lists (one element per image in batch) batch_results_confidence, batch_results_category = detector.segment( image, output_confidence_masks=request_confidence, - output_category_mask=request_category, + output_category_mask=request_category, running_mode=running_mode, - delegate_mode=delegate_mode + delegate_mode=delegate_mode, ) # Initialize output lists @@ -144,17 +177,17 @@ class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode): all_primary_masks = [] all_vis_images = [] all_category_masks = [] - default_vis_tensor = torch.zeros((h, w, 3), dtype=torch.float32, device=image.device) # HWC + default_vis_tensor = torch.zeros((h, w, 3), dtype=torch.float32, device=image.device) # HWC - # --- Process results based on node inputs --- + # --- Process results based on node inputs --- is_multiclass = self.is_multiclass_model(model_info) - + # Loop through batch results for i in range(batch_size): # Get results for the current image confidence_masks_hw_list = batch_results_confidence[i] if batch_results_confidence else None category_mask_hw = batch_results_category[i] if batch_results_category else None - + # Initialize outputs for this image current_primary_mask_hw = torch.zeros((h, w), dtype=torch.float32, device=image.device) current_vis_image_hwc = default_vis_tensor @@ -164,53 +197,54 @@ class MediaPipeImageSegmenterNode(BaseMediaPipeDetectorNode): # User wants confidence mask as primary output if confidence_masks_hw_list: # Take the first confidence mask (often foreground vs background) - conf_mask_hw = confidence_masks_hw_list[0] + conf_mask_hw = confidence_masks_hw_list[0] # Apply threshold conf_mask_hw_thresh = torch.where(conf_mask_hw >= threshold, conf_mask_hw, torch.zeros_like(conf_mask_hw)) - current_primary_mask_hw = conf_mask_hw_thresh # HW + current_primary_mask_hw = conf_mask_hw_thresh # HW # Visualization and multiclass output are disabled/invalid in confidence mode - generate_visualization_for_this = False - else: + generate_visualization_for_this = False + else: # User wants category mask (default or for visualization) if category_mask_hw is not None: - # Determine primary MASK output (HW Float) - if is_multiclass: - # Treat non-background as foreground for the primary mask - primary_mask_hw = (category_mask_hw > 0).float() - else: - # Binary models: Category 1 is foreground - primary_mask_hw = (category_mask_hw > 0).float() - current_primary_mask_hw = primary_mask_hw # HW - - # Prepare raw category mask output (HW Long Tensor) - current_category_mask_hw = category_mask_hw # Already Long - - # Generate visualization if requested (and category mask exists) - generate_visualization_for_this = generate_visualization - if generate_visualization_for_this: - # Convert to NumPy for OpenCV, then back to tensor - category_np = category_mask_hw.cpu().numpy().astype(np.uint8) - vis_np = self.create_visualization(category_np, is_multiclass) - if vis_np is not None: - # Convert back to tensor (HWC) - vis_tensor = torch.from_numpy(vis_np).to(dtype=torch.float32, device=image.device) / 255.0 - current_vis_image_hwc = vis_tensor # HWC - + # Determine primary MASK output (HW Float) + if is_multiclass: + # Treat non-background as foreground for the primary mask + primary_mask_hw = (category_mask_hw > 0).float() + else: + # Binary models: Category 1 is foreground + primary_mask_hw = (category_mask_hw > 0).float() + current_primary_mask_hw = primary_mask_hw # HW + + # Prepare raw category mask output (HW Long Tensor) + current_category_mask_hw = category_mask_hw # Already Long + + # Generate visualization if requested (and category mask exists) + generate_visualization_for_this = generate_visualization + if generate_visualization_for_this: + # Convert to NumPy for OpenCV, then back to tensor + category_np = category_mask_hw.cpu().numpy().astype(np.uint8) + vis_np = self.create_visualization(category_np, is_multiclass) + if vis_np is not None: + # Convert back to tensor (HWC) + vis_tensor = torch.from_numpy(vis_np).to(dtype=torch.float32, device=image.device) / 255.0 + current_vis_image_hwc = vis_tensor # HWC + # Add results for this image to the batch results - all_primary_masks.append(current_primary_mask_hw) # HW - all_vis_images.append(current_vis_image_hwc) # HWC - all_category_masks.append(current_category_mask_hw) # HW - + all_primary_masks.append(current_primary_mask_hw) # HW + all_vis_images.append(current_vis_image_hwc) # HWC + all_category_masks.append(current_category_mask_hw) # HW + # Stack all results to return as tensors - primary_mask_tensor = torch.stack(all_primary_masks, dim=0) # BHW - vis_image_tensor = torch.stack(all_vis_images, dim=0) # BHWC - category_mask_tensor = torch.stack(all_category_masks, dim=0) # BHW - + primary_mask_tensor = torch.stack(all_primary_masks, dim=0) # BHW + vis_image_tensor = torch.stack(all_vis_images, dim=0) # BHWC + category_mask_tensor = torch.stack(all_category_masks, dim=0) # BHW + # Ensure multiclass_segments is in the right format (BHW) for downstream nodes # Not adding an unsqueeze here, as we're handling both formats in the select_segment node - + return (primary_mask_tensor, vis_image_tensor, category_mask_tensor) + # Define mappings for ComfyUI NODE_CLASS_MAPPINGS = { "MediaPipeImageSegmenterModelLoader": MediaPipeImageSegmenterModelLoaderNode, @@ -220,4 +254,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeImageSegmenterModelLoader": "Load Image Segmenter Model (MediaPipe)", "MediaPipeImageSegmenter": "Image Segmenter (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py index dfa9dc7..ef9f72f 100644 --- a/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py +++ b/node_wrappers/mediapipe_vision/image_segmentation/image_segmentation_masking.py @@ -14,23 +14,29 @@ MULTICLASS_NAMES = { "Accessories/Other": 5, } + class SelectMediaPipeSegmentNode: """ Selects a specific segmentation mask from the raw multiclass output of the MediaPipeImageSegmenterNode. """ + @classmethod def INPUT_TYPES(cls): available_segments = list(MULTICLASS_NAMES.keys()) return { "required": { - "multiclass_segments": ("MP_INT_MASK", { - "tooltip": "Raw multiclass segmentation data (INT_MASK) from MediaPipe Image Segmenter node." - }), - "segment_name": (available_segments, { - "default": "Clothes" if "Clothes" in available_segments else available_segments[0], - "tooltip": "Select the specific segment to extract as a mask." - }), + "multiclass_segments": ( + "MP_INT_MASK", + {"tooltip": "Raw multiclass segmentation data (INT_MASK) from MediaPipe Image Segmenter node."}, + ), + "segment_name": ( + available_segments, + { + "default": "Clothes" if "Clothes" in available_segments else available_segments[0], + "tooltip": "Select the specific segment to extract as a mask.", + }, + ), } } @@ -43,14 +49,14 @@ class SelectMediaPipeSegmentNode: def select_segment(self, multiclass_segments: torch.Tensor, segment_name: str): if segment_name not in MULTICLASS_NAMES: raise ValueError(f"Invalid segment_name '{segment_name}'. Available: {list(MULTICLASS_NAMES.keys())}") - + if multiclass_segments is None or multiclass_segments.numel() == 0: - logger.warning("Input multiclass_segments is empty or None. Returning zero mask.") - # Need a shape to return a zero mask. How to get expected shape? - # For now, maybe raise error or return None (caller must handle) - # Returning None is safer if shape is unknown. - return (None,) # Or raise error - + logger.warning("Input multiclass_segments is empty or None. Returning zero mask.") + # Need a shape to return a zero mask. How to get expected shape? + # For now, maybe raise error or return None (caller must handle) + # Returning None is safer if shape is unknown. + return (None,) # Or raise error + # Check tensor format and reshape if needed # Input formats can be either BHW or B1HW (both are common in ComfyUI) if multiclass_segments.dim() == 3: # BHW format @@ -61,22 +67,24 @@ class SelectMediaPipeSegmentNode: elif multiclass_segments.dim() == 4 and multiclass_segments.shape[1] != 1: # If it's BCHW but not B1HW, raise error raise ValueError(f"Expected multiclass_segments to be B1HW or BHW tensor. Got shape {multiclass_segments.shape}") - + # Ensure input is a LongTensor if multiclass_segments.dtype != torch.int64: multiclass_segments = multiclass_segments.to(torch.int64) target_id = MULTICLASS_NAMES[segment_name] selected_mask = (multiclass_segments == target_id).float() - + # If result is B1HW, squeeze back to BHW for ComfyUI mask format compatibility if selected_mask.dim() == 4: selected_mask = selected_mask.squeeze(1) - - return (selected_mask,) -# --- End Select Segment Node --- -# --- Mappings --- + return (selected_mask,) + + +# --- End Select Segment Node --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "SelectMediaPipeSegment": SelectMediaPipeSegmentNode, } diff --git a/node_wrappers/mediapipe_vision/interactive_segmentation/__init__.py b/node_wrappers/mediapipe_vision/interactive_segmentation/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/interactive_segmentation/__init__.py +++ b/node_wrappers/mediapipe_vision/interactive_segmentation/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py b/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py index 93938c3..c0b5fd9 100644 --- a/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py +++ b/node_wrappers/mediapipe_vision/interactive_segmentation/interactive_segmentation.py @@ -6,22 +6,27 @@ import torch from ....src.mediapipe_vision.interactive_segmentation.detector import ( InteractiveSegmenterProcessor, ) -from ....src.mediapipe_vision. types import PointOfInterest +from ....src.mediapipe_vision.types import PointOfInterest # Import Base Loader and Detector from ..common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/InteractiveSegmentation" -# --- Model Loader --- + + +# --- Model Loader --- class MediaPipeInteractiveSegmenterModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Interactive Segmenter models.""" + # Usually uses a specific model like selfie_segmenter or a custom one - TASK_TYPE = "interactive_segmenter" # Need to define this task and models in model_loader.py + TASK_TYPE = "interactive_segmenter" # Need to define this task and models in model_loader.py RETURN_TYPES = ("INTERACTIVE_SEGMENTER_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category -# --- Segmenter Node --- + + +# --- Segmenter Node --- class MediaPipeInteractiveSegmenterNode: """ComfyUI node for MediaPipe Interactive Image Segmentation.""" @@ -36,32 +41,54 @@ class MediaPipeInteractiveSegmenterNode: "image": ("IMAGE",), "model_info": ("INTERACTIVE_SEGMENTER_MODEL_INFO",), # Input points as a JSON string list of lists/tuples: e.g., "[[0.5, 0.5], [0.2, 0.3]]" - "points_json": ("STRING", {"default": "[[0.5, 0.5]]", "multiline": True, - "tooltip": "Points of interest (normalized coords) as JSON: [[x1,y1], [x2,y2], ...]"}), - "output_category_mask": ("BOOLEAN", {"default": True, - "tooltip": "Output a category mask (0=background, 1=foreground). Priority over confidence."}), - "output_confidence_mask": ("BOOLEAN", {"default": False, - "tooltip": "Output a confidence mask (0.0-1.0). Used if category mask is False."}), + "points_json": ( + "STRING", + { + "default": "[[0.5, 0.5]]", + "multiline": True, + "tooltip": "Points of interest (normalized coords) as JSON: [[x1,y1], [x2,y2], ...]", + }, + ), + "output_category_mask": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Output a category mask (0=background, 1=foreground). Priority over confidence.", + }, + ), + "output_confidence_mask": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Output a confidence mask (0.0-1.0). Used if category mask is False.", + }, + ), } } - RETURN_TYPES = ("MASK",) # Output is a standard ComfyUI mask + RETURN_TYPES = ("MASK",) # Output is a standard ComfyUI mask RETURN_NAMES = ("mask",) FUNCTION = "segment_interactive" CATEGORY = _category - def segment_interactive(self, image: torch.Tensor, model_info: dict, points_json: str, - output_category_mask: bool, output_confidence_mask: bool): + def segment_interactive( + self, + image: torch.Tensor, + model_info: dict, + points_json: str, + output_category_mask: bool, + output_confidence_mask: bool, + ): """Performs interactive segmentation.""" - - task_type = model_info.get('task_type') + + task_type = model_info.get("task_type") expected_task_type = MediaPipeInteractiveSegmenterModelLoaderNode.TASK_TYPE if not isinstance(model_info, dict) or task_type != expected_task_type: - raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but got '{task_type}'.") - model_path = model_info.get('model_path') + raise ValueError(f"Invalid model_info. Expected task_type '{expected_task_type}' but got '{task_type}'.") + model_path = model_info.get("model_path") if not model_path: - raise ValueError("Model path not found or invalid in model_info.") - + raise ValueError("Model path not found or invalid in model_info.") + # Parse points from JSON try: points_list = json.loads(points_json) @@ -73,42 +100,47 @@ class MediaPipeInteractiveSegmenterNode: # Manage detector instance if self._processor is None or self._model_path != model_path: - if self._processor and hasattr(self._processor, 'close'): - try: self._processor.close() - except Exception as e: logger.warning(f"Error closing InteractiveSegmenterProcessor: {e}") + if self._processor and hasattr(self._processor, "close"): + try: + self._processor.close() + except Exception as e: + logger.warning(f"Error closing InteractiveSegmenterProcessor: {e}") logger.info(f"Creating new InteractiveSegmenterProcessor instance for {model_path}") self._processor = InteractiveSegmenterProcessor(model_path) self._model_path = model_path - + # Call the detector's segment method mask_batch = self._processor.segment( image, points_of_interest=points_of_interest, output_category_mask=output_category_mask, - output_confidence_mask=output_confidence_mask + output_confidence_mask=output_confidence_mask, ) - + # Convert HW mask list to B1HW tensor for ComfyUI MASK type # Assuming detector returns one mask per batch item comfy_masks = [] for mask_hw in mask_batch: - if mask_hw is not None: - comfy_masks.append(mask_hw.unsqueeze(0)) # Add channel dim - else: - # Handle None case - maybe return zero mask? Needs image shape - h, w = image.shape[1], image.shape[2] - comfy_masks.append(torch.zeros((1, h, w), dtype=torch.float32, device=image.device)) - - output_tensor = torch.stack(comfy_masks, dim=0) # Stack along batch dim + if mask_hw is not None: + comfy_masks.append(mask_hw.unsqueeze(0)) # Add channel dim + else: + # Handle None case - maybe return zero mask? Needs image shape + h, w = image.shape[1], image.shape[2] + comfy_masks.append(torch.zeros((1, h, w), dtype=torch.float32, device=image.device)) + + output_tensor = torch.stack(comfy_masks, dim=0) # Stack along batch dim return (output_tensor,) def __del__(self): - if hasattr(self, '_processor') and self._processor and hasattr(self._processor, 'close'): - try: self._processor.close() - except Exception as e: logger.warning(f"Error closing InteractiveSegmenterProcessor in __del__: {e}") - self._processor = None + if hasattr(self, "_processor") and self._processor and hasattr(self._processor, "close"): + try: + self._processor.close() + except Exception as e: + logger.warning(f"Error closing InteractiveSegmenterProcessor in __del__: {e}") + self._processor = None -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MediaPipeInteractiveSegmenterModelLoader": MediaPipeInteractiveSegmenterModelLoaderNode, "MediaPipeInteractiveSegmenter": MediaPipeInteractiveSegmenterNode, @@ -117,4 +149,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeInteractiveSegmenterModelLoader": "Load Interactive Segmenter Model (MediaPipe)", "MediaPipeInteractiveSegmenter": "Interactive Segmenter (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/object_detection/__init__.py b/node_wrappers/mediapipe_vision/object_detection/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/object_detection/__init__.py +++ b/node_wrappers/mediapipe_vision/object_detection/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/object_detection/object_detection.py b/node_wrappers/mediapipe_vision/object_detection/object_detection.py index ea25a92..92164c4 100644 --- a/node_wrappers/mediapipe_vision/object_detection/object_detection.py +++ b/node_wrappers/mediapipe_vision/object_detection/object_detection.py @@ -9,13 +9,16 @@ from ..common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ObjectDetection" + class MediaPipeObjectDetectorModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Object Detector models.""" + TASK_TYPE = "object_detector" RETURN_TYPES = ("OBJECT_DETECTOR_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category + class MediaPipeObjectDetectorNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Object Detection.""" @@ -32,38 +35,54 @@ class MediaPipeObjectDetectorNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add object detection specific parameters - inputs["required"].update({ - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, - "tooltip": "Minimum confidence score for detected objects"}), - "max_results": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1, - "tooltip": "Maximum number of objects to detect"}), - }) - + inputs["required"].update( + { + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Minimum confidence score for detected objects", + }, + ), + "max_results": ( + "INT", + {"default": 5, "min": 1, "max": 50, "step": 1, "tooltip": "Maximum number of objects to detect"}, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, min_confidence: float, - max_results: int, running_mode: str, delegate: str): + def detect( + self, + image: torch.Tensor, + model_info: dict, + min_confidence: float, + max_results: int, + running_mode: str, + delegate: str, + ): """Performs object detection with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Perform detection with all parameters batch_results = detector.detect( - image, - score_threshold=min_confidence, - max_results=max_results, - running_mode=running_mode, - delegate=delegate + image, score_threshold=min_confidence, max_results=max_results, running_mode=running_mode, delegate=delegate ) - + return (batch_results,) + # Define mappings for ComfyUI NODE_CLASS_MAPPINGS = { "MediaPipeObjectDetectorModelLoader": MediaPipeObjectDetectorModelLoaderNode, @@ -73,4 +92,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipeObjectDetectorModelLoader": "Load Object Detector Model (MediaPipe)", "MediaPipeObjectDetector": "Object Detector (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py index 0b50833..2225091 100644 --- a/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py +++ b/node_wrappers/mediapipe_vision/object_detection/object_detection_visualization.py @@ -8,60 +8,76 @@ from ..common.base_visualization_nodes import BaseDetectionVisualizationNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/ObjectDetection/Visualization" + class VisualizeObjectDetections(BaseDetectionVisualizationNode): """Draws object detection bounding boxes and labels onto images.""" + CATEGORY = _category @classmethod def INPUT_TYPES(cls): return { "required": { - "object_detections": ("OBJECT_DETECTIONS",), # Expect the specific type + "object_detections": ("OBJECT_DETECTIONS",), # Expect the specific type "image": ("IMAGE",), }, "optional": { - "box_color": ("STRING", {"default": "#FF0000"}), # Red + "box_color": ("STRING", {"default": "#FF0000"}), # Red "box_thickness": ("INT", {"default": 2, "min": 1, "max": 10}), - "label_color": ("STRING", {"default": "#FFFFFF"}), # White - "label_bg_color": ("STRING", {"default": "#FF0000"}), # Red BG + "label_color": ("STRING", {"default": "#FFFFFF"}), # White + "label_bg_color": ("STRING", {"default": "#FF0000"}), # Red BG # Font size might be tricky to implement well, skipping for now - } + }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "draw_visualizations" - def draw_visualizations(self, object_detections, image, - box_color="#FF0000", box_thickness=2, - label_color="#FFFFFF", label_bg_color="#FF0000"): + def draw_visualizations( + self, + object_detections, + image, + box_color="#FF0000", + box_thickness=2, + label_color="#FFFFFF", + label_bg_color="#FF0000", + ): """Main function that processes the batch of images and detections.""" return self.process_batch( - object_detections, + object_detections, image, self.draw_objects, box_color=box_color, box_thickness=box_thickness, label_color=label_color, - label_bg_color=label_bg_color + label_bg_color=label_bg_color, ) - - def draw_objects(self, image, width, height, detections, - box_color="#FF0000", box_thickness=2, - label_color="#FFFFFF", label_bg_color="#FF0000"): + + def draw_objects( + self, + image, + width, + height, + detections, + box_color="#FF0000", + box_thickness=2, + label_color="#FFFFFF", + label_bg_color="#FF0000", + ): """Callback function to draw objects on a single image.""" # Convert hex colors to BGR tuples for cv2 box_color_bgr = self._hex_to_bgr(box_color) - + for detection in detections: bbox = detection.bounding_box # Top-left corner x1, y1 = int(bbox.origin_x), int(bbox.origin_y) # Bottom-right corner x2, y2 = int(bbox.origin_x + bbox.width), int(bbox.origin_y + bbox.height) - + # Draw bounding box cv2.rectangle(image, (x1, y1), (x2, y2), box_color_bgr, box_thickness) - + # Prepare label text (use first category's display name if available) label = "N/A" score = 0.0 @@ -69,36 +85,27 @@ class VisualizeObjectDetections(BaseDetectionVisualizationNode): top_category = detection.categories[0] label = top_category.display_name or top_category.category_name or "N/A" score = top_category.score - + label_text = f"{label}: {score:.2f}" - + # Draw label background and text inside the box top-left try: - self.draw_text_with_background( - image, - label_text, - (x1, y1), - text_color=label_color, - bg_color=label_bg_color - ) + self.draw_text_with_background(image, label_text, (x1, y1), text_color=label_color, bg_color=label_bg_color) except Exception as e: logger.error(f"Error drawing text label: {e}") # Fallback: Draw text without background cv2.putText( - image, - label_text, - (x1 + 2, y1 + 15), - cv2.FONT_HERSHEY_SIMPLEX, - 0.5, - self._hex_to_bgr(label_color), - 1 + image, + label_text, + (x1 + 2, y1 + 15), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + self._hex_to_bgr(label_color), + 1, ) -# Node registration -NODE_CLASS_MAPPINGS = { - "VisualizeObjectDetections": VisualizeObjectDetections -} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizeObjectDetections": "Visualize Object Detections (MediaPipe)" -} \ No newline at end of file +# Node registration +NODE_CLASS_MAPPINGS = {"VisualizeObjectDetections": VisualizeObjectDetections} + +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizeObjectDetections": "Visualize Object Detections (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/pose/__init__.py b/node_wrappers/mediapipe_vision/pose/__init__.py index 0519ecb..e69de29 100644 --- a/node_wrappers/mediapipe_vision/pose/__init__.py +++ b/node_wrappers/mediapipe_vision/pose/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark.py b/node_wrappers/mediapipe_vision/pose/pose_landmark.py index 40397f6..54b0638 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark.py @@ -9,14 +9,17 @@ from ..common.model_loader import MediaPipeModelLoaderBaseNode logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Pose" + class MediaPipePoseLandmarkerModelLoaderNode(MediaPipeModelLoaderBaseNode): """ComfyUI node for loading MediaPipe Pose Landmarker models.""" + TASK_TYPE = "pose_landmarker" RETURN_TYPES = ("POSE_LANDMARKER_MODEL_INFO",) RETURN_NAMES = ("model_info",) CATEGORY = _category DESCRIPTION = "Loads a MediaPipe Pose Landmarker model for detecting human body poses and postures in images. Required before using the Pose Landmarker node." + class MediaPipePoseLandmarkerNode(BaseMediaPipeDetectorNode): """ComfyUI node for MediaPipe Pose Landmark Detection.""" @@ -24,7 +27,10 @@ class MediaPipePoseLandmarkerNode(BaseMediaPipeDetectorNode): DETECTOR_CLASS = PoseLandmarkDetector MODEL_INFO_TYPE = "POSE_LANDMARKER_MODEL_INFO" EXPECTED_TASK_TYPE = "pose_landmarker" - RETURN_TYPES = ("POSE_LANDMARKS", "MASK_LIST",) + RETURN_TYPES = ( + "POSE_LANDMARKS", + "MASK_LIST", + ) RETURN_NAMES = ("landmarks", "segmentation_masks") FUNCTION = "detect" CATEGORY = _category @@ -34,35 +40,82 @@ class MediaPipePoseLandmarkerNode(BaseMediaPipeDetectorNode): def INPUT_TYPES(cls): # Start with the base inputs from the parent class inputs = super().INPUT_TYPES() - + # Add pose-specific parameters - inputs["required"].update({ - "max_results": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1, - "tooltip": "Maximum number of people/poses to detect in the image - higher values can detect more people but use more processing power"}), - "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Minimum confidence threshold for pose detection - lower values detect more poses but may include false positives"}), - "min_tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence threshold for pose tracking between frames - lower values maintain tracking longer but may be less accurate"}), - "min_presence_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "(VIDEO mode) Minimum confidence threshold for pose presence in the image - lower values may include false negatives"}), - "output_segmentation_masks": ("BOOLEAN", {"default": False, - "tooltip": "When enabled, generates masks that separate people from the background"}), - }) - + inputs["required"].update( + { + "max_results": ( + "INT", + { + "default": 1, + "min": 1, + "max": 4, + "step": 1, + "tooltip": "Maximum number of people/poses to detect in the image - higher values can detect more people but use more processing power", + }, + ), + "min_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum confidence threshold for pose detection - lower values detect more poses but may include false positives", + }, + ), + "min_tracking_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence threshold for pose tracking between frames - lower values maintain tracking longer but may be less accurate", + }, + ), + "min_presence_confidence": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "(VIDEO mode) Minimum confidence threshold for pose presence in the image - lower values may include false negatives", + }, + ), + "output_segmentation_masks": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, generates masks that separate people from the background", + }, + ), + } + ) + return inputs - def detect(self, image: torch.Tensor, model_info: dict, max_results: int, - min_confidence: float, min_tracking_confidence: float, - min_presence_confidence: float, output_segmentation_masks: bool, - running_mode: str, delegate: str): + def detect( + self, + image: torch.Tensor, + model_info: dict, + max_results: int, + min_confidence: float, + min_tracking_confidence: float, + min_presence_confidence: float, + output_segmentation_masks: bool, + running_mode: str, + delegate: str, + ): """Performs pose landmark detection with the configured parameters.""" - + # Validate model_info and get model path model_path = self.validate_model_info(model_info) - + # Initialize or update detector detector = self.initialize_or_update_detector(model_path) - + # Perform detection with all parameters pose_landmarks_batch, segmentation_masks_batch = detector.detect( image, @@ -72,18 +125,19 @@ class MediaPipePoseLandmarkerNode(BaseMediaPipeDetectorNode): min_presence_confidence=min_presence_confidence, output_segmentation_masks=output_segmentation_masks, running_mode=running_mode, - delegate=delegate + delegate=delegate, ) - + # Process segmentation masks if needed flat_masks = [] if segmentation_masks_batch: for img_masks in segmentation_masks_batch: if img_masks: flat_masks.extend(img_masks) - + return (pose_landmarks_batch, flat_masks if output_segmentation_masks else []) + # Define mappings for ComfyUI NODE_CLASS_MAPPINGS = { "MediaPipePoseLandmarkerModelLoader": MediaPipePoseLandmarkerModelLoaderNode, @@ -93,4 +147,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "MediaPipePoseLandmarkerModelLoader": "Load Pose Landmarker Model (MediaPipe)", "MediaPipePoseLandmarker": "Pose Landmarker (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py index f5a308f..9f48113 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_masking.py @@ -15,22 +15,24 @@ _category = "Realtime Nodes/MediaPipe Vision/Pose/Masking" # --- Pose Part Definitions (Based on standard MediaPipe Pose Indices - BlazePose) --- # These indices need verification POSE_PART_INDICES = { - "HEAD": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], # Example - "TORSO": [11, 12, 23, 24], # Example - "LEFT_ARM": [11, 13, 15], # Upper arm only - "RIGHT_ARM": [12, 14, 16], # Upper arm only - "LEFT_FOREARM": [13, 15, 17, 19, 21], # Includes hand points? - "RIGHT_FOREARM": [14, 16, 18, 20, 22], # Includes hand points? - "LEFT_LEG": [23, 25, 27], # Upper leg - "RIGHT_LEG": [24, 26, 28], # Upper leg - "LEFT_LOWER_LEG": [25, 27, 29, 31], # Includes foot? - "RIGHT_LOWER_LEG": [26, 28, 30, 32], # Includes foot? + "HEAD": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], # Example + "TORSO": [11, 12, 23, 24], # Example + "LEFT_ARM": [11, 13, 15], # Upper arm only + "RIGHT_ARM": [12, 14, 16], # Upper arm only + "LEFT_FOREARM": [13, 15, 17, 19, 21], # Includes hand points? + "RIGHT_FOREARM": [14, 16, 18, 20, 22], # Includes hand points? + "LEFT_LEG": [23, 25, 27], # Upper leg + "RIGHT_LEG": [24, 26, 28], # Upper leg + "LEFT_LOWER_LEG": [25, 27, 29, 31], # Includes foot? + "RIGHT_LOWER_LEG": [26, 28, 30, 32], # Includes foot? "FULL_BODY": list(range(33)), } -# --- Node Definition --- + +# --- Node Definition --- class MaskFromPoseLandmarksNode: """Creates a mask for a specific body part from Pose Landmarks.""" + CATEGORY = _category DESCRIPTION = "Creates precise masks for specific body parts like head, torso, arms, or legs. Perfect for targeted editing or creating effects on specific body regions." @@ -39,22 +41,47 @@ class MaskFromPoseLandmarksNode: part_names = list(POSE_PART_INDICES.keys()) return { "required": { - "pose_landmarks": ("POSE_LANDMARKS", {"tooltip": "Pose landmark data from the Pose Landmarker node containing the precise locations of body joints"}), - "image_for_dimensions": ("IMAGE", {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}), - "part_name": (part_names, {"default": "FULL_BODY", - "tooltip": "Choose which body part to create a mask for - FULL_BODY (entire person), HEAD, TORSO, LEFT_ARM, etc."}), - "min_visibility": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, - "tooltip": "Only include landmarks that are at least this visible (0.0 = include all points regardless of visibility, 1.0 = only fully visible points)"}), + "pose_landmarks": ( + "POSE_LANDMARKS", + {"tooltip": "Pose landmark data from the Pose Landmarker node containing the precise locations of body joints"}, + ), + "image_for_dimensions": ( + "IMAGE", + {"tooltip": "Reference image to determine the size of the output mask - the mask will match these dimensions"}, + ), + "part_name": ( + part_names, + { + "default": "FULL_BODY", + "tooltip": "Choose which body part to create a mask for - FULL_BODY (entire person), HEAD, TORSO, LEFT_ARM, etc.", + }, + ), + "min_visibility": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Only include landmarks that are at least this visible (0.0 = include all points regardless of visibility, 1.0 = only fully visible points)", + }, + ), } } RETURN_TYPES = ("MASK",) FUNCTION = "generate_mask" - def generate_mask(self, pose_landmarks: List[List[PoseLandmarksResult]], image_for_dimensions: torch.Tensor, part_name: str, min_visibility: float): + def generate_mask( + self, + pose_landmarks: List[List[PoseLandmarksResult]], + image_for_dimensions: torch.Tensor, + part_name: str, + min_visibility: float, + ): if part_name not in POSE_PART_INDICES: raise ValueError(f"Unknown pose part name: '{part_name}'. Available: {list(POSE_PART_INDICES.keys())}") - + # Get dimensions from image if image_for_dimensions.dim() != 4: raise ValueError("image_for_dimensions must be BHWC format.") @@ -65,7 +92,9 @@ class MaskFromPoseLandmarksNode: # Validate batch sizes match if len(pose_landmarks) != input_batch_size: - logger.warning(f"Batch size mismatch between pose_landmarks ({len(pose_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size.") + logger.warning( + f"Batch size mismatch between pose_landmarks ({len(pose_landmarks)}) and image_for_dimensions ({input_batch_size}). Using landmark batch size." + ) batch_size = len(pose_landmarks) else: batch_size = input_batch_size @@ -74,7 +103,7 @@ class MaskFromPoseLandmarksNode: output_masks = [] for i in range(batch_size): - landmarks_for_image = pose_landmarks[i] # List[PoseLandmarksResult] + landmarks_for_image = pose_landmarks[i] # List[PoseLandmarksResult] combined_points_for_image = [] # Combine points from all detected poses for the selected part @@ -92,24 +121,25 @@ class MaskFromPoseLandmarksNode: part_points.append((px, py)) # No warning for missing index combined_points_for_image.extend(part_points) - + # Create mask from all collected points for this image mask_tensor = create_mask_from_points(height, width, combined_points_for_image, device=device) output_masks.append(mask_tensor) if not output_masks: - logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") - output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) + logger.warning("No landmarks found to generate any masks. Returning zero mask batch.") + output_batch = torch.zeros((input_batch_size, 1, height, width), dtype=torch.float32, device=device) else: - output_batch = torch.stack(output_masks, dim=0) - + output_batch = torch.stack(output_masks, dim=0) + return (output_batch,) -# --- Mappings --- + +# --- Mappings --- NODE_CLASS_MAPPINGS = { "MaskFromPoseLandmarks": MaskFromPoseLandmarksNode, } NODE_DISPLAY_NAME_MAPPINGS = { "MaskFromPoseLandmarks": "Mask From Pose Landmarks (MediaPipe)", -} \ No newline at end of file +} diff --git a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py index 630a7e8..5c0870c 100644 --- a/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py +++ b/node_wrappers/mediapipe_vision/pose/pose_landmark_visualizer.py @@ -11,9 +11,10 @@ _category = "Realtime Nodes/MediaPipe Vision/Pose/Visualization" # Define standard pose connections POSE_CONNECTIONS = mp.solutions.pose.POSE_CONNECTIONS + class VisualizePoseLandmarks(BaseLandmarkVisualizationNode): """Draws Pose landmarks and connections onto images.""" - + # Override base class properties LANDMARKS_TYPE = "POSE_LANDMARKS" CONNECTIONS = POSE_CONNECTIONS @@ -22,18 +23,14 @@ class VisualizePoseLandmarks(BaseLandmarkVisualizationNode): DEFAULT_POINT_RADIUS = 3 DEFAULT_CONNECTION_COLOR = "#FFFFFF" # White DEFAULT_CONNECTION_THICKNESS = 2 - SUPPORTS_VISIBILITY = True # Pose landmarks have visibility + SUPPORTS_VISIBILITY = True # Pose landmarks have visibility CATEGORY = _category - + def get_landmark_list(self, detection_result: Any) -> Optional[List[Any]]: return detection_result.landmarks # Node registration -NODE_CLASS_MAPPINGS = { - "VisualizePoseLandmarks": VisualizePoseLandmarks -} +NODE_CLASS_MAPPINGS = {"VisualizePoseLandmarks": VisualizePoseLandmarks} -NODE_DISPLAY_NAME_MAPPINGS = { - "VisualizePoseLandmarks": "Visualize Pose Landmarks (MediaPipe)" -} \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = {"VisualizePoseLandmarks": "Visualize Pose Landmarks (MediaPipe)"} diff --git a/node_wrappers/mediapipe_vision/pose/position.py b/node_wrappers/mediapipe_vision/pose/position.py index 589539d..01ae786 100644 --- a/node_wrappers/mediapipe_vision/pose/position.py +++ b/node_wrappers/mediapipe_vision/pose/position.py @@ -10,9 +10,11 @@ from ....src.mediapipe_vision.landmark_definitions import POSE_LANDMARK_TOOLTIP logger = logging.getLogger(__name__) _category = "Realtime Nodes/MediaPipe Vision/Pose/PoseLandmark/Position" + class PoseLandmarkPositionNode(LandmarkPositionBase): """Extracts position lists (x, y, z, vis, pres) for a specific landmark index - from POSE_LANDMARKS across the batch.""" + from POSE_LANDMARKS across the batch.""" + CATEGORY = _category # Define the specific input type string this node expects LANDMARKS_TYPE = "POSE_LANDMARKS" @@ -25,14 +27,21 @@ class PoseLandmarkPositionNode(LandmarkPositionBase): "required": { "landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), "landmark_index": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": POSE_LANDMARK_TOOLTIP}), - "result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the pose detection to use (0=first detected pose)"}), - "use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}), + "result_index": ( + "INT", + {"default": 0, "min": 0, "tooltip": "Index of the pose detection to use (0=first detected pose)"}, + ), + "use_world_coordinates": ( + "BOOLEAN", + {"default": False, "tooltip": "Use world coordinates if available"}, + ), } } # No need to override INPUT_TYPES or extract_positions, inheriting from base is sufficient. # The base class handles everything using self.LANDMARKS_TYPE. + # --- Node Mappings --- NODE_CLASS_MAPPINGS = { "PoseLandmarkPosition": PoseLandmarkPositionNode, @@ -40,4 +49,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "PoseLandmarkPosition": "Pose Landmark Position Extractor (Batch)", -} \ No newline at end of file +} diff --git a/node_wrappers/realtimenodes/__init__.py b/node_wrappers/realtimenodes/__init__.py index da22290..6e7633b 100644 --- a/node_wrappers/realtimenodes/__init__.py +++ b/node_wrappers/realtimenodes/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/node_wrappers/realtimenodes/controls/__init__.py b/node_wrappers/realtimenodes/controls/__init__.py index da22290..6e7633b 100644 --- a/node_wrappers/realtimenodes/controls/__init__.py +++ b/node_wrappers/realtimenodes/controls/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/node_wrappers/realtimenodes/controls/mask_controls.py b/node_wrappers/realtimenodes/controls/mask_controls.py index c35ef38..f48947f 100644 --- a/node_wrappers/realtimenodes/controls/mask_controls.py +++ b/node_wrappers/realtimenodes/controls/mask_controls.py @@ -7,44 +7,47 @@ from ....src.realtimenodes.mask_controls import MaskControlMixin class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin): """Node that maintains a mask that repulses or attracts to input masks""" + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "input_mask": ("MASK",), - "x_pos": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Initial X position (0-1)" - }), - "y_pos": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Initial Y position (0-1)" - }), - "size": ("FLOAT", { - "default": 0.2, - "min": 0.01, - "max": 1.0, - "step": 0.01, - "tooltip": "Size of the mask relative to image size" - }), - "interaction_type": (["repulse", "attract"], { - "default": "repulse", - "tooltip": "Whether to move away from or toward input masks" - }), - "interaction_strength": ("FLOAT", { - "default": 0.1, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Strength of the interaction effect" - }) - }) + inputs["required"].update( + { + "input_mask": ("MASK",), + "x_pos": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Initial X position (0-1)"}, + ), + "y_pos": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Initial Y position (0-1)"}, + ), + "size": ( + "FLOAT", + { + "default": 0.2, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Size of the mask relative to image size", + }, + ), + "interaction_type": ( + ["repulse", "attract"], + {"default": "repulse", "tooltip": "Whether to move away from or toward input masks"}, + ), + "interaction_strength": ( + "FLOAT", + { + "default": 0.1, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength of the interaction effect", + }, + ), + } + ) return inputs RETURN_TYPES = ("MASK",) @@ -54,45 +57,45 @@ class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin): def interact(self, mask, input_mask_np, state, strength, height, width, is_attract): if np.max(input_mask_np) < 0.1: return state - + y_indices, x_indices = np.nonzero(input_mask_np) if len(y_indices) == 0: return state - + weights = input_mask_np[y_indices, x_indices] com_y = np.average(y_indices, weights=weights) com_x = np.average(x_indices, weights=weights) - + center_y = state["y"] * height center_x = state["x"] * width - + dir_y = center_y - com_y dir_x = center_x - com_x - + dist = np.sqrt(dir_x**2 + dir_y**2) if dist < 1: return state - - scale = min(1.0, 20.0/dist) + + scale = min(1.0, 20.0 / dist) dir_y = (dir_y / dist) * scale dir_x = (dir_x / dist) * scale - + # Flip direction if attracting if is_attract: dir_y *= -1 dir_x *= -1 - + new_y = state["y"] + (dir_y * strength) new_x = state["x"] + (dir_x * strength) - + state["y"] = np.clip(new_y, 0.0, 1.0) state["x"] = np.clip(new_x, 0.0, 1.0) - + return state def update(self, input_mask, x_pos, y_pos, size, interaction_type, interaction_strength, always_execute=True): was_reset = len(self.state_manager._states) == 0 - + if was_reset: state = self.get_initial_state(x_pos, y_pos, size) else: @@ -100,71 +103,81 @@ class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin): batch_size, height, width = input_mask.shape masks = np.zeros((batch_size, height, width), dtype=np.float32) - + for b in range(batch_size): center_y, center_x = int(state["y"] * height), int(state["x"] * width) mask = self.create_circle_mask(height, width, center_y, center_x, state["size"]) - - state = self.interact(mask, input_mask[b].cpu().numpy(), state, - interaction_strength, height, width, - interaction_type == "attract") - + + state = self.interact( + mask, + input_mask[b].cpu().numpy(), + state, + interaction_strength, + height, + width, + interaction_type == "attract", + ) + center_y, center_x = int(state["y"] * height), int(state["x"] * width) masks[b] = self.create_circle_mask(height, width, center_y, center_x, state["size"]) - + self.set_state(state) return (torch.from_numpy(masks),) + class ResizeMaskNode(ControlNodeBase, MaskControlMixin): """Node that maintains a mask that resizes based on proximity to input masks""" + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "input_mask": ("MASK",), - "x_pos": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Initial X position (0-1)" - }), - "y_pos": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Initial Y position (0-1)" - }), - "size": ("FLOAT", { - "default": 0.2, - "min": 0.01, - "max": 1.0, - "step": 0.01, - "tooltip": "Size of the mask relative to image size" - }), - "min_size": ("FLOAT", { - "default": 0.1, - "min": 0.01, - "max": 1.0, - "step": 0.01, - "tooltip": "Minimum size the mask can shrink to" - }), - "max_size": ("FLOAT", { - "default": 0.4, - "min": 0.01, - "max": 1.0, - "step": 0.01, - "tooltip": "Maximum size the mask can grow to" - }), - "resize_strength": ("FLOAT", { - "default": 0.1, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Strength of the resize effect" - }) - }) + inputs["required"].update( + { + "input_mask": ("MASK",), + "x_pos": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Initial X position (0-1)"}, + ), + "y_pos": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Initial Y position (0-1)"}, + ), + "size": ( + "FLOAT", + { + "default": 0.2, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Size of the mask relative to image size", + }, + ), + "min_size": ( + "FLOAT", + { + "default": 0.1, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum size the mask can shrink to", + }, + ), + "max_size": ( + "FLOAT", + { + "default": 0.4, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Maximum size the mask can grow to", + }, + ), + "resize_strength": ( + "FLOAT", + {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Strength of the resize effect"}, + ), + } + ) return inputs RETURN_TYPES = ("MASK",) @@ -181,13 +194,13 @@ class ResizeMaskNode(ControlNodeBase, MaskControlMixin): center_y = int(state["y"] * height) center_x = int(state["x"] * width) y_indices, x_indices = np.nonzero(input_mask_np > 0.1) - + if len(y_indices) == 0: return state - - distances = np.sqrt((y_indices - center_y)**2 + (x_indices - center_x)**2) + + distances = np.sqrt((y_indices - center_y) ** 2 + (x_indices - center_x) ** 2) min_dist = np.min(distances) - + # Map distance directly to size # When dist=0, size=min_size # When dist=100, size=max_size @@ -195,12 +208,12 @@ class ResizeMaskNode(ControlNodeBase, MaskControlMixin): size_range = state["max_size"] - state["min_size"] size = state["min_size"] + (min_dist / dist_range) * size_range state["size"] = np.clip(size, state["min_size"], state["max_size"]) - + return state def update(self, input_mask, x_pos, y_pos, size, min_size, max_size, resize_strength, always_execute=True): was_reset = len(self.state_manager._states) == 0 - + if was_reset: state = self.get_initial_state(x_pos, y_pos, size, min_size, max_size) else: @@ -208,15 +221,15 @@ class ResizeMaskNode(ControlNodeBase, MaskControlMixin): batch_size, height, width = input_mask.shape masks = np.zeros((batch_size, height, width), dtype=np.float32) - + for b in range(batch_size): center_y, center_x = int(state["y"] * height), int(state["x"] * width) mask = self.create_circle_mask(height, width, center_y, center_x, state["size"]) - + state = self.resize(mask, input_mask[b].cpu().numpy(), state, resize_strength, height, width) - + center_y, center_x = int(state["y"] * height), int(state["x"] * width) masks[b] = self.create_circle_mask(height, width, center_y, center_x, state["size"]) - + self.set_state(state) - return (torch.from_numpy(masks),) \ No newline at end of file + return (torch.from_numpy(masks),) diff --git a/node_wrappers/realtimenodes/controls/motion_controls.py b/node_wrappers/realtimenodes/controls/motion_controls.py index b340329..9f4bd0a 100644 --- a/node_wrappers/realtimenodes/controls/motion_controls.py +++ b/node_wrappers/realtimenodes/controls/motion_controls.py @@ -10,35 +10,35 @@ from ....src.realtimenodes.control_base import ControlNodeBase class ROIAction(Enum): # Behavioral actions - TOGGLE = "toggle" # Toggles between min/max values + TOGGLE = "toggle" # Toggles between min/max values MOMENTARY = "momentary" # Outputs max while motion detected - TRIGGER = "trigger" # Triggers once per motion event - COUNTER = "counter" # Counts motion events + TRIGGER = "trigger" # Triggers once per motion event + COUNTER = "counter" # Counts motion events # Mathematical actions - ADD = "add" # Add value when motion detected - SUBTRACT = "subtract" # Subtract value when motion detected - MULTIPLY = "multiply" # Multiply by value when motion detected - DIVIDE = "divide" # Divide by value when motion detected - SET = "set" # Set to specific value when motion detected + ADD = "add" # Add value when motion detected + SUBTRACT = "subtract" # Subtract value when motion detected + MULTIPLY = "multiply" # Multiply by value when motion detected + DIVIDE = "divide" # Divide by value when motion detected + SET = "set" # Set to specific value when motion detected + class ROINode(ControlNodeBase): """Defines a single region of interest and its action""" - + @classmethod def INPUT_TYPES(cls): return { "required": { "mask": ("MASK",), # ROI defined by mask "action": (list(action.value for action in ROIAction),), - "value": ("FLOAT", { - "default": 0.1, - "step": 0.01, - "tooltip": "Value to use for mathematical operations" - }), + "value": ( + "FLOAT", + {"default": 0.1, "step": 0.01, "tooltip": "Value to use for mathematical operations"}, + ), }, "optional": { "next_roi": ("ROI",) # Chain to next ROI - } + }, } RETURN_TYPES = ("ROI",) @@ -48,7 +48,7 @@ class ROINode(ControlNodeBase): def define_roi(self, mask, action, value, next_roi=None): # Convert mask to numpy for bounding box calculation mask_np = mask[0].cpu().numpy() - + # Find non-zero coordinates in mask coords = np.nonzero(mask_np) if len(coords[0]) == 0: # Empty mask @@ -56,78 +56,73 @@ class ROINode(ControlNodeBase): else: y_min, y_max = coords[0].min(), coords[0].max() x_min, x_max = coords[1].min(), coords[1].max() - + roi_data = { "mask": mask_np, "bounds": (y_min, x_min, y_max, x_max), "action": action, "value": value, # Store the operation value - "next": next_roi + "next": next_roi, } return (roi_data,) def update(self, *args, **kwargs): return self.define_roi(*args, **kwargs) + class MotionController(ControlNodeBase): """Processes motion detection for a chain of ROIs and manages their states""" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "image": ("IMAGE",), - "roi_chain": ("ROI",), - "threshold": ("FLOAT", { - "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, - "tooltip": "Motion detection threshold" - }), - "blur_size": ("INT", { - "default": 5, "min": 1, "max": 21, "step": 2, - "tooltip": "Size of Gaussian blur kernel" - }), - "minimum_value": ("FLOAT", { - "default": 0.0, - "tooltip": "Minimum output value" - }), - "maximum_value": ("FLOAT", { - "default": 1.0, - "tooltip": "Maximum output value" - }), - "starting_value": ("FLOAT", { - "default": 0.0, - "tooltip": "Initial output value" - }) - }) + inputs["required"].update( + { + "image": ("IMAGE",), + "roi_chain": ("ROI",), + "threshold": ( + "FLOAT", + {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Motion detection threshold"}, + ), + "blur_size": ( + "INT", + {"default": 5, "min": 1, "max": 21, "step": 2, "tooltip": "Size of Gaussian blur kernel"}, + ), + "minimum_value": ("FLOAT", {"default": 0.0, "tooltip": "Minimum output value"}), + "maximum_value": ("FLOAT", {"default": 1.0, "tooltip": "Maximum output value"}), + "starting_value": ("FLOAT", {"default": 0.0, "tooltip": "Initial output value"}), + } + ) return inputs RETURN_TYPES = ("FLOAT", "MASK") FUNCTION = "process_motion" CATEGORY = "Realtime Nodes/control/motion" - def process_motion(self, image, roi_chain, threshold, blur_size, - minimum_value, maximum_value, starting_value, always_execute=True): + def process_motion(self, image, roi_chain, threshold, blur_size, minimum_value, maximum_value, starting_value, always_execute=True): # Get or initialize state with new caching fields - state = self.get_state({ - "prev_frame": None, - "prev_frame_blurred": None, # Cache blurred previous frame - "roi_states": {}, - "current_value": starting_value, - "last_cleanup": time.time() # Track last state cleanup - }) - + state = self.get_state( + { + "prev_frame": None, + "prev_frame_blurred": None, # Cache blurred previous frame + "roi_states": {}, + "current_value": starting_value, + "last_cleanup": time.time(), # Track last state cleanup + } + ) + # Convert image tensor to numpy array and prepare current frame current_frame = (image[0] * 255).cpu().numpy().astype(np.uint8) current_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) current_blurred = cv2.GaussianBlur(current_gray, (blur_size, blur_size), 0) - + # Initialize motion mask motion_mask = np.zeros_like(current_gray, dtype=np.float32) - + if state["prev_frame"] is not None: # Collect active ROIs for state cleanup active_rois = set() - + # Process each ROI in the chain current_roi = roi_chain while current_roi is not None: @@ -136,39 +131,36 @@ class MotionController(ControlNodeBase): action = current_roi["action"] roi_id = str(bounds) active_rois.add(roi_id) - + # Extract ROI regions once y_min, x_min, y_max, x_max = bounds - roi_region = slice(y_min, y_max+1), slice(x_min, x_max+1) - + roi_region = slice(y_min, y_max + 1), slice(x_min, x_max + 1) + # Process ROI using cached blurred frames roi_current = current_blurred[roi_region] roi_prev = state["prev_frame_blurred"][roi_region] roi_mask = mask[roi_region] - + # Detect motion in ROI diff = cv2.absdiff(roi_current, roi_prev) _, thresh = cv2.threshold(diff, threshold * 255, 255, cv2.THRESH_BINARY) thresh = thresh * (roi_mask > 0.5) - + # Update motion mask - motion_mask[roi_region] = np.maximum( - motion_mask[roi_region], - thresh / 255.0 - ) - + motion_mask[roi_region] = np.maximum(motion_mask[roi_region], thresh / 255.0) + # Calculate motion magnitude roi_area = np.sum(roi_mask > 0.5) motion_detected = np.sum(thresh) / (255 * roi_area) > threshold if roi_area > 0 else False - + # Get or initialize ROI state roi_state = state["roi_states"].get(roi_id, {"active": False, "count": 0}) value = current_roi["value"] - + # Process actions with optimized value updates if motion_detected and not roi_state["active"]: roi_state["active"] = True - + # Update current value based on action type if action == ROIAction.ADD.value: state["current_value"] = min(maximum_value, state["current_value"] + value) @@ -187,69 +179,59 @@ class MotionController(ControlNodeBase): state["current_value"] = maximum_value elif action == ROIAction.COUNTER.value: roi_state["count"] += 1 - state["current_value"] = minimum_value + ( - roi_state["count"] % (int((maximum_value - minimum_value) + 1)) - ) - + state["current_value"] = minimum_value + (roi_state["count"] % (int((maximum_value - minimum_value) + 1))) + elif action == ROIAction.MOMENTARY.value: state["current_value"] = maximum_value if motion_detected else minimum_value - + elif not motion_detected: roi_state["active"] = False if action == ROIAction.TRIGGER.value: state["current_value"] = minimum_value - + state["roi_states"][roi_id] = roi_state current_roi = current_roi["next"] - + # Cleanup inactive ROI states periodically (every 60 seconds) if time.time() - state["last_cleanup"] > 60: - state["roi_states"] = {k: v for k, v in state["roi_states"].items() - if k in active_rois} + state["roi_states"] = {k: v for k, v in state["roi_states"].items() if k in active_rois} state["last_cleanup"] = time.time() - + # Update state with current frame data state["prev_frame"] = current_gray state["prev_frame_blurred"] = current_blurred self.set_state(state) - + # Convert mask to tensor mask_tensor = torch.from_numpy(motion_mask).unsqueeze(0) - - return (state["current_value"], mask_tensor) + + return (state["current_value"], mask_tensor) def update(self, *args, **kwargs): - return self.process_motion(*args, **kwargs) + return self.process_motion(*args, **kwargs) + class RoundingMode(Enum): - ROUND = "round" # Standard rounding - FLOOR = "floor" # Round down - CEIL = "ceil" # Round up + ROUND = "round" # Standard rounding + FLOOR = "floor" # Round down + CEIL = "ceil" # Round up TRUNCATE = "truncate" # Remove decimal part + class IntegerMotionController(MotionController): """Processes motion detection for ROIs and outputs integer values""" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() # Convert float parameters to int - inputs["required"]["minimum_value"] = ("INT", { - "default": 0, - "tooltip": "Minimum output value" - }) - inputs["required"]["maximum_value"] = ("INT", { - "default": 100, - "tooltip": "Maximum output value" - }) - inputs["required"]["starting_value"] = ("INT", { - "default": 0, - "tooltip": "Initial output value" - }) + inputs["required"]["minimum_value"] = ("INT", {"default": 0, "tooltip": "Minimum output value"}) + inputs["required"]["maximum_value"] = ("INT", {"default": 100, "tooltip": "Maximum output value"}) + inputs["required"]["starting_value"] = ("INT", {"default": 0, "tooltip": "Initial output value"}) # Add rounding mode parameter inputs["required"]["rounding_mode"] = ( [mode.value for mode in RoundingMode], - {"default": RoundingMode.ROUND.value} + {"default": RoundingMode.ROUND.value}, ) return inputs @@ -268,20 +250,34 @@ class IntegerMotionController(MotionController): else: # TRUNCATE return int(value) - def process_motion(self, image, roi_chain, threshold, blur_size, - minimum_value, maximum_value, starting_value, - rounding_mode, always_execute=True): + def process_motion( + self, + image, + roi_chain, + threshold, + blur_size, + minimum_value, + maximum_value, + starting_value, + rounding_mode, + always_execute=True, + ): # Call parent class method with float conversion of bounds float_result, mask = super().process_motion( - image, roi_chain, threshold, blur_size, - float(minimum_value), float(maximum_value), float(starting_value), - always_execute + image, + roi_chain, + threshold, + blur_size, + float(minimum_value), + float(maximum_value), + float(starting_value), + always_execute, ) - + # Convert result to integer using specified rounding mode int_result = self._apply_rounding(float_result, rounding_mode) - + # Ensure result stays within bounds int_result = max(minimum_value, min(maximum_value, int_result)) - - return (int_result, mask) \ No newline at end of file + + return (int_result, mask) diff --git a/node_wrappers/realtimenodes/controls/sequence_controls.py b/node_wrappers/realtimenodes/controls/sequence_controls.py index 64f3c95..3fe22bf 100644 --- a/node_wrappers/realtimenodes/controls/sequence_controls.py +++ b/node_wrappers/realtimenodes/controls/sequence_controls.py @@ -3,43 +3,44 @@ from ....src.realtimenodes.control_base import ControlNodeBase class SequenceControlBase(ControlNodeBase): """Base class for sequence-based controls""" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "steps_per_item": ("INT", { - "default": 30, - "min": 1, - "max": 1000, - "step": 1, - "tooltip": "Number of frames to show each item" - }), - "sequence_mode": (["forward", "reverse", "pingpong", "random"], { - "default": "forward", - "tooltip": "How to move through the sequence" - }) - }) + inputs["required"].update( + { + "steps_per_item": ( + "INT", + {"default": 30, "min": 1, "max": 1000, "step": 1, "tooltip": "Number of frames to show each item"}, + ), + "sequence_mode": ( + ["forward", "reverse", "pingpong", "random"], + {"default": "forward", "tooltip": "How to move through the sequence"}, + ), + } + ) return inputs def update_sequence_base(self, values: list, steps_per_item: int, sequence_mode: str, always_execute: bool = True): # Initialize or get state - state = self.get_state({ - "current_index": 0, - "step_counter": 0, - "direction": 1 # 1 for forward, -1 for reverse (used in pingpong) - }) - + state = self.get_state( + { + "current_index": 0, + "step_counter": 0, + "direction": 1, # 1 for forward, -1 for reverse (used in pingpong) + } + ) + if not values: return (values[0] if values else None,) - + # Update step counter state["step_counter"] += 1 - + # Check if we should move to next item if state["step_counter"] >= steps_per_item: state["step_counter"] = 0 - + # Update index based on sequence mode if sequence_mode == "forward": state["current_index"] = (state["current_index"] + 1) % len(values) @@ -53,17 +54,19 @@ class SequenceControlBase(ControlNodeBase): state["current_index"] = next_index elif sequence_mode == "random": import random + state["current_index"] = random.randint(0, len(values) - 1) - + # Save state self.set_state(state) - + # Return current value return (values[state["current_index"]],) + class FloatSequence(SequenceControlBase): """Cycles through a sequence of float values""" - + RETURN_TYPES = ("FLOAT",) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/sequence" @@ -71,22 +74,18 @@ class FloatSequence(SequenceControlBase): @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "values": ("STRING", { - "default": "0.0, 0.5, 1.0", - "tooltip": "List of float values (comma separated)" - }) - }) + inputs["required"].update({"values": ("STRING", {"default": "0.0, 0.5, 1.0", "tooltip": "List of float values (comma separated)"})}) return inputs - + def update(self, values: str, steps_per_item: int, sequence_mode: str, always_execute: bool = True): # Parse comma-separated float values - float_values = [float(v.strip()) for v in values.split(',') if v.strip()] + float_values = [float(v.strip()) for v in values.split(",") if v.strip()] return self.update_sequence_base(float_values, steps_per_item, sequence_mode, always_execute) + class IntSequence(SequenceControlBase): """Cycles through a sequence of integer values""" - + RETURN_TYPES = ("INT",) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/sequence" @@ -94,22 +93,18 @@ class IntSequence(SequenceControlBase): @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "values": ("STRING", { - "default": "0, 5, 10", - "tooltip": "List of integer values (comma separated)" - }) - }) + inputs["required"].update({"values": ("STRING", {"default": "0, 5, 10", "tooltip": "List of integer values (comma separated)"})}) return inputs - + def update(self, values: str, steps_per_item: int, sequence_mode: str, always_execute: bool = True): # Parse comma-separated integer values - int_values = [int(v.strip()) for v in values.split(',') if v.strip()] + int_values = [int(v.strip()) for v in values.split(",") if v.strip()] return self.update_sequence_base(int_values, steps_per_item, sequence_mode, always_execute) + class StringSequence(SequenceControlBase): """Cycles through a sequence of strings""" - + RETURN_TYPES = ("STRING",) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/sequence" @@ -117,16 +112,17 @@ class StringSequence(SequenceControlBase): @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "values": ("STRING", { - "multiline": True, - "default": "first\nsecond\nthird", - "tooltip": "List of strings (one per line)" - }) - }) + inputs["required"].update( + { + "values": ( + "STRING", + {"multiline": True, "default": "first\nsecond\nthird", "tooltip": "List of strings (one per line)"}, + ) + } + ) return inputs - + def update(self, values: str, steps_per_item: int, sequence_mode: str, always_execute: bool = True): # Split into string values - string_values = [v.strip() for v in values.split('\n') if v.strip()] - return self.update_sequence_base(string_values, steps_per_item, sequence_mode, always_execute) \ No newline at end of file + string_values = [v.strip() for v in values.split("\n") if v.strip()] + return self.update_sequence_base(string_values, steps_per_item, sequence_mode, always_execute) diff --git a/node_wrappers/realtimenodes/controls/state_management_controls.py b/node_wrappers/realtimenodes/controls/state_management_controls.py index bad46d6..3ac2b67 100644 --- a/node_wrappers/realtimenodes/controls/state_management_controls.py +++ b/node_wrappers/realtimenodes/controls/state_management_controls.py @@ -6,18 +6,13 @@ from ....src.utils.general import AlwaysEqualProxy class StateResetNode(ControlNodeBase): """Node that resets all control node states when triggered""" - + OUTPUT_NODE = True - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "trigger": ("BOOLEAN", { - "default": False, - "tooltip": "Set to True to reset all states" - }) - }) + inputs["required"].update({"trigger": ("BOOLEAN", {"default": False, "tooltip": "Set to True to reset all states"})}) return inputs RETURN_TYPES = ("BOOLEAN",) @@ -30,21 +25,21 @@ class StateResetNode(ControlNodeBase): return (True,) return (False,) + class StateTestNode(ControlNodeBase): """Simple node that maintains a counter to test state management""" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "increment": ("INT", { - "default": 1, - "min": 1, - "max": 100, - "step": 1, - "tooltip": "Amount to increment counter by" - }) - }) + inputs["required"].update( + { + "increment": ( + "INT", + {"default": 1, "min": 1, "max": 100, "step": 1, "tooltip": "Amount to increment counter by"}, + ) + } + ) return inputs RETURN_TYPES = ("INT",) @@ -52,87 +47,112 @@ class StateTestNode(ControlNodeBase): CATEGORY = "Realtime Nodes/control/utility" def update(self, increment, always_execute=True): - state = self.get_state({ - "counter": 0 - }) - + state = self.get_state({"counter": 0}) + state["counter"] += increment self.set_state(state) - - return (state["counter"],) + return (state["counter"],) class GetStateNode(ControlNodeBase): """ Node that retrieves a value from the global state using a user-specified key. """ + CATEGORY = "utils" RETURN_TYPES = (AlwaysEqualProxy("*"),) RETURN_NAMES = ("value",) FUNCTION = "update" DESCRIPTION = "Retrieve a value from the global state using the given key. If the key is not found, return the default value." + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "key": ("STRING", {"default": "default_key", "tooltip": "The key to retrieve the value from. If not provided, the default value will be returned."}), - "default_value": (AlwaysEqualProxy("*"), {"tooltip": "The value to return if the key is not found."}), - "use_default": ("BOOLEAN", {"default": False, "tooltip": "If True, the default value will be returned if the key is not found."}), - }) + inputs["required"].update( + { + "key": ( + "STRING", + { + "default": "default_key", + "tooltip": "The key to retrieve the value from. If not provided, the default value will be returned.", + }, + ), + "default_value": (AlwaysEqualProxy("*"), {"tooltip": "The value to return if the key is not found."}), + "use_default": ( + "BOOLEAN", + { + "default": False, + "tooltip": "If True, the default value will be returned if the key is not found.", + }, + ), + } + ) return inputs - + def update(self, key: str, default_value, use_default: bool, always_execute=True): """ Retrieve a value from the global state using the given key. """ if not key or use_default: return (default_value,) - + # Get the shared state dictionary shared_state = self.state_manager.get_state("__shared_keys__", {}) - + # Check if the key exists if key in shared_state: return (shared_state[key],) - + # Return default value if key not found return (default_value,) - + + class SetStateNode(ControlNodeBase): """ Node that stores a value in the global state with a user-specified key. The value will be accessible in future workflow runs through GetStateNode. """ + CATEGORY = "utils" RETURN_TYPES = (AlwaysEqualProxy("*"),) RETURN_NAMES = ("value",) FUNCTION = "update" OUTPUT_NODE = True - DESCRIPTION = "Store a value in the global state with the given key. The value will be accessible in future workflow runs through GetStateNode." + DESCRIPTION = ( + "Store a value in the global state with the given key. The value will be accessible in future workflow runs through GetStateNode." + ) + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "key": ("STRING", {"default": "default_key", "tooltip": "The key to store the value under. If not provided, the value will not be stored."}), - "value": (AlwaysEqualProxy("*"), {"tooltip": "The value to store in the global state."}), - }) + inputs["required"].update( + { + "key": ( + "STRING", + { + "default": "default_key", + "tooltip": "The key to store the value under. If not provided, the value will not be stored.", + }, + ), + "value": (AlwaysEqualProxy("*"), {"tooltip": "The value to store in the global state."}), + } + ) return inputs - + def update(self, key: str, value, always_execute=True): """ Store a value in the global state with the given key. """ if not key: return (value,) - + try: shared_state = self.state_manager.get_state("__shared_keys__", {}) shared_state[key] = copy.deepcopy(value) self.state_manager.set_state("__shared_keys__", shared_state) except Exception as e: print(f"[State Node] Error storing value: {str(e)}") - - return (value,) + return (value,) diff --git a/node_wrappers/realtimenodes/controls/value_controls.py b/node_wrappers/realtimenodes/controls/value_controls.py index 90950c7..5800f5e 100644 --- a/node_wrappers/realtimenodes/controls/value_controls.py +++ b/node_wrappers/realtimenodes/controls/value_controls.py @@ -4,23 +4,28 @@ from ....src.realtimenodes.patterns.movement_patterns import MOVEMENT_PATTERNS class ValueControlBase(ControlNodeBase): """Base class for float and integer control nodes""" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "steps_per_cycle": ("INT", { - "default": 30, - "min": 1, - "max": 1000, - "step": 1, - "tooltip": "Number of steps to complete one full cycle" - }), - "movement_type": (list(MOVEMENT_PATTERNS.keys()), { - "default": "sine", - "tooltip": "Pattern of value changes over time" - }), - }) + inputs["required"].update( + { + "steps_per_cycle": ( + "INT", + { + "default": 30, + "min": 1, + "max": 1000, + "step": 1, + "tooltip": "Number of steps to complete one full cycle", + }, + ), + "movement_type": ( + list(MOVEMENT_PATTERNS.keys()), + {"default": "sine", "tooltip": "Pattern of value changes over time"}, + ), + } + ) return inputs def update(self, *args, **kwargs): @@ -28,11 +33,8 @@ class ValueControlBase(ControlNodeBase): return self.update_value_base(*args, **kwargs) def update_value_base(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True): - state = self.get_state({ - "current_value": None, - "phase": 0.0 - }) - + state = self.get_state({"current_value": None, "phase": 0.0}) + # Initialize if this is the first run if state["current_value"] is None: state["current_value"] = starting_value @@ -40,147 +42,175 @@ class ValueControlBase(ControlNodeBase): return (starting_value,) # Update phase - state["phase"] = (state["phase"] + 1/steps_per_cycle) % 1.0 - + state["phase"] = (state["phase"] + 1 / steps_per_cycle) % 1.0 + # Calculate new value using pattern pattern = MOVEMENT_PATTERNS[movement_type] new_value = pattern.calculate(state["phase"], minimum_value, maximum_value) - + # Update state state["current_value"] = new_value self.set_state(state) - + return (new_value,) + class FloatControl(ValueControlBase): - DESCRIPTION = "Generates a floating point value that changes over time according to various patterns." @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "maximum_value": ("FLOAT", { - "default": 1.0, - "min": -10000.0, - "max": 10000.0, - "step": 0.1, - "tooltip": "Maximum value that can be output" - }), - "minimum_value": ("FLOAT", { - "default": 0.0, - "min": -10000.0, - "max": 10000.0, - "step": 0.1, - "tooltip": "Minimum value that can be output" - }), - "starting_value": ("FLOAT", { - "default": 0.5, - "min": -10000.0, - "max": 10000.0, - "step": 0.1, - "tooltip": "Initial value when the node first executes" - }), - }) + inputs["required"].update( + { + "maximum_value": ( + "FLOAT", + { + "default": 1.0, + "min": -10000.0, + "max": 10000.0, + "step": 0.1, + "tooltip": "Maximum value that can be output", + }, + ), + "minimum_value": ( + "FLOAT", + { + "default": 0.0, + "min": -10000.0, + "max": 10000.0, + "step": 0.1, + "tooltip": "Minimum value that can be output", + }, + ), + "starting_value": ( + "FLOAT", + { + "default": 0.5, + "min": -10000.0, + "max": 10000.0, + "step": 0.1, + "tooltip": "Initial value when the node first executes", + }, + ), + } + ) return inputs RETURN_TYPES = ("FLOAT",) FUNCTION = "update_value" CATEGORY = "Realtime Nodes/control/value" - + def update_value(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True): return self.update_value_base(maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute) -class IntControl(ValueControlBase): - + +class IntControl(ValueControlBase): DESCRIPTION = "Generates an integer value that changes over time according to various patterns." @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "maximum_value": ("INT", { - "default": 100, - "min": -10000, - "max": 10000, - "step": 1, - "tooltip": "Maximum value that can be output" - }), - "minimum_value": ("INT", { - "default": 0, - "min": -10000, - "max": 10000, - "step": 1, - "tooltip": "Minimum value that can be output" - }), - "starting_value": ("INT", { - "default": 50, - "min": -10000, - "max": 10000, - "step": 1, - "tooltip": "Initial value when the node first executes" - }), - }) + inputs["required"].update( + { + "maximum_value": ( + "INT", + { + "default": 100, + "min": -10000, + "max": 10000, + "step": 1, + "tooltip": "Maximum value that can be output", + }, + ), + "minimum_value": ( + "INT", + { + "default": 0, + "min": -10000, + "max": 10000, + "step": 1, + "tooltip": "Minimum value that can be output", + }, + ), + "starting_value": ( + "INT", + { + "default": 50, + "min": -10000, + "max": 10000, + "step": 1, + "tooltip": "Initial value when the node first executes", + }, + ), + } + ) return inputs RETURN_TYPES = ("INT",) FUNCTION = "update_value" CATEGORY = "Realtime Nodes/control/value" - + def update_value(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True): result = self.update_value_base(maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute) return (int(round(result[0])),) + class StringControl(ControlNodeBase): - DESCRIPTION = "Generates string outputs that change over time according to various patterns." @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "strings": ("STRING", { - "multiline": True, - "default": "first string\nsecond string\nthird string", - "tooltip": "List of strings to cycle through (one per line)" - }), - "steps_per_cycle": ("INT", { - "default": 30, - "min": 1, - "max": 1000, - "step": 1, - "tooltip": "Number of steps to complete one full cycle" - }), - "movement_type": (list(MOVEMENT_PATTERNS.keys()), { - "default": "sine", - "tooltip": "Pattern of value changes over time" - }), - }) + inputs["required"].update( + { + "strings": ( + "STRING", + { + "multiline": True, + "default": "first string\nsecond string\nthird string", + "tooltip": "List of strings to cycle through (one per line)", + }, + ), + "steps_per_cycle": ( + "INT", + { + "default": 30, + "min": 1, + "max": 1000, + "step": 1, + "tooltip": "Number of steps to complete one full cycle", + }, + ), + "movement_type": ( + list(MOVEMENT_PATTERNS.keys()), + {"default": "sine", "tooltip": "Pattern of value changes over time"}, + ), + } + ) return inputs RETURN_TYPES = ("STRING",) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/value" - + def update(self, strings, steps_per_cycle, movement_type, always_execute=True): # Split the input strings into a list - string_list = [s.strip() for s in strings.split('\n') if s.strip()] + string_list = [s.strip() for s in strings.split("\n") if s.strip()] if not string_list: return ("",) - - state = self.get_state({ - "phase": 0.0 - }) - + + state = self.get_state({"phase": 0.0}) + # Update phase - state["phase"] = (state["phase"] + 1/steps_per_cycle) % 1.0 - + state["phase"] = (state["phase"] + 1 / steps_per_cycle) % 1.0 + # Calculate index using pattern pattern = MOVEMENT_PATTERNS[movement_type] index_float = pattern.calculate(state["phase"], 0, len(string_list) - 1) index = int(round(index_float)) % len(string_list) - + # Update state self.set_state(state) - - return (string_list[index],) \ No newline at end of file + + return (string_list[index],) diff --git a/node_wrappers/realtimenodes/logic/logic.py b/node_wrappers/realtimenodes/logic/logic.py index 2f56b42..778f319 100644 --- a/node_wrappers/realtimenodes/logic/logic.py +++ b/node_wrappers/realtimenodes/logic/logic.py @@ -3,30 +3,39 @@ from ....src.utils.general import AlwaysEqualProxy class LazyCondition(ControlNodeBase): - DESCRIPTION = "Uses lazy evaluation to truly skip execution of unused paths. Maintains state of the last value to circumvent feedback loops." - + DESCRIPTION = ( + "Uses lazy evaluation to truly skip execution of unused paths. Maintains state of the last value to circumvent feedback loops." + ) + @classmethod def INPUT_TYPES(s): return { "required": { - "condition": (AlwaysEqualProxy("*"), { - "tooltip": "When truthy (non-zero, non-empty, non-None), evaluates and returns if_true path. When falsy, returns either fallback or previous state of if_true.", - "forceInput": True, - }), - "if_true": (AlwaysEqualProxy("*"), { - "lazy": True, - "tooltip": "The path that should only be evaluated when condition is truthy" - }), - "fallback": (AlwaysEqualProxy("*"), { - "tooltip": "Alternative value to use when condition is falsy or no previous state of if_true" - }), - "use_fallback": ("BOOLEAN", { - "default": False, - "tooltip": "When False, uses last successful if_true result (if one exists). When True, uses fallback value" - }), + "condition": ( + AlwaysEqualProxy("*"), + { + "tooltip": "When truthy (non-zero, non-empty, non-None), evaluates and returns if_true path. When falsy, returns either fallback or previous state of if_true.", + "forceInput": True, + }, + ), + "if_true": ( + AlwaysEqualProxy("*"), + {"lazy": True, "tooltip": "The path that should only be evaluated when condition is truthy"}, + ), + "fallback": ( + AlwaysEqualProxy("*"), + {"tooltip": "Alternative value to use when condition is falsy or no previous state of if_true"}, + ), + "use_fallback": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When False, uses last successful if_true result (if one exists). When True, uses fallback value", + }, + ), } } - + RETURN_TYPES = (AlwaysEqualProxy("*"),) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/utility" @@ -40,14 +49,12 @@ class LazyCondition(ControlNodeBase): def update(self, condition, if_true, fallback, use_fallback): """Route to either if_true output or fallback value, maintaining state of last if_true.""" - state = self.get_state({ - "prev_output": None - }) + state = self.get_state({"prev_output": None}) if condition: # Let Python handle truthiness # Update last state when we run if_true path state["prev_output"] = if_true if if_true is not None else fallback - if hasattr(if_true, 'detach'): + if hasattr(if_true, "detach"): state["prev_output"] = if_true.detach().clone() self.set_state(state) return (if_true,) @@ -57,32 +64,42 @@ class LazyCondition(ControlNodeBase): else: return (state["prev_output"],) + class LogicOperator(ControlNodeBase): DESCRIPTION = "Performs logical operations (AND, OR, NOT, XOR) on inputs based on their truthiness" - + @classmethod def INPUT_TYPES(s): return { "required": { - "operation": (["AND", "OR", "NOT", "XOR"], { - "default": "AND", - "tooltip": "Logical operation to perform" - }), - "input_a": (AlwaysEqualProxy("*"), { - "tooltip": "First input to evaluate for truthiness", - "forceInput": True, - }), - "always_execute": ("BOOLEAN", { - "default": True, - }), + "operation": ( + ["AND", "OR", "NOT", "XOR"], + {"default": "AND", "tooltip": "Logical operation to perform"}, + ), + "input_a": ( + AlwaysEqualProxy("*"), + { + "tooltip": "First input to evaluate for truthiness", + "forceInput": True, + }, + ), + "always_execute": ( + "BOOLEAN", + { + "default": True, + }, + ), }, "optional": { - "input_b": (AlwaysEqualProxy("*"), { - "tooltip": "Second input to evaluate for truthiness (not used for NOT operation)", - }), - } + "input_b": ( + AlwaysEqualProxy("*"), + { + "tooltip": "Second input to evaluate for truthiness (not used for NOT operation)", + }, + ), + }, } - + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "update" CATEGORY = "Realtime Nodes/control/logic" @@ -90,19 +107,19 @@ class LogicOperator(ControlNodeBase): def update(self, operation, input_a, always_execute=True, input_b=None): """Perform the selected logical operation on inputs based on their truthiness.""" a = bool(input_a) - + if operation == "NOT": return (not a,) - + # For all other operations, we need input_b b = bool(input_b) - + if operation == "AND": return (a and b,) elif operation == "OR": return (a or b,) elif operation == "XOR": return (a != b,) - + # Should never get here, but just in case - return (False,) \ No newline at end of file + return (False,) diff --git a/node_wrappers/realtimenodes/media/__init__.py b/node_wrappers/realtimenodes/media/__init__.py index da22290..6e7633b 100644 --- a/node_wrappers/realtimenodes/media/__init__.py +++ b/node_wrappers/realtimenodes/media/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/node_wrappers/realtimenodes/media/media_pipe_nodes.py b/node_wrappers/realtimenodes/media/media_pipe_nodes.py index a860d75..b3d68aa 100644 --- a/node_wrappers/realtimenodes/media/media_pipe_nodes.py +++ b/node_wrappers/realtimenodes/media/media_pipe_nodes.py @@ -10,66 +10,51 @@ class HandTrackingNode: return { "required": { "image": ("IMAGE",), - "detection_confidence": ("FLOAT", { - "default": 0.5, - "min": 0.1, - "max": 1.0, - "step": 0.05 - }), - "tracking_confidence": ("FLOAT", { - "default": 0.5, - "min": 0.1, - "max": 1.0, - "step": 0.05 - }), + "detection_confidence": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}), + "tracking_confidence": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}), }, "optional": { - "max_num_hands": ("INT", { - "default": 2, - "min": 1, - "max": 4, - "step": 1 - }), - "static_image_mode": ("BOOLEAN", { - "default": False - }), - "draw_debug": ("BOOLEAN", { - "default": False - }) - } + "max_num_hands": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1}), + "static_image_mode": ("BOOLEAN", {"default": False}), + "draw_debug": ("BOOLEAN", {"default": False}), + }, } RETURN_TYPES = ("IMAGE", "HAND_DATA") RETURN_NAMES = ("debug_image", "hand_data") FUNCTION = "track_hands" CATEGORY = "image/processing" - DESCRIPTION = "(((EXPERIMENTAL))) Track hands in an image using MediaPipe. Returns a debug image with landmarks and a dictionary of hand data." + DESCRIPTION = ( + "(((EXPERIMENTAL))) Track hands in an image using MediaPipe. Returns a debug image with landmarks and a dictionary of hand data." + ) + def __init__(self): self.last_detection_conf = 0.5 self.last_tracking_conf = 0.5 self.last_max_hands = 2 self.last_static_mode = True - + self.mp_hands = mp.solutions.hands self.hands = self.mp_hands.Hands( static_image_mode=self.last_static_mode, max_num_hands=self.last_max_hands, min_detection_confidence=self.last_detection_conf, - min_tracking_confidence=self.last_tracking_conf + min_tracking_confidence=self.last_tracking_conf, ) def track_hands(self, image, detection_confidence, tracking_confidence, max_num_hands=2, static_image_mode=True, draw_debug=True): # Check if we need to reinitialize MediaPipe hands - if (detection_confidence != self.last_detection_conf or - tracking_confidence != self.last_tracking_conf or - max_num_hands != self.last_max_hands or - static_image_mode != self.last_static_mode): - + if ( + detection_confidence != self.last_detection_conf + or tracking_confidence != self.last_tracking_conf + or max_num_hands != self.last_max_hands + or static_image_mode != self.last_static_mode + ): self.hands = self.mp_hands.Hands( static_image_mode=static_image_mode, max_num_hands=max_num_hands, min_detection_confidence=detection_confidence, - min_tracking_confidence=tracking_confidence + min_tracking_confidence=tracking_confidence, ) self.last_detection_conf = detection_confidence self.last_tracking_conf = tracking_confidence @@ -84,61 +69,53 @@ class HandTrackingNode: # Process each image in batch processed_images = [] hand_data_batch = [] - + for i in range(batch_size): single_image = image_np[i] image_rgb = cv2.cvtColor(single_image, cv2.COLOR_BGR2RGB) results = self.hands.process(image_rgb) - - hand_data = { - "left_hand": None, - "right_hand": None, - "num_hands": 0, - "hands_present": False - } - + + hand_data = {"left_hand": None, "right_hand": None, "num_hands": 0, "hands_present": False} + if results.multi_hand_landmarks: hand_data["num_hands"] = len(results.multi_hand_landmarks) hand_data["hands_present"] = True - + # Create single debug image for all hands debug_image = single_image.copy() if draw_debug else single_image - + # Process all detected hands for hand_idx, (hand_landmarks, handedness) in enumerate(zip(results.multi_hand_landmarks, results.multi_handedness)): hand_type = handedness.classification[0].label.lower() landmarks = np.array([[lm.x, lm.y, lm.z] for lm in hand_landmarks.landmark]) - + # Store hand data based on type if hand_type == "left" and hand_data["left_hand"] is None: hand_data["left_hand"] = landmarks elif hand_type == "right" and hand_data["right_hand"] is None: hand_data["right_hand"] = landmarks - + # Draw landmarks for this hand on the shared debug image if draw_debug: - mp.solutions.drawing_utils.draw_landmarks( - debug_image, - hand_landmarks, - self.mp_hands.HAND_CONNECTIONS - ) - + mp.solutions.drawing_utils.draw_landmarks(debug_image, hand_landmarks, self.mp_hands.HAND_CONNECTIONS) + processed_images.append(debug_image) else: processed_images.append(single_image) - + hand_data_batch.append(hand_data) # Convert back to tensors processed_batch = np.stack(processed_images, axis=0) output_tensor = torch.from_numpy(processed_batch).float() / 255.0 - + return (output_tensor, hand_data_batch) def __del__(self): - if hasattr(self, 'hands'): + if hasattr(self, "hands"): self.hands.close() + class HandMaskNode: @classmethod def INPUT_TYPES(cls): @@ -157,57 +134,57 @@ class HandMaskNode: def create_mask(self, image, hand_data, mask_type): batch_size, height, width, channels = image.shape masks = np.zeros((batch_size, height, width), dtype=np.float32) - + # Finger indices in MediaPipe hand landmarks finger_indices = { - "thumb": [1,2,3,4], - "index": [5,6,7,8], - "middle": [9,10,11,12], - "ring": [13,14,15,16], - "pinky": [17,18,19,20] + "thumb": [1, 2, 3, 4], + "index": [5, 6, 7, 8], + "middle": [9, 10, 11, 12], + "ring": [13, 14, 15, 16], + "pinky": [17, 18, 19, 20], } - palm_indices = [0,1,5,9,13,17] # Wrist and base of each finger - + palm_indices = [0, 1, 5, 9, 13, 17] # Wrist and base of each finger + for i in range(batch_size): frame_hand_data = hand_data[i] if not frame_hand_data["hands_present"]: continue - + frame_mask = np.zeros((height, width), dtype=np.float32) - + for hand_key in ["right_hand", "left_hand"]: hand = frame_hand_data[hand_key] if hand is None: continue - + points = hand[:, :2] # Take only x,y points = (points * np.array([width, height])).astype(np.int32) - + if mask_type == "full_hand": hull = cv2.convexHull(points) cv2.fillConvexPoly(frame_mask, hull, 1.0) - + elif mask_type == "palm_only": palm_points = points[palm_indices] hull = cv2.convexHull(palm_points) cv2.fillConvexPoly(frame_mask, hull, 1.0) - + elif mask_type == "fingers_only": for finger in finger_indices.values(): finger_points = points[finger] cv2.fillConvexPoly(frame_mask, cv2.convexHull(finger_points), 1.0) - + elif mask_type in finger_indices: finger_points = points[finger_indices[mask_type]] cv2.fillConvexPoly(frame_mask, cv2.convexHull(finger_points), 1.0) - + masks[i] = frame_mask - + # # Debug output # if frame_hand_data["right_hand"] is not None: # print(f"Right hand points range: x({np.min(points[:, 0]):.1f}-{np.max(points[:, 0]):.1f}), y({np.min(points[:, 1]):.1f}-{np.max(points[:, 1]):.1f})") # if frame_hand_data["left_hand"] is not None: # print(f"Left hand points range: x({np.min(points[:, 0]):.1f}-{np.max(points[:, 0]):.1f}), y({np.min(points[:, 1]):.1f}-{np.max(points[:, 1]):.1f})") - + mask_tensor = torch.from_numpy(masks) - return (image, mask_tensor) \ No newline at end of file + return (image, mask_tensor) diff --git a/node_wrappers/realtimenodes/utils/__init__.py b/node_wrappers/realtimenodes/utils/__init__.py index da22290..6e7633b 100644 --- a/node_wrappers/realtimenodes/utils/__init__.py +++ b/node_wrappers/realtimenodes/utils/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/node_wrappers/realtimenodes/utils/utility_controls.py b/node_wrappers/realtimenodes/utils/utility_controls.py index 6bab29a..433cecb 100644 --- a/node_wrappers/realtimenodes/utils/utility_controls.py +++ b/node_wrappers/realtimenodes/utils/utility_controls.py @@ -12,152 +12,141 @@ from ....src.utils.similar_image_filter import SimilarImageFilter class FPSMonitor(ControlNodeBase): """Generates an FPS overlay as an image and mask""" - + RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "update" CATEGORY = "Realtime Nodes/control/utility" - + @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() - inputs["required"].update({ - "width": ("INT", { - "default": 512, - "min": 64, - "max": 4096, - "step": 1 - }), - "height": ("INT", { - "default": 512, - "min": 64, - "max": 4096, - "step": 1 - }), - "text_color": ("INT", { - "default": 255, - "min": 0, - "max": 255, - "step": 1 - }), - "text_size": ("FLOAT", { - "default": 1.0, - "min": 0.1, - "max": 5.0, - "step": 0.1 - }), - "window_size": ("INT", { - "default": 60, - "min": 1, - "max": 300, - "step": 1, - "tooltip": "Number of frames to average over" - }) - }) + inputs["required"].update( + { + "width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 1}), + "text_color": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}), + "text_size": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}), + "window_size": ( + "INT", + {"default": 60, "min": 1, "max": 300, "step": 1, "tooltip": "Number of frames to average over"}, + ), + } + ) return inputs def update(self, width: int, height: int, text_color: int, text_size: float, window_size: int, always_execute=True): current_time = time.time() - state = self.get_state({ - "last_time": current_time, - "frame_times": [], - "last_fps": 0, - "last_avg_fps": 0, - "cached_image": None, - "cached_mask": None, - "last_dims": (0, 0), - "last_color": 0, - "last_size": 0 - }) - + state = self.get_state( + { + "last_time": current_time, + "frame_times": [], + "last_fps": 0, + "last_avg_fps": 0, + "cached_image": None, + "cached_mask": None, + "last_dims": (0, 0), + "last_color": 0, + "last_size": 0, + } + ) + # Calculate FPS delta_time = current_time - state["last_time"] current_fps = 1.0 / delta_time if delta_time > 0 else 0 - + # Update rolling window state["frame_times"].append(delta_time) if len(state["frame_times"]) > window_size: state["frame_times"].pop(0) - + # Calculate average FPS avg_delta = sum(state["frame_times"]) / len(state["frame_times"]) average_fps = 1.0 / avg_delta if avg_delta > 0 else 0 - + # Check if we need to redraw needs_redraw = ( - abs(current_fps - state["last_fps"]) > 1.0 or # FPS changed significantly - abs(average_fps - state["last_avg_fps"]) > 0.5 or # Avg FPS changed significantly - state["last_dims"] != (width, height) or # Dimensions changed - state["last_color"] != text_color or # Color changed - state["last_size"] != text_size or # Text size changed - state["cached_image"] is None # No cache exists + abs(current_fps - state["last_fps"]) > 1.0 # FPS changed significantly + or abs(average_fps - state["last_avg_fps"]) > 0.5 # Avg FPS changed significantly + or state["last_dims"] != (width, height) # Dimensions changed + or state["last_color"] != text_color # Color changed + or state["last_size"] != text_size # Text size changed + or state["cached_image"] is None # No cache exists ) - + if needs_redraw: # Create text fps_text = f"FPS: {current_fps:.1f} (avg: {average_fps:.1f})" - + # Create image and mask using numpy/cv2 image = np.zeros((height, width, 3), dtype=np.uint8) mask = np.zeros((height, width), dtype=np.uint8) - + # Position text in top-left with padding padding = 10 pos = (padding, padding + 20) # Add vertical offset for cv2's text positioning - + # Draw text - much simpler with cv2! cv2.putText(image, fps_text, pos, cv2.FONT_HERSHEY_SIMPLEX, text_size, (text_color, text_color, text_color), 1) cv2.putText(mask, fps_text, pos, cv2.FONT_HERSHEY_SIMPLEX, text_size, 255, 1) - + # Convert to float32 and normalize image = image.astype(np.float32) / 255.0 mask = mask.astype(np.float32) / 255.0 - + # Convert to torch tensors with correct shapes state["cached_image"] = torch.from_numpy(image[None, ...]) # BHWC - state["cached_mask"] = torch.from_numpy(mask[None, ...]) # BHW - + state["cached_mask"] = torch.from_numpy(mask[None, ...]) # BHW + # Update cache parameters state["last_fps"] = current_fps state["last_avg_fps"] = average_fps state["last_dims"] = (width, height) state["last_color"] = text_color state["last_size"] = text_size - + state["last_time"] = current_time self.set_state(state) - + return (state["cached_image"], state["cached_mask"]) + class SimilarityFilter(ControlNodeBase): DESCRIPTION = "Filters out similar consecutive images and outputs a signal to control downstream execution." - + @classmethod def INPUT_TYPES(s): return { "required": { - "image": ("IMAGE", { - "tooltip": "Input image to compare with previous frame" - }), - "always_execute": ("BOOLEAN", { - "default": False, - }), - "threshold": ("FLOAT", { - "default": 0.98, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Similarity threshold (0-1). Higher values mean more frames are considered similar" - }), - "max_skip_frames": ("INT", { - "default": 10, - "min": 1, - "max": 100, - "step": 1, - "tooltip": "Maximum number of consecutive frames to skip before forcing execution" - }), + "image": ("IMAGE", {"tooltip": "Input image to compare with previous frame"}), + "always_execute": ( + "BOOLEAN", + { + "default": False, + }, + ), + "threshold": ( + "FLOAT", + { + "default": 0.98, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Similarity threshold (0-1). Higher values mean more frames are considered similar", + }, + ), + "max_skip_frames": ( + "INT", + { + "default": 10, + "min": 1, + "max": 100, + "step": 1, + "tooltip": "Maximum number of consecutive frames to skip before forcing execution", + }, + ), } } - + RETURN_TYPES = ("IMAGE", "BOOLEAN") RETURN_NAMES = ("image", "should_execute") FUNCTION = "update" @@ -169,12 +158,9 @@ class SimilarityFilter(ControlNodeBase): def update(self, image, threshold=0.98, max_skip_frames=10, always_execute=False): print(f"[DEBUG] Input image object: {hex(id(image))}, shape: {image.shape}, device: {image.device}") - + # Get state with defaults - state = self.get_state({ - "prev_image": None, - "skip_count": 0 - }) + state = self.get_state({"prev_image": None, "skip_count": 0}) # First frame case if state["prev_image"] is None: @@ -186,7 +172,7 @@ class SimilarityFilter(ControlNodeBase): # Update filter parameters self._similarity_filter.set_threshold(threshold) self._similarity_filter.set_max_skip_frame(max_skip_frames) - + # Use filter to check similarity result = self._similarity_filter(image) should_execute = result is not None @@ -199,7 +185,7 @@ class SimilarityFilter(ControlNodeBase): state["skip_count"] = 0 self.set_state(state) return (image, True) # Force execution after max skips - + # Skip frame - return previous image and False for execution state["skip_count"] += 1 self.set_state(state) @@ -210,12 +196,3 @@ class SimilarityFilter(ControlNodeBase): state["skip_count"] = 0 self.set_state(state) return (image, True) - - - - - - - - - diff --git a/node_wrappers/utils/misc_nodes.py b/node_wrappers/utils/misc_nodes.py index c9a60a3..153819d 100644 --- a/node_wrappers/utils/misc_nodes.py +++ b/node_wrappers/utils/misc_nodes.py @@ -13,9 +13,10 @@ from ...src.utils.general import AlwaysEqualProxy MAX_RESOLUTION = nodes.MAX_RESOLUTION # Get the same max resolution as core nodes + class DTypeConverter: """Converts masks to specified data types""" - + @classmethod def INPUT_TYPES(cls): return { @@ -29,16 +30,11 @@ class DTypeConverter: RETURN_TYPES = ("MASK",) # Return only MASK type FUNCTION = "convert_dtype" - DTYPE_MAP = { - "uint8": torch.uint8, - "float16": torch.float16, - "float32": torch.float32, - "float64": torch.float64 - } + DTYPE_MAP = {"uint8": torch.uint8, "float16": torch.float16, "float32": torch.float32, "float64": torch.float64} def convert_dtype(self, mask, dtype): target_dtype = self.DTYPE_MAP[dtype] - + if target_dtype == torch.uint8: if mask.is_floating_point(): converted = (mask * 255.0).round().to(torch.uint8) @@ -49,8 +45,9 @@ class DTypeConverter: converted = (mask > 0).to(target_dtype) else: converted = mask.to(target_dtype) - - return (converted,) + + return (converted,) + class FastWebcamCapture: @classmethod @@ -63,55 +60,127 @@ class FastWebcamCapture: "capture_on_queue": ("BOOLEAN", {"default": True}), } } - + RETURN_TYPES = ("IMAGE",) FUNCTION = "process_capture" - + CATEGORY = "image" def process_capture(self, image, width, height, capture_on_queue): # Check if we got a data URL - if isinstance(image, str) and image.startswith('data:image/'): + if isinstance(image, str) and image.startswith("data:image/"): # Extract the base64 data after the comma - base64_data = re.sub('^data:image/.+;base64,', '', image) - + base64_data = re.sub("^data:image/.+;base64,", "", image) + # Convert base64 to PIL Image buffer = BytesIO(base64.b64decode(base64_data)) pil_image = Image.open(buffer).convert("RGB") - + # Convert PIL to numpy array image = np.array(pil_image) - + # Handle resize if requested if width > 0 and height > 0: import cv2 + image = cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA) - + # Convert to float32 and normalize to 0-1 range image = image.astype(np.float32) / 255.0 - + # Add batch dimension and convert to torch tensor # ComfyUI expects BHWC format - image = torch.from_numpy(image)[None,...] - + image = torch.from_numpy(image)[None, ...] + return (image,) else: raise ValueError("Invalid image format received from webcam") + coco_classes = [ - 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', - 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', - 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', - 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', - 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', - 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', - 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', - 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', - 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', - 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', - 'toothbrush' + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", ] + class YOLOSimilarityCompare: @classmethod def INPUT_TYPES(s): @@ -119,58 +188,80 @@ class YOLOSimilarityCompare: "required": { "ULTRALYTICS_RESULTS1": ("ULTRALYTICS_RESULTS",), "ULTRALYTICS_RESULTS2": ("ULTRALYTICS_RESULTS",), - "class_weight": ("FLOAT", { - "default": 0.3, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Weight for class similarity - how much to consider matching object types" - }), - "spatial_weight": ("FLOAT", { - "default": 0.2, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Weight for spatial similarity - how much to consider object positions" - }), - "confidence_weight": ("FLOAT", { - "default": 0.2, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Weight for confidence similarity - how much to consider detection confidence" - }), - "size_weight": ("FLOAT", { - "default": 0.15, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Weight for size similarity - how much to consider matching object sizes" - }), - "relationship_weight": ("FLOAT", { - "default": 0.15, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Weight for relationship similarity - how much to consider distances between objects" - }), - "threshold": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "display": "slider", - "tooltip": "Threshold for similarity - scores above this value will return True" - }), + "class_weight": ( + "FLOAT", + { + "default": 0.3, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Weight for class similarity - how much to consider matching object types", + }, + ), + "spatial_weight": ( + "FLOAT", + { + "default": 0.2, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Weight for spatial similarity - how much to consider object positions", + }, + ), + "confidence_weight": ( + "FLOAT", + { + "default": 0.2, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Weight for confidence similarity - how much to consider detection confidence", + }, + ), + "size_weight": ( + "FLOAT", + { + "default": 0.15, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Weight for size similarity - how much to consider matching object sizes", + }, + ), + "relationship_weight": ( + "FLOAT", + { + "default": 0.15, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Weight for relationship similarity - how much to consider distances between objects", + }, + ), + "threshold": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Threshold for similarity - scores above this value will return True", + }, + ), } } - - RETURN_TYPES = ("FLOAT", "BOOLEAN", "STRING",) + + RETURN_TYPES = ( + "FLOAT", + "BOOLEAN", + "STRING", + ) RETURN_NAMES = ("similarity_score", "above_threshold", "explanation") FUNCTION = "compare_images" CATEGORY = "ScavengerHunt" @@ -178,25 +269,25 @@ class YOLOSimilarityCompare: def compute_size_similarity(self, boxes1, boxes2): if len(boxes1) == 0 or len(boxes2) == 0: return 0.0 - + # Compute areas areas1 = (boxes1[:, 2] * boxes1[:, 3]).cpu().numpy() # width * height areas2 = (boxes2[:, 2] * boxes2[:, 3]).cpu().numpy() - + # Normalize areas by image size norm_areas1 = areas1 / np.max(areas1) if len(areas1) > 0 else areas1 norm_areas2 = areas2 / np.max(areas2) if len(areas2) > 0 else areas2 - + # Compare average sizes avg_size1 = np.mean(norm_areas1) avg_size2 = np.mean(norm_areas2) - + return 1.0 - abs(avg_size1 - avg_size2) def compute_relationship_similarity(self, boxes1, boxes2): if len(boxes1) < 2 or len(boxes2) < 2: return 0.0 - + def compute_pairwise_distances(boxes): centers = boxes[:, :2] # Get centers (x, y) distances = [] @@ -205,82 +296,96 @@ class YOLOSimilarityCompare: dist = np.linalg.norm(centers[i] - centers[j]) distances.append(dist) return np.array(distances) - + # Compute normalized pairwise distances dist1 = compute_pairwise_distances(boxes1.cpu().numpy()) dist2 = compute_pairwise_distances(boxes2.cpu().numpy()) - + # Normalize distances if len(dist1) > 0: dist1 = dist1 / np.max(dist1) if len(dist2) > 0: dist2 = dist2 / np.max(dist2) - + # Compare average normalized distances avg_dist1 = np.mean(dist1) if len(dist1) > 0 else 0 avg_dist2 = np.mean(dist2) if len(dist2) > 0 else 0 - + return 1.0 - abs(avg_dist1 - avg_dist2) - def compare_images(self, ULTRALYTICS_RESULTS1, ULTRALYTICS_RESULTS2, class_weight=0.3, - spatial_weight=0.2, confidence_weight=0.2, size_weight=0.15, - relationship_weight=0.15, threshold=0.5): + def compare_images( + self, + ULTRALYTICS_RESULTS1, + ULTRALYTICS_RESULTS2, + class_weight=0.3, + spatial_weight=0.2, + confidence_weight=0.2, + size_weight=0.15, + relationship_weight=0.15, + threshold=0.5, + ): # Get detections from first image boxes1 = ULTRALYTICS_RESULTS1[0].boxes classes1 = boxes1.cls.cpu().numpy() conf1 = boxes1.conf.cpu().numpy() xyxy1 = boxes1.xyxy.cpu().numpy() - + # Get detections from second image boxes2 = ULTRALYTICS_RESULTS2[0].boxes classes2 = boxes2.cls.cpu().numpy() conf2 = boxes2.conf.cpu().numpy() xyxy2 = boxes2.xyxy.cpu().numpy() - + # 1. Class similarity - what objects appear in both images unique_classes1 = set(classes1) unique_classes2 = set(classes2) - class_overlap = len(unique_classes1.intersection(unique_classes2)) / len(unique_classes1.union(unique_classes2)) if unique_classes1 or unique_classes2 else 0 - + class_overlap = ( + len(unique_classes1.intersection(unique_classes2)) / len(unique_classes1.union(unique_classes2)) + if unique_classes1 or unique_classes2 + else 0 + ) + # 2. Spatial similarity - compare object locations spatial_sim = 0 if len(xyxy1) > 0 and len(xyxy2) > 0: # Normalize coordinates to 0-1 range img_size1 = ULTRALYTICS_RESULTS1[0].boxes.orig_shape img_size2 = ULTRALYTICS_RESULTS2[0].boxes.orig_shape - + norm_boxes1 = xyxy1.copy() - norm_boxes1[:,[0,2]] /= img_size1[1] # normalize x by width - norm_boxes1[:,[1,3]] /= img_size1[0] # normalize y by height - + norm_boxes1[:, [0, 2]] /= img_size1[1] # normalize x by width + norm_boxes1[:, [1, 3]] /= img_size1[0] # normalize y by height + norm_boxes2 = xyxy2.copy() - norm_boxes2[:,[0,2]] /= img_size2[1] - norm_boxes2[:,[1,3]] /= img_size2[0] - + norm_boxes2[:, [0, 2]] /= img_size2[1] + norm_boxes2[:, [1, 3]] /= img_size2[0] + # Compare average box positions center1 = np.mean(norm_boxes1, axis=0) center2 = np.mean(norm_boxes2, axis=0) spatial_sim = 1 - np.mean(np.abs(center1 - center2)) - + # 3. Confidence similarity - compare detection confidences conf_sim = 1 - abs(np.mean(conf1) - np.mean(conf2)) if len(conf1) > 0 and len(conf2) > 0 else 0 - + # 4. Size similarity - compare object sizes size_sim = self.compute_size_similarity(boxes1.xywh, boxes2.xywh) - + # 5. Relationship similarity - compare distances between objects relationship_sim = self.compute_relationship_similarity(boxes1.xywh, boxes2.xywh) - + # Weighted combination - similarity = (class_weight * class_overlap + - spatial_weight * spatial_sim + - confidence_weight * conf_sim + - size_weight * size_sim + - relationship_weight * relationship_sim) - + similarity = ( + class_weight * class_overlap + + spatial_weight * spatial_sim + + confidence_weight * conf_sim + + size_weight * size_sim + + relationship_weight * relationship_sim + ) + # Check if above threshold above_threshold = similarity >= threshold - + # Generate explanation explanation = f"""Similarity Score: {similarity:.3f} (Threshold: {threshold}) Above Threshold: {above_threshold} @@ -293,118 +398,112 @@ Relationship Similarity: {relationship_sim:.3f} (weight: {relationship_weight}) Detected Classes in Image 1: {[coco_classes[int(c)] for c in unique_classes1]} Detected Classes in Image 2: {[coco_classes[int(c)] for c in unique_classes2]}""" - - return (float(similarity), above_threshold, explanation,) + + return ( + float(similarity), + above_threshold, + explanation, + ) + class QuickShapeMask: """A node that quickly generates basic shape masks""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "shape": (["circle", "square"], { - "default": "circle", - "tooltip": "The shape of the mask to generate" - }), - "width": ("INT", { - "default": 64, - "min": 1, - "max": 4096, - "step": 1, - "tooltip": "Width of the shape in pixels" - }), - "height": ("INT", { - "default": 64, - "min": 1, - "max": 4096, - "step": 1, - "tooltip": "Height of the shape in pixels" - }), - "x": ("INT", { - "default": 256, - "min": 0, - "max": 4096, - "step": 1, - "tooltip": "X position of the shape center (0 = left edge)" - }), - "y": ("INT", { - "default": 256, - "min": 0, - "max": 4096, - "step": 1, - "tooltip": "Y position of the shape center (0 = top edge)" - }), - "canvas_width": ("INT", { - "default": 512, - "min": 1, - "max": 4096, - "step": 1, - "tooltip": "Width of the output mask" - }), - "canvas_height": ("INT", { - "default": 512, - "min": 1, - "max": 4096, - "step": 1, - "tooltip": "Height of the output mask" - }), - "batch_size": ("INT", { - "default": 1, - "min": 1, - "max": 64, - "step": 1, - "tooltip": "Number of identical masks to generate" - }) + "shape": (["circle", "square"], {"default": "circle", "tooltip": "The shape of the mask to generate"}), + "width": ( + "INT", + {"default": 64, "min": 1, "max": 4096, "step": 1, "tooltip": "Width of the shape in pixels"}, + ), + "height": ( + "INT", + {"default": 64, "min": 1, "max": 4096, "step": 1, "tooltip": "Height of the shape in pixels"}, + ), + "x": ( + "INT", + { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "X position of the shape center (0 = left edge)", + }, + ), + "y": ( + "INT", + { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "Y position of the shape center (0 = top edge)", + }, + ), + "canvas_width": ( + "INT", + {"default": 512, "min": 1, "max": 4096, "step": 1, "tooltip": "Width of the output mask"}, + ), + "canvas_height": ( + "INT", + {"default": 512, "min": 1, "max": 4096, "step": 1, "tooltip": "Height of the output mask"}, + ), + "batch_size": ( + "INT", + {"default": 1, "min": 1, "max": 64, "step": 1, "tooltip": "Number of identical masks to generate"}, + ), } } RETURN_TYPES = ("MASK",) FUNCTION = "generate_mask" CATEGORY = "mask" - + DESCRIPTION = "Generates a mask containing a basic shape (circle or square) with high performance" def generate_mask(self, shape, width, height, x, y, canvas_width, canvas_height, batch_size): # Create empty mask mask = np.zeros((canvas_height, canvas_width), dtype=np.float32) - + # Calculate boundaries half_width = width // 2 half_height = height // 2 - + # Calculate shape boundaries left = max(0, x - half_width) right = min(canvas_width, x + half_width) top = max(0, y - half_height) bottom = min(canvas_height, y + half_height) - + if shape == "square": # Simple square mask mask[top:bottom, left:right] = 1.0 - + else: # circle # Create coordinate grids for the region of interest Y, X = np.ogrid[top:bottom, left:right] - + # Calculate distances from center for the region - dist_x = (X - x) - dist_y = (Y - y) - + dist_x = X - x + dist_y = Y - y + # Create circle mask using distance formula - circle_mask = (dist_x**2 / (width/2)**2 + dist_y**2 / (height/2)**2) <= 1 - + circle_mask = (dist_x**2 / (width / 2) ** 2 + dist_y**2 / (height / 2) ** 2) <= 1 + # Apply circle to the region mask[top:bottom, left:right][circle_mask] = 1.0 - + # Convert to torch tensor and add batch dimension (BHW format) mask_tensor = torch.from_numpy(mask) - + # Expand to requested batch size mask_tensor = mask_tensor.unsqueeze(0).expand(batch_size, -1, -1) - + return (mask_tensor,) + class TextRenderer: @classmethod def INPUT_TYPES(s): @@ -419,54 +518,62 @@ class TextRenderer: "x_offset": ("INT", {"default": 0, "min": -4096, "max": 4096, "step": 1}), "y_offset": ("INT", {"default": 0, "min": -4096, "max": 4096, "step": 1}), "align": (["left", "center", "right"], {"default": "center"}), - "wrap_width": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8, - "tooltip": "Width to wrap text at (0 = no wrapping)"}), + "wrap_width": ( + "INT", + { + "default": 0, + "min": 0, + "max": MAX_RESOLUTION, + "step": 8, + "tooltip": "Width to wrap text at (0 = no wrapping)", + }, + ), } } - + RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "render_text" CATEGORY = "real-time" - def render_text(self, any, width, height, font_size, font_color, background_color, - x_offset, y_offset, align, wrap_width): + def render_text(self, any, width, height, font_size, font_color, background_color, x_offset, y_offset, align, wrap_width): # Convert input to string if it isn't already any = str(any) - + # Create image with background - image = Image.new('RGBA', (width, height), background_color) - mask = Image.new('L', (width, height), 0) # Black mask initially - + image = Image.new("RGBA", (width, height), background_color) + mask = Image.new("L", (width, height), 0) # Black mask initially + # Load font from assets directory import os + # Construct the path relative to the main package directory - base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) # Go up 3 levels + base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) # Go up 3 levels font_path = os.path.join(base_path, "fonts", "dejavu-sans", "DejaVuSans.ttf") try: pil_font = ImageFont.truetype(font_path, font_size) except Exception as e: print(f"Failed to load font from {font_path}, falling back to default. Error: {str(e)}") pil_font = ImageFont.load_default() - + # Create drawing objects draw = ImageDraw.Draw(image) mask_draw = ImageDraw.Draw(mask) - + import textwrap - + # Handle text wrapping if wrap_width > 0: # Convert wrap_width from pixels to approximate character count # Assuming average character width is font_size/2 pixels - char_width = font_size/2 + char_width = font_size / 2 char_count = max(1, int(wrap_width / char_width)) any = "\n".join(textwrap.wrap(any, width=char_count)) - + # Calculate text size bbox = draw.textbbox((0, 0), any, font=pil_font) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] - + # Calculate position if align == "left": x = 0 + x_offset @@ -474,79 +581,85 @@ class TextRenderer: x = width - text_width + x_offset else: # center x = (width - text_width) // 2 + x_offset - + y = (height - text_height) // 2 + y_offset - + # Draw text draw.multiline_text((x, y), any, font=pil_font, fill=font_color, align=align) mask_draw.multiline_text((x, y), any, font=pil_font, fill=255, align=align) - + # Convert PIL to tensor image_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,] mask_tensor = torch.from_numpy(np.array(mask).astype(np.float32) / 255.0)[None,] - + # If image has alpha channel, use it to modify the mask if image_tensor.shape[-1] == 4: mask_tensor = mask_tensor * image_tensor[..., 3] image_tensor = image_tensor[..., :3] # Keep only RGB channels - + return (image_tensor, mask_tensor) + class MultilineText: """A simple node for creating and formatting multiline text""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "text": ("STRING", { - "multiline": True, - "default": "Line 1\nLine 2\nLine 3", - "tooltip": "Enter text (use \\n for new lines)" - }), - "strip_whitespace": ("BOOLEAN", { - "default": True, - "tooltip": "Remove leading/trailing whitespace from each line" - }), - "remove_empty_lines": ("BOOLEAN", { - "default": False, - "tooltip": "Remove empty or whitespace-only lines" - }) + "text": ( + "STRING", + { + "multiline": True, + "default": "Line 1\nLine 2\nLine 3", + "tooltip": "Enter text (use \\n for new lines)", + }, + ), + "strip_whitespace": ( + "BOOLEAN", + {"default": True, "tooltip": "Remove leading/trailing whitespace from each line"}, + ), + "remove_empty_lines": ( + "BOOLEAN", + {"default": False, "tooltip": "Remove empty or whitespace-only lines"}, + ), } } RETURN_TYPES = ("STRING",) FUNCTION = "process_text" CATEGORY = "text" - + def process_text(self, text, strip_whitespace, remove_empty_lines): # Split text into lines - lines = text.split('\n') - + lines = text.split("\n") + # Process lines according to settings if strip_whitespace: lines = [line.strip() for line in lines] - + if remove_empty_lines: lines = [line for line in lines if line.strip()] - + # Rejoin lines - result = '\n'.join(lines) - + result = "\n".join(lines) + return (result,) + class LoadImageFromPath_: """A simple node that loads an image from a given file path""" - + @classmethod def INPUT_TYPES(s): - return {"required": { - "image_path": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Full path to the image file to load" - }) - }} + return { + "required": { + "image_path": ( + "STRING", + {"default": "", "multiline": False, "tooltip": "Full path to the image file to load"}, + ) + } + } RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "load_image" @@ -556,153 +669,150 @@ class LoadImageFromPath_: # Open and process the image img = Image.open(image_path) img = ImageOps.exif_transpose(img) - + # Convert to RGB if needed - if img.mode == 'I': + if img.mode == "I": img = img.point(lambda i: i * (1 / 255)) image = img.convert("RGB") - + # Convert to tensor image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] - + # Create mask from alpha channel if it exists - if 'A' in img.getbands(): - mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0 - mask = 1. - torch.from_numpy(mask) + if "A" in img.getbands(): + mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0 + mask = 1.0 - torch.from_numpy(mask) else: - mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") - + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + mask = mask.unsqueeze(0) - + return (image, mask) + class FeedbackEffect: """ Ultra-fast real-time visual feedback processor designed for SetState/GetState loops. Creates stunning, hypnotic visuals with optimized GPU-accelerated transformations. - + For best results: connect to a SetState/GetState loop and adjust parameters in real-time. """ + CATEGORY = "image/effects" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "apply_effect" - + @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE", {"tooltip": "Input image (typically from GetStateNode)"}), - "effect": (["fractal_zoom", "kaleidoscope", "plasma_field", "quantum_wave", - "neural_dream", "digital_flow", "hypershock"], { - "default": "fractal_zoom", - "tooltip": "Type of real-time effect to apply" - }), - "intensity": ("FLOAT", { - "default": 0.65, - "min": 0.1, - "max": 1.0, - "step": 0.05, - "tooltip": "Effect strength/intensity" - }), - "variation": ("INT", { - "default": 0, - "min": 0, - "max": 5, - "step": 1, - "tooltip": "Variation/mode of the selected effect" - }), - "speed": ("FLOAT", { - "default": 0.5, - "min": 0.05, - "max": 1.0, - "step": 0.05, - "tooltip": "Animation speed factor" - }), - "feedback_blend": ("FLOAT", { - "default": 0.95, - "min": 0.5, - "max": 0.999, - "step": 0.005, - "tooltip": "How much of previous frame to retain (higher = more trail effect)" - }), - "color_shift": ("FLOAT", { - "default": 0.2, - "min": 0.0, - "max": 1.0, - "step": 0.05, - "tooltip": "Amount of color transformation to apply" - }), - "reset": ("BOOLEAN", { - "default": False, - "tooltip": "Reset the feedback loop" - }), + "effect": ( + [ + "fractal_zoom", + "kaleidoscope", + "plasma_field", + "quantum_wave", + "neural_dream", + "digital_flow", + "hypershock", + ], + {"default": "fractal_zoom", "tooltip": "Type of real-time effect to apply"}, + ), + "intensity": ( + "FLOAT", + {"default": 0.65, "min": 0.1, "max": 1.0, "step": 0.05, "tooltip": "Effect strength/intensity"}, + ), + "variation": ( + "INT", + {"default": 0, "min": 0, "max": 5, "step": 1, "tooltip": "Variation/mode of the selected effect"}, + ), + "speed": ( + "FLOAT", + {"default": 0.5, "min": 0.05, "max": 1.0, "step": 0.05, "tooltip": "Animation speed factor"}, + ), + "feedback_blend": ( + "FLOAT", + { + "default": 0.95, + "min": 0.5, + "max": 0.999, + "step": 0.005, + "tooltip": "How much of previous frame to retain (higher = more trail effect)", + }, + ), + "color_shift": ( + "FLOAT", + { + "default": 0.2, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Amount of color transformation to apply", + }, + ), + "reset": ("BOOLEAN", {"default": False, "tooltip": "Reset the feedback loop"}), } } - + def __init__(self): self.frame_count = 0 self.time = 0 self.last_image = None self.last_effect = None self.effect_buffers = {} - + def apply_effect(self, image, effect, intensity, variation, speed, feedback_blend, color_shift, reset): batch_size, height, width, channels = image.shape device = image.device - + # Reset state if requested or if effect changes if reset or effect != self.last_effect: self.frame_count = 0 self.time = 0 self.effect_buffers = {} - + # Update time and frame counter self.frame_count += 1 frame_time = self.frame_count * speed * 0.03 self.time = frame_time self.last_effect = effect - + # Fast path for batch processing - we'll process each image separately result_images = [] - + for b in range(batch_size): img = image[b] result = self._process_single_image( - img, effect, intensity, variation, - feedback_blend, color_shift, device, - width, height, frame_time, speed + img, effect, intensity, variation, feedback_blend, color_shift, device, width, height, frame_time, speed ) result_images.append(result) return (torch.stack(result_images),) - - def _process_single_image(self, img, effect, intensity, variation, - feedback_blend, color_shift, device, - width, height, time, speed): + + def _process_single_image(self, img, effect, intensity, variation, feedback_blend, color_shift, device, width, height, time, speed): # Highly optimized processing path - calculate grid transformations in one go # Coordinates grid (only calculate once and reuse) - if not hasattr(self, 'grid_cache') or self.grid_cache['size'] != (height, width): + if not hasattr(self, "grid_cache") or self.grid_cache["size"] != (height, width): y_norm, x_norm = torch.meshgrid( - torch.linspace(-1, 1, height, device=device), - torch.linspace(-1, 1, width, device=device), - indexing='ij' + torch.linspace(-1, 1, height, device=device), torch.linspace(-1, 1, width, device=device), indexing="ij" ) self.grid_cache = { - 'size': (height, width), - 'x_norm': x_norm, - 'y_norm': y_norm, - 'r': torch.sqrt(x_norm**2 + y_norm**2), - 'theta': torch.atan2(y_norm, x_norm) + "size": (height, width), + "x_norm": x_norm, + "y_norm": y_norm, + "r": torch.sqrt(x_norm**2 + y_norm**2), + "theta": torch.atan2(y_norm, x_norm), } - + # Get cached grids - x_norm = self.grid_cache['x_norm'] - y_norm = self.grid_cache['y_norm'] - r = self.grid_cache['r'] - theta = self.grid_cache['theta'] - + x_norm = self.grid_cache["x_norm"] + y_norm = self.grid_cache["y_norm"] + r = self.grid_cache["r"] + theta = self.grid_cache["theta"] + # Initialize source coordinates to identity source_x = x_norm.clone() source_y = y_norm.clone() @@ -711,339 +821,347 @@ class FeedbackEffect: if effect == "fractal_zoom": # Fractal zoom with organic movement zoom = 0.98 + math.sin(time * 2.5) * 0.02 - + # Variation affects spiral pattern spiral = variation * 0.2 * intensity angle = theta + r * spiral + time * speed - + # Complex transform that creates evolving spiral patterns dx = torch.sin(r * (3 + variation) + time * 3) * 0.05 * intensity dy = torch.cos(r * (3 + variation) + time * 3.5) * 0.05 * intensity - + # Apply zoom source_x = (x_norm * zoom + dx) * (1 + 0.03 * intensity * math.sin(time)) source_y = (y_norm * zoom + dy) * (1 + 0.03 * intensity * math.cos(time)) - + # Add swirling effect twist = r * (4.0 * intensity) + time * (1.5 + variation * 0.3) source_x += torch.sin(twist + theta) * 0.05 * intensity source_y += torch.cos(twist - theta) * 0.05 * intensity - + elif effect == "kaleidoscope": # Kaleidoscope effect with dynamic symmetry segments = 4 + variation * 2 segment_angle = 2 * torch.pi / segments - + # Rotate over time - rot_theta = theta + time * (0.2 + 0.1 * variation) - + rot_theta = theta + time * (0.2 + 0.1 * variation) + # Mirror segments segment_id = torch.floor(rot_theta / segment_angle) in_segment = rot_theta - segment_id * segment_angle - flip_mask = (segment_id % 2 == 0) + flip_mask = segment_id % 2 == 0 mirror_theta = torch.where(flip_mask, in_segment, segment_angle - in_segment) - + # Apply transform with zoom breathing and ripples zoom = 0.8 + 0.2 * math.sin(time * 1.5) + r * 0.1 * intensity ripple = torch.sin(r * (8 + variation) - time * 3) * 0.06 * intensity - + source_x = torch.cos(mirror_theta) * r * (zoom + ripple) source_y = torch.sin(mirror_theta) * r * (zoom + ripple) - + elif effect == "plasma_field": # Plasma with flow dynamics scale = 5.0 + variation * 2.0 speed_factor = 3.0 + variation * 0.5 - + # Create plasma-like distortion fields - plasma_x = torch.sin(x_norm * scale + time * speed_factor) * torch.cos(y_norm * (scale*0.8) + time) - plasma_y = torch.sin(y_norm * scale - time * (speed_factor*0.7)) * torch.cos(x_norm * (scale*0.9) - time * 1.3) - plasma_xy = torch.sin((x_norm + y_norm) * (scale*0.4) + time * 1.5) * 0.3 - + plasma_x = torch.sin(x_norm * scale + time * speed_factor) * torch.cos(y_norm * (scale * 0.8) + time) + plasma_y = torch.sin(y_norm * scale - time * (speed_factor * 0.7)) * torch.cos(x_norm * (scale * 0.9) - time * 1.3) + plasma_xy = torch.sin((x_norm + y_norm) * (scale * 0.4) + time * 1.5) * 0.3 + # Combine fields for rich movement flow_x = (plasma_x + plasma_xy) * 0.15 * intensity flow_y = (plasma_y + plasma_xy) * 0.15 * intensity - + # Apply field distortion source_x = x_norm + flow_x source_y = y_norm + flow_y - + elif effect == "quantum_wave": # Quantum wave effect with interference patterns time_factor = time * (2.0 + variation * 0.4) - + # Dynamic wave parameters wave_scale = (4.0 + variation) + math.sin(time) * 1.5 wave_speed = 1.5 + variation * 0.2 - + # Create interference pattern wave1 = torch.sin(x_norm * wave_scale + time_factor * wave_speed) wave2 = torch.sin(y_norm * wave_scale - time_factor * wave_speed) wave3 = torch.sin((x_norm + y_norm) * wave_scale * 0.5 + time_factor * 1.3) wave4 = torch.sin(r * wave_scale * 0.8 - time_factor) - + # Combine wave patterns interference = (wave1 + wave2 + wave3 + wave4) * 0.25 - + # Apply distortion with intensity scaling dist_scale = 0.12 * intensity * (1.0 + 0.3 * math.sin(time * 0.5)) source_x = x_norm + interference * dist_scale * (1.0 + 0.5 * x_norm) source_y = y_norm + interference * dist_scale * (1.0 + 0.5 * y_norm) - + # Add subtle spiral rotation = time * 0.3 * (1 + variation * 0.1) rot_x = x_norm * torch.cos(torch.tensor(rotation, device=device)) - y_norm * torch.sin(torch.tensor(rotation, device=device)) rot_y = x_norm * torch.sin(torch.tensor(rotation, device=device)) + y_norm * torch.cos(torch.tensor(rotation, device=device)) - + # Blend rotation into final coordinates source_x = source_x * 0.85 + rot_x * 0.15 source_y = source_y * 0.85 + rot_y * 0.15 - + elif effect == "neural_dream": # Neural network-like patterns that evolve # Create layered activation patterns scale_factor = 3.0 + variation * 0.5 time_offset = time * (1.2 + variation * 0.2) - + # First activation layer - positional encodings act1 = torch.sin(x_norm * scale_factor + time_offset) act2 = torch.sin(y_norm * scale_factor * 1.1 - time_offset * 0.7) - act3 = torch.sin((x_norm*0.5 + y_norm*1.5) * scale_factor * 0.7 + time_offset * 0.5) - + act3 = torch.sin((x_norm * 0.5 + y_norm * 1.5) * scale_factor * 0.7 + time_offset * 0.5) + # Second layer - non-linear combinations layer2_1 = torch.tanh(act1 * act2 * 0.5 + act3 * 0.3) layer2_2 = torch.tanh(act2 * act3 * 0.5 + act1 * 0.3) - + # Output layer with residual connection dx = layer2_1 * 0.15 * intensity dy = layer2_2 * 0.15 * intensity - + # Apply distortion with subtle spiral angle = r * (variation * 0.2) + time * (0.2 + variation * 0.05) source_x = x_norm + dx + torch.sin(angle) * r * 0.02 * intensity source_y = y_norm + dy + torch.cos(angle) * r * 0.02 * intensity - + elif effect == "digital_flow": # Digital fluid-like motion with emergent patterns # Base flow field scale = 4.0 + variation time_factor = time * (1.0 + variation * 0.2) - + # Create flow vectors using simplex-like patterns - flow_x1 = torch.sin(x_norm * scale + time_factor) * torch.cos(y_norm * (scale*0.5) + time_factor * 0.6) - flow_y1 = torch.cos(x_norm * (scale*0.6) - time_factor * 0.7) * torch.sin(y_norm * scale - time_factor) - + flow_x1 = torch.sin(x_norm * scale + time_factor) * torch.cos(y_norm * (scale * 0.5) + time_factor * 0.6) + flow_y1 = torch.cos(x_norm * (scale * 0.6) - time_factor * 0.7) * torch.sin(y_norm * scale - time_factor) + # Second flow field at different scale - flow_x2 = torch.sin((x_norm + y_norm) * (scale*0.4) + time_factor * 1.3) - flow_y2 = torch.cos((y_norm - x_norm) * (scale*0.4) - time_factor * 1.1) - + flow_x2 = torch.sin((x_norm + y_norm) * (scale * 0.4) + time_factor * 1.3) + flow_y2 = torch.cos((y_norm - x_norm) * (scale * 0.4) - time_factor * 1.1) + # Combine flow fields flow_x = (flow_x1 * 0.6 + flow_x2 * 0.4) * 0.12 * intensity flow_y = (flow_y1 * 0.6 + flow_y2 * 0.4) * 0.12 * intensity - + # Apply flow with subtle zoom pulse zoom = 1.0 + math.sin(time * 1.5) * 0.01 * intensity source_x = x_norm * zoom + flow_x source_y = y_norm * zoom + flow_y - + # Add digital quantization if higher variation if variation >= 3: # Quantize stronger for higher variations quant_level = 15 + (5 - variation) * 5 # More divisions for lower variation source_x = torch.round(source_x * quant_level) / quant_level source_y = torch.round(source_y * quant_level) / quant_level - + elif effect == "hypershock": # High energy effect with shockwaves and energy fields # Base frequencies freq = 5.0 + variation time_factor = time * (2.0 + variation * 0.3) - + # Create pulsing shockwaves shock1_radius = (time_factor * 0.7) % 2.0 - 1.0 - shock1 = torch.exp(-((r - shock1_radius) * 10)**2) * 0.5 - + shock1 = torch.exp(-(((r - shock1_radius) * 10) ** 2)) * 0.5 + shock2_radius = (time_factor * 0.5 + 1.0) % 2.0 - 1.0 - shock2 = torch.exp(-((r - shock2_radius) * 10)**2) * 0.4 - + shock2 = torch.exp(-(((r - shock2_radius) * 10) ** 2)) * 0.4 + # Energy field energy_x = torch.sin(x_norm * freq + time_factor * 1.5) * 0.1 energy_y = torch.sin(y_norm * freq - time_factor * 1.3) * 0.1 - + # Combine effects dx = (shock1 + shock2 * 0.7) * x_norm * 0.3 + energy_x dy = (shock1 + shock2 * 0.7) * y_norm * 0.3 + energy_y - + # Apply distortion with high energy scaling source_x = x_norm + dx * intensity source_y = y_norm + dy * intensity - + # Add spiral warp for higher variations if variation >= 2: warp = torch.sin(r * (8 - variation) + time_factor) * 0.1 * intensity source_x += torch.sin(theta + torch.tensor(time, device=device)) * warp source_y += torch.cos(theta + torch.tensor(time, device=device)) * warp - + # Ensure coordinates stay within bounds source_x = torch.clamp(source_x, -1.98, 1.98) source_y = torch.clamp(source_y, -1.98, 1.98) - + # Convert from normalized coords to actual pixel sampling grid [-1, 1] grid = torch.stack([source_x, source_y], dim=-1) grid = grid.unsqueeze(0) # Add batch dimension - + # Sample the image (fast tensor operation) img_rgb = img[..., :3].permute(2, 0, 1).unsqueeze(0) # [1, C, H, W] - sampled = torch.nn.functional.grid_sample( - img_rgb, grid, - mode='bilinear', - padding_mode='reflection', - align_corners=False - ).squeeze(0).permute(1, 2, 0) # Back to [H, W, C] - + sampled = ( + torch.nn.functional.grid_sample(img_rgb, grid, mode="bilinear", padding_mode="reflection", align_corners=False) + .squeeze(0) + .permute(1, 2, 0) + ) # Back to [H, W, C] + # Apply color transformations based on selected effect if color_shift > 0: # Optimize color processing by pushing calculations to the GPU if effect == "fractal_zoom": # Vibrant color amplification color_factor = time * 0.1 % 1.0 - hsv_rotation = torch.tensor([ - [math.cos(color_factor * 6.28), -math.sin(color_factor * 6.28), 0], - [math.sin(color_factor * 6.28), math.cos(color_factor * 6.28), 0], - [0, 0, 1] - ], device=device) - + hsv_rotation = torch.tensor( + [ + [math.cos(color_factor * 6.28), -math.sin(color_factor * 6.28), 0], + [math.sin(color_factor * 6.28), math.cos(color_factor * 6.28), 0], + [0, 0, 1], + ], + device=device, + ) + flat_rgb = sampled.reshape(-1, 3) rotated_rgb = torch.matmul(flat_rgb, hsv_rotation) sampled = rotated_rgb.reshape(sampled.shape) - + # Vibrance enhancement luma = sampled[..., 0] * 0.299 + sampled[..., 1] * 0.587 + sampled[..., 2] * 0.114 chroma = sampled - luma.unsqueeze(-1) sampled = luma.unsqueeze(-1) + chroma * (1 + color_shift * 0.5) - + elif effect in ["kaleidoscope", "hypershock"]: # High contrast with glow # Add bloom/glow to bright areas bloom_threshold = 0.7 bloom_mask = (sampled.mean(dim=-1) > bloom_threshold).float().unsqueeze(-1) bloom = sampled * bloom_mask * color_shift * 0.5 - + # Apply bloom with additive blending sampled = torch.clamp(sampled + bloom, 0, 1) - + # Edge enhancement for kaleidoscope if effect == "kaleidoscope": edge_x = torch.abs(sampled[:, 1:] - sampled[:, :-1]).mean(-1, keepdim=True) edge_x = torch.cat([edge_x, edge_x[:, -1:]], dim=1) - + edge_y = torch.abs(sampled[1:, :] - sampled[:-1, :]).mean(-1, keepdim=True) edge_y = torch.cat([edge_y, edge_y[-1:, :]], dim=0) - + edge = (edge_x + edge_y) * 0.5 sampled = torch.clamp(sampled + edge * color_shift, 0, 1) - + elif effect in ["plasma_field", "neural_dream"]: # Hue rotation with time hue_shift = (time * 0.05) % 1.0 - + # RGB to HLS-like transform (simplified) cmax, _ = torch.max(sampled, dim=-1, keepdim=True) cmin, _ = torch.min(sampled, dim=-1, keepdim=True) delta = cmax - cmin - + # Calculate luminance luma = (cmax + cmin) * 0.5 - + # Calculate saturation sat = torch.where(delta > 0, delta / (1 - torch.abs(2 * luma - 1) + 1e-6), torch.zeros_like(delta)) - + # Enhance saturation sat = torch.clamp(sat * (1 + color_shift), 0, 1) - + # Apply new saturation while preserving luminance normalized = (sampled - cmin) / (delta + 1e-6) shifted = luma + (normalized - 0.5) * sat * (1 - torch.abs(2 * luma - 1)) sampled = torch.clamp(shifted, 0, 1) - + elif effect == "quantum_wave": # Quantum interference colors # Shift RGB channels with interference pattern shift = torch.sin(r * 4 + time * 2) * 0.1 * color_shift - + # RGB channel shifting - r_shift = torch.roll(sampled[..., 0], shifts=int(shift[0,0].item() * width), dims=0) - r_shift = torch.roll(r_shift, shifts=int(shift[0,0].item() * height), dims=1) - - b_shift = torch.roll(sampled[..., 2], shifts=-int(shift[0,0].item() * width), dims=0) - b_shift = torch.roll(b_shift, shifts=-int(shift[0,0].item() * height), dims=1) - + r_shift = torch.roll(sampled[..., 0], shifts=int(shift[0, 0].item() * width), dims=0) + r_shift = torch.roll(r_shift, shifts=int(shift[0, 0].item() * height), dims=1) + + b_shift = torch.roll(sampled[..., 2], shifts=-int(shift[0, 0].item() * width), dims=0) + b_shift = torch.roll(b_shift, shifts=-int(shift[0, 0].item() * height), dims=1) + sampled = torch.stack([r_shift, sampled[..., 1], b_shift], dim=-1) - + elif effect == "digital_flow": # Digital color quantization if color_shift > 0.3: quant_levels = max(3, int(10 * (1 - color_shift))) # Fewer levels with higher color_shift sampled = torch.floor(sampled * quant_levels) / quant_levels - + # Alpha handling if img.shape[-1] == 4: # Keep original alpha channel - alpha = torch.nn.functional.grid_sample( - img[..., 3:].permute(2, 0, 1).unsqueeze(0), - grid, - mode='bilinear', - padding_mode='reflection', - align_corners=False - ).squeeze(0).permute(1, 2, 0) - + alpha = ( + torch.nn.functional.grid_sample( + img[..., 3:].permute(2, 0, 1).unsqueeze(0), + grid, + mode="bilinear", + padding_mode="reflection", + align_corners=False, + ) + .squeeze(0) + .permute(1, 2, 0) + ) + # Recombine with RGB sampled = torch.cat([sampled, alpha], dim=-1) - + # Apply feedback blend if we have a last image if self.last_image is not None and img.shape == self.last_image.shape: # Optimize blend operation result = sampled * (1.0 - feedback_blend) + self.last_image * feedback_blend else: result = sampled - + # Store for next frame self.last_image = result - + return result - + + class ToString: """A node that converts any input to a string.""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "any": (AlwaysEqualProxy("*"),{}), # Accept any input type - "format_string": ("STRING", { - "default": "{}", - "multiline": False, - "tooltip": "Optional format string (e.g., '{:.2f}' for 2 decimal place float)" - }) + "any": (AlwaysEqualProxy("*"), {}), # Accept any input type + "format_string": ( + "STRING", + { + "default": "{}", + "multiline": False, + "tooltip": "Optional format string (e.g., '{:.2f}' for 2 decimal place float)", + }, + ), }, "optional": { - "prepend": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Text to add before the converted string" - }), - "append": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Text to add after the converted string" - }) - } + "prepend": ( + "STRING", + {"default": "", "multiline": False, "tooltip": "Text to add before the converted string"}, + ), + "append": ( + "STRING", + {"default": "", "multiline": False, "tooltip": "Text to add after the converted string"}, + ), + }, } RETURN_TYPES = ("STRING",) FUNCTION = "to_string" CATEGORY = "utils" - + def to_string(self, any, format_string="{}", prepend="", append=""): try: # Try to use the format string if provided @@ -1051,30 +1169,25 @@ class ToString: except (ValueError, TypeError): # Fall back to simple string conversion if format fails result = str(any) - + # Add prepend and append text result = prepend + result + append - + return (result,) + class RoundNode: """A node that rounds a float to a specified number of decimal places.""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "value": ("FLOAT", { - "default": 0.0, - "tooltip": "Float value to round" - }), - "decimal_places": ("INT", { - "default": 2, - "min": 0, - "max": 10, - "step": 1, - "tooltip": "Number of decimal places to round to" - }) + "value": ("FLOAT", {"default": 0.0, "tooltip": "Float value to round"}), + "decimal_places": ( + "INT", + {"default": 2, "min": 0, "max": 10, "step": 1, "tooltip": "Number of decimal places to round to"}, + ), } } @@ -1082,7 +1195,7 @@ class RoundNode: RETURN_NAMES = ("rounded_value",) FUNCTION = "round_value" CATEGORY = "utils" - + def round_value(self, value, decimal_places): if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)): rounded = [round(v, decimal_places) for v in value] @@ -1090,42 +1203,20 @@ class RoundNode: rounded = round(value, decimal_places) return (rounded,) + class RenormalizeFloat: """Renormalizes a float value from one range to another""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "value": ("FLOAT", { - "default": 0.5, - "forceInput": True, - "tooltip": "Float value to renormalize" - }), - "input_min": ("FLOAT", { - "default": 0.0, - "step": 0.01, - "tooltip": "Minimum value of the input range" - }), - "input_max": ("FLOAT", { - "default": 1.0, - "step": 0.01, - "tooltip": "Maximum value of the input range" - }), - "output_min": ("FLOAT", { - "default": 0.0, - "step": 0.01, - "tooltip": "Minimum value of the output range" - }), - "output_max": ("FLOAT", { - "default": 1.0, - "step": 0.01, - "tooltip": "Maximum value of the output range" - }), - "clamp": ("BOOLEAN", { - "default": True, - "tooltip": "Clamp output to the output range" - }) + "value": ("FLOAT", {"default": 0.5, "forceInput": True, "tooltip": "Float value to renormalize"}), + "input_min": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Minimum value of the input range"}), + "input_max": ("FLOAT", {"default": 1.0, "step": 0.01, "tooltip": "Maximum value of the input range"}), + "output_min": ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Minimum value of the output range"}), + "output_max": ("FLOAT", {"default": 1.0, "step": 0.01, "tooltip": "Maximum value of the output range"}), + "clamp": ("BOOLEAN", {"default": True, "tooltip": "Clamp output to the output range"}), } } @@ -1133,7 +1224,7 @@ class RenormalizeFloat: RETURN_NAMES = ("renormalized",) FUNCTION = "renormalize" CATEGORY = "utils" - + def renormalize(self, value, input_min, input_max, output_min, output_max, clamp): """Rescales a value from an input range to an output range""" # Check if value is a list, tuple, numpy array, or torch tensor @@ -1147,60 +1238,42 @@ class RenormalizeFloat: # Process single value result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp) return (result,) - + def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp): """Helper method to renormalize a single value""" # Handle cases where input_min equals input_max if abs(input_max - input_min) < 1e-6: return output_min - + # Calculate normalized position in input range normalized = (value - input_min) / (input_max - input_min) - + # Map to output range result = output_min + normalized * (output_max - output_min) - + # Apply clamping if requested if clamp: if output_min <= output_max: result = max(output_min, min(output_max, result)) else: result = max(output_max, min(output_min, result)) - + return result + class RenormalizeInt: """Renormalizes an integer value from one range to another""" - + @classmethod def INPUT_TYPES(cls): return { "required": { - "value": ("INT", { - "default": 50, - "forceInput": True, - "tooltip": "Integer value to renormalize" - }), - "input_min": ("INT", { - "default": 0, - "tooltip": "Minimum value of the input range" - }), - "input_max": ("INT", { - "default": 100, - "tooltip": "Maximum value of the input range" - }), - "output_min": ("INT", { - "default": 0, - "tooltip": "Minimum value of the output range" - }), - "output_max": ("INT", { - "default": 100, - "tooltip": "Maximum value of the output range" - }), - "clamp": ("BOOLEAN", { - "default": True, - "tooltip": "Clamp output to the output range" - }) + "value": ("INT", {"default": 50, "forceInput": True, "tooltip": "Integer value to renormalize"}), + "input_min": ("INT", {"default": 0, "tooltip": "Minimum value of the input range"}), + "input_max": ("INT", {"default": 100, "tooltip": "Maximum value of the input range"}), + "output_min": ("INT", {"default": 0, "tooltip": "Minimum value of the output range"}), + "output_max": ("INT", {"default": 100, "tooltip": "Maximum value of the output range"}), + "clamp": ("BOOLEAN", {"default": True, "tooltip": "Clamp output to the output range"}), } } @@ -1208,7 +1281,7 @@ class RenormalizeInt: RETURN_NAMES = ("renormalized",) FUNCTION = "renormalize" CATEGORY = "utils" - + def renormalize(self, value, input_min, input_max, output_min, output_max, clamp): """Rescales an integer value from an input range to an output range""" # Check if value is a list, tuple, numpy array, or torch tensor @@ -1222,24 +1295,24 @@ class RenormalizeInt: # Process single value result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp) return (result,) - + def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp): """Helper method to renormalize a single integer value""" # Handle cases where input_min equals input_max if input_max == input_min: return output_min - + # Calculate normalized position in input range normalized = (value - input_min) / (input_max - input_min) - + # Map to output range and round to integer result = round(output_min + normalized * (output_max - output_min)) - + # Apply clamping if requested if clamp: if output_min <= output_max: result = max(output_min, min(output_max, result)) else: result = max(output_max, min(output_min, result)) - + return result diff --git a/pyproject.toml b/pyproject.toml index 5f78415..afd31b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,3 +15,6 @@ Repository = "https://github.com/ryanontheinside/ComfyUI_RealTimeNodes" PublisherId = "ryanontheinside" DisplayName = "ComfyUI_RealTimeNodes" Icon = "⏱️" + +[tool.ruff] +line-length = 140 \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py index da22290..6e7633b 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/src/coordinates/__init__.py b/src/coordinates/__init__.py index 2503eb2..d23a798 100644 --- a/src/coordinates/__init__.py +++ b/src/coordinates/__init__.py @@ -9,4 +9,4 @@ with support for single values, batches, and point collections. from .coordinate_system import CoordinateSystem, Point, PointList from .drawing_engine import DrawingEngine -__all__ = ['CoordinateSystem', 'Point', 'PointList', 'DrawingEngine'] \ No newline at end of file +__all__ = ["CoordinateSystem", "Point", "PointList", "DrawingEngine"] diff --git a/src/coordinates/coordinate_delta.py b/src/coordinates/coordinate_delta.py index c306819..e1ac4c2 100644 --- a/src/coordinates/coordinate_delta.py +++ b/src/coordinates/coordinate_delta.py @@ -4,21 +4,22 @@ import math import typing # Add typing for Dict hint from typing import Optional -#TODO: consider moving base nodes to src +# TODO: consider moving base nodes to src # Import from the new consolidated utilities logger = logging.getLogger(__name__) -MAX_POSITION_HISTORY = 50 # Matches MAX_LANDMARK_HISTORY for consistency -FLOAT_EQUALITY_TOLERANCE = 1e-6 # Same tolerance as base_delta_nodes +MAX_POSITION_HISTORY = 50 # Matches MAX_LANDMARK_HISTORY for consistency +FLOAT_EQUALITY_TOLERANCE = 1e-6 # Same tolerance as base_delta_nodes # --- Base Position Delta Logic (Handling Multiple Histories) --- class BaseCoordinateDelta: """Base class for nodes calculating delta from position history. - Handles single floats or lists of coordinates. + Handles single floats or lists of coordinates. """ + FUNCTION = "execute" def __init__(self): @@ -35,23 +36,21 @@ class BaseCoordinateDelta: def _get_position_delta(self, x: float, y: float, z: float, window_size: int, item_index: int = 0) -> Optional[float]: """Calculates delta using the history deque for the given item_index.""" history_deque = self._get_or_create_history(item_index) - + window_size = max(1, min(window_size, MAX_POSITION_HISTORY - 1)) current_pos = (x, y, z) history_deque.append(current_pos) - + if len(history_deque) > window_size: pos_now = history_deque[-1] - pos_past = history_deque[-(window_size + 1)] # Correct indexing + pos_past = history_deque[-(window_size + 1)] # Correct indexing if pos_now is None or pos_past is None: - logger.warning(f"{self.__class__.__name__}[{item_index}]: Missing position data in history window.") - return None - delta = math.sqrt((pos_now[0] - pos_past[0])**2 + - (pos_now[1] - pos_past[1])**2 + - (pos_now[2] - pos_past[2])**2) + logger.warning(f"{self.__class__.__name__}[{item_index}]: Missing position data in history window.") + return None + delta = math.sqrt((pos_now[0] - pos_past[0]) ** 2 + (pos_now[1] - pos_past[1]) ** 2 + (pos_now[2] - pos_past[2]) ** 2) return delta else: # logger.debug(f"{self.__class__.__name__}[{item_index}]: Not enough history ({len(history_deque)}/{window_size+1}).") - return None # Not enough history yet - + return None # Not enough history yet + # Remove _handle_input_coords, logic moved to execute diff --git a/src/coordinates/coordinate_system.py b/src/coordinates/coordinate_system.py index a21120b..9e539ff 100644 --- a/src/coordinates/coordinate_system.py +++ b/src/coordinates/coordinate_system.py @@ -12,37 +12,40 @@ import torch logger = logging.getLogger(__name__) + class CoordinateSystem: """Unified coordinate system handling for coordinate space conversions.""" - + # Constants for coordinate spaces NORMALIZED = "normalized" # 0.0-1.0 range - PIXEL = "pixel" # Pixel coordinates - + PIXEL = "pixel" # Pixel coordinates + @staticmethod - def normalize(coords: Union[float, List[float], torch.Tensor], - dimensions: Union[int, Tuple[int, ...]], - input_space: str = PIXEL) -> Union[float, List[float], torch.Tensor]: + def normalize( + coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + input_space: str = PIXEL, + ) -> Union[float, List[float], torch.Tensor]: """ Normalize coordinates to 0.0-1.0 range. - + Args: coords: Coordinates to normalize dimensions: Dimensions to normalize by (width, height) input_space: Current coordinate space - + Returns: Normalized coordinates (0.0-1.0) """ if input_space == CoordinateSystem.NORMALIZED: return coords - + # Handle different dimension specifications if isinstance(dimensions, tuple): dim_size = dimensions[0] if len(dimensions) == 1 else dimensions else: dim_size = dimensions - + # Handle different input types if isinstance(coords, list): if isinstance(dim_size, tuple): @@ -67,31 +70,33 @@ class CoordinateSystem: return coords.float() / dim_size else: return float(coords) / dim_size - + @staticmethod - def denormalize(coords: Union[float, List[float], torch.Tensor], - dimensions: Union[int, Tuple[int, ...]], - output_space: str = PIXEL) -> Union[float, List[float], torch.Tensor]: + def denormalize( + coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + output_space: str = PIXEL, + ) -> Union[float, List[float], torch.Tensor]: """ Convert normalized coordinates to target space. - + Args: coords: Normalized coordinates (0.0-1.0) dimensions: Target dimensions output_space: Target coordinate space - + Returns: Denormalized coordinates in target space """ if output_space == CoordinateSystem.NORMALIZED: return coords - + # Handle different dimension specifications if isinstance(dimensions, tuple): dim_size = dimensions[0] if len(dimensions) == 1 else dimensions else: dim_size = dimensions - + # Handle different input types if isinstance(coords, list): if isinstance(dim_size, tuple): @@ -117,69 +122,68 @@ class CoordinateSystem: return torch.clamp(coords * dim_size, 0, dim_size - 1) else: return float(max(0, min(dim_size - 1, coords * dim_size))) - + @staticmethod - def convert(coords: Union[float, List[float], torch.Tensor], - dimensions: Union[int, Tuple[int, ...]], - from_space: str, - to_space: str) -> Union[float, List[float], torch.Tensor]: + def convert( + coords: Union[float, List[float], torch.Tensor], + dimensions: Union[int, Tuple[int, ...]], + from_space: str, + to_space: str, + ) -> Union[float, List[float], torch.Tensor]: """ Convert between coordinate spaces. - + Args: coords: Coordinates to convert dimensions: Target dimensions from_space: Source coordinate space to_space: Target coordinate space - + Returns: Converted coordinates """ if from_space == to_space: return coords - + # First normalize if needed if from_space != CoordinateSystem.NORMALIZED: normalized = CoordinateSystem.normalize(coords, dimensions, from_space) else: normalized = coords - + # Then denormalize to target space if needed if to_space != CoordinateSystem.NORMALIZED: return CoordinateSystem.denormalize(normalized, dimensions, to_space) else: return normalized - + @staticmethod def get_dimensions_from_tensor(tensor: torch.Tensor) -> Tuple[int, int]: """ Extract dimensions from tensor in BHWC format. - + Args: tensor: Input tensor in BHWC format - + Returns: Tuple of (width, height) """ if tensor.dim() != 4: raise ValueError(f"Expected 4D tensor (BHWC), got shape {tensor.shape}") - + _, height, width, _ = tensor.shape return (width, height) + class Point: """ Represents a 2D or 3D point with coordinate space awareness. """ - - def __init__(self, - x: float, - y: float, - z: Optional[float] = None, - space: str = CoordinateSystem.NORMALIZED): + + def __init__(self, x: float, y: float, z: Optional[float] = None, space: str = CoordinateSystem.NORMALIZED): """ Initialize a point with coordinates and space information. - + Args: x: X coordinate y: Y coordinate @@ -191,85 +195,85 @@ class Point: self.z = z self.space = space self.is_3d = z is not None - - def to_normalized(self, dimensions: Tuple[int, ...]) -> 'Point': + + def to_normalized(self, dimensions: Tuple[int, ...]) -> "Point": """ Convert point to normalized space. - + Args: dimensions: Dimensions for conversion - + Returns: New point in normalized space """ if self.space == CoordinateSystem.NORMALIZED: return self - + norm_x = CoordinateSystem.normalize(self.x, dimensions[0], self.space) norm_y = CoordinateSystem.normalize(self.y, dimensions[1], self.space) norm_z = None - + if self.is_3d and len(dimensions) > 2: norm_z = CoordinateSystem.normalize(self.z, dimensions[2], self.space) - + return Point(norm_x, norm_y, norm_z, CoordinateSystem.NORMALIZED) - - def to_pixel(self, dimensions: Tuple[int, ...]) -> 'Point': + + def to_pixel(self, dimensions: Tuple[int, ...]) -> "Point": """ Convert point to pixel space. - + Args: dimensions: Dimensions for conversion - + Returns: New point in pixel space """ if self.space == CoordinateSystem.PIXEL: return self - + # First normalize if not already if self.space != CoordinateSystem.NORMALIZED: normalized = self.to_normalized(dimensions) else: normalized = self - + pixel_x = CoordinateSystem.denormalize(normalized.x, dimensions[0], CoordinateSystem.PIXEL) pixel_y = CoordinateSystem.denormalize(normalized.y, dimensions[1], CoordinateSystem.PIXEL) pixel_z = None - + if self.is_3d and len(dimensions) > 2: pixel_z = CoordinateSystem.denormalize(normalized.z, dimensions[2], CoordinateSystem.PIXEL) - + return Point(pixel_x, pixel_y, pixel_z, CoordinateSystem.PIXEL) - - def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> 'Point': + + def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> "Point": """ Convert point to target space. - + Args: dimensions: Dimensions for conversion target_space: Target coordinate space - + Returns: New point in target space """ if self.space == target_space: return self - + if target_space == CoordinateSystem.NORMALIZED: return self.to_normalized(dimensions) elif target_space == CoordinateSystem.PIXEL: return self.to_pixel(dimensions) else: raise ValueError(f"Unsupported target space: {target_space}") - + def to_tuple(self, include_z: bool = False) -> Tuple[float, ...]: """ Convert point to tuple representation. - + Args: include_z: Whether to include Z coordinate if available - + Returns: Tuple of coordinates """ @@ -277,14 +281,14 @@ class Point: return (self.x, self.y, self.z) else: return (self.x, self.y) - + def to_list(self, include_z: bool = False) -> List[float]: """ Convert point to list representation. - + Args: include_z: Whether to include Z coordinate if available - + Returns: List of coordinates """ @@ -297,126 +301,127 @@ class Point: else: return f"Point({self.x}, {self.y}, space='{self.space}')" + class PointList: """ Represents a list of 2D or 3D points with coordinate space awareness. Useful for handling batches of points or landmark lists. """ - - def __init__(self, - points: List[Point] = None, - space: str = CoordinateSystem.NORMALIZED): + + def __init__(self, points: List[Point] = None, space: str = CoordinateSystem.NORMALIZED): """ Initialize a point list. - + Args: points: List of Point objects space: Coordinate space for all points """ self.points = points or [] self.space = space - + @classmethod - def from_coordinates(cls, - x_coords: List[float], - y_coords: List[float], - z_coords: Optional[List[float]] = None, - space: str = CoordinateSystem.NORMALIZED) -> 'PointList': + def from_coordinates( + cls, + x_coords: List[float], + y_coords: List[float], + z_coords: Optional[List[float]] = None, + space: str = CoordinateSystem.NORMALIZED, + ) -> "PointList": """ Create a PointList from separate coordinate lists. - + Args: x_coords: List of X coordinates y_coords: List of Y coordinates z_coords: List of Z coordinates (optional) space: Coordinate space - + Returns: New PointList object """ if len(x_coords) != len(y_coords): raise ValueError(f"Coordinate lists must have same length: x={len(x_coords)}, y={len(y_coords)}") - + has_z = z_coords is not None if has_z and len(z_coords) != len(x_coords): raise ValueError(f"Z coordinates must match length of X/Y: x={len(x_coords)}, z={len(z_coords)}") - + points = [] for i in range(len(x_coords)): z = z_coords[i] if has_z else None points.append(Point(x_coords[i], y_coords[i], z, space)) - + return cls(points, space) - - def to_normalized(self, dimensions: Tuple[int, ...]) -> 'PointList': + + def to_normalized(self, dimensions: Tuple[int, ...]) -> "PointList": """ Convert all points to normalized space. - + Args: dimensions: Dimensions for conversion - + Returns: New PointList in normalized space """ if self.space == CoordinateSystem.NORMALIZED: return self - + normalized_points = [point.to_normalized(dimensions) for point in self.points] return PointList(normalized_points, CoordinateSystem.NORMALIZED) - - def to_pixel(self, dimensions: Tuple[int, ...]) -> 'PointList': + + def to_pixel(self, dimensions: Tuple[int, ...]) -> "PointList": """ Convert all points to pixel space. - + Args: dimensions: Dimensions for conversion - + Returns: New PointList in pixel space """ if self.space == CoordinateSystem.PIXEL: return self - + pixel_points = [point.to_pixel(dimensions) for point in self.points] return PointList(pixel_points, CoordinateSystem.PIXEL) - - def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> 'PointList': + + def to_space(self, dimensions: Tuple[int, ...], target_space: str) -> "PointList": """ Convert all points to target space. - + Args: dimensions: Dimensions for conversion target_space: Target coordinate space - + Returns: New PointList in target space """ if self.space == target_space: return self - + if target_space == CoordinateSystem.NORMALIZED: return self.to_normalized(dimensions) elif target_space == CoordinateSystem.PIXEL: return self.to_pixel(dimensions) else: raise ValueError(f"Unsupported target space: {target_space}") - + def get_x_list(self) -> List[float]: """Get list of X coordinates.""" return [point.x for point in self.points] - + def get_y_list(self) -> List[float]: """Get list of Y coordinates.""" return [point.y for point in self.points] - + def get_z_list(self) -> List[Optional[float]]: """Get list of Z coordinates (may contain None values).""" return [point.z for point in self.points] - + def __len__(self) -> int: """Get number of points in the list.""" return len(self.points) - + def __getitem__(self, idx: int) -> Point: """Access point by index.""" - return self.points[idx] \ No newline at end of file + return self.points[idx] diff --git a/src/coordinates/drawing_engine.py b/src/coordinates/drawing_engine.py index 39ed0d4..69bfcbd 100644 --- a/src/coordinates/drawing_engine.py +++ b/src/coordinates/drawing_engine.py @@ -9,19 +9,20 @@ from .coordinate_system import CoordinateSystem logger = logging.getLogger(__name__) + class DrawingEngine: """High-performance drawing engine for real-time applications. Optimized for minimal tensor-numpy conversions and efficient batch operations. """ - + @staticmethod def hex_to_bgr(hex_color: str) -> Tuple[int, int, int]: """Converts hex color to BGR tuple with validation.""" - hex_color = hex_color.lstrip('#') + hex_color = hex_color.lstrip("#") if len(hex_color) != 6: return (0, 0, 0) # Default black try: - rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + rgb = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) return (rgb[2], rgb[1], rgb[0]) # RGB to BGR except ValueError: return (0, 0, 0) @@ -31,7 +32,7 @@ class DrawingEngine: """Fast preparation of single image from batch for drawing - minimal conversion.""" if image.dim() != 4: raise ValueError("Input image must be BHWC format") - + device, dtype = image.device, image.dtype # Get single image, convert to float32 for processing img_tensor = image[batch_idx].to(torch.float32) @@ -45,17 +46,19 @@ class DrawingEngine: """Fast conversion from numpy back to tensor with proper device/dtype.""" return torch.from_numpy(img_np.astype(np.float32) / 255.0).to(device=device, dtype=dtype) - def draw_points(self, - image: torch.Tensor, - points_x: Union[float, List[float]], - points_y: Union[float, List[float]], - radius: int = 5, - color: str = "#FF0000", - is_normalized: bool = True, - batch_mapping: str = "broadcast") -> torch.Tensor: + def draw_points( + self, + image: torch.Tensor, + points_x: Union[float, List[float]], + points_y: Union[float, List[float]], + radius: int = 5, + color: str = "#FF0000", + is_normalized: bool = True, + batch_mapping: str = "broadcast", + ) -> torch.Tensor: """ Draw points on images efficiently. - + Args: image: BHWC tensor points_x/y: Coordinate values (normalized 0-1 or pixel) @@ -63,30 +66,30 @@ class DrawingEngine: color: Hex color is_normalized: If True, coords are 0-1 normalized batch_mapping: How to map points to batch ("broadcast", "one-to-one", "all-on-first") - + Returns: Modified image tensor """ batch_size, height, width, _ = image.shape output_batch = image.clone() # Avoid modifying input bgr_color = self.hex_to_bgr(color) - + # Standardize input to lists for unified processing x_list = points_x if isinstance(points_x, list) else [points_x] y_list = points_y if isinstance(points_y, list) else [points_y] - + # Determine batch mapping strategy points_len = len(x_list) if batch_mapping == "one-to-one" and points_len != batch_size and points_len != 1: logger.warning(f"One-to-one mapping requested but point list length ({points_len}) != batch size ({batch_size})") batch_mapping = "all-on-first" - + # Process each image in batch for b in range(batch_size): # Skip processing based on mapping strategy if batch_mapping == "all-on-first" and b > 0: continue - + # Get points for this batch item based on mapping strategy if batch_mapping == "one-to-one" and points_len == batch_size: # One point per image @@ -100,19 +103,19 @@ class DrawingEngine: # All points on first image only curr_x = x_list curr_y = y_list - + # Skip if no points to draw if not curr_x: continue - + # Single efficient conversion to numpy img_np, device, dtype = self.prep_batch_image(output_batch, b) - + # Draw all points for this image for xi, yi in zip(curr_x, curr_y): if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)): continue - + # Convert to pixel coordinates if needed if is_normalized: px = int(CoordinateSystem.denormalize(xi, width, CoordinateSystem.PIXEL)) @@ -120,34 +123,36 @@ class DrawingEngine: else: px = int(max(0, min(width - 1, xi))) py = int(max(0, min(height - 1, yi))) - + # Draw circle cv2.circle(img_np, (px, py), radius, bgr_color, -1) - + # Convert back to tensor and update batch output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) - + return output_batch - - def draw_lines(self, - image: torch.Tensor, - x1: Union[float, List[float]], - y1: Union[float, List[float]], - x2: Union[float, List[float]], - y2: Union[float, List[float]], - thickness: int = 2, - color: str = "#00FF00", - is_normalized: bool = True, - batch_mapping: str = "broadcast", - draw_endpoints: bool = False, - point_radius: int = 3, - draw_label: bool = False, - label_text: str = "", - label_position: str = "Midpoint", - font_scale: float = 0.5) -> torch.Tensor: + + def draw_lines( + self, + image: torch.Tensor, + x1: Union[float, List[float]], + y1: Union[float, List[float]], + x2: Union[float, List[float]], + y2: Union[float, List[float]], + thickness: int = 2, + color: str = "#00FF00", + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_endpoints: bool = False, + point_radius: int = 3, + draw_label: bool = False, + label_text: str = "", + label_position: str = "Midpoint", + font_scale: float = 0.5, + ) -> torch.Tensor: """ Draw lines on images efficiently. - + Args: image: BHWC tensor x1/y1/x2/y2: Line endpoint coordinates @@ -157,36 +162,36 @@ class DrawingEngine: batch_mapping: How to map lines to batch draw_endpoints: Whether to draw circles at endpoints Other params: Styling options - + Returns: Modified image tensor """ batch_size, height, width, _ = image.shape output_batch = image.clone() bgr_color = self.hex_to_bgr(color) - + # Standardize input to lists x1_list = x1 if isinstance(x1, list) else [x1] y1_list = y1 if isinstance(y1, list) else [y1] x2_list = x2 if isinstance(x2, list) else [x2] y2_list = y2 if isinstance(y2, list) else [y2] - + # Validate input lengths line_len = len(x1_list) if not all(len(coord) == line_len for coord in [y1_list, x2_list, y2_list]): raise ValueError("All coordinate lists must have the same length") - + # Determine batch mapping if batch_mapping == "one-to-one" and line_len != batch_size and line_len != 1: logger.warning(f"One-to-one mapping requested but line list length ({line_len}) != batch size ({batch_size})") batch_mapping = "all-on-first" - + # Process each image in batch for b in range(batch_size): # Skip processing based on mapping strategy if batch_mapping == "all-on-first" and b > 0: continue - + # Get lines for this batch item if batch_mapping == "one-to-one" and line_len == batch_size: indices = [b] @@ -194,22 +199,22 @@ class DrawingEngine: indices = range(line_len) else: # "all-on-first" indices = range(line_len) - + # Skip if no lines to draw if not indices: continue - + # Single efficient conversion to numpy img_np, device, dtype = self.prep_batch_image(output_batch, b) - + # Draw all lines for this image for i in indices: _x1, _y1 = x1_list[i], y1_list[i] _x2, _y2 = x2_list[i], y2_list[i] - + if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _x2, _y2]): continue - + # Convert to pixel coordinates if needed if is_normalized: px1 = int(CoordinateSystem.denormalize(_x1, width, CoordinateSystem.PIXEL)) @@ -221,15 +226,15 @@ class DrawingEngine: py1 = int(max(0, min(height - 1, _y1))) px2 = int(max(0, min(width - 1, _x2))) py2 = int(max(0, min(height - 1, _y2))) - + # Draw line cv2.line(img_np, (px1, py1), (px2, py2), bgr_color, thickness) - + # Optional features if draw_endpoints: cv2.circle(img_np, (px1, py1), point_radius, bgr_color, -1) cv2.circle(img_np, (px2, py2), point_radius, bgr_color, -1) - + if draw_label and label_text: # Calculate label position if label_position == "Midpoint": @@ -239,43 +244,47 @@ class DrawingEngine: label_x, label_y = px1, py1 else: # "End" label_x, label_y = px2, py2 - + # Get text size text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) - + # Draw background rectangle - cv2.rectangle(img_np, - (label_x - 5, label_y - text_size[1] - 5), - (label_x + text_size[0] + 5, label_y + 5), - (50, 50, 50), -1) - + cv2.rectangle( + img_np, + (label_x - 5, label_y - text_size[1] - 5), + (label_x + text_size[0] + 5, label_y + 5), + (50, 50, 50), + -1, + ) + # Draw text - cv2.putText(img_np, label_text, (label_x, label_y), - cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 1) - + cv2.putText(img_np, label_text, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 1) + # Convert back to tensor and update batch output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) - + return output_batch - - def draw_polygon(self, - image: torch.Tensor, - points_x: List[float], - points_y: List[float], - color: str = "#0000FF", - thickness: int = 2, - fill: bool = False, - is_normalized: bool = True, - batch_mapping: str = "broadcast", - draw_vertices: bool = False, - vertex_radius: int = 3, - draw_label: bool = False, - label_text: str = "", - label_position: str = "Center", - font_scale: float = 0.5) -> torch.Tensor: + + def draw_polygon( + self, + image: torch.Tensor, + points_x: List[float], + points_y: List[float], + color: str = "#0000FF", + thickness: int = 2, + fill: bool = False, + is_normalized: bool = True, + batch_mapping: str = "broadcast", + draw_vertices: bool = False, + vertex_radius: int = 3, + draw_label: bool = False, + label_text: str = "", + label_position: str = "Center", + font_scale: float = 0.5, + ) -> torch.Tensor: """ Draw polygon on images efficiently. - + Args: image: BHWC tensor points_x/y: Lists of polygon vertex coordinates @@ -285,46 +294,46 @@ class DrawingEngine: is_normalized: If True, coords are 0-1 normalized batch_mapping: How to map polygons to batch Other params: Styling options - + Returns: Modified image tensor """ batch_size, height, width, _ = image.shape output_batch = image.clone() bgr_color = self.hex_to_bgr(color) - + # Validate input lengths if len(points_x) != len(points_y): raise ValueError("X and Y coordinate lists must have the same length") - + # Ensure lists have at least 3 points for a polygon if len(points_x) < 3: logger.warning("Polygon needs at least 3 points. Drawing nothing.") return output_batch - + # Determine batch mapping if batch_mapping == "one-to-one" and batch_size > 1: logger.warning(f"One-to-one mapping not supported for polygons. Using broadcast.") batch_mapping = "broadcast" - + # Process each image in batch for b in range(batch_size): # Skip processing based on mapping strategy if batch_mapping == "all-on-first" and b > 0: continue - + # Single efficient conversion to numpy img_np, device, dtype = self.prep_batch_image(output_batch, b) - + # Convert polygon points to pixel coordinates pixel_points = [] for i in range(len(points_x)): x_coord = points_x[i] y_coord = points_y[i] - + if not isinstance(x_coord, (float, int)) or not isinstance(y_coord, (float, int)): continue - + # Convert to pixel coordinates if needed if is_normalized: px = int(CoordinateSystem.denormalize(x_coord, width, CoordinateSystem.PIXEL)) @@ -332,47 +341,56 @@ class DrawingEngine: else: px = int(max(0, min(width - 1, x_coord))) py = int(max(0, min(height - 1, y_coord))) - + pixel_points.append((px, py)) - + # Draw vertices if requested if draw_vertices: cv2.circle(img_np, (px, py), vertex_radius, bgr_color, -1) - + # Skip if no valid points if len(pixel_points) < 3: continue - + # Convert to numpy array for OpenCV pts = np.array(pixel_points, dtype=np.int32) - + # Draw filled or outline polygon if fill: cv2.fillPoly(img_np, [pts], bgr_color) else: cv2.polylines(img_np, [pts], True, bgr_color, thickness) - + # Draw label if requested if draw_label and label_text: # Calculate label position (center of polygon) center_x = int(sum(p[0] for p in pixel_points) / len(pixel_points)) center_y = int(sum(p[1] for p in pixel_points) / len(pixel_points)) - + # Get text size text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) - + # Draw background rectangle - cv2.rectangle(img_np, - (center_x - text_size[0]//2 - 5, center_y - text_size[1]//2 - 5), - (center_x + text_size[0]//2 + 5, center_y + text_size[1]//2 + 5), - (50, 50, 50), -1) - + cv2.rectangle( + img_np, + (center_x - text_size[0] // 2 - 5, center_y - text_size[1] // 2 - 5), + (center_x + text_size[0] // 2 + 5, center_y + text_size[1] // 2 + 5), + (50, 50, 50), + -1, + ) + # Draw text - cv2.putText(img_np, label_text, - (center_x - text_size[0]//2, center_y + text_size[1]//2), - cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 1) - + cv2.putText( + img_np, + label_text, + (center_x - text_size[0] // 2, center_y + text_size[1] // 2), + cv2.FONT_HERSHEY_SIMPLEX, + font_scale, + (255, 255, 255), + 1, + ) + # Convert back to tensor and update batch output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) - - return output_batch \ No newline at end of file + + return output_batch diff --git a/src/mediapipe_vision/common/position.py b/src/mediapipe_vision/common/position.py index 323811c..f3fda94 100644 --- a/src/mediapipe_vision/common/position.py +++ b/src/mediapipe_vision/common/position.py @@ -1,7 +1,7 @@ import logging from typing import List, Optional, Tuple -#TODO: consider moving base nodes to src +# TODO: consider moving base nodes to src # Assuming types are correctly imported relative to this file's location from ..types import LandmarkPoint @@ -9,11 +9,13 @@ from ..types import LandmarkPoint logger = logging.getLogger(__name__) -MAX_POSITION_HISTORY = 50 # Matches MAX_LANDMARK_HISTORY for consistency -FLOAT_EQUALITY_TOLERANCE = 1e-6 # Same tolerance as base_delta_nodes -#TODO refactor to use CoordinateSystem +MAX_POSITION_HISTORY = 50 # Matches MAX_LANDMARK_HISTORY for consistency +FLOAT_EQUALITY_TOLERANCE = 1e-6 # Same tolerance as base_delta_nodes + + +# TODO refactor to use CoordinateSystem class LandmarkPositionBase: - #NOTE: use x,y,and z for the landmark position rather than LandmarkPoint object for max flexibility + # NOTE: use x,y,and z for the landmark position rather than LandmarkPoint object for max flexibility """ Base class for extracting position lists from specific landmark types. Outputs values corresponding to the input batch size, using defaults for missing items. @@ -21,27 +23,30 @@ class LandmarkPositionBase: Processes only the detected object at the specified result_index per batch item. RETURN_TYPES declare the semantic type (FLOAT/BOOLEAN); the actual return value may be a list. """ + FUNCTION = "extract_positions" # Use base types for semantic declaration RETURN_TYPES = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "BOOLEAN") - RETURN_NAMES = ("x", "y", "z", "visibility", "presence", "is_valid") # Singular names - CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Utils" + RETURN_NAMES = ("x", "y", "z", "visibility", "presence", "is_valid") # Singular names + CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Utils" DEFAULT_VALUES = (0.0, 0.0, 0.0, 0.0, 0.0) - LANDMARKS_TYPE = "ANY_LANDMARKS" # Placeholder + LANDMARKS_TYPE = "ANY_LANDMARKS" # Placeholder # INPUT_TYPES are defined in subclasses - def extract_positions(self, landmarks: list, landmark_index: int, result_index: int = 0, use_world_coordinates: bool = False) -> Tuple[List[float], List[float], List[float], List[float], List[float], List[bool]]: + def extract_positions( + self, landmarks: list, landmark_index: int, result_index: int = 0, use_world_coordinates: bool = False + ) -> Tuple[List[float], List[float], List[float], List[float], List[float], List[bool]]: """Extracts position and properties for the specified landmark across the entire batch, - padding with defaults and providing validity flags for missing items.""" + padding with defaults and providing validity flags for missing items.""" x_coords = [] y_coords = [] z_coords = [] vis_coords = [] pres_coords = [] - is_valid_list = [] # New list to track validity + is_valid_list = [] # New list to track validity if not isinstance(landmarks, list): raise TypeError(f"{self.__class__.__name__}: Input landmarks is not a list. Received type: {type(landmarks)}.") @@ -51,12 +56,12 @@ class LandmarkPositionBase: # Iterate through the batch (List[...]) for batch_idx in range(batch_size): item_list = landmarks[batch_idx] - + current_x, current_y, current_z, current_vis, current_pres = self.DEFAULT_VALUES - is_valid = False # Assume invalid until proven otherwise + is_valid = False # Assume invalid until proven otherwise landmark_point: Optional[LandmarkPoint] = None target_landmarks_list: Optional[List[LandmarkPoint]] = None - + # Check if item_list is valid and contains at least one result item if isinstance(item_list, list) and item_list: # Get the result at the specified index if it exists @@ -65,21 +70,25 @@ class LandmarkPositionBase: if result_item is not None: try: - # --- Determine Structure and Extract Landmark List --- + # --- Determine Structure and Extract Landmark List --- # Case 1: result_item is the list itself (FACE_LANDMARKS) if isinstance(result_item, list): target_landmarks_list = result_item if use_world_coordinates and self.LANDMARKS_TYPE == "FACE_LANDMARKS": - logger.warning(f"World coordinates requested but not available for FACE_LANDMARKS. Using image coordinates.") + logger.warning( + f"World coordinates requested but not available for FACE_LANDMARKS. Using image coordinates." + ) # Case 2: result_item is an object with .landmarks (HAND/POSE_LANDMARKS) - elif hasattr(result_item, 'landmarks'): - landmarks_attr = getattr(result_item, 'landmarks', None) + elif hasattr(result_item, "landmarks"): + landmarks_attr = getattr(result_item, "landmarks", None) if landmarks_attr is not None: - use_world = (use_world_coordinates and - self.LANDMARKS_TYPE != "FACE_LANDMARKS" and - hasattr(result_item, 'world_landmarks') and - getattr(result_item, 'world_landmarks', None)) + use_world = ( + use_world_coordinates + and self.LANDMARKS_TYPE != "FACE_LANDMARKS" + and hasattr(result_item, "world_landmarks") + and getattr(result_item, "world_landmarks", None) + ) if use_world: target_landmarks_list = result_item.world_landmarks else: @@ -89,22 +98,27 @@ class LandmarkPositionBase: # --- Find Landmark Point by Index --- if isinstance(target_landmarks_list, list) and target_landmarks_list: landmark_point = next((lm for lm in target_landmarks_list if lm.index == landmark_index), None) - + except Exception as e: - logger.error(f"{self.__class__.__name__}: Unexpected error processing batch item {batch_idx}: {e}", exc_info=True) - # Keep landmark_point as None, is_valid as False + logger.error( + f"{self.__class__.__name__}: Unexpected error processing batch item {batch_idx}: {e}", + exc_info=True, + ) + # Keep landmark_point as None, is_valid as False else: - logger.debug(f"{self.__class__.__name__}: Result index {result_index} exceeds available detections for batch item {batch_idx} (has {len(item_list)})") + logger.debug( + f"{self.__class__.__name__}: Result index {result_index} exceeds available detections for batch item {batch_idx} (has {len(item_list)})" + ) # else: item_list was not a list or was empty, keep defaults/invalid # --- Extract Values or Use Defaults, Set Validity --- if landmark_point: - is_valid = True # Mark as valid since we found the point + is_valid = True # Mark as valid since we found the point current_x = landmark_point.x current_y = landmark_point.y current_z = landmark_point.z - current_vis = getattr(landmark_point, 'visibility', None) - current_pres = getattr(landmark_point, 'presence', None) + current_vis = getattr(landmark_point, "visibility", None) + current_pres = getattr(landmark_point, "presence", None) current_vis = current_vis if current_vis is not None else 0.0 current_pres = current_pres if current_pres is not None else 0.0 # else: defaults are used, is_valid remains False @@ -120,5 +134,5 @@ class LandmarkPositionBase: # Check if each list has only one element, if so return single values instead of lists if len(x_coords) == 1: return (x_coords[0], y_coords[0], z_coords[0], vis_coords[0], pres_coords[0], is_valid_list[0]) - + return (x_coords, y_coords, z_coords, vis_coords, pres_coords, is_valid_list) diff --git a/src/mediapipe_vision/face_detection/__init__.py b/src/mediapipe_vision/face_detection/__init__.py index fd9b0f9..cc80784 100644 --- a/src/mediapipe_vision/face_detection/__init__.py +++ b/src/mediapipe_vision/face_detection/__init__.py @@ -1 +1 @@ -"""Face Detection implementation using MediaPipe.""" \ No newline at end of file +"""Face Detection implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/face_detection/detector.py b/src/mediapipe_vision/face_detection/detector.py index 75272ce..f9919c3 100644 --- a/src/mediapipe_vision/face_detection/detector.py +++ b/src/mediapipe_vision/face_detection/detector.py @@ -32,28 +32,33 @@ class FaceDetector: self.model_path = model_path self._detector_instance = None - self._current_options = None # Store current options for comparison - self._timestamp_provider = None # Initialize TimestampProvider later + self._current_options = None # Store current options for comparison + self._timestamp_provider = None # Initialize TimestampProvider later def _initialize_detector(self, min_detection_confidence: float, running_mode: str, delegate: str): """Initializes the FaceDetector.""" # Map string mode to MediaPipe enum mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU - + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU + base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.FaceDetectorOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode - min_detection_confidence=min_detection_confidence + running_mode=mode_enum, # Use enum mode + min_detection_confidence=min_detection_confidence, ) - self._current_options = options # Store for comparison + self._current_options = options # Store for comparison # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.FaceDetector.create_from_options(options) - def detect(self, image: torch.Tensor, min_detection_confidence: float = 0.5, - running_mode: str = "video", delegate: str = 'cpu') -> List[List[FaceDetectionResult]]: # Added running_mode + def detect( + self, + image: torch.Tensor, + min_detection_confidence: float = 0.5, + running_mode: str = "video", + delegate: str = "cpu", + ) -> List[List[FaceDetectionResult]]: # Added running_mode """Detects faces in the input image tensor. Returns: @@ -67,15 +72,15 @@ class FaceDetector: new_options_tuple = (min_detection_confidence, running_mode, delegate.lower()) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.min_detection_confidence, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' + current_options_tuple = ( + self._current_options.min_detection_confidence, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", ) if self._detector_instance is None or new_options_tuple != current_options_tuple: # Close previous instance if it exists and has a close method - if self._detector_instance and hasattr(self._detector_instance, 'close'): + if self._detector_instance and hasattr(self._detector_instance, "close"): self._detector_instance.close() self._detector_instance = self._initialize_detector(min_detection_confidence, running_mode, delegate) @@ -86,7 +91,7 @@ class FaceDetector: detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -97,39 +102,23 @@ class FaceDetector: else: # Call image mode function detection_result = detect_func(mp_image) - + current_image_detections = [] - if detection_result and detection_result.detections: # Check if result is not None + if detection_result and detection_result.detections: # Check if result is not None for detection in detection_result.detections: bbox_mp = detection.bounding_box keypoints_mp = detection.keypoints - score = detection.categories[0].score if detection.categories else None # Face detection usually has one score - - bbox = BoundingBox( - origin_x=bbox_mp.origin_x, - origin_y=bbox_mp.origin_y, - width=bbox_mp.width, - height=bbox_mp.height - ) - + score = detection.categories[0].score if detection.categories else None # Face detection usually has one score + + bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) + keypoints = None if keypoints_mp: - keypoints = [ - FaceKeypoint( - label=kp.label, - x=kp.x, - y=kp.y - ) - for kp in keypoints_mp - ] + keypoints = [FaceKeypoint(label=kp.label, x=kp.x, y=kp.y) for kp in keypoints_mp] + + current_image_detections.append(FaceDetectionResult(bounding_box=bbox, keypoints=keypoints, score=score)) - current_image_detections.append(FaceDetectionResult( - bounding_box=bbox, - keypoints=keypoints, - score=score - )) - batch_results.append(current_image_detections) # Check MediaPipe documentation if resource cleanup is required here. - return batch_results \ No newline at end of file + return batch_results diff --git a/src/mediapipe_vision/face_landmark/__init__.py b/src/mediapipe_vision/face_landmark/__init__.py index 7830f06..9d9dcfb 100644 --- a/src/mediapipe_vision/face_landmark/__init__.py +++ b/src/mediapipe_vision/face_landmark/__init__.py @@ -1 +1 @@ -"""Face Landmark Detection implementation using MediaPipe.""" \ No newline at end of file +"""Face Landmark Detection implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/face_landmark/detector.py b/src/mediapipe_vision/face_landmark/detector.py index f04601b..c828bcb 100644 --- a/src/mediapipe_vision/face_landmark/detector.py +++ b/src/mediapipe_vision/face_landmark/detector.py @@ -26,38 +26,53 @@ class FaceLandmarkDetector: raise ValueError("A valid model_path must be provided.") self.model_path = model_path - self._detector_instance = None - self._current_options = None # Store current options for comparison - self._timestamp_provider = None # Added + self._detector_instance = None + self._current_options = None # Store current options for comparison + self._timestamp_provider = None # Added - def _initialize_detector(self, num_faces: int, min_detection_confidence: float, - min_presence_confidence: float, min_tracking_confidence: float, - output_blendshapes: bool, output_transform_matrix: bool, - running_mode: str, delegate: str): # Added running_mode + def _initialize_detector( + self, + num_faces: int, + min_detection_confidence: float, + min_presence_confidence: float, + min_tracking_confidence: float, + output_blendshapes: bool, + output_transform_matrix: bool, + running_mode: str, + delegate: str, + ): # Added running_mode """Initializes the FaceLandmarker detector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.FaceLandmarkerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode num_faces=num_faces, min_face_detection_confidence=min_detection_confidence, # min_face_presence_confidence and min_tracking_confidence are for VIDEO/LIVE_STREAM - min_face_presence_confidence=min_presence_confidence, - min_tracking_confidence=min_tracking_confidence, + min_face_presence_confidence=min_presence_confidence, + min_tracking_confidence=min_tracking_confidence, output_face_blendshapes=output_blendshapes, - output_facial_transformation_matrixes=output_transform_matrix + output_facial_transformation_matrixes=output_transform_matrix, ) - self._current_options = options # Store for comparison + self._current_options = options # Store for comparison # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.FaceLandmarker.create_from_options(options) - def detect(self, image: torch.Tensor, num_faces: int = 1, min_detection_confidence: float = 0.5, - min_presence_confidence: float = 0.5, min_tracking_confidence: float = 0.5, - output_blendshapes: bool = False, output_transform_matrix: bool = False, - running_mode: str = "video", delegate: str = 'cpu') -> Tuple[List[List[LandmarkPoint]], Optional[List[Dict[str, float]]], Optional[List[np.ndarray]]]: # Added running_mode + def detect( + self, + image: torch.Tensor, + num_faces: int = 1, + min_detection_confidence: float = 0.5, + min_presence_confidence: float = 0.5, + min_tracking_confidence: float = 0.5, + output_blendshapes: bool = False, + output_transform_matrix: bool = False, + running_mode: str = "video", + delegate: str = "cpu", + ) -> Tuple[List[List[LandmarkPoint]], Optional[List[Dict[str, float]]], Optional[List[np.ndarray]]]: # Added running_mode """Detects face landmarks in the input image tensor. Returns: @@ -68,32 +83,46 @@ class FaceLandmarkDetector: """ if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - - # Check if options changed + + # Check if options changed # Add running_mode to tuple check - new_options_tuple = (num_faces, min_detection_confidence, min_presence_confidence, min_tracking_confidence, - output_blendshapes, output_transform_matrix, running_mode, delegate.lower()) + new_options_tuple = ( + num_faces, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + output_blendshapes, + output_transform_matrix, + running_mode, + delegate.lower(), + ) current_options_tuple = None if self._current_options: current_options_tuple = ( self._current_options.num_faces, self._current_options.min_face_detection_confidence, - self._current_options.min_face_presence_confidence, - self._current_options.min_tracking_confidence, + self._current_options.min_face_presence_confidence, + self._current_options.min_tracking_confidence, self._current_options.output_face_blendshapes, self._current_options.output_facial_transformation_matrixes, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", ) if self._detector_instance is None or new_options_tuple != current_options_tuple: - if self._detector_instance and hasattr(self._detector_instance, 'close'): # Close previous instance if options changed - self._detector_instance.close() + if self._detector_instance and hasattr(self._detector_instance, "close"): # Close previous instance if options changed + self._detector_instance.close() # Pass new params to initializer - self._detector_instance = self._initialize_detector(num_faces, min_detection_confidence, - min_presence_confidence, min_tracking_confidence, - output_blendshapes, output_transform_matrix, - running_mode, delegate) # Pass running_mode + self._detector_instance = self._initialize_detector( + num_faces, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + output_blendshapes, + output_transform_matrix, + running_mode, + delegate, + ) # Pass running_mode batch_size = image.shape[0] batch_results_landmarks = [] @@ -103,9 +132,9 @@ class FaceLandmarkDetector: # Determine the correct detection function based on mode is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect - + for i in range(batch_size): - img_tensor = image[i] + img_tensor = image[i] np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -114,24 +143,30 @@ class FaceLandmarkDetector: detection_result = detect_func(mp_image, timestamp_ms) else: detection_result = detect_func(mp_image) - + current_image_landmarks = [] current_image_blendshapes = [] if output_blendshapes else None current_image_matrices = [] if output_transform_matrix else None - if detection_result and detection_result.face_landmarks: # Added check for detection_result + if detection_result and detection_result.face_landmarks: # Added check for detection_result for face_idx, face_landmarks_mp in enumerate(detection_result.face_landmarks): landmarks = [LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(face_landmarks_mp)] current_image_landmarks.append(landmarks) if output_blendshapes and detection_result.face_blendshapes and face_idx < len(detection_result.face_blendshapes): - blendshapes = {bs.category_name: bs.score for bs in detection_result.face_blendshapes[face_idx]} - current_image_blendshapes.append(blendshapes) - if output_transform_matrix and detection_result.facial_transformation_matrixes and face_idx < len(detection_result.facial_transformation_matrixes): - current_image_matrices.append(detection_result.facial_transformation_matrixes[face_idx]) + blendshapes = {bs.category_name: bs.score for bs in detection_result.face_blendshapes[face_idx]} + current_image_blendshapes.append(blendshapes) + if ( + output_transform_matrix + and detection_result.facial_transformation_matrixes + and face_idx < len(detection_result.facial_transformation_matrixes) + ): + current_image_matrices.append(detection_result.facial_transformation_matrixes[face_idx]) batch_results_landmarks.append(current_image_landmarks) - if output_blendshapes: batch_results_blendshapes.append(current_image_blendshapes) - if output_transform_matrix: batch_results_matrices.append(current_image_matrices) - + if output_blendshapes: + batch_results_blendshapes.append(current_image_blendshapes) + if output_transform_matrix: + batch_results_matrices.append(current_image_matrices) + # Close detector after batch? Landmarker might need closing per batch? Check docs. # For IMAGE mode, closing after batch seems okay. # self._detector_instance.close() # Close is handled if re-initialized - return batch_results_landmarks, batch_results_blendshapes, batch_results_matrices \ No newline at end of file + return batch_results_landmarks, batch_results_blendshapes, batch_results_matrices diff --git a/src/mediapipe_vision/face_stylization/__init__.py b/src/mediapipe_vision/face_stylization/__init__.py index aa53844..7de2b07 100644 --- a/src/mediapipe_vision/face_stylization/__init__.py +++ b/src/mediapipe_vision/face_stylization/__init__.py @@ -1 +1 @@ -"""Face Stylization implementation using MediaPipe.""" \ No newline at end of file +"""Face Stylization implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/face_stylization/detector.py b/src/mediapipe_vision/face_stylization/detector.py index 03fdfaa..8681dba 100644 --- a/src/mediapipe_vision/face_stylization/detector.py +++ b/src/mediapipe_vision/face_stylization/detector.py @@ -28,12 +28,12 @@ class FaceStylizer: options = vision.FaceStylizerOptions( base_options=base_options # running_mode=vision.RunningMode.VIDEO # Not supported - ) + ) self.current_options = () self._timestamp_provider = TimestampProvider() return vision.FaceStylizer.create_from_options(options) - def stylize(self, image: torch.Tensor) -> torch.Tensor: # Output is IMAGE + def stylize(self, image: torch.Tensor) -> torch.Tensor: # Output is IMAGE """Applies face stylization to the input image tensor(s).""" if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") @@ -45,24 +45,24 @@ class FaceStylizer: batch_results = [] for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) stylization_result = self.detector_instance.stylize(mp_image) - + output_image_tensor = None if stylization_result: - # Convert MediaPipe Image back to NumPy (HWC, uint8) - stylized_np = stylization_result.numpy_view() - # Convert to float tensor (HWC, 0-1) - output_image_tensor = torch.from_numpy(stylized_np).float().div(255.0) + # Convert MediaPipe Image back to NumPy (HWC, uint8) + stylized_np = stylization_result.numpy_view() + # Convert to float tensor (HWC, 0-1) + output_image_tensor = torch.from_numpy(stylized_np).float().div(255.0) else: - # Handle cases where stylization fails (e.g., no face detected) - # Return original image tensor for this batch item - output_image_tensor = img_tensor - + # Handle cases where stylization fails (e.g., no face detected) + # Return original image tensor for this batch item + output_image_tensor = img_tensor + batch_results.append(output_image_tensor) output_batch_tensor = torch.stack(batch_results, dim=0) - return output_batch_tensor \ No newline at end of file + return output_batch_tensor diff --git a/src/mediapipe_vision/gesture_recognition/__init__.py b/src/mediapipe_vision/gesture_recognition/__init__.py index 81fdcae..996e4e5 100644 --- a/src/mediapipe_vision/gesture_recognition/__init__.py +++ b/src/mediapipe_vision/gesture_recognition/__init__.py @@ -1 +1 @@ -"""Gesture Recognition implementation using MediaPipe.""" \ No newline at end of file +"""Gesture Recognition implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/gesture_recognition/detector.py b/src/mediapipe_vision/gesture_recognition/detector.py index 3257720..64c6338 100644 --- a/src/mediapipe_vision/gesture_recognition/detector.py +++ b/src/mediapipe_vision/gesture_recognition/detector.py @@ -22,57 +22,81 @@ class GestureRecognizer: self.model_path = model_path self._detector_instance = None self._current_options = None - self._timestamp_provider = None # Added + self._timestamp_provider = None # Added - def _initialize_detector(self, num_hands: int, min_detection_confidence: float, - min_tracking_confidence: float, min_presence_confidence: float, - running_mode: str, delegate: str): # Added running_mode - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + def _initialize_detector( + self, + num_hands: int, + min_detection_confidence: float, + min_tracking_confidence: float, + min_presence_confidence: float, + running_mode: str, + delegate: str, + ): # Added running_mode + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) # Note: GestureRecognizerOptions has min_tracking_confidence, useful for VIDEO mode options = vision.GestureRecognizerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode num_hands=num_hands, min_hand_detection_confidence=min_detection_confidence, - min_tracking_confidence=min_tracking_confidence, - min_hand_presence_confidence=min_presence_confidence + min_tracking_confidence=min_tracking_confidence, + min_hand_presence_confidence=min_presence_confidence, # Add other relevant VIDEO mode parameters if needed (e.g., canned_gestures_classifier_options) ) self._current_options = options # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.GestureRecognizer.create_from_options(options) - def recognize(self, image: torch.Tensor, num_hands: int = 2, min_detection_confidence: float = 0.5, - min_tracking_confidence: float = 0.5, min_presence_confidence: float = 0.5, - running_mode: str = "video", delegate: str = 'cpu') -> List[List[GestureRecognitionResult]]: # Added running_mode + def recognize( + self, + image: torch.Tensor, + num_hands: int = 2, + min_detection_confidence: float = 0.5, + min_tracking_confidence: float = 0.5, + min_presence_confidence: float = 0.5, + running_mode: str = "video", + delegate: str = "cpu", + ) -> List[List[GestureRecognitionResult]]: # Added running_mode """Recognizes hand gestures in the input image tensor.""" if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - + # Add running_mode to tuple check - new_options_tuple = (num_hands, min_detection_confidence, min_tracking_confidence, - min_presence_confidence, running_mode, delegate.lower()) + new_options_tuple = ( + num_hands, + min_detection_confidence, + min_tracking_confidence, + min_presence_confidence, + running_mode, + delegate.lower(), + ) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.num_hands, - self._current_options.min_hand_detection_confidence, - self._current_options.min_tracking_confidence, - self._current_options.min_hand_presence_confidence, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' - ) + current_options_tuple = ( + self._current_options.num_hands, + self._current_options.min_hand_detection_confidence, + self._current_options.min_tracking_confidence, + self._current_options.min_hand_presence_confidence, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", + ) if self._detector_instance is None or new_options_tuple != current_options_tuple: - if self._detector_instance and hasattr(self._detector_instance, 'close'): - self._detector_instance.close() + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() # Pass new params to initializer - self._detector_instance = self._initialize_detector(num_hands, min_detection_confidence, - min_tracking_confidence, min_presence_confidence, - running_mode, delegate) # Pass running_mode + self._detector_instance = self._initialize_detector( + num_hands, + min_detection_confidence, + min_tracking_confidence, + min_presence_confidence, + running_mode, + delegate, + ) # Pass running_mode batch_size = image.shape[0] batch_results = [] @@ -81,7 +105,7 @@ class GestureRecognizer: recognize_func = self._detector_instance.recognize_for_video if is_video_mode else self._detector_instance.recognize for i in range(batch_size): - img_tensor = image[i] + img_tensor = image[i] np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -90,17 +114,26 @@ class GestureRecognizer: recognition_result = recognize_func(mp_image, timestamp_ms) else: recognition_result = recognize_func(mp_image) - + current_image_gestures = [] - if recognition_result and recognition_result.gestures: # Added check for recognition_result + if recognition_result and recognition_result.gestures: # Added check for recognition_result for hand_idx, hand_gestures in enumerate(recognition_result.gestures): - gesture_categories = [GestureCategory(index=cat.index, score=cat.score, display_name=cat.display_name, category_name=cat.category_name) for cat in hand_gestures] + gesture_categories = [ + GestureCategory( + index=cat.index, + score=cat.score, + display_name=cat.display_name, + category_name=cat.category_name, + ) + for cat in hand_gestures + ] handedness = None if recognition_result.handedness and hand_idx < len(recognition_result.handedness): handedness_cats = recognition_result.handedness[hand_idx] - if handedness_cats: handedness = handedness_cats[0].display_name + if handedness_cats: + handedness = handedness_cats[0].display_name current_image_gestures.append(GestureRecognitionResult(gestures=gesture_categories, handedness=handedness)) batch_results.append(current_image_gestures) # self._detector_instance.close() # Close handled on re-init - return batch_results \ No newline at end of file + return batch_results diff --git a/src/mediapipe_vision/hand_landmark/__init__.py b/src/mediapipe_vision/hand_landmark/__init__.py index fd70627..699a703 100644 --- a/src/mediapipe_vision/hand_landmark/__init__.py +++ b/src/mediapipe_vision/hand_landmark/__init__.py @@ -1 +1 @@ -"""Hand Landmark Detection implementation using MediaPipe.""" \ No newline at end of file +"""Hand Landmark Detection implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/hand_landmark/detector.py b/src/mediapipe_vision/hand_landmark/detector.py index d577dec..420d996 100644 --- a/src/mediapipe_vision/hand_landmark/detector.py +++ b/src/mediapipe_vision/hand_landmark/detector.py @@ -15,7 +15,7 @@ from ..types import HandLandmarksResult, LandmarkPoint class HandLandmarkDetector: """Detects hand landmarks using MediaPipe HandLandmarker.""" - + def __init__(self, model_path: str): """Initialize the detector with the model path. @@ -25,36 +25,49 @@ class HandLandmarkDetector: if not model_path: raise ValueError("A valid model_path must be provided.") # Verify file exists? Optional, MediaPipe might handle this. - + self.model_path = model_path self._detector_instance = None self._current_options = None - self._timestamp_provider = None # Added + self._timestamp_provider = None # Added - def _initialize_detector(self, num_hands: int, min_detection_confidence: float, - min_presence_confidence: float, min_tracking_confidence: float, - running_mode: str, delegate: str): # Added running_mode + def _initialize_detector( + self, + num_hands: int, + min_detection_confidence: float, + min_presence_confidence: float, + min_tracking_confidence: float, + running_mode: str, + delegate: str, + ): # Added running_mode """Initializes the HandLandmarker detector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.HandLandmarkerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode num_hands=num_hands, min_hand_detection_confidence=min_detection_confidence, # min_hand_presence_confidence and min_tracking_confidence are for VIDEO/LIVE_STREAM min_hand_presence_confidence=min_presence_confidence, - min_tracking_confidence=min_tracking_confidence + min_tracking_confidence=min_tracking_confidence, ) self._current_options = options # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.HandLandmarker.create_from_options(options) - def detect(self, image: torch.Tensor, num_hands: int = 2, min_detection_confidence: float = 0.5, - min_presence_confidence: float = 0.5, min_tracking_confidence: float = 0.5, - running_mode: str = "video", delegate: str = 'cpu') -> List[List[HandLandmarksResult]]: # Added running_mode + def detect( + self, + image: torch.Tensor, + num_hands: int = 2, + min_detection_confidence: float = 0.5, + min_presence_confidence: float = 0.5, + min_tracking_confidence: float = 0.5, + running_mode: str = "video", + delegate: str = "cpu", + ) -> List[List[HandLandmarksResult]]: # Added running_mode """Detects hand landmarks in the input image tensor. Returns: @@ -64,72 +77,78 @@ class HandLandmarkDetector: raise ValueError("Input tensor must be in BHWC format.") # Add running_mode to tuple check - new_options_tuple = (num_hands, min_detection_confidence, - min_presence_confidence, min_tracking_confidence, running_mode, delegate.lower()) + new_options_tuple = ( + num_hands, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + running_mode, + delegate.lower(), + ) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.num_hands, - self._current_options.min_hand_detection_confidence, - self._current_options.min_hand_presence_confidence, - self._current_options.min_tracking_confidence, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' - ) + current_options_tuple = ( + self._current_options.num_hands, + self._current_options.min_hand_detection_confidence, + self._current_options.min_hand_presence_confidence, + self._current_options.min_tracking_confidence, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", + ) if self._detector_instance is None or new_options_tuple != current_options_tuple: - if self._detector_instance and hasattr(self._detector_instance, 'close'): - self._detector_instance.close() - # Pass new params to initializer - self._detector_instance = self._initialize_detector(num_hands, min_detection_confidence, - min_presence_confidence, min_tracking_confidence, - running_mode, delegate) # Pass running_mode + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() + # Pass new params to initializer + self._detector_instance = self._initialize_detector( + num_hands, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + running_mode, + delegate, + ) # Pass running_mode batch_size = image.shape[0] batch_results = [] # Determine the correct detection function based on mode is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect - + for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) - + if is_video_mode: timestamp_ms = self._timestamp_provider.next() detection_result = detect_func(mp_image, timestamp_ms) else: detection_result = detect_func(mp_image) - + current_image_results = [] - if detection_result and detection_result.hand_landmarks: # Added check for detection_result + if detection_result and detection_result.hand_landmarks: # Added check for detection_result for hand_idx, hand_landmarks_mp in enumerate(detection_result.hand_landmarks): - - landmarks = [ - LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) - for lm_idx, lm in enumerate(hand_landmarks_mp) - ] - + landmarks = [LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(hand_landmarks_mp)] + world_landmarks = None if detection_result.hand_world_landmarks and hand_idx < len(detection_result.hand_world_landmarks): world_landmarks = [ LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z) for lm_idx, lm in enumerate(detection_result.hand_world_landmarks[hand_idx]) ] - + handedness = None if detection_result.handedness and hand_idx < len(detection_result.handedness): handedness_cats = detection_result.handedness[hand_idx] - if handedness_cats: handedness = handedness_cats[0].display_name - - current_image_results.append(HandLandmarksResult( - landmarks=landmarks, - world_landmarks=world_landmarks, - handedness=handedness - )) - + if handedness_cats: + handedness = handedness_cats[0].display_name + + current_image_results.append( + HandLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks, handedness=handedness) + ) + batch_results.append(current_image_results) # self._detector_instance.close() # Close handled on re-init - return batch_results \ No newline at end of file + return batch_results diff --git a/src/mediapipe_vision/holistic_landmark/detector.py b/src/mediapipe_vision/holistic_landmark/detector.py index 8f7176f..3c12f6c 100644 --- a/src/mediapipe_vision/holistic_landmark/detector.py +++ b/src/mediapipe_vision/holistic_landmark/detector.py @@ -19,14 +19,14 @@ from ..types import HolisticLandmarksResult, LandmarkPoint class HolisticLandmarkDetector: """Detects holistic landmarks combining face, pose and hand tracking. - + This detector uses the legacy MediaPipe Solutions API rather than the Task API since Holistic is not yet available in the Task API format. """ - + def __init__(self, model_path: Optional[str] = None): """Initialize the detector. - + Args: model_path: Unused parameter to maintain API compatibility with other detectors. The holistic solution doesn't use a model path as it uses the built-in models. @@ -36,15 +36,17 @@ class HolisticLandmarkDetector: self._current_options = None self._timestamp_provider = TimestampProvider() - def _initialize_detector(self, - min_detection_confidence: float, - min_tracking_confidence: float, - model_complexity: int, - enable_segmentation: bool, - refine_face_landmarks: bool, - static_image_mode: bool) -> mp.solutions.holistic.Holistic: + def _initialize_detector( + self, + min_detection_confidence: float, + min_tracking_confidence: float, + model_complexity: int, + enable_segmentation: bool, + refine_face_landmarks: bool, + static_image_mode: bool, + ) -> mp.solutions.holistic.Holistic: """Initializes the Holistic detector with the appropriate parameters. - + Args: min_detection_confidence: Minimum confidence for detection to be considered min_tracking_confidence: Minimum confidence to maintain tracking @@ -52,20 +54,20 @@ class HolisticLandmarkDetector: enable_segmentation: Whether to generate segmentation masks refine_face_landmarks: Whether to use refined face landmarks (more details around eyes/lips) static_image_mode: If True, treats input as independent images (slower but more accurate) - + Returns: Initialized Holistic detector instance """ # Store options as a dictionary for parameter comparison self._current_options = { - 'min_detection_confidence': min_detection_confidence, - 'min_tracking_confidence': min_tracking_confidence, - 'model_complexity': model_complexity, - 'enable_segmentation': enable_segmentation, - 'refine_face_landmarks': refine_face_landmarks, - 'static_image_mode': static_image_mode + "min_detection_confidence": min_detection_confidence, + "min_tracking_confidence": min_tracking_confidence, + "model_complexity": model_complexity, + "enable_segmentation": enable_segmentation, + "refine_face_landmarks": refine_face_landmarks, + "static_image_mode": static_image_mode, } - + # Create the holistic solution instance with specified options return mp.solutions.holistic.Holistic( static_image_mode=static_image_mode, @@ -73,18 +75,21 @@ class HolisticLandmarkDetector: enable_segmentation=enable_segmentation, refine_face_landmarks=refine_face_landmarks, min_detection_confidence=min_detection_confidence, - min_tracking_confidence=min_tracking_confidence + min_tracking_confidence=min_tracking_confidence, ) - def detect(self, image: torch.Tensor, - min_detection_confidence: float = 0.5, - min_tracking_confidence: float = 0.5, - model_complexity: int = 1, - enable_segmentation: bool = False, - refine_face_landmarks: bool = False, - static_image_mode: bool = False) -> Tuple[List[HolisticLandmarksResult], Optional[List[torch.Tensor]]]: + def detect( + self, + image: torch.Tensor, + min_detection_confidence: float = 0.5, + min_tracking_confidence: float = 0.5, + model_complexity: int = 1, + enable_segmentation: bool = False, + refine_face_landmarks: bool = False, + static_image_mode: bool = False, + ) -> Tuple[List[HolisticLandmarksResult], Optional[List[torch.Tensor]]]: """Detects holistic landmarks in the input image tensor. - + Args: image: Input image tensor in BHWC format min_detection_confidence: Minimum confidence for detection (0.0-1.0) @@ -93,7 +98,7 @@ class HolisticLandmarkDetector: enable_segmentation: Whether to output segmentation masks refine_face_landmarks: Whether to use refined face landmarks static_image_mode: If True, treats input as independent images - + Returns: Tuple containing: - holistic_landmarks_batch: List of HolisticLandmarksResult objects for each image @@ -104,17 +109,17 @@ class HolisticLandmarkDetector: # Create a parameters tuple for comparison new_options = { - 'min_detection_confidence': min_detection_confidence, - 'min_tracking_confidence': min_tracking_confidence, - 'model_complexity': model_complexity, - 'enable_segmentation': enable_segmentation, - 'refine_face_landmarks': refine_face_landmarks, - 'static_image_mode': static_image_mode + "min_detection_confidence": min_detection_confidence, + "min_tracking_confidence": min_tracking_confidence, + "model_complexity": model_complexity, + "enable_segmentation": enable_segmentation, + "refine_face_landmarks": refine_face_landmarks, + "static_image_mode": static_image_mode, } # Initialize or update detector if needed if self._detector_instance is None or self._current_options != new_options: - if self._detector_instance and hasattr(self._detector_instance, 'close'): + if self._detector_instance and hasattr(self._detector_instance, "close"): self._detector_instance.close() self._detector_instance = self._initialize_detector( min_detection_confidence, @@ -122,42 +127,40 @@ class HolisticLandmarkDetector: model_complexity, enable_segmentation, refine_face_landmarks, - static_image_mode + static_image_mode, ) batch_size = image.shape[0] batch_holistic_landmarks = [] batch_segmentation_masks = [] if enable_segmentation else None - + for i in range(batch_size): img_tensor = image[i] # HWC # Convert to numpy with proper scaling and datatype np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) - + # Process the image with MediaPipe Holistic - holistic_results = self._detector_instance.process( - cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB) - ) - + holistic_results = self._detector_instance.process(cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)) + # Convert results to our standardized format holistic_result = self._convert_to_landmark_result(holistic_results) batch_holistic_landmarks.append(holistic_result) - + # Handle segmentation mask if requested if enable_segmentation and holistic_results.segmentation_mask is not None: mask_np = holistic_results.segmentation_mask mask_tensor = torch.from_numpy(mask_np).float() # HW if batch_segmentation_masks is not None: batch_segmentation_masks.append(mask_tensor) - + return batch_holistic_landmarks, batch_segmentation_masks def _convert_to_landmark_result(self, holistic_results: Any) -> HolisticLandmarksResult: """Converts MediaPipe Holistic results to standardized HolisticLandmarksResult. - + Args: holistic_results: Raw results from mp.solutions.holistic - + Returns: Standardized HolisticLandmarksResult containing all landmarks """ @@ -170,75 +173,53 @@ class HolisticLandmarkDetector: x=landmark.x, y=landmark.y, z=landmark.z, - visibility=getattr(landmark, 'visibility', None) + visibility=getattr(landmark, "visibility", None), ) for idx, landmark in enumerate(holistic_results.face_landmarks.landmark) ] - + # Convert pose landmarks pose_landmarks = None if holistic_results.pose_landmarks: pose_landmarks = [ - LandmarkPoint( - index=idx, - x=landmark.x, - y=landmark.y, - z=landmark.z, - visibility=landmark.visibility - ) + LandmarkPoint(index=idx, x=landmark.x, y=landmark.y, z=landmark.z, visibility=landmark.visibility) for idx, landmark in enumerate(holistic_results.pose_landmarks.landmark) ] - + # Convert pose world landmarks pose_world_landmarks = None if holistic_results.pose_world_landmarks: pose_world_landmarks = [ - LandmarkPoint( - index=idx, - x=landmark.x, - y=landmark.y, - z=landmark.z, - visibility=landmark.visibility - ) + LandmarkPoint(index=idx, x=landmark.x, y=landmark.y, z=landmark.z, visibility=landmark.visibility) for idx, landmark in enumerate(holistic_results.pose_world_landmarks.landmark) ] - + # Convert left hand landmarks left_hand_landmarks = None if holistic_results.left_hand_landmarks: left_hand_landmarks = [ - LandmarkPoint( - index=idx, - x=landmark.x, - y=landmark.y, - z=landmark.z - ) + LandmarkPoint(index=idx, x=landmark.x, y=landmark.y, z=landmark.z) for idx, landmark in enumerate(holistic_results.left_hand_landmarks.landmark) ] - + # Convert right hand landmarks right_hand_landmarks = None if holistic_results.right_hand_landmarks: right_hand_landmarks = [ - LandmarkPoint( - index=idx, - x=landmark.x, - y=landmark.y, - z=landmark.z - ) + LandmarkPoint(index=idx, x=landmark.x, y=landmark.y, z=landmark.z) for idx, landmark in enumerate(holistic_results.right_hand_landmarks.landmark) ] - + return HolisticLandmarksResult( face_landmarks=face_landmarks, pose_landmarks=pose_landmarks, pose_world_landmarks=pose_world_landmarks, left_hand_landmarks=left_hand_landmarks, - right_hand_landmarks=right_hand_landmarks + right_hand_landmarks=right_hand_landmarks, ) - + def close(self): """Releases resources used by the detector.""" - if self._detector_instance and hasattr(self._detector_instance, 'close'): + if self._detector_instance and hasattr(self._detector_instance, "close"): self._detector_instance.close() - self._detector_instance = None \ No newline at end of file + self._detector_instance = None diff --git a/src/mediapipe_vision/image_classification/__init__.py b/src/mediapipe_vision/image_classification/__init__.py index b839c10..70c9d96 100644 --- a/src/mediapipe_vision/image_classification/__init__.py +++ b/src/mediapipe_vision/image_classification/__init__.py @@ -1 +1 @@ -"""Image Classification implementation using MediaPipe.""" \ No newline at end of file +"""Image Classification implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/image_embedding/__init__.py b/src/mediapipe_vision/image_embedding/__init__.py index cf29562..6091720 100644 --- a/src/mediapipe_vision/image_embedding/__init__.py +++ b/src/mediapipe_vision/image_embedding/__init__.py @@ -1 +1 @@ -"""Image Embedding implementation using MediaPipe.""" \ No newline at end of file +"""Image Embedding implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/image_embedding/detector.py b/src/mediapipe_vision/image_embedding/detector.py index 721e34c..bf87f51 100644 --- a/src/mediapipe_vision/image_embedding/detector.py +++ b/src/mediapipe_vision/image_embedding/detector.py @@ -23,45 +23,50 @@ class ImageEmbedder: self.model_path = model_path self._detector_instance = None self._current_options = None - self._timestamp_provider = None # Added + self._timestamp_provider = None # Added - def _initialize_detector(self, l2_normalize: bool, quantize: bool, - running_mode: str, delegate: str): # Added running_mode - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + def _initialize_detector(self, l2_normalize: bool, quantize: bool, running_mode: str, delegate: str): # Added running_mode + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.ImageEmbedderOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode l2_normalize=l2_normalize, - quantize=quantize + quantize=quantize, ) self._current_options = options # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.ImageEmbedder.create_from_options(options) - def embed(self, image: torch.Tensor, l2_normalize: bool = True, quantize: bool = False, - running_mode: str = "video", delegate: str = 'cpu') -> List[List[ImageEmbedderResult]]: # Added running_mode + def embed( + self, + image: torch.Tensor, + l2_normalize: bool = True, + quantize: bool = False, + running_mode: str = "video", + delegate: str = "cpu", + ) -> List[List[ImageEmbedderResult]]: # Added running_mode """Generates embeddings for the input image tensor(s).""" if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") - + # Add running_mode to tuple check new_options_tuple = (l2_normalize, quantize, running_mode, delegate.lower()) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.l2_normalize, - self._current_options.quantize, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' - ) + current_options_tuple = ( + self._current_options.l2_normalize, + self._current_options.quantize, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", + ) if self._detector_instance is None or new_options_tuple != current_options_tuple: # ImageEmbedder doesn't seem to have close(), so just replace instance - if self._detector_instance and hasattr(self._detector_instance, 'close'): - self._detector_instance.close() + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() self._detector_instance = self._initialize_detector(l2_normalize, quantize, running_mode, delegate) batch_size = image.shape[0] @@ -71,7 +76,7 @@ class ImageEmbedder: embed_func = self._detector_instance.embed_for_video if is_video_mode else self._detector_instance.embed for i in range(batch_size): - img_tensor = image[i] + img_tensor = image[i] np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -80,12 +85,17 @@ class ImageEmbedder: embedding_result = embed_func(mp_image, timestamp_ms) else: embedding_result = embed_func(mp_image) - + current_image_embeddings = [] - if embedding_result and embedding_result.embeddings: # Added check for embedding_result + if embedding_result and embedding_result.embeddings: # Added check for embedding_result for emb in embedding_result.embeddings: - result = ImageEmbedderResult(float_embedding=list(emb.embedding) if not quantize else None, quantized_embedding=bytes(emb.embedding) if quantize else None, head_index=emb.head_index, head_name=emb.head_name) + result = ImageEmbedderResult( + float_embedding=list(emb.embedding) if not quantize else None, + quantized_embedding=bytes(emb.embedding) if quantize else None, + head_index=emb.head_index, + head_name=emb.head_name, + ) current_image_embeddings.append(result) batch_results.append(current_image_embeddings) - - return batch_results \ No newline at end of file + + return batch_results diff --git a/src/mediapipe_vision/image_segmentation/__init__.py b/src/mediapipe_vision/image_segmentation/__init__.py index 87f49bb..58bf4b5 100644 --- a/src/mediapipe_vision/image_segmentation/__init__.py +++ b/src/mediapipe_vision/image_segmentation/__init__.py @@ -1 +1 @@ -"""Image Segmentation implementation using MediaPipe.""" \ No newline at end of file +"""Image Segmentation implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/image_segmentation/segmenter.py b/src/mediapipe_vision/image_segmentation/segmenter.py index 8b8e318..cc3a833 100644 --- a/src/mediapipe_vision/image_segmentation/segmenter.py +++ b/src/mediapipe_vision/image_segmentation/segmenter.py @@ -17,9 +17,10 @@ from ...utils.timing import TimestampProvider logger = logging.getLogger(__name__) + class ImageSegmenter: """Performs image segmentation using MediaPipe ImageSegmenter. - Adapted from ComfyUI-Stream-Pack implementation. + Adapted from ComfyUI-Stream-Pack implementation. """ def __init__(self, model_path: str): @@ -35,49 +36,47 @@ class ImageSegmenter: # Keep track of the initialized segmenter instance and its config self._detector_instance = None self.current_config = None - self.delegate = None # Store delegate used - self._timestamp_provider = None # Added + self.delegate = None # Store delegate used + self._timestamp_provider = None # Added - def _initialize_segmenter( - self, - output_confidence_masks: bool, - output_category_mask: bool, - running_mode: str, delegate_mode="cpu"): + def _initialize_segmenter(self, output_confidence_masks: bool, output_category_mask: bool, running_mode: str, delegate_mode="cpu"): """Initializes the ImageSegmenter or returns existing if config matches. - Handles GPU delegate selection. + Handles GPU delegate selection. """ # Return type removed from signature as it's handled internally new_config = (output_confidence_masks, output_category_mask, running_mode, delegate_mode) - + # Re-initialize if config changed or not initialized if self._detector_instance is None or self.current_config != new_config: if self._detector_instance: - logger.info("Closing existing ImageSegmenter due to config change.") - try: - self._detector_instance.close() - except Exception as e: - logger.warning(f"Error closing previous segmenter: {e}") - self._detector_instance = None - - logger.info(f"Initializing ImageSegmenter: model='{self.model_path}', conf_mask={output_confidence_masks}, cat_mask={output_category_mask}, mode={running_mode}, delegate={delegate_mode}") - + logger.info("Closing existing ImageSegmenter due to config change.") + try: + self._detector_instance.close() + except Exception as e: + logger.warning(f"Error closing previous segmenter: {e}") + self._detector_instance = None + + logger.info( + f"Initializing ImageSegmenter: model='{self.model_path}', conf_mask={output_confidence_masks}, cat_mask={output_category_mask}, mode={running_mode}, delegate={delegate_mode}" + ) + # Configure BaseOptions with delegate - use_gpu = (delegate_mode.lower() == "gpu") + use_gpu = delegate_mode.lower() == "gpu" if use_gpu and platform.system().lower() == "windows": logger.warning("GPU delegate is not supported on Windows for Python MediaPipe Tasks. Falling back to CPU.") use_gpu = False delegate_mode = "cpu" - + delegate = python.BaseOptions.Delegate.GPU if use_gpu else python.BaseOptions.Delegate.CPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate) - + try: mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO options = vision.ImageSegmenterOptions( base_options=base_options, running_mode=mode_enum, output_confidence_masks=output_confidence_masks, - output_category_mask=output_category_mask + output_category_mask=output_category_mask, ) self._detector_instance = vision.ImageSegmenter.create_from_options(options) self.current_config = new_config @@ -91,15 +90,21 @@ class ImageSegmenter: else: # logger.debug("Reusing existing ImageSegmenter instance.") # Optional: for debugging pass - + # No explicit return needed, instance is stored in self._detector_instance - def segment(self, image: torch.Tensor, output_confidence_masks: bool = True, output_category_mask: bool = False, - running_mode: str = "video", delegate_mode="cpu"): + def segment( + self, + image: torch.Tensor, + output_confidence_masks: bool = True, + output_category_mask: bool = False, + running_mode: str = "video", + delegate_mode="cpu", + ): """Segments the input image tensor. Args: - image: Input image tensor (BHWC, RGB, 0-1 float). + image: Input image tensor (BHWC, RGB, 0-1 float). NOTE: This implementation currently only supports batch size 1 for simplicity. output_confidence_masks: Whether to output confidence masks (float, 0-1). output_category_mask: Whether to output a category mask (uint8). @@ -115,18 +120,18 @@ class ImageSegmenter: raise ValueError("Input tensor must be in BHWC format.") # if image.shape[0] != 1: # Temporarily removed batch size 1 restriction # logger.warning("ImageSegmenter currently processes batch > 1 sequentially. Performance may vary.") - # raise ValueError("Input tensor must be in BHWC format with batch size 1.") + # raise ValueError("Input tensor must be in BHWC format with batch size 1.") if not output_confidence_masks and not output_category_mask: - raise ValueError("At least one output type (confidence or category mask) must be enabled.") + raise ValueError("At least one output type (confidence or category mask) must be enabled.") # Initialize (or get existing) segmenter instance # Initialization now happens within the function if needed self._initialize_segmenter(output_confidence_masks, output_category_mask, running_mode, delegate_mode) if not self._detector_instance: - # Initialization failed, return None or raise error - logger.error("Segmenter instance is not available after initialization attempt.") - # Return structure matching expected output format (list per batch item) - return [None] * image.shape[0], [None] * image.shape[0] + # Initialization failed, return None or raise error + logger.error("Segmenter instance is not available after initialization attempt.") + # Return structure matching expected output format (list per batch item) + return [None] * image.shape[0], [None] * image.shape[0] batch_size = image.shape[0] batch_results_confidence = [] if output_confidence_masks else None @@ -137,7 +142,7 @@ class ImageSegmenter: segment_func = self._detector_instance.segment_for_video if is_video_mode else self._detector_instance.segment for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC # Convert image format np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -151,11 +156,13 @@ class ImageSegmenter: else: segmentation_result = segment_func(mp_image) except Exception as e: - logger.error(f"Error during segmentation for image {i}: {e}") - # Append None results for this image - if output_confidence_masks: batch_results_confidence.append(None) - if output_category_mask: batch_results_category.append(None) - continue # Skip to the next image + logger.error(f"Error during segmentation for image {i}: {e}") + # Append None results for this image + if output_confidence_masks: + batch_results_confidence.append(None) + if output_category_mask: + batch_results_category.append(None) + continue # Skip to the next image # Process results for the current image current_image_confidence_masks = None @@ -166,7 +173,7 @@ class ImageSegmenter: current_image_confidence_masks = [] for mask in segmentation_result.confidence_masks: mask_np = mask.numpy_view() - mask_tensor = torch.from_numpy(mask_np).float() # Keep as HW + mask_tensor = torch.from_numpy(mask_np).float() # Keep as HW current_image_confidence_masks.append(mask_tensor) batch_results_confidence.append(current_image_confidence_masks) @@ -174,23 +181,23 @@ class ImageSegmenter: if segmentation_result and segmentation_result.category_mask: mask_np = segmentation_result.category_mask.numpy_view() # Category mask is uint8, convert to long tensor - current_image_category_mask = torch.from_numpy(mask_np).long() # Keep as HW + current_image_category_mask = torch.from_numpy(mask_np).long() # Keep as HW batch_results_category.append(current_image_category_mask) return batch_results_confidence, batch_results_category def close(self): - """Closes the segmenter instance if it exists.""" - if self._detector_instance: + """Closes the segmenter instance if it exists.""" + if self._detector_instance: logger.info("Closing ImageSegmenter instance.") try: - self._detector_instance.close() + self._detector_instance.close() except Exception as e: - logger.warning(f"Error closing segmenter instance: {e}") + logger.warning(f"Error closing segmenter instance: {e}") self._detector_instance = None self.current_config = None self.delegate = None def __del__(self): # Ensure resources are released when the object is destroyed - self.close() \ No newline at end of file + self.close() diff --git a/src/mediapipe_vision/interactive_segmentation/__init__.py b/src/mediapipe_vision/interactive_segmentation/__init__.py index 0519ecb..e69de29 100644 --- a/src/mediapipe_vision/interactive_segmentation/__init__.py +++ b/src/mediapipe_vision/interactive_segmentation/__init__.py @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/src/mediapipe_vision/interactive_segmentation/detector.py b/src/mediapipe_vision/interactive_segmentation/detector.py index cab2ec8..fe6deec 100644 --- a/src/mediapipe_vision/interactive_segmentation/detector.py +++ b/src/mediapipe_vision/interactive_segmentation/detector.py @@ -15,6 +15,7 @@ logger = logging.getLogger(__name__) NormalizedKeypoint = keypoint.NormalizedKeypoint + class InteractiveSegmenterProcessor: """Performs interactive image segmentation using MediaPipe.""" @@ -29,16 +30,21 @@ class InteractiveSegmenterProcessor: def _initialize_processor(self, output_category_mask: bool, output_confidence_mask: bool): base_options = python.BaseOptions(model_asset_path=self.model_path) options = interactive_segmenter.InteractiveSegmenterOptions( - base_options=base_options, - output_category_mask=output_category_mask, - output_confidence_masks=output_confidence_mask - ) + base_options=base_options, + output_category_mask=output_category_mask, + output_confidence_masks=output_confidence_mask, + ) self.current_options = (output_category_mask, output_confidence_mask) self._timestamp_provider = TimestampProvider() return interactive_segmenter.InteractiveSegmenter.create_from_options(options) - def segment(self, image: torch.Tensor, points_of_interest: List[PointOfInterest], - output_category_mask: bool = True, output_confidence_mask: bool = False) -> List[Optional[torch.Tensor]]: + def segment( + self, + image: torch.Tensor, + points_of_interest: List[PointOfInterest], + output_category_mask: bool = True, + output_confidence_mask: bool = False, + ) -> List[Optional[torch.Tensor]]: """Performs interactive segmentation based on points of interest. Iterates through points, generates mask for each, and combines them. @@ -55,19 +61,19 @@ class InteractiveSegmenterProcessor: if image.dim() != 4: raise ValueError("Input tensor must be in BHWC format.") if not points_of_interest: - raise ValueError("At least one point_of_interest must be provided.") + raise ValueError("At least one point_of_interest must be provided.") batch_size = image.shape[0] batch_results_combined = [] - + new_options = (output_category_mask, output_confidence_mask) if self.detector_instance is None or self.current_options != new_options: - if self.detector_instance and hasattr(self.detector_instance, 'close'): - self.detector_instance.close() - self.detector_instance = self._initialize_processor(output_category_mask, output_confidence_mask) + if self.detector_instance and hasattr(self.detector_instance, "close"): + self.detector_instance.close() + self.detector_instance = self._initialize_processor(output_category_mask, output_confidence_mask) for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -76,12 +82,11 @@ class InteractiveSegmenterProcessor: for point in points_of_interest: roi_keypoint_single = NormalizedKeypoint(x=point.x, y=point.y, label=point.label) roi = interactive_segmenter.RegionOfInterest( - format=interactive_segmenter.RegionOfInterest.Format.KEYPOINT, - keypoint=roi_keypoint_single + format=interactive_segmenter.RegionOfInterest.Format.KEYPOINT, keypoint=roi_keypoint_single ) - + segmentation_result = self.detector_instance.segment(mp_image, roi) - + current_point_mask = None if output_category_mask and segmentation_result.category_mask: mask_np = segmentation_result.category_mask.numpy_view() @@ -90,7 +95,7 @@ class InteractiveSegmenterProcessor: if segmentation_result.confidence_masks: mask_np = segmentation_result.confidence_masks[0].numpy_view() current_point_mask = torch.from_numpy(mask_np).float() - + if current_point_mask is not None: individual_masks_for_image.append(current_point_mask) else: @@ -105,6 +110,6 @@ class InteractiveSegmenterProcessor: else: mask_stack = torch.stack(individual_masks_for_image, dim=0) combined_mask_for_image = torch.max(mask_stack, dim=0).values - + batch_results_combined.append(combined_mask_for_image) - return batch_results_combined \ No newline at end of file + return batch_results_combined diff --git a/src/mediapipe_vision/landmark_definitions.py b/src/mediapipe_vision/landmark_definitions.py index edd0f5f..24ac09a 100644 --- a/src/mediapipe_vision/landmark_definitions.py +++ b/src/mediapipe_vision/landmark_definitions.py @@ -44,7 +44,6 @@ FACE_LANDMARK_DESCRIPTIONS = { 86: "Bottom lip bottom center", 14: "Upper teeth", 13: "Lower teeth", - # General regions (for tooltip summary) # 0-17: Jaw/face outline # 18-35: Eyebrows @@ -75,7 +74,7 @@ Notable: # Hand landmarks (21 points) HAND_LANDMARK_DESCRIPTIONS = { 0: "Wrist", - 1: "Thumb CMC", + 1: "Thumb CMC", 2: "Thumb MCP", 3: "Thumb IP", 4: "Thumb tip", @@ -94,7 +93,7 @@ HAND_LANDMARK_DESCRIPTIONS = { 17: "Pinky MCP", 18: "Pinky PIP", 19: "Pinky DIP", - 20: "Pinky tip" + 20: "Pinky tip", } # Tooltip for hand landmarks @@ -143,7 +142,7 @@ POSE_LANDMARK_DESCRIPTIONS = { 29: "Left heel", 30: "Right heel", 31: "Left foot index", - 32: "Right foot index" + 32: "Right foot index", } # Tooltip for pose landmarks @@ -158,4 +157,4 @@ POSE_LANDMARK_TOOLTIP = """MediaPipe Pose Landmarks (0-32): 23-24: Hips 25-26: Knees 27-28: Ankles - 29-32: Feet""" \ No newline at end of file + 29-32: Feet""" diff --git a/src/mediapipe_vision/model_loader.py b/src/mediapipe_vision/model_loader.py index c1cb9cb..8327f4e 100644 --- a/src/mediapipe_vision/model_loader.py +++ b/src/mediapipe_vision/model_loader.py @@ -15,7 +15,7 @@ import folder_paths # ComfyUI path helper # logging.basicConfig(level=logging.INFO) # Avoid reconfiguring root logger logger = logging.getLogger(__name__) -# --- Model Definitions --- +# --- Model Definitions --- # Central dictionary mapping task types and variants to model details # Structure: { task_type: { variant_name: { url: ..., filename: ... } } } # Task types should match the keys used elsewhere (e.g., detector class validation) @@ -28,99 +28,99 @@ MEDIAPIPE_MODELS = { "short_range": { "url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite", "filename": "blaze_face_short_range.tflite", - "description": "Optimized for detecting faces at close range, ideal for selfie-style images and near-field face detection." + "description": "Optimized for detecting faces at close range, ideal for selfie-style images and near-field face detection.", }, "full_range_sparse": { "url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_back_camera/float16/latest/blaze_face_back_camera.tflite", "filename": "blaze_face_full_range_sparse.tflite", - "description": "Designed for detecting faces at various distances, suitable for group photos and wider-angle shots." - } + "description": "Designed for detecting faces at various distances, suitable for group photos and wider-angle shots.", + }, }, "face_landmarker": { "default": { "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task", "filename": "face_landmarker.task", - "description": "Standard face landmark detection model that locates key facial points for tracking and analysis." + "description": "Standard face landmark detection model that locates key facial points for tracking and analysis.", }, - "with_blendshapes": { # If a separate model exists + "with_blendshapes": { # If a separate model exists "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker_v2_with_blendshapes/float16/latest/face_landmarker_v2_with_blendshapes.task", "filename": "face_landmarker_v2_with_blendshapes.task", - "description": "Advanced face landmark model that includes blendshape coefficients for facial expression analysis." - } + "description": "Advanced face landmark model that includes blendshape coefficients for facial expression analysis.", + }, }, "hand_landmarker": { - "default": { # Often includes world landmarks + "default": { # Often includes world landmarks "url": "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/latest/hand_landmarker.task", "filename": "hand_landmarker.task", - "description": "Detects hand landmarks for tracking finger positions and hand gestures in 2D and 3D space." + "description": "Detects hand landmarks for tracking finger positions and hand gestures in 2D and 3D space.", } }, "pose_landmarker": { "lite": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/latest/pose_landmarker_lite.task", "filename": "pose_landmarker_lite.task", - "description": "Lightweight pose landmark model for faster inference with slightly reduced accuracy." + "description": "Lightweight pose landmark model for faster inference with slightly reduced accuracy.", }, "full": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/latest/pose_landmarker_full.task", "filename": "pose_landmarker_full.task", - "description": "Balance of performance and accuracy for detecting body pose landmarks in typical use cases." + "description": "Balance of performance and accuracy for detecting body pose landmarks in typical use cases.", }, "heavy": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/latest/pose_landmarker_heavy.task", "filename": "pose_landmarker_heavy.task", - "description": "Higher accuracy pose detection model for more precise landmark tracking at the cost of performance." - } + "description": "Higher accuracy pose detection model for more precise landmark tracking at the cost of performance.", + }, }, "object_detector": { "efficientdet_lite0": { "url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite0/float16/latest/efficientdet_lite0.tflite", "filename": "efficientdet_lite0.tflite", - "description": "Lightweight object detection model that balances speed and accuracy for real-time applications." + "description": "Lightweight object detection model that balances speed and accuracy for real-time applications.", }, "efficientdet_lite2": { - "url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite2/float16/latest/efficientdet_lite2.tflite", - "filename": "efficientdet_lite2.tflite", - "description": "Medium-sized object detection model with improved accuracy over lite0 while maintaining reasonable performance." - }, - "ssd_mobilenet_v2": { - "url": "https://storage.googleapis.com/mediapipe-models/object_detector/ssd_mobilenet_v2/float32/latest/ssd_mobilenet_v2.tflite", - "filename": "ssd_mobilenet_v2.tflite", - "description": "Single Shot Detector with MobileNetV2 backbone for efficient object detection tasks." - } + "url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite2/float16/latest/efficientdet_lite2.tflite", + "filename": "efficientdet_lite2.tflite", + "description": "Medium-sized object detection model with improved accuracy over lite0 while maintaining reasonable performance.", + }, + "ssd_mobilenet_v2": { + "url": "https://storage.googleapis.com/mediapipe-models/object_detector/ssd_mobilenet_v2/float32/latest/ssd_mobilenet_v2.tflite", + "filename": "ssd_mobilenet_v2.tflite", + "description": "Single Shot Detector with MobileNetV2 backbone for efficient object detection tasks.", + }, }, "image_segmenter": { "deeplab_v3": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/deeplab_v3/float32/latest/deeplab_v3.tflite", "filename": "deeplab_v3.tflite", - "description": "General-purpose semantic segmentation model that can segment images into different object categories." + "description": "General-purpose semantic segmentation model that can segment images into different object categories.", }, "selfie_segmenter_landscape": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter_landscape/float16/latest/selfie_segmenter_landscape.tflite", "filename": "selfie_segmenter_landscape.tflite", - "description": "Optimized for landscape-oriented images to accurately separate people from backgrounds." + "description": "Optimized for landscape-oriented images to accurately separate people from backgrounds.", }, "selfie_segmenter_general": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite", "filename": "selfie_segmenter.tflite", - "description": "General-purpose person segmentation model that works across different orientations and scenarios." + "description": "General-purpose person segmentation model that works across different orientations and scenarios.", }, "hair_segmenter": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/hair_segmenter/float32/latest/hair_segmenter.tflite", "filename": "hair_segmenter.tflite", - "description": "Specialized segmentation model for isolating hair in portrait images with high precision." + "description": "Specialized segmentation model for isolating hair in portrait images with high precision.", }, "selfie_multiclass": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite", "filename": "selfie_multiclass_256x256.tflite", - "description": "Multi-class segmentation model that can separate people into different body parts and clothing items." - } + "description": "Multi-class segmentation model that can separate people into different body parts and clothing items.", + }, }, "gesture_recognizer": { "default": { "url": "https://storage.googleapis.com/mediapipe-models/gesture_recognizer/gesture_recognizer/float16/latest/gesture_recognizer.task", "filename": "gesture_recognizer.task", - "description": "Detects and classifies hand gestures such as thumbs up, victory sign, and pointing gestures." + "description": "Detects and classifies hand gestures such as thumbs up, victory sign, and pointing gestures.", } }, # "holistic_landmarker": { @@ -135,74 +135,75 @@ MEDIAPIPE_MODELS = { # } # Add other tasks like image_classification etc. here "image_embedder": { - "mobilenet_v3_small": { + "mobilenet_v3_small": { "url": "https://storage.googleapis.com/mediapipe-models/image_embedder/mobilenet_v3_small/float32/1/mobilenet_v3_small.tflite", "filename": "mobilenet_v3_small.tflite", - "description": "Converts images into feature vectors (embeddings) for similarity comparison and classification tasks." + "description": "Converts images into feature vectors (embeddings) for similarity comparison and classification tasks.", }, "efficientnet_lite0": { "url": "https://storage.googleapis.com/mediapipe-models/image_embedder/efficientnet_lite0/float32/latest/efficientnet_lite0.tflite", "filename": "efficientnet_lite0.tflite", - "description": "EfficientNet Lite0-based embedder providing robust image embeddings with a balance of accuracy and efficiency." - } + "description": "EfficientNet Lite0-based embedder providing robust image embeddings with a balance of accuracy and efficiency.", + }, }, # Placeholder for Interactive Segmenter - Typically uses a general segmentation model "interactive_segmenter": { "magic_touch": { "url": "https://storage.googleapis.com/mediapipe-models/interactive_segmenter/magic_touch/float32/latest/magic_touch.tflite", "filename": "magic_touch.tflite", - "description": "Interactive segmentation model that segments objects in an image based on a user-provided point of interest." + "description": "Interactive segmentation model that segments objects in an image based on a user-provided point of interest.", }, "default": { - "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite", - "filename": "selfie_segmenter_interactive.tflite", - "description": "Enables user-guided segmentation by providing points or regions to indicate foreground objects." - } + "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite", + "filename": "selfie_segmenter_interactive.tflite", + "description": "Enables user-guided segmentation by providing points or regions to indicate foreground objects.", + }, }, "face_stylizer": { "color_sketch": { "url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_color_sketch.task", "filename": "face_stylizer_color_sketch.task", - "description": "Transforms faces into colorful sketch-like artistic renderings with clean lines and vibrant colors." + "description": "Transforms faces into colorful sketch-like artistic renderings with clean lines and vibrant colors.", }, "color_ink": { "url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_color_ink.task", "filename": "face_stylizer_color_ink.task", - "description": "Applies a color ink effect to faces with flowing brush-like strokes and dynamic color blending." + "description": "Applies a color ink effect to faces with flowing brush-like strokes and dynamic color blending.", }, "oil_painting": { "url": "https://storage.googleapis.com/mediapipe-models/face_stylizer/blaze_face_stylizer/float32/latest/face_stylizer_oil_painting.task", "filename": "face_stylizer_oil_painting.task", - "description": "Renders faces in classic oil painting style with textured brush strokes and rich color tones." - } + "description": "Renders faces in classic oil painting style with textured brush strokes and rich color tones.", + }, }, "image_classifier": { "efficientnet_lite0": { "url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite0/float32/latest/efficientnet_lite0.tflite", "filename": "efficientnet_lite0.tflite", - "description": "EfficientNet-Lite0 model for general-purpose image classification tasks with optimized performance." + "description": "EfficientNet-Lite0 model for general-purpose image classification tasks with optimized performance.", }, "efficientnet_lite2": { "url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite2/float32/latest/efficientnet_lite2.tflite", "filename": "efficientnet_lite2.tflite", - "description": "EfficientNet-Lite2 model offering higher accuracy for image classification tasks at the cost of increased computation." - } + "description": "EfficientNet-Lite2 model offering higher accuracy for image classification tasks at the cost of increased computation.", + }, }, "face_mesh": { "default": { "url": "https://storage.googleapis.com/mediapipe-models/face_mesh/face_mesh/float16/latest/face_mesh.tflite", "filename": "face_mesh.tflite", - "description": "High-fidelity 3D face landmark model estimating 468 landmarks for detailed facial geometry analysis." + "description": "High-fidelity 3D face landmark model estimating 468 landmarks for detailed facial geometry analysis.", }, "attention": { "url": "https://storage.googleapis.com/mediapipe-models/face_mesh/face_mesh_attention/float16/latest/face_mesh_attention.tflite", "filename": "face_mesh_attention.tflite", - "description": "Enhanced face mesh model with attention mechanism for improved accuracy around eyes and lips." - } - } + "description": "Enhanced face mesh model with attention mechanism for improved accuracy around eyes and lips.", + }, + }, } -# --- Utility Functions --- +# --- Utility Functions --- + def get_mediapipe_model_dir(task_type: str) -> str: """Gets the dedicated directory for a specific MediaPipe task type within ComfyUI models.""" @@ -214,7 +215,7 @@ def get_mediapipe_model_dir(task_type: str) -> str: base_path = mediapipe_paths[0] logger.debug(f"Using specific MediaPipe model path: {base_path}") else: - logger.debug("No valid specific 'mediapipe' path found or defined.") + logger.debug("No valid specific 'mediapipe' path found or defined.") except (KeyError, IndexError, TypeError) as e: logger.debug(f"Could not retrieve specific 'mediapipe' path ({e}). Falling back to default.") @@ -222,11 +223,11 @@ def get_mediapipe_model_dir(task_type: str) -> str: if base_path is None: try: # Ensure models_dir exists - if hasattr(folder_paths, 'models_dir') and os.path.isdir(folder_paths.models_dir): + if hasattr(folder_paths, "models_dir") and os.path.isdir(folder_paths.models_dir): base_path = os.path.join(folder_paths.models_dir, "mediapipe") logger.debug(f"Using fallback MediaPipe model path: {base_path}") else: - raise RuntimeError("ComfyUI models directory ('models_dir') not found or invalid.") + raise RuntimeError("ComfyUI models directory ('models_dir') not found or invalid.") except Exception as e: logger.error(f"Error determining fallback MediaPipe directory: {e}") raise RuntimeError("Could not determine a valid directory for MediaPipe models.") from e @@ -244,17 +245,18 @@ def get_mediapipe_model_dir(task_type: str) -> str: return task_dir + def _download_url_to_file(url: str, save_path: str): """Downloads a file with progress indication.""" logger.info(f"Downloading model from {url} to {save_path}...") try: - with urllib.request.urlopen(url) as response, open(save_path, 'wb') as out_file: - total_size_header = response.getheader('Content-Length') + with urllib.request.urlopen(url) as response, open(save_path, "wb") as out_file: + total_size_header = response.getheader("Content-Length") total_size = int(total_size_header) if total_size_header else None bytes_downloaded = 0 block_size = 8192 last_print_bytes = 0 - print_interval = 1 * 1024 * 1024 # Print every 1MB + print_interval = 1 * 1024 * 1024 # Print every 1MB if total_size: logger.info(f" File size: {total_size / (1024*1024):.2f} MB") @@ -275,10 +277,10 @@ def _download_url_to_file(url: str, save_path: str): progress = (bytes_downloaded / total_size) * 100 progress_str += f" / {total_size / (1024*1024):.2f} MB ({progress:.1f}%)" # Use print with carriage return for inline update - print(progress_str, end=' \r') + print(progress_str, end=" \r") last_print_bytes = bytes_downloaded - print() # Newline after download finishes + print() # Newline after download finishes logger.info(f"Download complete: {save_path}") return True except Exception as e: @@ -293,6 +295,7 @@ def _download_url_to_file(url: str, save_path: str): logger.error(f"Error removing incomplete download {save_path}: {rm_err}") return False + def get_model_path(task_type: str, model_variant: str) -> Optional[str]: """Gets the full path to the model file, downloading if necessary.""" if task_type not in MEDIAPIPE_MODELS: @@ -308,32 +311,35 @@ def get_model_path(task_type: str, model_variant: str) -> Optional[str]: model_details = MEDIAPIPE_MODELS[task_type][model_variant] model_filename = model_details["filename"] model_url = model_details["url"] - + # Get the specific directory for this task try: model_dir = get_mediapipe_model_dir(task_type) except Exception as e: - logger.error(f"Cannot determine model directory for {task_type}: {e}") - return None - + logger.error(f"Cannot determine model directory for {task_type}: {e}") + return None + model_path = os.path.join(model_dir, model_filename) # Check and download if not os.path.exists(model_path): logger.info(f"Model not found locally at {model_path}. Attempting download.") if not _download_url_to_file(model_url, model_path): - return None # Download failed + return None # Download failed else: logger.debug(f"Found model locally: {model_path}") return model_path -# --- Public Functions for Node Wrapper --- + +# --- Public Functions for Node Wrapper --- + def get_available_tasks() -> List[str]: """Get a list of available MediaPipe Vision task types from the definition.""" return list(MEDIAPIPE_MODELS.keys()) + def get_available_models(task_type: str) -> List[str]: """Get a list of available model variants for a specific task.""" if task_type in MEDIAPIPE_MODELS: @@ -342,50 +348,51 @@ def get_available_models(task_type: str) -> List[str]: logger.warning(f"Task type '{task_type}' not found in model definitions.") return [] + def get_model_description(task_type: str, model_variant: str) -> str: """Get the description for a specific model variant. - + Args: task_type: The MediaPipe task type model_variant: The specific model variant - + Returns: The description string if found, or a default message if not available """ if task_type in MEDIAPIPE_MODELS and model_variant in MEDIAPIPE_MODELS[task_type]: - return MEDIAPIPE_MODELS[task_type][model_variant].get( - "description", "No description available" - ) + return MEDIAPIPE_MODELS[task_type][model_variant].get("description", "No description available") return "Model description not available" -# --- Model Loader Class (Simplified) --- + +# --- Model Loader Class (Simplified) --- + class MediaPipeModelLoader: """Helper class to manage loading state and info. - Used temporarily by the loader node to get model info. + Used temporarily by the loader node to get model info. """ - + def __init__(self): self.model_path = None self.task_type = None self.model_variant = None self.model_loaded = False self.model_info = {} - + def load_model(self, task_type: str, model_variant: str) -> bool: """Sets model info and triggers path retrieval (incl. download).""" logger.debug(f"MediaPipeModelLoader: Loading {task_type}/{model_variant}") self.task_type = task_type self.model_variant = model_variant self.model_path = get_model_path(task_type, model_variant) - + if self.model_path and os.path.exists(self.model_path): self.model_loaded = True self.model_info = { - 'task_type': self.task_type, - 'model_variant': self.model_variant, - 'model_path': self.model_path, - 'description': get_model_description(task_type, model_variant) + "task_type": self.task_type, + "model_variant": self.model_variant, + "model_path": self.model_path, + "description": get_model_description(task_type, model_variant), } logger.debug(f"Model info set: {self.model_info}") return True @@ -394,23 +401,24 @@ class MediaPipeModelLoader: self.model_loaded = False self.model_info = {} return False - + def get_model_info(self) -> Dict[str, Any]: """Get information about the loaded model.""" # Ensure info is consistent with loaded state if self.model_loaded: - # Add the description to the model info - if self.task_type and self.model_variant: - description = get_model_description(self.task_type, self.model_variant) - self.model_info['description'] = description - return self.model_info + # Add the description to the model info + if self.task_type and self.model_variant: + description = get_model_description(self.task_type, self.model_variant) + self.model_info["description"] = description + return self.model_info else: - # Return minimal info or empty dict if not loaded - return { - 'task_type': self.task_type, - 'model_variant': self.model_variant, - 'model_path': None # Indicate path is not valid/loaded - } + # Return minimal info or empty dict if not loaded + return { + "task_type": self.task_type, + "model_variant": self.model_variant, + "model_path": None, # Indicate path is not valid/loaded + } + # Removed original get_model_path, download_model etc. from global scope -# They are now integrated into the new get_model_path function. \ No newline at end of file +# They are now integrated into the new get_model_path function. diff --git a/src/mediapipe_vision/object_detection/__init__.py b/src/mediapipe_vision/object_detection/__init__.py index f1dde9e..e3d9c67 100644 --- a/src/mediapipe_vision/object_detection/__init__.py +++ b/src/mediapipe_vision/object_detection/__init__.py @@ -1 +1 @@ -"""Object Detection implementation using MediaPipe.""" \ No newline at end of file +"""Object Detection implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/object_detection/detector.py b/src/mediapipe_vision/object_detection/detector.py index b822ffc..032e9a3 100644 --- a/src/mediapipe_vision/object_detection/detector.py +++ b/src/mediapipe_vision/object_detection/detector.py @@ -28,27 +28,32 @@ class ObjectDetector: self.model_path = model_path self._detector_instance = None self._current_options = None - self._timestamp_provider = None # Added + self._timestamp_provider = None # Added - def _initialize_detector(self, score_threshold: float, max_results: int, - running_mode: str, delegate: str): # Added running_mode + def _initialize_detector(self, score_threshold: float, max_results: int, running_mode: str, delegate: str): # Added running_mode """Initializes the ObjectDetector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.ObjectDetectorOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode score_threshold=score_threshold, - max_results=max_results + max_results=max_results, ) self._current_options = options # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.ObjectDetector.create_from_options(options) - def detect(self, image: torch.Tensor, score_threshold: float = 0.5, max_results: int = 5, - running_mode: str = "video", delegate: str = 'cpu') -> List[List[ObjectDetectionResult]]: # Added running_mode + def detect( + self, + image: torch.Tensor, + score_threshold: float = 0.5, + max_results: int = 5, + running_mode: str = "video", + delegate: str = "cpu", + ) -> List[List[ObjectDetectionResult]]: # Added running_mode """Detects objects in the input image tensor. Returns: @@ -61,16 +66,16 @@ class ObjectDetector: new_options_tuple = (score_threshold, max_results, running_mode, delegate.lower()) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.score_threshold, - self._current_options.max_results, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' - ) + current_options_tuple = ( + self._current_options.score_threshold, + self._current_options.max_results, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", + ) if self._detector_instance is None or new_options_tuple != current_options_tuple: - if self._detector_instance and hasattr(self._detector_instance, 'close'): - self._detector_instance.close() + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() self._detector_instance = self._initialize_detector(score_threshold, max_results, running_mode, delegate) batch_size = image.shape[0] @@ -78,9 +83,9 @@ class ObjectDetector: # Determine the correct detection function based on mode is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect - + for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -89,36 +94,28 @@ class ObjectDetector: detection_result = detect_func(mp_image, timestamp_ms) else: detection_result = detect_func(mp_image) - + current_image_detections = [] - if detection_result and detection_result.detections: # Added check for detection_result + if detection_result and detection_result.detections: # Added check for detection_result for detection in detection_result.detections: bbox_mp = detection.bounding_box categories_mp = detection.categories - - bbox = BoundingBox( - origin_x=bbox_mp.origin_x, - origin_y=bbox_mp.origin_y, - width=bbox_mp.width, - height=bbox_mp.height - ) - + + bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height) + categories = [ ObjectDetectionCategory( - index=cat.index, + index=cat.index, score=cat.score, - display_name=cat.display_name, - category_name=cat.category_name + display_name=cat.display_name, + category_name=cat.category_name, ) for cat in categories_mp ] - - current_image_detections.append(ObjectDetectionResult( - bounding_box=bbox, - categories=categories - )) - + + current_image_detections.append(ObjectDetectionResult(bounding_box=bbox, categories=categories)) + batch_results.append(current_image_detections) # self._detector_instance.close() # Close handled on re-init - return batch_results \ No newline at end of file + return batch_results diff --git a/src/mediapipe_vision/pose_landmark/__init__.py b/src/mediapipe_vision/pose_landmark/__init__.py index 2816aed..4c275a9 100644 --- a/src/mediapipe_vision/pose_landmark/__init__.py +++ b/src/mediapipe_vision/pose_landmark/__init__.py @@ -1 +1 @@ -"""Pose Landmark Detection implementation using MediaPipe.""" \ No newline at end of file +"""Pose Landmark Detection implementation using MediaPipe.""" diff --git a/src/mediapipe_vision/pose_landmark/detector.py b/src/mediapipe_vision/pose_landmark/detector.py index 1eb5c53..405b685 100644 --- a/src/mediapipe_vision/pose_landmark/detector.py +++ b/src/mediapipe_vision/pose_landmark/detector.py @@ -15,7 +15,7 @@ from ..types import LandmarkPoint, PoseLandmarksResult class PoseLandmarkDetector: """Detects pose landmarks in an image using MediaPipe PoseLandmarker.""" - + def __init__(self, model_path: str): """Initialize the detector with the model path. @@ -24,38 +24,51 @@ class PoseLandmarkDetector: """ if not model_path: raise ValueError("A valid model_path must be provided.") - + self.model_path = model_path self._detector_instance = None self._current_options = None - self._timestamp_provider = None # Added + self._timestamp_provider = None # Added - def _initialize_detector(self, num_poses: int, min_detection_confidence: float, - min_presence_confidence: float, min_tracking_confidence: float, - output_segmentation_masks: bool, - running_mode: str, delegate: str) -> vision.PoseLandmarker: # Added running_mode + def _initialize_detector( + self, + num_poses: int, + min_detection_confidence: float, + min_presence_confidence: float, + min_tracking_confidence: float, + output_segmentation_masks: bool, + running_mode: str, + delegate: str, + ) -> vision.PoseLandmarker: # Added running_mode """Initializes the PoseLandmarker detector.""" - mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added - delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == 'cpu' else BaseOptions.Delegate.GPU + mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO # Added + delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum) options = vision.PoseLandmarkerOptions( base_options=base_options, - running_mode=mode_enum, # Use enum mode + running_mode=mode_enum, # Use enum mode num_poses=num_poses, min_pose_detection_confidence=min_detection_confidence, - min_pose_presence_confidence=min_presence_confidence, + min_pose_presence_confidence=min_presence_confidence, min_tracking_confidence=min_tracking_confidence, - output_segmentation_masks=output_segmentation_masks + output_segmentation_masks=output_segmentation_masks, ) self._current_options = options # Only create timestamp provider for VIDEO mode - self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None + self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None return vision.PoseLandmarker.create_from_options(options) - def detect(self, image: torch.Tensor, num_poses: int = 1, min_detection_confidence: float = 0.5, - min_presence_confidence: float = 0.5, min_tracking_confidence: float = 0.5, - output_segmentation_masks: bool = False, - running_mode: str = "video", delegate: str = 'cpu') -> Tuple[List[List[PoseLandmarksResult]], Optional[List[List[torch.Tensor]]]]: # Added running_mode + def detect( + self, + image: torch.Tensor, + num_poses: int = 1, + min_detection_confidence: float = 0.5, + min_presence_confidence: float = 0.5, + min_tracking_confidence: float = 0.5, + output_segmentation_masks: bool = False, + running_mode: str = "video", + delegate: str = "cpu", + ) -> Tuple[List[List[PoseLandmarksResult]], Optional[List[List[torch.Tensor]]]]: # Added running_mode """Detects pose landmarks in the input image tensor. Returns: @@ -67,38 +80,51 @@ class PoseLandmarkDetector: raise ValueError("Input tensor must be in BHWC format.") # Add running_mode to tuple check - new_options_tuple = (num_poses, min_detection_confidence, min_presence_confidence, - min_tracking_confidence, output_segmentation_masks, running_mode, delegate.lower()) + new_options_tuple = ( + num_poses, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + output_segmentation_masks, + running_mode, + delegate.lower(), + ) current_options_tuple = None if self._current_options: - current_options_tuple = ( - self._current_options.num_poses, - self._current_options.min_pose_detection_confidence, - self._current_options.min_pose_presence_confidence, - self._current_options.min_tracking_confidence, - self._current_options.output_segmentation_masks, - "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string - 'cpu' if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else 'gpu' - ) + current_options_tuple = ( + self._current_options.num_poses, + self._current_options.min_pose_detection_confidence, + self._current_options.min_pose_presence_confidence, + self._current_options.min_tracking_confidence, + self._current_options.output_segmentation_masks, + "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video", # Map enum back to string + "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu", + ) if self._detector_instance is None or new_options_tuple != current_options_tuple: - if self._detector_instance and hasattr(self._detector_instance, 'close'): - self._detector_instance.close() + if self._detector_instance and hasattr(self._detector_instance, "close"): + self._detector_instance.close() # Pass new params to initializer - self._detector_instance = self._initialize_detector(num_poses, min_detection_confidence, - min_presence_confidence, min_tracking_confidence, - output_segmentation_masks, running_mode, delegate) # Pass running_mode + self._detector_instance = self._initialize_detector( + num_poses, + min_detection_confidence, + min_presence_confidence, + min_tracking_confidence, + output_segmentation_masks, + running_mode, + delegate, + ) # Pass running_mode batch_size = image.shape[0] batch_results_landmarks = [] batch_results_masks = [] if output_segmentation_masks else None - + # Determine the correct detection function based on mode is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect - + for i in range(batch_size): - img_tensor = image[i] # HWC + img_tensor = image[i] # HWC np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image) @@ -111,42 +137,49 @@ class PoseLandmarkDetector: current_image_landmarks = [] current_image_masks = [] if output_segmentation_masks else None - if detection_result and detection_result.pose_landmarks: # Added check for detection_result + if detection_result and detection_result.pose_landmarks: # Added check for detection_result for pose_idx, pose_landmarks_mp in enumerate(detection_result.pose_landmarks): landmarks = [ - LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z, - visibility=getattr(lm, 'visibility', None), # Use getattr for safety - presence=getattr(lm, 'presence', None)) + LandmarkPoint( + index=lm_idx, + x=lm.x, + y=lm.y, + z=lm.z, + visibility=getattr(lm, "visibility", None), # Use getattr for safety + presence=getattr(lm, "presence", None), + ) for lm_idx, lm in enumerate(pose_landmarks_mp) ] - + world_landmarks = None if detection_result.pose_world_landmarks and pose_idx < len(detection_result.pose_world_landmarks): - world_landmarks = [ - LandmarkPoint(index=lm_idx, x=lm.x, y=lm.y, z=lm.z, - visibility=getattr(lm, 'visibility', None), - presence=getattr(lm, 'presence', None)) - for lm_idx, lm in enumerate(detection_result.pose_world_landmarks[pose_idx]) - ] - - current_image_landmarks.append(PoseLandmarksResult( - landmarks=landmarks, - world_landmarks=world_landmarks - )) + world_landmarks = [ + LandmarkPoint( + index=lm_idx, + x=lm.x, + y=lm.y, + z=lm.z, + visibility=getattr(lm, "visibility", None), + presence=getattr(lm, "presence", None), + ) + for lm_idx, lm in enumerate(detection_result.pose_world_landmarks[pose_idx]) + ] + + current_image_landmarks.append(PoseLandmarksResult(landmarks=landmarks, world_landmarks=world_landmarks)) # Handle segmentation masks separately as they are not indexed per pose in the result if output_segmentation_masks and detection_result.segmentation_masks: - # MediaPipe returns one mask per image if output_segmentation_masks is True - # We store it as a list associated with the image, not per pose - # Assuming only one mask is relevant per image, might need adjustment if multiple masks are returned - mask_mp = detection_result.segmentation_masks[0] # Get the first mask - mask_np = mask_mp.numpy_view() - mask_tensor = torch.from_numpy(mask_np).float() # HW - current_image_masks = [mask_tensor] # Store as a list for consistency with return type - + # MediaPipe returns one mask per image if output_segmentation_masks is True + # We store it as a list associated with the image, not per pose + # Assuming only one mask is relevant per image, might need adjustment if multiple masks are returned + mask_mp = detection_result.segmentation_masks[0] # Get the first mask + mask_np = mask_mp.numpy_view() + mask_tensor = torch.from_numpy(mask_np).float() # HW + current_image_masks = [mask_tensor] # Store as a list for consistency with return type + batch_results_landmarks.append(current_image_landmarks) if output_segmentation_masks: batch_results_masks.append(current_image_masks) # self._detector_instance.close() # Close handled on re-init - return batch_results_landmarks, batch_results_masks \ No newline at end of file + return batch_results_landmarks, batch_results_masks diff --git a/src/mediapipe_vision/types.py b/src/mediapipe_vision/types.py index 171a629..f5b88dc 100644 --- a/src/mediapipe_vision/types.py +++ b/src/mediapipe_vision/types.py @@ -9,93 +9,109 @@ class LandmarkPoint(NamedTuple): x: float y: float z: float - visibility: Optional[float] = None # Available in Pose landmarks - presence: Optional[float] = None # Available in Pose landmarks + visibility: Optional[float] = None # Available in Pose landmarks + presence: Optional[float] = None # Available in Pose landmarks + # Define container structures for the output of each landmark task # These will be returned by the respective detector nodes + class FaceLandmarksResult(NamedTuple): landmarks: List[LandmarkPoint] # Add other face-specific results if needed later (e.g., bounding box) + class HandLandmarksResult(NamedTuple): landmarks: List[LandmarkPoint] - world_landmarks: Optional[List[LandmarkPoint]] # Keep world landmarks optional for now - handedness: Optional[str] # 'Left', 'Right' + world_landmarks: Optional[List[LandmarkPoint]] # Keep world landmarks optional for now + handedness: Optional[str] # 'Left', 'Right' + class PoseLandmarksResult(NamedTuple): landmarks: List[LandmarkPoint] world_landmarks: Optional[List[LandmarkPoint]] # Segmentation mask info could potentially go here if tightly coupled + # We might add more types later for detections, segmentations etc. # For now, focus is on landmarks. # Define the list types that will be used for batching # These will be the types used in node RETURN_TYPES/INPUT_TYPES strings -FACE_LANDMARKS_LIST = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] -HAND_LANDMARKS_LIST = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] -POSE_LANDMARKS_LIST = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] +FACE_LANDMARKS_LIST = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] +HAND_LANDMARKS_LIST = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] +POSE_LANDMARKS_LIST = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] -# --- Blendshape Result --- + +# --- Blendshape Result --- class Blendshape(NamedTuple): index: Optional[int] score: float - category_name: str # e.g., mouthSmileLeft, eyeBlinkRight + category_name: str # e.g., mouthSmileLeft, eyeBlinkRight display_name: Optional[str] -# --- Common Types --- + +# --- Common Types --- class BoundingBox(NamedTuple): origin_x: int origin_y: int width: int height: int -# --- Face Detection Result --- + +# --- Face Detection Result --- class FaceKeypoint(NamedTuple): - label: Optional[str] # e.g., 'nose_tip', 'left_eye' + label: Optional[str] # e.g., 'nose_tip', 'left_eye' x: float y: float + class FaceDetectionResult(NamedTuple): bounding_box: BoundingBox keypoints: Optional[List[FaceKeypoint]] - score: Optional[float] # Confidence score for the detection + score: Optional[float] # Confidence score for the detection -# --- Object Detection Result --- + +# --- Object Detection Result --- class ObjectDetectionCategory(NamedTuple): index: Optional[int] score: float display_name: Optional[str] category_name: Optional[str] + class ObjectDetectionResult(NamedTuple): bounding_box: BoundingBox categories: List[ObjectDetectionCategory] -# --- Batch List Types for Nodes --- + +# --- Batch List Types for Nodes --- # Define the list types that will be used for batching # These will be the types used in node RETURN_TYPES/INPUT_TYPES strings -FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] -HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] -POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] -OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]] -FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]] +FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] +HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] +POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] +OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]] +FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]] -# --- Gesture Recognition Result --- + +# --- Gesture Recognition Result --- class GestureCategory(NamedTuple): index: Optional[int] score: float display_name: Optional[str] category_name: Optional[str] + class GestureRecognitionResult(NamedTuple): - gestures: List[GestureCategory] # Top gesture(s) recognized for a hand - handedness: Optional[str] # Which hand this gesture belongs to + gestures: List[GestureCategory] # Top gesture(s) recognized for a hand + handedness: Optional[str] # Which hand this gesture belongs to + # Define Batch List Type -GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]] +GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]] + # --- Image Embedding Result --- class EmbeddingEntry(NamedTuple): @@ -104,49 +120,55 @@ class EmbeddingEntry(NamedTuple): display_name: Optional[str] category_name: Optional[str] + class ImageEmbedderResult(NamedTuple): float_embedding: Optional[List[float]] quantized_embedding: Optional[bytes] head_index: int head_name: Optional[str] + # Define Batch List Type -IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]] +IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]] + # --- Interactive Segmentation Input/Result --- class PointOfInterest(NamedTuple): x: float y: float - label: Optional[str] = None # Optional label for the point + label: Optional[str] = None # Optional label for the point + # Result is typically a mask, potentially multiple if categories are involved # Using existing MASK type might suffice, or define a specific one if needed. -# --- Holistic Landmark Result --- + +# --- Holistic Landmark Result --- class HolisticLandmarksResult(NamedTuple): - face_landmarks: Optional[List[LandmarkPoint]] # Only one face per person + face_landmarks: Optional[List[LandmarkPoint]] # Only one face per person pose_landmarks: Optional[List[LandmarkPoint]] pose_world_landmarks: Optional[List[LandmarkPoint]] left_hand_landmarks: Optional[List[LandmarkPoint]] right_hand_landmarks: Optional[List[LandmarkPoint]] # Pose segmentation mask could also be included here + # Define Batch List Type -HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson] +HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson] # ======================================= -# --- Batch List Types for Node I/O --- +# --- Batch List Types for Node I/O --- # ======================================= -FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] -HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] -POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] -OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]] -FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]] -GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]] -IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]] -HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson] +FACE_LANDMARKS = List[List[FaceLandmarksResult]] # Batch[FacesPerImage[Result]] +HAND_LANDMARKS = List[List[HandLandmarksResult]] # Batch[HandsPerImage[Result]] +POSE_LANDMARKS = List[List[PoseLandmarksResult]] # Batch[PosesPerImage[Result]] +OBJECT_DETECTIONS = List[List[ObjectDetectionResult]] # Batch[DetectionsPerImage[Result]] +FACE_DETECTIONS = List[List[FaceDetectionResult]] # Batch[DetectionsPerImage[Result]] +GESTURE_RECOGNITIONS = List[List[GestureRecognitionResult]] # Batch[Hands[Gestures]] +IMAGE_EMBEDDINGS = List[List[ImageEmbedderResult]] # Batch[EmbeddingsPerHead[Result]] +HOLISTIC_LANDMARKS = List[HolisticLandmarksResult] # Batch[ResultPerPerson] # New Types - Defined AFTER Blendshape class # Updated names to match FaceLandmarker output convention -BLENDSHAPES_LIST = List[List[List[Blendshape]]] # Batch[FacesPerImage[BlendshapeList[Blendshape]]] -TRANSFORM_MATRIX_LIST = List[List[np.ndarray]] # Batch[FacesPerImage[Matrix]] \ No newline at end of file +BLENDSHAPES_LIST = List[List[List[Blendshape]]] # Batch[FacesPerImage[BlendshapeList[Blendshape]]] +TRANSFORM_MATRIX_LIST = List[List[np.ndarray]] # Batch[FacesPerImage[Matrix]] diff --git a/src/mediapipe_vision/utils/__init__.py b/src/mediapipe_vision/utils/__init__.py index 82af41d..efeabcc 100644 --- a/src/mediapipe_vision/utils/__init__.py +++ b/src/mediapipe_vision/utils/__init__.py @@ -1,7 +1,7 @@ """Utility functions for MediaPipe Vision.""" # DEPRECATED: These imports are maintained only for backward compatibility. -# For new code, please import directly from src.utils +# For new code, please import directly from src.utils from ...utils.image import convert_to_cv2, convert_to_tensor from ...utils.math import scale_value @@ -10,10 +10,4 @@ from .delta_utils import calculate_euclidean_delta # Keep only the functions specific to mediapipe_vision from .download_utils import download_model -__all__ = [ - "convert_to_cv2", - "convert_to_tensor", - "calculate_euclidean_delta", - "scale_value", - "download_model" -] \ No newline at end of file +__all__ = ["convert_to_cv2", "convert_to_tensor", "calculate_euclidean_delta", "scale_value", "download_model"] diff --git a/src/mediapipe_vision/utils/blendshape_utils.py b/src/mediapipe_vision/utils/blendshape_utils.py index e3b564e..eb03839 100644 --- a/src/mediapipe_vision/utils/blendshape_utils.py +++ b/src/mediapipe_vision/utils/blendshape_utils.py @@ -4,25 +4,64 @@ from typing import List logger = logging.getLogger(__name__) + def get_blendshape_categories() -> List[str]: """Returns the list of standard blendshape category names from MediaPipe.""" - # Define blendshape categories directly + # Define blendshape categories directly # This is based on the MediaPipe standard blendshapes standard_blendshapes = [ - "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", - "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight", - "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight", - "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight", - "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight", - "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", - "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", - "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft", - "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft", - "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", - "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", - "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", "mouthUpperUpLeft", - "mouthUpperUpRight", "noseSneerLeft", "noseSneerRight" + "browDownLeft", + "browDownRight", + "browInnerUp", + "browOuterUpLeft", + "browOuterUpRight", + "cheekPuff", + "cheekSquintLeft", + "cheekSquintRight", + "eyeBlinkLeft", + "eyeBlinkRight", + "eyeLookDownLeft", + "eyeLookDownRight", + "eyeLookInLeft", + "eyeLookInRight", + "eyeLookOutLeft", + "eyeLookOutRight", + "eyeLookUpLeft", + "eyeLookUpRight", + "eyeSquintLeft", + "eyeSquintRight", + "eyeWideLeft", + "eyeWideRight", + "jawForward", + "jawLeft", + "jawOpen", + "jawRight", + "mouthClose", + "mouthDimpleLeft", + "mouthDimpleRight", + "mouthFrownLeft", + "mouthFrownRight", + "mouthFunnel", + "mouthLeft", + "mouthLowerDownLeft", + "mouthLowerDownRight", + "mouthPressLeft", + "mouthPressRight", + "mouthPucker", + "mouthRight", + "mouthRollLower", + "mouthRollUpper", + "mouthShrugLower", + "mouthShrugUpper", + "mouthSmileLeft", + "mouthSmileRight", + "mouthStretchLeft", + "mouthStretchRight", + "mouthUpperUpLeft", + "mouthUpperUpRight", + "noseSneerLeft", + "noseSneerRight", ] - + logger.info("Using predefined list of standard blendshape categories.") - return standard_blendshapes \ No newline at end of file + return standard_blendshapes diff --git a/src/mediapipe_vision/utils/delta_utils.py b/src/mediapipe_vision/utils/delta_utils.py index c3d03b0..7e396ea 100644 --- a/src/mediapipe_vision/utils/delta_utils.py +++ b/src/mediapipe_vision/utils/delta_utils.py @@ -11,17 +11,18 @@ from typing import Optional from ..types import LandmarkPoint # Assuming LandmarkPoint is in src/types.py -#TODO: refactor to common after unified coordinate system +# TODO: refactor to common after unified coordinate system def calculate_euclidean_delta(p1: Optional[LandmarkPoint], p2: Optional[LandmarkPoint]) -> Optional[float]: """Calculates the Euclidean distance between two LandmarkPoint objects.""" if p1 is None or p2 is None: - return None # Cannot calculate delta if one point is missing - + return None # Cannot calculate delta if one point is missing + delta_x = p1.x - p2.x delta_y = p1.y - p2.y delta_z = p1.z - p2.z - + distance = math.sqrt(delta_x**2 + delta_y**2 + delta_z**2) return distance -# scale_value is now imported from utils.math \ No newline at end of file + +# scale_value is now imported from utils.math diff --git a/src/mediapipe_vision/utils/download_utils.py b/src/mediapipe_vision/utils/download_utils.py index d65c19f..7c9815b 100644 --- a/src/mediapipe_vision/utils/download_utils.py +++ b/src/mediapipe_vision/utils/download_utils.py @@ -16,11 +16,11 @@ logger = logging.getLogger(__name__) def download_model(task_type: str, model_variant: str) -> Optional[str]: """Download a MediaPipe model for a specific task and variant. - + Args: task_type: The task category (e.g., 'face_mesh', 'pose_detection') model_variant: The specific model variant (e.g., 'standard', 'lite') - + Returns: Path to the downloaded model file, or None if download failed """ @@ -28,85 +28,85 @@ def download_model(task_type: str, model_variant: str) -> Optional[str]: if task_type not in MEDIAPIPE_MODELS: logger.error(f"Task '{task_type}' not found in available models.") return None - + if model_variant not in MEDIAPIPE_MODELS[task_type]: logger.error(f"Model variant '{model_variant}' not found for task '{task_type}'.") return None - + # Get model details model_info = MEDIAPIPE_MODELS[task_type][model_variant] url = model_info["url"] filename = model_info["filename"] - + # Get the full path for the model output_path = get_model_path(task_type, filename) - + # If model already exists, return the path if os.path.exists(output_path): logger.info(f"Using existing model: {output_path}") return output_path - + # Download the model logger.info(f"Downloading {task_type}/{model_variant} model from {url}...") try: # Ensure parent directory exists os.makedirs(os.path.dirname(output_path), exist_ok=True) - + # Stream the download with progress bar response = requests.get(url, stream=True) response.raise_for_status() - + # Get content length for progress bar if available - total_size = int(response.headers.get('content-length', 0)) + total_size = int(response.headers.get("content-length", 0)) block_size = 8192 # 8 KB - - with open(output_path, 'wb') as f: - with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"{task_type}/{model_variant}") as pbar: + + with open(output_path, "wb") as f: + with tqdm(total=total_size, unit="B", unit_scale=True, desc=f"{task_type}/{model_variant}") as pbar: for chunk in response.iter_content(chunk_size=block_size): if chunk: f.write(chunk) pbar.update(len(chunk)) - + logger.info(f"Successfully downloaded {task_type}/{model_variant} to {output_path}") return output_path - + except Exception as e: logger.error(f"Error downloading {task_type}/{model_variant}: {str(e)}") - + # Clean up partial download if it exists if os.path.exists(output_path): try: os.remove(output_path) except: pass - + return None def get_model_options(task_type: Optional[str] = None) -> Dict[str, Any]: """Get available model options for a specific task or all tasks. - + Args: task_type: The task category to get options for, or None for all tasks - + Returns: Dictionary of available models and variants """ if task_type is None: # Return all tasks and their variants return {task: list(variants.keys()) for task, variants in MEDIAPIPE_MODELS.items()} - + # Return variants for specific task if task_type in MEDIAPIPE_MODELS: return {task_type: list(MEDIAPIPE_MODELS[task_type].keys())} - + # Task not found return {} def get_all_task_types() -> list: """Get a list of all available task types. - + Returns: List of task type strings """ @@ -115,10 +115,10 @@ def get_all_task_types() -> list: def get_model_variants(task_type: str) -> list: """Get a list of all available model variants for a specific task. - + Args: task_type: The task category - + Returns: List of model variant strings, or empty list if task not found """ @@ -129,11 +129,11 @@ def get_model_variants(task_type: str) -> list: def get_model_info(task_type: str, model_variant: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: """Get information about a specific model. - + Args: task_type: The task category model_variant: The specific model variant - + Returns: Tuple containing: - Path to the model file (downloaded if needed) @@ -142,16 +142,16 @@ def get_model_info(task_type: str, model_variant: str) -> Tuple[Optional[str], O # Check if the task and variant exist if task_type not in MEDIAPIPE_MODELS or model_variant not in MEDIAPIPE_MODELS[task_type]: return None, None - + # Get model details model_info = MEDIAPIPE_MODELS[task_type][model_variant].copy() - + # Download the model if needed and get the path model_path = download_model(task_type, model_variant) if not model_path: return None, None - + # Add the model path to the info model_info["model_path"] = model_path - - return model_path, model_info \ No newline at end of file + + return model_path, model_info diff --git a/src/mediapipe_vision/utils/model_utils.py b/src/mediapipe_vision/utils/model_utils.py index dd6fa4a..474ad5e 100644 --- a/src/mediapipe_vision/utils/model_utils.py +++ b/src/mediapipe_vision/utils/model_utils.py @@ -10,6 +10,7 @@ logger = logging.getLogger(__name__) # Import ComfyUI's folder_paths to use the models directory try: import folder_paths + HAS_COMFY_PATHS = True except ImportError: logger.warning("ComfyUI folder_paths module not available. Using local paths.") @@ -18,9 +19,9 @@ except ImportError: def get_mediapipe_models_directory() -> str: """Get the directory where MediaPipe models are stored. - + Uses ComfyUI's models directory structure when available. - + Returns: Path to the MediaPipe models directory """ @@ -31,7 +32,7 @@ def get_mediapipe_models_directory() -> str: # Fallback to local path base_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) mediapipe_dir = os.path.join(base_path, "models") - + # Create the directory if it doesn't exist os.makedirs(mediapipe_dir, exist_ok=True) return mediapipe_dir @@ -39,18 +40,18 @@ def get_mediapipe_models_directory() -> str: def get_model_path(model_type: str, model_name: str) -> str: """Get the path for a specific MediaPipe Vision model. - + Args: model_type: Type/category of model (e.g., 'face_mesh', 'pose') model_name: Name of the specific model file - + Returns: Full path to the model file """ models_dir = get_mediapipe_models_directory() model_type_dir = os.path.join(models_dir, model_type) os.makedirs(model_type_dir, exist_ok=True) - + return os.path.join(model_type_dir, model_name) @@ -60,88 +61,88 @@ MEDIAPIPE_MODELS = { "face_detection": { "short_range": { "url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/1/blaze_face_short_range.task", - "filename": "blaze_face_short_range.task" + "filename": "blaze_face_short_range.task", }, "full_range": { "url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face/float16/1/blaze_face.task", - "filename": "blaze_face.task" - } + "filename": "blaze_face.task", + }, }, "face_mesh": { "standard": { "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task", "filename": "face_landmarker.task", - "description": "Detects face landmarks and facial features (468 points)." + "description": "Detects face landmarks and facial features (468 points).", }, "attention": { "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/2/face_landmarker.task", "filename": "face_landmarker_v2.task", - "description": "Improved version with better accuracy for face landmarks." - } + "description": "Improved version with better accuracy for face landmarks.", + }, }, "hand_tracking": { "standard": { "url": "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task", "filename": "hand_landmarker.task", - "description": "Detects hand landmarks (21 points per hand)." + "description": "Detects hand landmarks (21 points per hand).", } }, "pose_detection": { "lite": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/1/pose_landmarker_lite.task", "filename": "pose_landmarker_lite.task", - "description": "Lightweight pose detection model (33 points)." + "description": "Lightweight pose detection model (33 points).", }, "full": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/1/pose_landmarker_full.task", "filename": "pose_landmarker_full.task", - "description": "Full pose detection model with high accuracy." + "description": "Full pose detection model with high accuracy.", }, "heavy": { "url": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task", "filename": "pose_landmarker_heavy.task", - "description": "Most accurate pose detection model." - } + "description": "Most accurate pose detection model.", + }, }, "image_segmentation": { "selfie": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/1/selfie_segmenter.task", "filename": "selfie_segmenter.task", - "description": "Segments people from background." + "description": "Segments people from background.", }, "selfie_landscape": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter_landscape/float16/1/selfie_segmenter_landscape.task", "filename": "selfie_segmenter_landscape.task", - "description": "Optimized for landscape orientation." + "description": "Optimized for landscape orientation.", }, "multiclass": { "url": "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/1/selfie_multiclass_256x256.task", "filename": "selfie_multiclass_256x256.task", - "description": "Segments multiple body parts (face, hair, clothes, etc.)." - } + "description": "Segments multiple body parts (face, hair, clothes, etc.).", + }, }, "object_detection": { "efficientdet_lite0": { "url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite0/float16/1/efficientdet_lite0.task", "filename": "efficientdet_lite0.task", - "description": "General object detection (80 classes)." + "description": "General object detection (80 classes).", }, "efficientdet_lite2": { "url": "https://storage.googleapis.com/mediapipe-models/object_detector/efficientdet_lite2/float16/1/efficientdet_lite2.task", "filename": "efficientdet_lite2.task", - "description": "Higher accuracy object detection." - } + "description": "Higher accuracy object detection.", + }, }, "image_classification": { "efficientnet_lite0": { "url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite0/float32/1/efficientnet_lite0.task", "filename": "efficientnet_lite0.task", - "description": "General image classification (1000 classes)." + "description": "General image classification (1000 classes).", }, "efficientnet_lite4": { "url": "https://storage.googleapis.com/mediapipe-models/image_classifier/efficientnet_lite4/float32/2/efficientnet_lite4.task", "filename": "efficientnet_lite4.task", - "description": "Higher accuracy image classification." - } - } -} \ No newline at end of file + "description": "Higher accuracy image classification.", + }, + }, +} diff --git a/src/realtimenodes/__init__.py b/src/realtimenodes/__init__.py index da22290..6e7633b 100644 --- a/src/realtimenodes/__init__.py +++ b/src/realtimenodes/__init__.py @@ -1 +1 @@ -# This file intentionally left blank \ No newline at end of file +# This file intentionally left blank diff --git a/src/realtimenodes/control_base.py b/src/realtimenodes/control_base.py index 9e29c82..197f7ea 100644 --- a/src/realtimenodes/control_base.py +++ b/src/realtimenodes/control_base.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod class Pattern(ABC): """Base class for all patterns (movement, timing, etc)""" + @abstractmethod def calculate(self, phase: float, min_val: float, max_val: float) -> float: """Calculate pattern value at given phase (0-1)""" @@ -14,24 +15,26 @@ class Pattern(ABC): """Pattern name for UI""" pass + class StateManager: """Handles persistent state between executions""" + def __init__(self): - self._states = {} - + self._states = {} + def get_state(self, node_id: str, default=None): return self._states.get(node_id, default) - + def set_state(self, node_id: str, state): self._states[node_id] = state - + def clear_state(self, node_id: str) -> bool: """Clear state for a specific node. Returns True if state was found and cleared.""" if node_id in self._states: del self._states[node_id] return True return False - + def clear_all_states(self): """Clear all states""" self._states.clear() @@ -39,7 +42,7 @@ class StateManager: def get_all_node_ids(self): """Get list of all node IDs with state""" return list(self._states.keys()) - + def cleanup_orphaned_states(self, active_node_ids: set[str]): """Remove states for nodes that no longer exist""" orphaned = [nid for nid in self._states.keys() if nid not in active_node_ids] @@ -47,36 +50,37 @@ class StateManager: del self._states[nid] return len(orphaned) + class ControlNodeBase(ABC): """Base class for all control nodes""" - + # Shared state manager state_manager = StateManager() - + @classmethod def INPUT_TYPES(cls): """Base input types common to all control nodes""" return { "required": { - "always_execute": ("BOOLEAN", { - "default": True, - "tooltip": "When enabled, the node updates every execution" - }), + "always_execute": ( + "BOOLEAN", + {"default": True, "tooltip": "When enabled, the node updates every execution"}, + ), }, } - + @classmethod def IS_CHANGED(cls, **kwargs): - if kwargs.get('always_execute', True): + if kwargs.get("always_execute", True): return float(time.time()) return False - + def __init__(self): self.node_id = str(id(self)) - + def get_state(self, default=None): return self.state_manager.get_state(self.node_id, default) - + def set_state(self, state): self.state_manager.set_state(self.node_id, state) @@ -90,4 +94,4 @@ class ControlNodeBase(ABC): @abstractmethod def update(self, *args, **kwargs): """Main update method to be implemented by control nodes""" - pass \ No newline at end of file + pass diff --git a/src/realtimenodes/mask_controls.py b/src/realtimenodes/mask_controls.py index e03c579..93b68a7 100644 --- a/src/realtimenodes/mask_controls.py +++ b/src/realtimenodes/mask_controls.py @@ -3,24 +3,20 @@ import numpy as np class MaskControlMixin: """Mixin providing common mask functionality""" + def create_circle_mask(self, height, width, center_y, center_x, size): """Create a circular mask with anti-aliased edges""" y, x = np.ogrid[:height, :width] radius = size * min(height, width) / 2 - dist = np.sqrt((x - center_x)**2 + (y - center_y)**2) + dist = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2) mask = np.clip(radius + 1 - dist, 0, 1).astype(np.float32) return mask def get_initial_state(self, x_pos, y_pos, size, min_size=None, max_size=None): """Get initial state dictionary""" - state = { - "x": x_pos, - "y": y_pos, - "size": size, - "initialized": True - } + state = {"x": x_pos, "y": y_pos, "size": size, "initialized": True} if min_size is not None: state["min_size"] = min_size if max_size is not None: state["max_size"] = max_size - return state \ No newline at end of file + return state diff --git a/src/realtimenodes/patterns/movement_patterns.py b/src/realtimenodes/patterns/movement_patterns.py index 0444053..7551877 100644 --- a/src/realtimenodes/patterns/movement_patterns.py +++ b/src/realtimenodes/patterns/movement_patterns.py @@ -7,79 +7,88 @@ from ..control_base import Pattern class StaticPattern(Pattern): def get_name(self) -> str: return "static" - + def calculate(self, phase, min_val, max_val): return min_val + class SinePattern(Pattern): def get_name(self) -> str: return "sine" - + def calculate(self, phase, min_val, max_val): amplitude = (max_val - min_val) / 2 center = min_val + amplitude return center + amplitude * math.sin(2 * math.pi * phase) + class TrianglePattern(Pattern): def get_name(self) -> str: return "triangle" - + def calculate(self, phase, min_val, max_val): if phase < 0.5: return min_val + (max_val - min_val) * (2 * phase) return max_val - (max_val - min_val) * (2 * (phase - 0.5)) + class SawtoothPattern(Pattern): def get_name(self) -> str: return "sawtooth" - + def calculate(self, phase, min_val, max_val): return min_val + (max_val - min_val) * phase + class SquarePattern(Pattern): def get_name(self) -> str: return "square" - + def calculate(self, phase, min_val, max_val): return max_val if phase < 0.5 else min_val + class BouncePattern(Pattern): def get_name(self) -> str: return "bounce" - + def calculate(self, phase, min_val, max_val): bounce = abs(math.sin(math.pi * phase)) return min_val + (max_val - min_val) * bounce + class ExponentialPattern(Pattern): def get_name(self) -> str: return "exponential" - + def calculate(self, phase, min_val, max_val): exp_phase = math.exp(4 * phase) - 1 exp_max = math.exp(4) - 1 return min_val + (max_val - min_val) * (exp_phase / exp_max) + class LogarithmicPattern(Pattern): def get_name(self) -> str: return "logarithmic" - + def calculate(self, phase, min_val, max_val): log_phase = math.log(1 + 99 * phase) / math.log(100) return min_val + (max_val - min_val) * log_phase + class PulsePattern(Pattern): def get_name(self) -> str: return "pulse" - + def calculate(self, phase, min_val, max_val): pulse = math.exp(-10 * ((phase - 0.5) ** 2)) return min_val + (max_val - min_val) * pulse + class RandomWalkPattern(Pattern): def get_name(self) -> str: return "random_walk" - + def __init__(self): self.last_delta = 0 self.current_pos = 0.5 @@ -88,24 +97,27 @@ class RandomWalkPattern(Pattern): momentum = 0.7 random_component = (2 * random.random() - 1) * 0.1 self.last_delta = momentum * self.last_delta + random_component - + self.current_pos += self.last_delta self.current_pos = max(0, min(1, self.current_pos)) - + return min_val + (max_val - min_val) * self.current_pos + class SmoothNoisePattern(Pattern): def get_name(self) -> str: return "smooth_noise" - + def calculate(self, phase, min_val, max_val): t = phase * 10 # Scale phase for more interesting variation - noise = (math.sin(t) + math.sin(2.2*t + 5.52) + math.sin(3.6*t + 4.12)) / 3 + noise = (math.sin(t) + math.sin(2.2 * t + 5.52) + math.sin(3.6 * t + 4.12)) / 3 return min_val + (max_val - min_val) * (noise + 1) / 2 + # Pattern registry MOVEMENT_PATTERNS = { - pattern().get_name(): pattern() for pattern in [ + pattern().get_name(): pattern() + for pattern in [ StaticPattern, SinePattern, TrianglePattern, @@ -116,6 +128,6 @@ MOVEMENT_PATTERNS = { LogarithmicPattern, PulsePattern, RandomWalkPattern, - SmoothNoisePattern + SmoothNoisePattern, ] -} \ No newline at end of file +} diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 4f93b20..22c9785 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -21,4 +21,4 @@ from .timing import TimestampProvider # For backward compatibility, expose the original names # This allows existing code to continue working while we transition image_to_cv2 = convert_to_cv2 -cv2_to_image = convert_to_tensor \ No newline at end of file +cv2_to_image = convert_to_tensor diff --git a/src/utils/general.py b/src/utils/general.py index e911761..d8153f5 100644 --- a/src/utils/general.py +++ b/src/utils/general.py @@ -2,12 +2,14 @@ General utilities for ComfyUI RealTimeNodes. """ + class AlwaysEqualProxy(str): """Helper class to allow connecting multiple specific input types to a single input node. - Originally borrowed from https://github.com/theUpsider/ComfyUI-Logic + Originally borrowed from https://github.com/theUpsider/ComfyUI-Logic """ + def __eq__(self, _): return True def __ne__(self, _): - return False \ No newline at end of file + return False diff --git a/src/utils/image.py b/src/utils/image.py index 90d7f24..911c2f5 100644 --- a/src/utils/image.py +++ b/src/utils/image.py @@ -16,55 +16,57 @@ def convert_to_cv2(tensor: torch.Tensor) -> list: """Converts a ComfyUI IMAGE tensor (BHWC, float32, 0-1) to a list of cv2 images (BGR, uint8).""" if tensor.ndim != 4 or tensor.shape[3] != 3: raise ValueError(f"Input tensor must be BHWC with 3 channels, got shape: {tensor.shape}") - + # Convert entire batch to numpy, scale to 0-255, and convert to uint8 images_np = tensor.cpu().numpy() images_np = (images_np * 255).clip(0, 255).astype(np.uint8) - + # Convert RGB to BGR for cv2 cv2_images = [cv2.cvtColor(img, cv2.COLOR_RGB2BGR) for img in images_np] return cv2_images + def convert_to_tensor(cv2_images) -> torch.Tensor: """Converts cv2 image(s) (BGR, uint8) to a ComfyUI IMAGE tensor (BHWC, float32, 0-1). - + Args: cv2_images: A single cv2 image or a list of cv2 images """ # Handle both single image and list of images if not isinstance(cv2_images, list): cv2_images = [cv2_images] - + # Process each image np_images = [] for cv2_image in cv2_images: # Convert BGR to RGB rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) - + # Scale to 0-1 image_np = rgb_image.astype(np.float32) / 255.0 np_images.append(image_np) - + # Stack all images into a batch batch_np = np.stack(np_images, axis=0) - + # Convert to tensor (already has batch dimension) tensor = torch.from_numpy(batch_np) - + return tensor -def create_mask_from_points(height: int, width: int, points: List[Tuple[int, int]], device='cpu') -> torch.Tensor: + +def create_mask_from_points(height: int, width: int, points: List[Tuple[int, int]], device="cpu") -> torch.Tensor: """Creates a binary mask tensor from a list of points using convex hull.""" mask = np.zeros((height, width), dtype=np.uint8) if len(points) < 3: return torch.from_numpy(mask).float().unsqueeze(0).to(device) - + try: np_points = np.array(points, dtype=np.int32) hull = cv2.convexHull(np_points) cv2.fillPoly(mask, [hull], 1) except Exception: return torch.from_numpy(mask).float().unsqueeze(0).to(device) - + mask_tensor = torch.from_numpy(mask).float().unsqueeze(0).to(device) - return mask_tensor \ No newline at end of file + return mask_tensor diff --git a/src/utils/math.py b/src/utils/math.py index 60bffb2..e88ec18 100644 --- a/src/utils/math.py +++ b/src/utils/math.py @@ -11,13 +11,12 @@ import math from typing import Optional, Tuple -def scale_value(value: Optional[float], - input_min: float, input_max: float, - output_min: float, output_max: float, - clamp: bool = True) -> Optional[float]: +def scale_value( + value: Optional[float], input_min: float, input_max: float, output_min: float, output_max: float, clamp: bool = True +) -> Optional[float]: """ Linearly scales a value from one range to another. - + Args: value: The value to scale input_min: Minimum value of input range @@ -25,7 +24,7 @@ def scale_value(value: Optional[float], output_min: Minimum value of output range output_max: Maximum value of output range clamp: Whether to clamp the output to the output range - + Returns: Scaled value, or None if input was None """ @@ -33,49 +32,49 @@ def scale_value(value: Optional[float], return None if input_max == input_min: # Avoid division by zero return output_min if value <= input_min else output_max - + # Scale scaled = output_min + (value - input_min) * (output_max - output_min) / (input_max - input_min) - + # Clamp if clamp: scaled = max(output_min, min(output_max, scaled)) - + return scaled -def calculate_euclidean_distance(point1: Tuple[float, float, float], - point2: Tuple[float, float, float]) -> float: + +def calculate_euclidean_distance(point1: Tuple[float, float, float], point2: Tuple[float, float, float]) -> float: """ Calculates the Euclidean distance between two 3D points. - + Args: point1: First point as (x, y, z) point2: Second point as (x, y, z) - + Returns: Euclidean distance between the points """ delta_x = point1[0] - point2[0] delta_y = point1[1] - point2[1] delta_z = point1[2] - point2[2] - + distance = math.sqrt(delta_x**2 + delta_y**2 + delta_z**2) return distance -def calculate_euclidean_distance_2d(point1: Tuple[float, float], - point2: Tuple[float, float]) -> float: + +def calculate_euclidean_distance_2d(point1: Tuple[float, float], point2: Tuple[float, float]) -> float: """ Calculates the Euclidean distance between two 2D points. - + Args: point1: First point as (x, y) point2: Second point as (x, y) - + Returns: Euclidean distance between the points """ delta_x = point1[0] - point2[0] delta_y = point1[1] - point2[1] - + distance = math.sqrt(delta_x**2 + delta_y**2) - return distance \ No newline at end of file + return distance diff --git a/src/utils/similar_image_filter.py b/src/utils/similar_image_filter.py index 2d9bb25..a5261cb 100644 --- a/src/utils/similar_image_filter.py +++ b/src/utils/similar_image_filter.py @@ -41,6 +41,6 @@ class SimilarImageFilter: def set_threshold(self, threshold: float) -> None: self.threshold = threshold - + def set_max_skip_frame(self, max_skip_frame: float) -> None: - self.max_skip_frame = max_skip_frame \ No newline at end of file + self.max_skip_frame = max_skip_frame diff --git a/src/utils/timing.py b/src/utils/timing.py index 2460e69..5b136c9 100644 --- a/src/utils/timing.py +++ b/src/utils/timing.py @@ -12,6 +12,7 @@ class TimestampProvider: This ensures timestamps are always increasing while staying aligned with real elapsed time. Useful for any process that requires monotonically increasing time values. """ + def __init__(self): self._start_time = time.time() self._last_timestamp = 0 @@ -22,4 +23,4 @@ class TimestampProvider: if now <= self._last_timestamp: now = self._last_timestamp + 1 self._last_timestamp = now - return now \ No newline at end of file + return now From 509a8eb215d3e7f880e935c62f4d5a48b0487212 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 26 Apr 2025 20:55:56 -0400 Subject: [PATCH 20/20] pyproject --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index afd31b3..506628d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "comfyui_realtimenodes" description = "These nodes are for real-time applications of ComfyUI." -version = "1.3.0" +version = "1.3.1" license = {file = "LICENSE"} dependencies = ["opencv-python"]