MdONeilsl
b6a2d1dd79
Introduce a complete GIMP/GEGL-inspired color module for ComfyUI. Adds: node declarations (nodes/*) and matching operator implementations (py/*) for many color/image operations (balance, temperature, hue/chroma, hue/saturation, saturation, exposure, levels, invert, auto adjustments, maps, posterize, desaturate, tone mapping, component extract/mix, etc.). Includes support libs (py/libs: color, process, blur), layer compositing (layers_blending), VRAM‑aware chunking, mask handling and alpha preservation. Adds web UI extension and widgets (web/*), package metadata, dynamic __init__.py loader, README update and .gitignore tweaks.
163 lines
4.4 KiB
Python
163 lines
4.4 KiB
Python
import torch
|
|
from typing import Tuple, Optional
|
|
|
|
from .libs.process import compute_chunk_size
|
|
|
|
|
|
def exec_operator_color_threshold(
|
|
image: torch.Tensor,
|
|
channel: str,
|
|
low: float,
|
|
high: float,
|
|
mask: Optional[torch.Tensor] = None,
|
|
) -> Tuple[torch.Tensor]:
|
|
"""
|
|
Apply a threshold to the input image, producing a binary output.
|
|
|
|
Supported channels:
|
|
Value, Red, Green, Blue, Alpha, Luminance, RGB
|
|
|
|
Image:
|
|
[B,H,W,C] float32 RGB/RGBA in [0,1]
|
|
|
|
Mask:
|
|
[H,W] or [B,H,W]
|
|
"""
|
|
try:
|
|
if not isinstance(image, torch.Tensor):
|
|
raise TypeError("image must be a torch.Tensor")
|
|
|
|
if image.ndim != 4:
|
|
raise ValueError("image must be [B,H,W,C]")
|
|
|
|
B, H, W, C = image.shape
|
|
|
|
if C not in (3, 4):
|
|
raise ValueError("image must be RGB or RGBA")
|
|
|
|
low = max(0.0, min(255.0, float(low))) / 255.0
|
|
high = max(0.0, min(255.0, float(high))) / 255.0
|
|
|
|
if low > high:
|
|
low, high = high, low
|
|
|
|
if mask is not None:
|
|
if not isinstance(mask, torch.Tensor):
|
|
raise TypeError("mask must be a torch.Tensor")
|
|
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
elif mask.ndim != 3:
|
|
raise ValueError("mask must be [H,W] or [B,H,W]")
|
|
|
|
if mask.shape[-2:] != (H, W):
|
|
raise ValueError("mask size mismatch")
|
|
|
|
if mask.shape[0] not in (1, B):
|
|
raise ValueError("mask batch mismatch")
|
|
|
|
mask = mask.to(dtype=image.dtype, device=image.device).clamp_(0, 1)
|
|
|
|
if torch.count_nonzero(mask) == 0:
|
|
return (image,)
|
|
|
|
if torch.all(mask == 1):
|
|
mask = None
|
|
|
|
device = image.device
|
|
|
|
chunk_size = (
|
|
max(1, compute_chunk_size(image, mask))
|
|
if device.type == "cuda"
|
|
else B
|
|
)
|
|
|
|
def threshold_chunk(
|
|
img: torch.Tensor,
|
|
msk: Optional[torch.Tensor],
|
|
) -> torch.Tensor:
|
|
|
|
rgb = img[..., :3]
|
|
|
|
if channel == "RGB":
|
|
inside = (
|
|
(rgb[..., 0] >= low)
|
|
& (rgb[..., 0] <= high)
|
|
& (rgb[..., 1] >= low)
|
|
& (rgb[..., 1] <= high)
|
|
& (rgb[..., 2] >= low)
|
|
& (rgb[..., 2] <= high)
|
|
)
|
|
|
|
else:
|
|
if channel == "Value":
|
|
intensity = rgb.amax(dim=-1)
|
|
|
|
elif channel == "Red":
|
|
intensity = rgb[..., 0]
|
|
|
|
elif channel == "Green":
|
|
intensity = rgb[..., 1]
|
|
|
|
elif channel == "Blue":
|
|
intensity = rgb[..., 2]
|
|
|
|
elif channel == "Alpha":
|
|
if C == 4:
|
|
intensity = img[..., 3]
|
|
else:
|
|
intensity = torch.ones(
|
|
img.shape[:-1],
|
|
device=img.device,
|
|
dtype=img.dtype,
|
|
)
|
|
|
|
elif channel == "Luminance":
|
|
intensity = (
|
|
rgb[..., 0] * 0.299
|
|
+ rgb[..., 1] * 0.587
|
|
+ rgb[..., 2] * 0.114
|
|
)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown channel '{channel}'")
|
|
|
|
inside = (intensity >= low) & (intensity <= high)
|
|
|
|
bw = inside.to(img.dtype)[..., None]
|
|
|
|
out_rgb = bw.expand(-1, -1, -1, 3).clone()
|
|
|
|
if C == 4:
|
|
out = torch.cat((out_rgb, img[..., 3:4]), dim=-1)
|
|
else:
|
|
out = out_rgb
|
|
|
|
if msk is not None:
|
|
out = out * msk[..., None] + img * (1.0 - msk[..., None])
|
|
|
|
return out
|
|
|
|
chunks = []
|
|
|
|
for start in range(0, B, chunk_size):
|
|
|
|
end = min(start + chunk_size, B)
|
|
|
|
img_chunk = image[start:end]
|
|
|
|
if mask is None:
|
|
mask_chunk = None
|
|
elif mask.shape[0] == 1:
|
|
mask_chunk = mask
|
|
else:
|
|
mask_chunk = mask[start:end]
|
|
|
|
chunks.append(threshold_chunk(img_chunk, mask_chunk))
|
|
|
|
return (torch.cat(chunks, dim=0),)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"Color threshold operation failed: {e}"
|
|
) from e |