1030 lines
36 KiB
Python
1030 lines
36 KiB
Python
"""
|
||
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",
|
||
}
|
||
|