From 73f9164fd19d46e48eef06cee6ffa89d89dcea8f Mon Sep 17 00:00:00 2001 From: ryanontheinstide Date: Sun, 8 Sep 2024 16:08:38 -0400 Subject: [PATCH] DEPTH --- __init__.py | 9 +++- node_configs.py | 15 ++++++ nodes/flex/feature_externals.py | 86 ++++++++++++++++++++++++++++++++ nodes/flex/feature_extractors.py | 28 ++++++++++- nodes/flex/feature_modulation.py | 2 +- nodes/flex/features.py | 61 +++++++++++++++++++++- pyproject.toml | 2 +- readme.md | 3 +- requirements.txt | 3 +- 9 files changed, 200 insertions(+), 9 deletions(-) diff --git a/__init__.py b/__init__.py index 4860673..51dfd87 100644 --- a/__init__.py +++ b/__init__.py @@ -65,6 +65,7 @@ from .nodes.flex.feature_extractors import( ColorFeatureNode, BrightnessFeatureNode, MotionFeatureNode, + AreaFeatureNode, ) from .nodes.flex.midi_feature_extractor import( @@ -160,6 +161,7 @@ from .nodes.images.flex_images import ( from .nodes.flex.feature_externals import ( FeatureToWeightsStrategy, + DepthShapeModifier, ) from .nodes.flex.feature_modulation import ( @@ -247,13 +249,16 @@ NODE_CLASS_MAPPINGS = { "ColorFeatureNode": ColorFeatureNode, "BrightnessFeatureNode": BrightnessFeatureNode, "MotionFeatureNode": MotionFeatureNode, - "FeatureToWeightsStrategy": FeatureToWeightsStrategy, "LocationFromMask": LocationFromMask, "ProximityFeatureNode": ProximityFeatureNode, "ProximityVisualizer": ProximityVisualizer, "LocationFromPoint": LocationFromPoint, "LocationTransform": LocationTransform, - + "AreaFeatureNode": AreaFeatureNode, + + "FeatureToWeightsStrategy": FeatureToWeightsStrategy, + "DepthShapeModifier": DepthShapeModifier, + #feature modulation "FeatureMixer": FeatureMixer, "FeatureCombine": FeatureCombine, diff --git a/node_configs.py b/node_configs.py index 3e713ab..9a7df1a 100644 --- a/node_configs.py +++ b/node_configs.py @@ -593,6 +593,21 @@ Analyzes the input depth maps to extract the specified depth-related feature. Th """ }) +add_node_config("AreaFeatureNode", { + "TOP_DESCRIPTION": "Extracts area-related features from mask sequences for mask modulation.", + "ADDITIONAL_INFO": """ +- `masks`: Input mask sequence to analyze (MASK type) +- `feature_type`: Type of area feature to extract + - Options: + - "total_area" (sum of pixels above threshold) + - "largest_contour" (area of the largest contiguous region) + - "bounding_box" (area of the bounding box containing the largest region) +- `threshold`: Threshold value for considering pixels as part of the area (0.0 to 1.0) + +This node analyzes the input mask sequence to extract the specified area-related feature. The resulting feature can be used to modulate masks based on changes in area over time, allowing for effects that respond to the size or extent of masked regions in the scene. +""" +}) + add_node_config("ProximityFeatureNode", { "TOP_DESCRIPTION": "Calculates a proximity feature based on the distance between anchor and query locations in video frames.", "ADDITIONAL_INFO": """ diff --git a/nodes/flex/feature_externals.py b/nodes/flex/feature_externals.py index 93df7ca..9348c19 100644 --- a/nodes/flex/feature_externals.py +++ b/nodes/flex/feature_externals.py @@ -1,3 +1,7 @@ +import torch +import torch.nn.functional as F +import numpy as np +from ... import RyanOnTheInside from ... import RyanOnTheInside class FlexExternalModulator(RyanOnTheInside): @@ -36,3 +40,85 @@ class FeatureToWeightsStrategy(FlexExternalModulator): return (weights_strategy,) + + + +import torch +import torch.nn.functional as F +import cv2 + +class DepthShapeModifier(RyanOnTheInside): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "depth_map": ("IMAGE",), + "mask": ("MASK",), + "gradient_steepness": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "depth_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "depth_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "modify_depth" + CATEGORY = "RyanOnTheInside/DepthModifiers" + + def modify_depth(self, depth_map, mask, gradient_steepness, depth_min, depth_max, strength): + device = depth_map.device + mask = mask.to(device).bool() + + b, h, w, c = depth_map.shape + + modified_depths = [] + for i in range(b): + mask_i = mask[i].cpu().numpy().astype(np.uint8) + + # Use cv2 to find contours (separate shapes) + contours, _ = cv2.findContours(mask_i, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + sphere_gradients = torch.zeros((h, w), device=device) + + for contour in contours: + # Create a mask for this contour + component_mask = np.zeros_like(mask_i) + cv2.drawContours(component_mask, [contour], 0, 1, -1) + component_mask = torch.from_numpy(component_mask).to(device).bool() + + # Find center and radius of the circular component + M = cv2.moments(contour) + if M["m00"] != 0: + center_x = int(M["m10"] / M["m00"]) + center_y = int(M["m01"] / M["m00"]) + else: + continue # Skip this contour if we can't find its center + + # Calculate radius (approximate) + radius = np.sqrt(cv2.contourArea(contour) / np.pi) + + # Generate spherical gradient + y, x = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing='ij') + distances = torch.sqrt((y - center_y)**2 + (x - center_x)**2) + sphere_gradient = torch.clamp(1 - (distances / radius) ** gradient_steepness, 0, 1) + + # Apply gradient only to this component + sphere_gradients += sphere_gradient * component_mask.float() + + # Scale gradient to depth range + depth_gradient = depth_min + sphere_gradients * (depth_max - depth_min) + + # Apply gradient to depth map + modified_depth = depth_map[i].clone() + depth_gradient = depth_gradient.unsqueeze(-1).repeat(1, 1, c) + modified_depth = torch.where(mask[i].unsqueeze(-1).repeat(1, 1, c), + depth_gradient, + modified_depth) + + # Blend modified depth with original depth + blend_mask = (mask[i].unsqueeze(-1) * strength).repeat(1, 1, c) + modified_depth = depth_map[i] * (1 - blend_mask) + modified_depth * blend_mask + modified_depths.append(modified_depth) + + result = torch.stack(modified_depths, dim=0) + return (result,) \ No newline at end of file diff --git a/nodes/flex/feature_extractors.py b/nodes/flex/feature_extractors.py index 1105f32..fa46237 100644 --- a/nodes/flex/feature_extractors.py +++ b/nodes/flex/feature_extractors.py @@ -1,6 +1,6 @@ from .feature_pipe import FeaturePipe from ... import RyanOnTheInside -from .features import AudioFeature, TimeFeature, DepthFeature, ColorFeature, BrightnessFeature, MotionFeature +from .features import AudioFeature, TimeFeature, DepthFeature, ColorFeature, BrightnessFeature, MotionFeature, AreaFeature from .proximity_feature import Location from tqdm import tqdm from comfy.utils import ProgressBar @@ -182,4 +182,28 @@ class MotionFeatureNode(FirstFeature): motion_feature.extract() self.end_progress() - return (motion_feature, feature_pipe) \ No newline at end of file + return (motion_feature, feature_pipe) + + +class AreaFeatureNode(FirstFeature): + @classmethod + def INPUT_TYPES(cls): + return { + **super().INPUT_TYPES(), + "required": { + **super().INPUT_TYPES()["required"], + "masks": ("MASK",), + "feature_type": (["total_area", "largest_contour", "bounding_box"],), + "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("FEATURE", "FEATURE_PIPE") + FUNCTION = "create_feature" + + def create_feature(self, masks, video_frames, frame_rate, feature_type, threshold): + feature_pipe = FeaturePipe(frame_rate, video_frames) + area_feature = AreaFeature("area_feature", feature_pipe.frame_rate, feature_pipe.frame_count, + masks, feature_type=feature_type, threshold=threshold) + area_feature.extract() + return (area_feature, feature_pipe) \ No newline at end of file diff --git a/nodes/flex/feature_modulation.py b/nodes/flex/feature_modulation.py index 262f4c9..ab55b78 100644 --- a/nodes/flex/feature_modulation.py +++ b/nodes/flex/feature_modulation.py @@ -6,7 +6,7 @@ from io import BytesIO from PIL import Image class FeatureModulationBase(RyanOnTheInside): - CATEGORY = "RyanOnTheInside/FlexFeatures" + CATEGORY = "RyanOnTheInside/FlexFeatures/FeatureModulators" FUNCTION = "modulate" @classmethod diff --git a/nodes/flex/features.py b/nodes/flex/features.py index 37e5c1e..28eedab 100644 --- a/nodes/flex/features.py +++ b/nodes/flex/features.py @@ -428,6 +428,65 @@ class MotionFeature(BaseFeature): self.feature_name = feature_name else: raise ValueError(f"Invalid feature name. Available features are: {', '.join(self.available_features)}") + +class AreaFeature(BaseFeature): + def __init__(self, name, frame_rate, frame_count, masks, feature_type='total_area', threshold=0.5): + self.masks = masks + self.feature_type = feature_type + self.threshold = threshold + self.available_features = ["total_area", "largest_contour", "bounding_box"] + super().__init__(name, "area", frame_rate, frame_count) + + def extract(self): + self.data = [] -#TODO area feature + for mask in self.masks: + mask_np = mask.cpu().numpy() + binary_mask = (mask_np > self.threshold).astype(np.uint8) + + if self.feature_type == 'total_area': + area = np.sum(binary_mask) + elif self.feature_type == 'largest_contour': + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if contours: + largest_contour = max(contours, key=cv2.contourArea) + area = cv2.contourArea(largest_contour) + else: + area = 0 + elif self.feature_type == 'bounding_box': + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if contours: + x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea)) + area = w * h + else: + area = 0 + else: + raise ValueError(f"Unsupported feature type: {self.feature_type}") + + self.data.append(area) + + return self.normalize() + + def normalize(self): + if self.data: + min_val = min(self.data) + max_val = max(self.data) + if max_val > min_val: + self.data = [(v - min_val) / (max_val - min_val) for v in self.data] + else: + self.data = [0] * len(self.data) + return self + + def get_value_at_frame(self, frame_index): + if self.data is not None and 0 <= frame_index < len(self.data): + return self.data[frame_index] + else: + raise ValueError("Invalid frame index or no data available") + + def set_active_feature(self, feature_name): + if feature_name in self.available_features: + self.feature_type = feature_name + self.extract() # Re-extract the data with the new feature type + else: + raise ValueError(f"Invalid feature name. Available features are: {', '.join(self.available_features)}") #TODO volume feature diff --git a/pyproject.toml b/pyproject.toml index 17b37fd..ad3cc4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "comfyui_ryanonyheinside" description = "Custom nodes introducing particle simulations, optical flow, audio manipulation & reactivity, and temporal masks" version = "1.6.0" license = {file = "LICENSE"} -dependencies = ["opencv-python==4.10.0","scipy","pillow","librosa","pymunk==6.8.1","matplotlib","openunmix","mido"] +dependencies = ["opencv-python==4.10.0","scipy","torchaudio", "pillow","librosa","pymunk==6.8.1","matplotlib","openunmix","mido"] [project.urls] Repository = "https://github.com/ryanontheinside/ComfyUI_RyanOnTheInside" diff --git a/readme.md b/readme.md index c952172..96979e7 100644 --- a/readme.md +++ b/readme.md @@ -33,7 +33,8 @@ *Examples showcasing various effects using particle emitters, vortices, and other node features* ##### 🆕 Recent Updates -- 9/10/24 - **Proximity Feature**: Adds proximity as a driving reactivity feature! Allows for the distince of objects from one another to control other nodes. +- 9/8/24 - **Area Feature**: Adds area as a driving reactivity feature! +- 9/7/24 - **Proximity Feature**: Adds proximity as a driving reactivity feature! Allows for the distince of objects from one another to control other nodes. - 9/5/24 - **FlexImageParallax**: Add illusory depth. - 9/3/24 - **FlexMaskDepthChamber**: Mask anything within a given depth range - 9/3/24 - **FeatureFade**: another feature modifier to allow mixing of features, controled by features. Think mixing depth and motion features according to a kick drum or somthing. Ridiculous. diff --git a/requirements.txt b/requirements.txt index d5e8409..4f9d321 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ librosa pymunk==6.8.1 matplotlib openunmix -mido \ No newline at end of file +mido +torchaudio \ No newline at end of file