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.
357 lines
11 KiB
Python
357 lines
11 KiB
Python
import torch
|
||
from typing import Optional, Tuple
|
||
|
||
from .libs.process import compute_chunk_size
|
||
|
||
# ============================================================
|
||
# Utility: safe clamp
|
||
# ============================================================
|
||
|
||
def _clamp01(x):
|
||
return torch.clamp(x, 0.0, 1.0)
|
||
|
||
# ============================================================
|
||
# Core GIMP-like Levels transfer (per-channel)
|
||
# ============================================================
|
||
|
||
def _levels_transfer(
|
||
v: torch.Tensor,
|
||
in_low: float,
|
||
in_high: float,
|
||
gamma: float,
|
||
out_low: float,
|
||
out_high: float,
|
||
clamp_input: bool,
|
||
clamp_output: bool,
|
||
):
|
||
denom = (in_high - in_low)
|
||
if abs(denom) < 1e-8:
|
||
denom = 1e-8
|
||
|
||
v = (v - in_low) / denom
|
||
|
||
if clamp_input:
|
||
v = _clamp01(v)
|
||
|
||
gamma = max(gamma, 1e-6)
|
||
v = torch.pow(torch.clamp(v, min=0.0), 1.0 / gamma)
|
||
|
||
v = out_low + v * (out_high - out_low)
|
||
|
||
if clamp_output:
|
||
v = _clamp01(v)
|
||
|
||
return v
|
||
|
||
# ============================================================
|
||
# HSV conversion (needed for Value channel mode like GIMP)
|
||
# ============================================================
|
||
|
||
def _rgb_to_hsv(img: torch.Tensor):
|
||
r, g, b = img[..., 0], img[..., 1], img[..., 2]
|
||
maxc, _ = torch.max(img[..., :3], dim=-1)
|
||
minc, _ = torch.min(img[..., :3], dim=-1)
|
||
v = maxc
|
||
delta = maxc - minc
|
||
s = torch.where(maxc > 0, delta / (maxc + 1e-8), torch.zeros_like(maxc))
|
||
hue = torch.zeros_like(maxc)
|
||
mask = delta > 1e-8
|
||
rc = (maxc - r) / (delta + 1e-8)
|
||
gc = (maxc - g) / (delta + 1e-8)
|
||
bc = (maxc - b) / (delta + 1e-8)
|
||
hue = torch.where(
|
||
(maxc == r) & mask,
|
||
(bc - gc),
|
||
hue
|
||
)
|
||
hue = torch.where(
|
||
(maxc == g) & mask,
|
||
2.0 + (rc - bc),
|
||
hue
|
||
)
|
||
hue = torch.where(
|
||
(maxc == b) & mask,
|
||
4.0 + (gc - rc),
|
||
hue
|
||
)
|
||
hue = (hue / 6.0) % 1.0
|
||
return torch.stack([hue, s, v], dim=-1)
|
||
|
||
def _hsv_to_rgb(hsv: torch.Tensor):
|
||
h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
|
||
h = (h % 1.0) * 6.0
|
||
i = torch.floor(h)
|
||
f = h - i
|
||
p = v * (1.0 - s)
|
||
q = v * (1.0 - s * f)
|
||
t = v * (1.0 - s * (1.0 - f))
|
||
i = i.long() % 6
|
||
r = torch.zeros_like(v)
|
||
g = torch.zeros_like(v)
|
||
b = torch.zeros_like(v)
|
||
cond = i == 0
|
||
r = torch.where(cond, v, r)
|
||
g = torch.where(cond, t, g)
|
||
b = torch.where(cond, p, b)
|
||
cond = i == 1
|
||
r = torch.where(cond, q, r)
|
||
g = torch.where(cond, v, g)
|
||
b = torch.where(cond, p, b)
|
||
cond = i == 2
|
||
r = torch.where(cond, p, r)
|
||
g = torch.where(cond, v, g)
|
||
b = torch.where(cond, t, b)
|
||
cond = i == 3
|
||
r = torch.where(cond, p, r)
|
||
g = torch.where(cond, q, g)
|
||
b = torch.where(cond, v, b)
|
||
cond = i == 4
|
||
r = torch.where(cond, t, r)
|
||
g = torch.where(cond, p, g)
|
||
b = torch.where(cond, v, b)
|
||
cond = i == 5
|
||
r = torch.where(cond, v, r)
|
||
g = torch.where(cond, p, g)
|
||
b = torch.where(cond, q, b)
|
||
return torch.stack([r, g, b], dim=-1)
|
||
|
||
# ============================================================
|
||
# Histogram (GIMP-style endpoint detection)
|
||
# ============================================================
|
||
|
||
def _histogram_bounds(channel: torch.Tensor):
|
||
flat = channel.reshape(-1)
|
||
flat = flat[torch.isfinite(flat)]
|
||
if flat.numel() == 0:
|
||
return 0.0, 1.0
|
||
lo = torch.min(flat)
|
||
hi = torch.max(flat)
|
||
return lo.item(), hi.item()
|
||
|
||
# ============================================================
|
||
# Mask blending (GEGL-style)
|
||
# ============================================================
|
||
|
||
def _apply_mask(src: torch.Tensor, dst: torch.Tensor, mask: Optional[torch.Tensor]):
|
||
if mask is None:
|
||
return dst
|
||
if mask.dim() == 3:
|
||
mask = mask.unsqueeze(-1)
|
||
mask = mask.to(dtype=src.dtype, device=src.device)
|
||
return src * (1.0 - mask) + dst * mask
|
||
|
||
# ============================================================
|
||
# Apply per-channel Levels
|
||
# ============================================================
|
||
|
||
def _apply_levels(
|
||
img: torch.Tensor,
|
||
channel: str,
|
||
in_low: float,
|
||
in_high: float,
|
||
gamma: float,
|
||
out_low: float,
|
||
out_high: float,
|
||
clamp_input: bool,
|
||
clamp_output: bool,
|
||
auto_input: bool,
|
||
mask: Optional[torch.Tensor],
|
||
):
|
||
B, H, W, C = img.shape
|
||
|
||
if auto_input:
|
||
if channel == "Value":
|
||
v = img[..., :3].max(dim=-1).values
|
||
in_low, in_high = _histogram_bounds(v)
|
||
elif channel == "Red":
|
||
in_low, in_high = _histogram_bounds(img[..., 0])
|
||
elif channel == "Green":
|
||
in_low, in_high = _histogram_bounds(img[..., 1])
|
||
elif channel == "Blue":
|
||
in_low, in_high = _histogram_bounds(img[..., 2])
|
||
elif channel == "Alpha" and C == 4:
|
||
in_low, in_high = _histogram_bounds(img[..., 3])
|
||
|
||
if channel == "Value":
|
||
hsv = _rgb_to_hsv(img[..., :3])
|
||
v = hsv[..., 2]
|
||
v2 = _levels_transfer(
|
||
v,
|
||
in_low, in_high,
|
||
gamma,
|
||
out_low, out_high,
|
||
clamp_input,
|
||
clamp_output,
|
||
)
|
||
hsv[..., 2] = v2
|
||
rgb = _hsv_to_rgb(hsv)
|
||
out = img.clone()
|
||
out[..., :3] = rgb
|
||
if C == 4:
|
||
out[..., 3] = img[..., 3]
|
||
return out
|
||
|
||
out = img.clone()
|
||
channel_map = {
|
||
"Red": 0,
|
||
"Green": 1,
|
||
"Blue": 2,
|
||
"Alpha": 3,
|
||
}
|
||
if channel in channel_map:
|
||
c = channel_map[channel]
|
||
if c < C:
|
||
processed = _levels_transfer(
|
||
img[..., c],
|
||
in_low, in_high,
|
||
gamma,
|
||
out_low, out_high,
|
||
clamp_input,
|
||
clamp_output,
|
||
)
|
||
out[..., c] = processed
|
||
|
||
return _apply_mask(img, out, mask)
|
||
|
||
# ============================================================
|
||
# FINAL OPERATOR ENTRY
|
||
# ============================================================
|
||
|
||
def exec_operator_color_levels(
|
||
image: torch.Tensor,
|
||
channel: str,
|
||
interpolation: str, # kept for API compatibility (ignored like GIMP Levels)
|
||
input_low: float,
|
||
input_high: float,
|
||
input_gamma: float,
|
||
input_clamp: bool,
|
||
output_low: float,
|
||
output_high: float,
|
||
output_clamp: bool,
|
||
auto_input: bool,
|
||
mask: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
"""
|
||
Apply a GIMP‑style Levels adjustment to an image tensor.
|
||
|
||
The adjustment remaps the selected channel (or the Value component of HSV)
|
||
using input low/high, gamma, and output low/high. The operation supports
|
||
batch processing with an optional per‑image mask. Input and output tensors
|
||
must be in RGB or RGBA format with values in [0,1].
|
||
|
||
Args:
|
||
image: Input tensor of shape (B, H, W, C) with C=3 (RGB) or C=4 (RGBA).
|
||
channel: Channel to modify: "Value", "Red", "Green", "Blue", or "Alpha".
|
||
interpolation: Unused, kept for API compatibility.
|
||
input_low, input_high: Input range in UI units (0‑255).
|
||
input_gamma: Gamma correction factor (>0).
|
||
input_clamp: Clamp normalized values to [0,1] before gamma.
|
||
output_low, output_high: Output range in UI units (0‑255).
|
||
output_clamp: Clamp final values to [0,1].
|
||
auto_input: If True, compute input low/high from the image histogram.
|
||
mask: Optional mask tensor of shape (B, H, W) or (B, H, W, 1) to blend
|
||
the original and adjusted images.
|
||
|
||
Returns:
|
||
Adjusted image tensor with the same shape and dtype as the input.
|
||
"""
|
||
# ------------------------------------------------------------------
|
||
# Input validation
|
||
# ------------------------------------------------------------------
|
||
if not isinstance(image, torch.Tensor):
|
||
raise TypeError("image must be a torch.Tensor")
|
||
if image.dim() != 4:
|
||
raise ValueError("image must have 4 dimensions [B, H, W, C]")
|
||
B, H, W, C = image.shape
|
||
if C not in (3, 4):
|
||
raise ValueError("image must have 3 or 4 channels (RGB or RGBA)")
|
||
|
||
if mask is not None:
|
||
if not isinstance(mask, torch.Tensor):
|
||
raise TypeError("mask must be a torch.Tensor")
|
||
if mask.dim() == 3:
|
||
mask = mask.unsqueeze(-1) # [B, H, W, 1]
|
||
elif mask.dim() == 4 and mask.shape[-1] == 1:
|
||
pass
|
||
else:
|
||
raise ValueError("mask must have shape [B, H, W] or [B, H, W, 1]")
|
||
if mask.shape[0] != B:
|
||
raise ValueError(
|
||
f"mask batch size {mask.shape[0]} does not match image batch size {B}"
|
||
)
|
||
mask = mask.to(device=image.device, dtype=image.dtype)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Early exit: identity transformation
|
||
# ------------------------------------------------------------------
|
||
if (not auto_input and
|
||
input_low == 0 and input_high == 255 and input_gamma == 1.0 and
|
||
output_low == 0 and output_high == 255 and
|
||
not input_clamp and not output_clamp):
|
||
return image.clone()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Normalize parameters to [0,1]
|
||
# ------------------------------------------------------------------
|
||
in_low = input_low / 255.0
|
||
in_high = input_high / 255.0
|
||
out_low = output_low / 255.0
|
||
out_high = output_high / 255.0
|
||
gamma = max(float(input_gamma), 1e-6)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Auto‑input: compute global histogram bounds on the whole batch
|
||
# ------------------------------------------------------------------
|
||
if auto_input:
|
||
if channel == "Value":
|
||
v = image[..., :3].max(dim=-1).values
|
||
in_low, in_high = _histogram_bounds(v)
|
||
elif channel == "Red":
|
||
in_low, in_high = _histogram_bounds(image[..., 0])
|
||
elif channel == "Green":
|
||
in_low, in_high = _histogram_bounds(image[..., 1])
|
||
elif channel == "Blue":
|
||
in_low, in_high = _histogram_bounds(image[..., 2])
|
||
elif channel == "Alpha" and C == 4:
|
||
in_low, in_high = _histogram_bounds(image[..., 3])
|
||
else:
|
||
# Channel not available: keep default [0,1] (no effect)
|
||
in_low, in_high = 0.0, 1.0
|
||
|
||
# ------------------------------------------------------------------
|
||
# Chunked processing for VRAM management
|
||
# ------------------------------------------------------------------
|
||
chunk_size = compute_chunk_size(image, mask)
|
||
if chunk_size is None or chunk_size <= 0:
|
||
chunk_size = B
|
||
|
||
outputs = []
|
||
for start in range(0, B, chunk_size):
|
||
end = min(start + chunk_size, B)
|
||
img_chunk = image[start:end]
|
||
mask_chunk = mask[start:end] if mask is not None else None
|
||
|
||
out_chunk = _apply_levels(
|
||
img_chunk,
|
||
channel=channel,
|
||
in_low=in_low,
|
||
in_high=in_high,
|
||
gamma=gamma,
|
||
out_low=out_low,
|
||
out_high=out_high,
|
||
clamp_input=input_clamp,
|
||
clamp_output=output_clamp,
|
||
auto_input=False, # already computed globally
|
||
mask=mask_chunk,
|
||
)
|
||
outputs.append(out_chunk)
|
||
|
||
if outputs:
|
||
result = torch.cat(outputs, dim=0)
|
||
else:
|
||
result = image # should not happen for non‑empty batch
|
||
|
||
return result.to(image.dtype).clamp(0.0, 1.0)
|
||
|
||
|