143 lines
5.2 KiB
Python
143 lines
5.2 KiB
Python
import torch
|
|
|
|
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons.
|
|
Used to create a generic type that can accept any input without type constraints."""
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any_type = AnyType("*")
|
|
|
|
|
|
class MaskAreaCondition:
|
|
"""
|
|
A ComfyUI node that analyzes the size of a mask relative to the total tensor area
|
|
and outputs whether it falls below a specified threshold percentage.
|
|
|
|
This is primarily used for conditional workflows, allowing different processing
|
|
based on the size of the mask (e.g., only run face detailing if the face is small).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"threshold_percent": ("FLOAT", {
|
|
"default": 10.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BOOLEAN", "FLOAT", "MASK")
|
|
RETURN_NAMES = ("is_below_threshold", "mask_area_percent", "mask_passthrough")
|
|
FUNCTION = "check_mask_area"
|
|
CATEGORY = "mask/conditional"
|
|
|
|
def check_mask_area(self, mask: torch.Tensor, threshold_percent: float):
|
|
"""
|
|
Checks the area of the mask relative to the total tensor area.
|
|
|
|
Args:
|
|
mask (torch.Tensor): The input mask tensor (batch, height, width).
|
|
threshold_percent (float): The threshold percentage (0.0 - 100.0).
|
|
|
|
Returns:
|
|
tuple[bool, float, torch.Tensor]:
|
|
- is_below_threshold: Boolean indicating if the mask area is below the threshold.
|
|
- mask_area_percent: The area of the mask as a percentage of the total area.
|
|
- mask_passthrough: The original mask tensor passed through for convenience.
|
|
"""
|
|
print(f"[MaskAreaCondition] Executing with threshold: {threshold_percent}%")
|
|
|
|
if mask.numel() == 0:
|
|
# Handle empty mask case
|
|
mask_area_percent = 0.0
|
|
# Empty mask = nothing detected = no need to run inpainting
|
|
is_below_threshold = False
|
|
print("[MaskAreaCondition] Input mask is empty. Returning False for is_below_threshold.")
|
|
else:
|
|
# Count non-zero pixels (considered part of the mask)
|
|
# sum() works for float masks (e.g., from blurring) as well as binary masks
|
|
mask_pixels = torch.sum(mask > 0).item()
|
|
|
|
# Total number of elements in the tensor
|
|
total_pixels = mask.numel()
|
|
|
|
# Calculate mask percentage
|
|
mask_area_percent = (mask_pixels / total_pixels) * 100.0
|
|
|
|
print(f"[MaskAreaCondition] Mask Pixels: {mask_pixels}, Total Pixels: {total_pixels}, Calculated Area: {mask_area_percent:.2f}%")
|
|
|
|
# Check if the mask area is below the threshold
|
|
is_below_threshold = mask_area_percent < threshold_percent
|
|
|
|
print(f"[MaskAreaCondition] Mask area {mask_area_percent:.2f}% is {'SMALL' if is_below_threshold else 'LARGE'} (threshold: {threshold_percent}%)")
|
|
|
|
# Return results, including the original mask for convenience in workflows
|
|
return (is_below_threshold, mask_area_percent, mask)
|
|
|
|
|
|
class SelectData:
|
|
"""
|
|
Selects one of two data inputs based on a boolean condition.
|
|
|
|
This node is essential for building conditional workflows, allowing different
|
|
processing paths based on conditions like mask size. It works with any data type
|
|
through the use of the AnyType helper.
|
|
|
|
Example use:
|
|
- When paired with MaskAreaCondition, can choose between different processing
|
|
paths based on the size of a detected feature.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"data_if_true": (any_type,),
|
|
"data_if_false": (any_type,),
|
|
"condition": ("BOOLEAN", {"forceInput": True})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (any_type,)
|
|
RETURN_NAMES = ("selected_data",)
|
|
FUNCTION = "select_data"
|
|
CATEGORY = "mask/conditional"
|
|
|
|
def select_data(self, data_if_true, data_if_false, condition):
|
|
"""
|
|
Selects between two inputs based on a boolean condition.
|
|
|
|
Args:
|
|
data_if_true: Data to return if condition is True
|
|
data_if_false: Data to return if condition is False
|
|
condition: Boolean value determining which input to select
|
|
|
|
Returns:
|
|
The selected data based on the condition
|
|
"""
|
|
print(f"[SelectData] Condition is {condition}. Selecting data_if_{str(condition).lower()}.")
|
|
|
|
if condition:
|
|
# If condition is True, select data_if_true
|
|
return (data_if_true,)
|
|
else:
|
|
# If condition is False, select data_if_false
|
|
return (data_if_false,)
|
|
|
|
# Node registration dictionary
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MaskAreaCondition": MaskAreaCondition,
|
|
"SelectData": SelectData
|
|
}
|
|
|
|
# Human-readable names dictionary
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MaskAreaCondition": "Mask Area Condition",
|
|
"SelectData": "Select Data based on Condition"
|
|
} |