""" Mask Generation Utilities for TextureAlchemy Advanced mask creation for weathering, selection, and compositing """ import torch import torch.nn.functional as F import math def _rgb_to_hsv_tensor(rgb): """Convert RGB tensor (..., 3) to HSV in 0-1. Used for color-to-mask logic.""" r = rgb[..., 0] g = rgb[..., 1] b = rgb[..., 2] max_c = torch.max(torch.max(r, g), b) min_c = torch.min(torch.min(r, g), b) diff = max_c - min_c h = torch.zeros_like(max_c) mask = diff != 0 h = torch.where((max_c == r) & mask, ((g - b) / (diff + 1e-7)) % 6.0, h) h = torch.where((max_c == g) & mask, ((b - r) / (diff + 1e-7)) + 2.0, h) h = torch.where((max_c == b) & mask, ((r - g) / (diff + 1e-7)) + 4.0, h) h = h / 6.0 s = torch.where(max_c != 0, diff / (max_c + 1e-7), torch.zeros_like(max_c)) v = max_c return torch.stack([h, s, v], dim=-1) class EdgeWearMaskGenerator: """ Generate edge wear masks from normal maps and curvature Perfect for realistic paint chipping and edge damage """ @classmethod def INPUT_TYPES(cls): return { "required": { "width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "edge_width": ("FLOAT", { "default": 0.1, "min": 0.01, "max": 0.5, "step": 0.01, "display": "number", "tooltip": "Width of edge wear (0.0-1.0)" }), "edge_sharpness": ("FLOAT", { "default": 2.0, "min": 0.1, "max": 10.0, "step": 0.1, "display": "number", "tooltip": "Edge falloff sharpness" }), "noise_scale": ("FLOAT", { "default": 10.0, "min": 1.0, "max": 100.0, "step": 1.0, "display": "number", "tooltip": "Variation noise scale" }), "noise_strength": ("FLOAT", { "default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Amount of edge variation" }), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), }, "optional": { "normal_map": ("IMAGE",), "curvature_map": ("IMAGE",), "ao_map": ("IMAGE",), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("mask",) FUNCTION = "generate" CATEGORY = "Texture Alchemist/Masks" def generate(self, width, height, edge_width, edge_sharpness, noise_scale, noise_strength, seed, normal_map=None, curvature_map=None, ao_map=None): """Generate edge wear mask""" print(f"\n{'='*60}") print(f"Edge Wear Mask Generator") print(f"{'='*60}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 # Start with procedural edge detection or use provided maps if curvature_map is not None: # Use curvature directly mask = curvature_map.clone() if mask.shape[1:3] != (height, width): mask = F.interpolate( mask.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) elif normal_map is not None: # Derive curvature from normal map mask = self._normal_to_curvature(normal_map, height, width) else: # Create procedural edge pattern mask = torch.zeros((1, height, width, 1), device=device, dtype=dtype) # Add noise variation if noise_strength > 0.0: noise = self._generate_noise(height, width, noise_scale, seed, device, dtype) mask = mask + noise * noise_strength # Apply edge width and sharpness mask = torch.pow(torch.clamp(mask, 0.0, 1.0), 1.0 / edge_sharpness) mask = torch.clamp(mask / edge_width, 0.0, 1.0) # Optionally modulate by AO (wear accumulates in crevices) if ao_map is not None: ao = ao_map.clone() if ao.shape[1:3] != (height, width): ao = F.interpolate( ao.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) if ao.shape[3] > 1: ao = ao.mean(dim=3, keepdim=True) # More wear in exposed areas (high AO) mask = mask * (1.0 - ao * 0.3) # Convert to 3-channel for compatibility mask = mask.mean(dim=3, keepdim=True).repeat(1, 1, 1, 3) mask = torch.clamp(mask, 0.0, 1.0) print(f"✓ Edge wear mask generated") print(f" Resolution: {width}×{height}") print(f" Edge width: {edge_width:.2f}") print(f" Noise: {noise_strength:.2f}") print(f"{'='*60}\n") return (mask,) def _normal_to_curvature(self, normal_map, height, width): """Derive curvature from normal map""" if normal_map.shape[1:3] != (height, width): normal_map = F.interpolate( normal_map.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) # Simple curvature: how much does normal change dx = torch.abs(normal_map[:, :, 1:, :] - normal_map[:, :, :-1, :]) dy = torch.abs(normal_map[:, 1:, :, :] - normal_map[:, :-1, :, :]) dx = F.pad(dx, (0, 0, 0, 1), mode='replicate') dy = F.pad(dy, (0, 0, 0, 0, 0, 1), mode='replicate') curv = (dx + dy).mean(dim=3, keepdim=True) return curv def _generate_noise(self, height, width, scale, seed, device, dtype): """Generate Perlin-like noise""" torch.manual_seed(seed) # Simple multi-octave noise noise = torch.zeros((1, height, width, 1), device=device, dtype=dtype) for octave in range(3): freq = scale * (2 ** octave) amp = 1.0 / (2 ** octave) grid_h = int(height / freq) + 2 grid_w = int(width / freq) + 2 grid = torch.rand((1, grid_h, grid_w, 1), device=device, dtype=dtype) upsampled = F.interpolate( grid.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) noise = noise + upsampled * amp noise = (noise - noise.min()) / (noise.max() - noise.min() + 1e-7) return noise class DirtGrimeMaskGenerator: """ Generate procedural dirt and grime accumulation masks Based on AO and crevice detection """ @classmethod def INPUT_TYPES(cls): return { "required": { "width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "density": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Overall dirt coverage" }), "scale": ("FLOAT", { "default": 20.0, "min": 1.0, "max": 100.0, "step": 1.0, "display": "number", "tooltip": "Size of dirt patches" }), "crevice_strength": ("FLOAT", { "default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Dirt accumulation in crevices" }), "contrast": ("FLOAT", { "default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1, "display": "number", "tooltip": "Dirt pattern contrast" }), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), }, "optional": { "ao_map": ("IMAGE",), "height_map": ("IMAGE",), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("mask",) FUNCTION = "generate" CATEGORY = "Texture Alchemist/Masks" def generate(self, width, height, density, scale, crevice_strength, contrast, seed, ao_map=None, height_map=None): """Generate dirt/grime mask""" print(f"\n{'='*60}") print(f"Dirt/Grime Mask Generator") print(f"{'='*60}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 torch.manual_seed(seed) # Generate base noise pattern mask = torch.zeros((1, height, width, 1), device=device, dtype=dtype) # Multi-scale noise for natural dirt patterns for octave in range(4): freq = scale * (2 ** octave) amp = 1.0 / (2 ** (octave + 1)) grid_h = max(int(height / freq), 2) + 2 grid_w = max(int(width / freq), 2) + 2 grid = torch.rand((1, grid_h, grid_w, 1), device=device, dtype=dtype) upsampled = F.interpolate( grid.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) mask = mask + upsampled * amp # Normalize mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-7) # Apply contrast mask = torch.pow(mask, 1.0 / contrast) # Apply density threshold mask = torch.clamp((mask - (1.0 - density)) / density, 0.0, 1.0) # Modulate by AO (dirt accumulates in crevices) if ao_map is not None: ao = ao_map.clone() if ao.shape[1:3] != (height, width): ao = F.interpolate( ao.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) if ao.shape[3] > 1: ao = ao.mean(dim=3, keepdim=True) # More dirt in dark areas (low AO = crevices) ao_factor = 1.0 - ao mask = mask + ao_factor * crevice_strength # Modulate by height (dirt in low areas) if height_map is not None: hm = height_map.clone() if hm.shape[1:3] != (height, width): hm = F.interpolate( hm.permute(0, 3, 1, 2), size=(height, width), mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) if hm.shape[3] > 1: hm = hm.mean(dim=3, keepdim=True) # More dirt in low areas height_factor = 1.0 - hm mask = mask + height_factor * crevice_strength * 0.5 # Convert to 3-channel mask = mask.repeat(1, 1, 1, 3) mask = torch.clamp(mask, 0.0, 1.0) print(f"✓ Dirt/grime mask generated") print(f" Resolution: {width}×{height}") print(f" Density: {density:.2f}") print(f" Scale: {scale:.1f}") print(f"{'='*60}\n") return (mask,) class GradientMaskGenerator: """ Generate gradient masks with multiple styles and noise overlay """ @classmethod def INPUT_TYPES(cls): return { "required": { "width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), "gradient_type": (["linear", "radial", "angle", "diamond", "square"],), "angle": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0, "display": "number", "tooltip": "Gradient angle (degrees)" }), "center_x": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Center X position (0-1)" }), "center_y": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Center Y position (0-1)" }), "noise_amount": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Add noise variation" }), "noise_scale": ("FLOAT", { "default": 20.0, "min": 1.0, "max": 100.0, "step": 1.0, "display": "number" }), "invert": ("BOOLEAN", {"default": False}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("mask",) FUNCTION = "generate" CATEGORY = "Texture Alchemist/Masks" def generate(self, width, height, gradient_type, angle, center_x, center_y, noise_amount, noise_scale, invert, seed): """Generate gradient mask""" print(f"\n{'='*60}") print(f"Gradient Mask Generator - {gradient_type.upper()}") print(f"{'='*60}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 # Create coordinate grid y_coords = torch.linspace(0, 1, height, device=device, dtype=dtype).view(-1, 1).repeat(1, width) x_coords = torch.linspace(0, 1, width, device=device, dtype=dtype).view(1, -1).repeat(height, 1) cx = center_x cy = center_y if gradient_type == "linear": # Linear gradient based on angle rad = math.radians(angle) mask = x_coords * math.cos(rad) + y_coords * math.sin(rad) elif gradient_type == "radial": # Radial from center dx = x_coords - cx dy = y_coords - cy mask = torch.sqrt(dx * dx + dy * dy) elif gradient_type == "angle": # Angular/conical gradient dx = x_coords - cx dy = y_coords - cy mask = (torch.atan2(dy, dx) + math.pi) / (2 * math.pi) # Rotate by angle mask = (mask + angle / 360.0) % 1.0 elif gradient_type == "diamond": # Diamond shape dx = torch.abs(x_coords - cx) dy = torch.abs(y_coords - cy) mask = dx + dy else: # square # Square shape dx = torch.abs(x_coords - cx) dy = torch.abs(y_coords - cy) mask = torch.max(dx, dy) # Normalize to 0-1 mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-7) # Add noise if requested if noise_amount > 0.0: torch.manual_seed(seed) noise = torch.rand((height, width), device=device, dtype=dtype) # Multi-scale noise for octave in range(2): freq = noise_scale * (2 ** octave) amp = 1.0 / (2 ** octave) grid_h = max(int(height / freq), 2) + 2 grid_w = max(int(width / freq), 2) + 2 grid = torch.rand((grid_h, grid_w), device=device, dtype=dtype) upsampled = F.interpolate( grid.unsqueeze(0).unsqueeze(0), size=(height, width), mode='bilinear', align_corners=False ).squeeze() noise = noise + upsampled * amp noise = (noise - noise.min()) / (noise.max() - noise.min() + 1e-7) mask = mask * (1.0 - noise_amount) + noise * noise_amount # Invert if requested if invert: mask = 1.0 - mask # Convert to image format mask = mask.unsqueeze(0).unsqueeze(3).repeat(1, 1, 1, 3) mask = torch.clamp(mask, 0.0, 1.0) print(f"✓ Gradient mask generated") print(f" Type: {gradient_type}") print(f" Resolution: {width}×{height}") print(f" Noise: {noise_amount:.2f}") print(f"{'='*60}\n") return (mask,) class ColorSelectionMask: """ Create masks by selecting color ranges (like Photoshop's Color Range) Perfect for isolating specific colors for editing """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "target_color": ("STRING", { "default": "#FF0000", "tooltip": "Target color in hex format (#RRGGBB)" }), "tolerance": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Color matching tolerance" }), "feather": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 0.5, "step": 0.01, "display": "number", "tooltip": "Edge softness" }), "mode": (["rgb", "hsv", "luminance"],), "invert": ("BOOLEAN", {"default": False}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("mask",) FUNCTION = "select_color" CATEGORY = "Texture Alchemist/Masks" def select_color(self, image, target_color, tolerance, feather, mode, invert): """Create mask from color selection""" print(f"\n{'='*60}") print(f"Color Selection Mask") print(f"{'='*60}") print(f" Target color: {target_color}") print(f" Mode: {mode}") print(f" Tolerance: {tolerance:.2f}") # Parse hex color target_color = target_color.lstrip('#') r = int(target_color[0:2], 16) / 255.0 g = int(target_color[2:4], 16) / 255.0 b = int(target_color[4:6], 16) / 255.0 target = torch.tensor([r, g, b], device=image.device, dtype=image.dtype) if mode == "rgb": # RGB distance diff = image - target.view(1, 1, 1, 3) distance = torch.sqrt((diff * diff).sum(dim=3, keepdim=True)) max_dist = math.sqrt(3.0) # Maximum RGB distance distance = distance / max_dist elif mode == "hsv": # HSV distance (more perceptual) hsv_image = self._rgb_to_hsv(image) hsv_target = self._rgb_to_hsv(target.view(1, 1, 1, 3)) # Hue is circular, handle wrap-around h_diff = torch.abs(hsv_image[:, :, :, 0:1] - hsv_target[:, :, :, 0:1]) h_diff = torch.min(h_diff, 1.0 - h_diff) * 2.0 # Wrap-around s_diff = torch.abs(hsv_image[:, :, :, 1:2] - hsv_target[:, :, :, 1:2]) v_diff = torch.abs(hsv_image[:, :, :, 2:3] - hsv_target[:, :, :, 2:3]) distance = torch.sqrt(h_diff * h_diff + s_diff * s_diff + v_diff * v_diff) distance = distance / math.sqrt(3.0) else: # luminance # Luminance distance lum_image = 0.299 * image[:, :, :, 0:1] + 0.587 * image[:, :, :, 1:2] + 0.114 * image[:, :, :, 2:3] lum_target = 0.299 * r + 0.587 * g + 0.114 * b distance = torch.abs(lum_image - lum_target) # Convert distance to mask with tolerance and feather mask = 1.0 - torch.clamp(distance / (tolerance + 1e-7), 0.0, 1.0) # Apply feather (smooth edges) if feather > 0.0: feather_range = feather mask = torch.where( mask > (1.0 - feather_range), mask, torch.where( mask < feather_range, torch.zeros_like(mask), (mask - feather_range) / (1.0 - 2.0 * feather_range) ) ) # Invert if requested if invert: mask = 1.0 - mask # Convert to 3-channel mask = mask.repeat(1, 1, 1, 3) mask = torch.clamp(mask, 0.0, 1.0) print(f"✓ Color selection mask created") print(f" Pixels selected: {(mask.mean() * 100):.1f}%") print(f"{'='*60}\n") return (mask,) def _rgb_to_hsv(self, rgb): """Convert RGB to HSV""" r = rgb[:, :, :, 0] g = rgb[:, :, :, 1] b = rgb[:, :, :, 2] max_c = torch.max(torch.max(r, g), b) min_c = torch.min(torch.min(r, g), b) diff = max_c - min_c # Hue h = torch.zeros_like(max_c) mask = diff != 0 h = torch.where((max_c == r) & mask, ((g - b) / (diff + 1e-7)) % 6.0, h) h = torch.where((max_c == g) & mask, ((b - r) / (diff + 1e-7)) + 2.0, h) h = torch.where((max_c == b) & mask, ((r - g) / (diff + 1e-7)) + 4.0, h) h = h / 6.0 # Saturation s = torch.where(max_c != 0, diff / (max_c + 1e-7), torch.zeros_like(max_c)) # Value v = max_c return torch.stack([h, s, v], dim=3) class ImageMaskCombiner: """ Combine image with mask from any source Use with LoadImage (with mask painting) or MaskEditor nodes Outputs both together for easy workflow routing """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "mask": ("MASK",), } } RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") RETURN_NAMES = ("image", "mask", "mask_preview") FUNCTION = "combine" CATEGORY = "Texture Alchemist/Masks" def combine(self, image, mask): """ Combine image and mask, creating mask visualization """ print("\n" + "="*60) print("Image + Mask Combiner") print("="*60) print(f"Image shape: {image.shape}") print(f"Mask shape: {mask.shape}") batch, height, width, channels = image.shape # Process mask if mask is not None: # Ensure mask matches image dimensions if len(mask.shape) == 3: # (B, H, W) if mask.shape[1:3] != (height, width): print(f" Resizing mask from {mask.shape[1:3]} to {(height, width)}") mask_reshaped = mask.unsqueeze(1) mask_resized = F.interpolate( mask_reshaped, size=(height, width), mode='bilinear', align_corners=False ) mask = mask_resized.squeeze(1) # Ensure batch size matches if mask.shape[0] != batch: print(f" Adjusting batch size from {mask.shape[0]} to {batch}") if mask.shape[0] == 1: mask = mask.repeat(batch, 1, 1) if len(mask.shape) == 3 else mask.repeat(batch, 1, 1, 1) else: mask = mask[:batch] # Ensure mask is 3D (B, H, W) if len(mask.shape) == 4: # Convert to grayscale if needed if mask.shape[-1] > 1: weights = torch.tensor([0.299, 0.587, 0.114], device=mask.device, dtype=mask.dtype) mask = torch.sum(mask[..., :3] * weights, dim=-1) else: mask = mask.squeeze(-1) # Create visualization of mask as RGB image mask_preview = mask.unsqueeze(-1).repeat(1, 1, 1, 3) # Convert to RGB # Print mask statistics mask_min = mask.min().item() mask_max = mask.max().item() mask_mean = mask.mean().item() nonzero_percent = (mask > 0.01).float().mean().item() * 100 print(f"\n📊 Mask Statistics:") print(f" Range: [{mask_min:.3f}, {mask_max:.3f}]") print(f" Mean: {mask_mean:.3f}") print(f" Coverage: {nonzero_percent:.1f}% of image") print(f"\n✓ Outputs:") print(f" Image: {image.shape} (passthrough)") print(f" Mask: {mask.shape}") print(f" Mask Preview: {mask_preview.shape}") print("="*60 + "\n") return (image, mask, mask_preview) class MaskCompositor: """ Interactive mask compositor with canvas editor Edit masks directly over images with painting tools Click 'Open Mask Editor' button to paint, then 'Save to Node' """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), }, "hidden": { "mask_data": "STRING", } } RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") RETURN_NAMES = ("image", "mask", "mask_preview") FUNCTION = "composite" CATEGORY = "Texture Alchemist/Masks" def composite(self, image, mask_data="", **kwargs): """ Interactive mask editing over image with painted data from widget """ print("\n" + "="*60) print("Mask Compositor") print("="*60) print(f"Image shape: {image.shape}") batch, height, width, channels = image.shape device = image.device dtype = image.dtype # Parse mask data from widget mask = None if mask_data and mask_data.strip(): try: import json import numpy as np print("📥 Parsing saved mask data...") saved_data = json.loads(mask_data) mask_width = saved_data['width'] mask_height = saved_data['height'] mask_array = np.array(saved_data['data'], dtype=np.uint8) # Reshape to (H, W, 4) RGBA mask_rgba = mask_array.reshape((mask_height, mask_width, 4)) # Convert to grayscale (use R channel since we painted in grayscale) mask_gray = mask_rgba[:, :, 0].astype(np.float32) / 255.0 # Convert to torch tensor mask = torch.from_numpy(mask_gray).to(device=device, dtype=dtype) # Add batch dimension mask = mask.unsqueeze(0) # (1, H, W) print(f"✓ Loaded painted mask: {mask_width}×{mask_height}") print(f" Mask range: [{mask.min():.3f}, {mask.max():.3f}]") # Resize to match image dimensions if needed if (mask_height, mask_width) != (height, width): print(f" Resizing mask from {mask_width}×{mask_height} to {width}×{height}") mask = mask.unsqueeze(1) # (1, 1, H, W) mask = F.interpolate( mask, size=(height, width), mode='bilinear', align_corners=False ) mask = mask.squeeze(1) # (1, H, W) # Match batch size if batch > 1: mask = mask.repeat(batch, 1, 1) except Exception as e: print(f"⚠️ Error parsing mask data: {e}") print(" Creating empty mask instead") mask = None # Create empty mask if no valid data if mask is None: print("⚠️ No painted mask - creating empty mask") print("💡 Click 'Open Mask Editor' button to paint mask") mask = torch.zeros((batch, height, width), device=device, dtype=dtype) # Create visualization mask_preview = mask.unsqueeze(-1).repeat(1, 1, 1, 3) # Print statistics mask_min = mask.min().item() mask_max = mask.max().item() mask_mean = mask.mean().item() coverage = (mask > 0.01).float().mean().item() * 100 print(f"\n📊 Mask Statistics:") print(f" Range: [{mask_min:.3f}, {mask_max:.3f}]") print(f" Mean: {mask_mean:.3f}") print(f" Coverage: {coverage:.1f}%") print(f"\n✓ Outputs:") print(f" Image: {image.shape}") print(f" Mask: {mask.shape}") print(f" Preview: {mask_preview.shape}") print("="*60 + "\n") return (image, mask, mask_preview) class CustomColorToMask: """ Turn the selected color (hex) into a mask. Use the color picker or eyedropper widget to pick a color from your image or screen. """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "color": ("STRING", { "default": "#FF0000", "tooltip": "Target color in hex (#RRGGBB). Use 🎨 Pick Color or 💧 Eyedropper below to select.", }), "tolerance": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "Color matching tolerance", }), "feather": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 0.5, "step": 0.01, "display": "number", "tooltip": "Edge softness", }), "mode": (["rgb", "hsv", "luminance"], { "default": "rgb", "tooltip": "Color distance mode", }), } } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("mask",) FUNCTION = "color_to_mask" CATEGORY = "Texture Alchemist/Masks" def color_to_mask(self, image, color, tolerance, feather, mode): batch, height, width, channels = image.shape device = image.device dtype = image.dtype color = str(color).strip().lstrip("#") if len(color) < 6: return (torch.zeros((batch, height, width), device=device, dtype=dtype),) try: r = int(color[0:2], 16) / 255.0 g = int(color[2:4], 16) / 255.0 b = int(color[4:6], 16) / 255.0 except ValueError: return (torch.zeros((batch, height, width), device=device, dtype=dtype),) target = torch.tensor([r, g, b], device=device, dtype=dtype).view(1, 1, 1, 3) if image.shape[-1] == 1: image = image.repeat(1, 1, 1, 3) img_rgb = image[:, :, :, :3] if mode == "rgb": diff = img_rgb - target distance = torch.sqrt((diff * diff).sum(dim=3, keepdim=True)) / math.sqrt(3.0) elif mode == "hsv": hsv_image = _rgb_to_hsv_tensor(img_rgb) hsv_target = _rgb_to_hsv_tensor(target) h_diff = torch.abs(hsv_image[:, :, :, 0:1] - hsv_target[:, :, :, 0:1]) h_diff = torch.min(h_diff, 1.0 - h_diff) * 2.0 s_diff = torch.abs(hsv_image[:, :, :, 1:2] - hsv_target[:, :, :, 1:2]) v_diff = torch.abs(hsv_image[:, :, :, 2:3] - hsv_target[:, :, :, 2:3]) distance = torch.sqrt(h_diff * h_diff + s_diff * s_diff + v_diff * v_diff) / math.sqrt(3.0) else: lum_img = 0.299 * img_rgb[:, :, :, 0:1] + 0.587 * img_rgb[:, :, :, 1:2] + 0.114 * img_rgb[:, :, :, 2:3] lum_t = 0.299 * r + 0.587 * g + 0.114 * b distance = torch.abs(lum_img - lum_t) mask = 1.0 - torch.clamp(distance / (tolerance + 1e-7), 0.0, 1.0) if feather > 0.0: mask = torch.where( mask > (1.0 - feather), mask, torch.where( mask < feather, torch.zeros_like(mask), (mask - feather) / (1.0 - 2.0 * feather), ), ) mask = torch.clamp(mask.squeeze(-1), 0.0, 1.0) return (mask,) class ShuffleMask: """ Extract one channel (red, green, blue, or alpha) as a black-and-white mask. The channel can be chosen via the dropdown or overridden by connecting an integer: 0=red, 1=green, 2=blue, 3=alpha (e.g. from an LLM or logic node). """ CHANNELS = ["red", "green", "blue", "alpha"] @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "channel": (cls.CHANNELS, { "default": "red", "tooltip": "Channel to use as mask. Overridden when channel_index is connected.", }), }, "optional": { "channel_index": ("INT", { "default": 0, "min": 0, "max": 3, "forceInput": True, "tooltip": "Connect an integer to override the dropdown: 0=red, 1=green, 2=blue, 3=alpha (e.g. from LLM).", }), } } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("mask",) FUNCTION = "extract" CATEGORY = "Texture Alchemist/Masks" def extract(self, image, channel, channel_index=None): batch, height, width, channels = image.shape device = image.device dtype = image.dtype # Resolve channel index: connected INT overrides dropdown if channel_index is not None: idx = int(channel_index) if not isinstance(channel_index, (list, tuple)) else int(channel_index[0]) idx = max(0, min(3, idx)) else: name_to_idx = {"red": 0, "green": 1, "blue": 2, "alpha": 3} idx = name_to_idx.get(channel, 0) # Extract channel; if alpha requested but image has no alpha, use ones (opaque) if idx < channels: mask = image[:, :, :, idx : idx + 1].squeeze(-1) # (B, H, W) else: mask = torch.ones((batch, height, width), device=device, dtype=dtype) return (mask,) # Node registration NODE_CLASS_MAPPINGS = { "EdgeWearMaskGenerator": EdgeWearMaskGenerator, "DirtGrimeMaskGenerator": DirtGrimeMaskGenerator, "GradientMaskGenerator": GradientMaskGenerator, "ColorSelectionMask": ColorSelectionMask, "CustomColorToMask": CustomColorToMask, "ImageMaskCombiner": ImageMaskCombiner, "MaskCompositor": MaskCompositor, "ShuffleMask": ShuffleMask, } NODE_DISPLAY_NAME_MAPPINGS = { "EdgeWearMaskGenerator": "Edge Wear Mask Generator", "DirtGrimeMaskGenerator": "Dirt/Grime Mask Generator", "GradientMaskGenerator": "Gradient Mask Generator", "ColorSelectionMask": "Color Selection Mask", "CustomColorToMask": "Custom Color to Mask", "ImageMaskCombiner": "Image + Mask Combiner", "MaskCompositor": "Mask Compositor (Interactive)", "ShuffleMask": "Shuffle Mask", }