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.
440 lines
14 KiB
Python
440 lines
14 KiB
Python
|
||
import torch
|
||
import math
|
||
|
||
#===================================================================================
|
||
def srgb_to_linear(c: float) -> float:
|
||
if c <= 0.04045: return c / 12.92
|
||
return ((c + 0.055) / 1.055) ** 2.4
|
||
|
||
#===================================================================================
|
||
def srgb_to_linear_vec(x: torch.Tensor) -> torch.Tensor:
|
||
"""Vectorised sRGB to linear (matches standard IEC 61966‑2‑1)."""
|
||
mask = x <= 0.04045
|
||
out = torch.zeros_like(x)
|
||
out[mask] = x[mask] / 12.92
|
||
out[~mask] = ((x[~mask] + 0.055) / 1.055) ** 2.4
|
||
return out
|
||
|
||
#===================================================================================
|
||
def linear_to_srgb_vec(x: torch.Tensor) -> torch.Tensor:
|
||
"""Vectorised linear to sRGB (inverse of above)."""
|
||
mask = x <= 0.0031308
|
||
out = torch.zeros_like(x)
|
||
out[mask] = x[mask] * 12.92
|
||
out[~mask] = 1.055 * (x[~mask] ** (1.0/2.4)) - 0.055
|
||
return torch.clamp(out, 0.0, 1.0)
|
||
|
||
#===================================================================================
|
||
def rgb_to_lch_torch(image: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Convert RGB [0,1] to CIE LCH(ab) using D65 white point.
|
||
Uses srgb_to_linear_vec (or your own) for gamma handling.
|
||
"""
|
||
# 1. sRGB → linear RGB
|
||
rgb_lin = srgb_to_linear_vec(image)
|
||
|
||
# 2. linear RGB → XYZ (D65)
|
||
M_rgb2xyz = torch.tensor([
|
||
[0.4124, 0.3576, 0.1805],
|
||
[0.2126, 0.7152, 0.0722],
|
||
[0.0193, 0.1192, 0.9505]
|
||
], device=image.device, dtype=image.dtype)
|
||
xyz = torch.einsum('ij,...j->...i', M_rgb2xyz, rgb_lin)
|
||
|
||
# 3. XYZ → CIE L*a*b* (D65 white point)
|
||
white = torch.tensor([0.95047, 1.0, 0.95883], device=image.device, dtype=image.dtype)
|
||
xyz_n = xyz / white
|
||
delta = 6.0 / 29.0
|
||
delta2 = delta * delta
|
||
delta3 = delta * delta * delta
|
||
|
||
f = torch.zeros_like(xyz_n)
|
||
mask_f = xyz_n > delta3
|
||
f[mask_f] = xyz_n[mask_f] ** (1.0/3.0)
|
||
f[~mask_f] = (xyz_n[~mask_f] / (3.0 * delta2)) + (4.0 / 29.0)
|
||
|
||
L = 116.0 * f[..., 1] - 16.0
|
||
a = 500.0 * (f[..., 0] - f[..., 1])
|
||
b = 200.0 * (f[..., 1] - f[..., 2])
|
||
|
||
# 4. Lab → LCH
|
||
C = torch.sqrt(a*a + b*b)
|
||
H = torch.atan2(b, a) * (180.0 / math.pi)
|
||
H = H % 360.0
|
||
H = torch.where(H >= 180.0, H - 360.0, H) # map to [-180, 180)
|
||
|
||
return torch.stack([L, C, H], dim=-1)
|
||
|
||
#===================================================================================
|
||
def lch_to_rgb_torch(lch: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Convert CIE LCH(ab) back to RGB [0,1].
|
||
Uses linear_to_srgb_vec for final gamma compression.
|
||
"""
|
||
L, C, H = lch[..., 0], lch[..., 1], lch[..., 2]
|
||
H_rad = (H % 360.0) * (math.pi / 180.0)
|
||
|
||
# LCH → Lab
|
||
a = C * torch.cos(H_rad)
|
||
b = C * torch.sin(H_rad)
|
||
|
||
# Lab → XYZ
|
||
delta = 6.0 / 29.0
|
||
delta2 = delta * delta
|
||
delta3 = delta * delta * delta
|
||
|
||
fy = (L + 16.0) / 116.0
|
||
fx = fy + a / 500.0
|
||
fz = fy - b / 200.0
|
||
|
||
xyz = torch.zeros_like(lch)
|
||
# Inverse f(t)
|
||
mask = fx > delta
|
||
xyz[..., 0][mask] = fx[mask] ** 3
|
||
xyz[..., 0][~mask] = 3.0 * delta2 * (fx[~mask] - 4.0/29.0)
|
||
|
||
mask = fy > delta
|
||
xyz[..., 1][mask] = fy[mask] ** 3
|
||
xyz[..., 1][~mask] = 3.0 * delta2 * (fy[~mask] - 4.0/29.0)
|
||
|
||
mask = fz > delta
|
||
xyz[..., 2][mask] = fz[mask] ** 3
|
||
xyz[..., 2][~mask] = 3.0 * delta2 * (fz[~mask] - 4.0/29.0)
|
||
|
||
white = torch.tensor([0.95047, 1.0, 0.95883], device=lch.device, dtype=lch.dtype)
|
||
xyz = xyz * white
|
||
|
||
# XYZ → linear RGB
|
||
M_xyz2rgb = torch.tensor([
|
||
[ 3.2406, -1.5372, -0.4986],
|
||
[-0.9689, 1.8758, 0.0415],
|
||
[ 0.0557, -0.2040, 1.0570]
|
||
], device=lch.device, dtype=lch.dtype)
|
||
rgb_lin = torch.einsum('ij,...j->...i', M_xyz2rgb, xyz)
|
||
|
||
# linear → sRGB
|
||
rgb_srgb = linear_to_srgb_vec(rgb_lin)
|
||
|
||
return rgb_srgb
|
||
|
||
#===================================================================================
|
||
def hsl_to_rgb(h, s, l):
|
||
hp = h / 60.0
|
||
c = (1.0 - torch.abs(2.0 * l - 1.0)) * s
|
||
x = c * (1.0 - torch.abs(hp % 2.0 - 1.0))
|
||
m = l - 0.5 * c
|
||
|
||
r = torch.zeros_like(l)
|
||
g = torch.zeros_like(l)
|
||
b = torch.zeros_like(l)
|
||
|
||
m0 = (hp >= 0) & (hp < 1)
|
||
m1 = (hp >= 1) & (hp < 2)
|
||
m2 = (hp >= 2) & (hp < 3)
|
||
m3 = (hp >= 3) & (hp < 4)
|
||
m4 = (hp >= 4) & (hp < 5)
|
||
m5 = (hp >= 5) & (hp < 6)
|
||
|
||
r[m0], g[m0], b[m0] = c[m0], x[m0], 0
|
||
r[m1], g[m1], b[m1] = x[m1], c[m1], 0
|
||
r[m2], g[m2], b[m2] = 0, c[m2], x[m2]
|
||
r[m3], g[m3], b[m3] = 0, x[m3], c[m3]
|
||
r[m4], g[m4], b[m4] = x[m4], 0, c[m4]
|
||
r[m5], g[m5], b[m5] = c[m5], 0, x[m5]
|
||
|
||
return torch.clamp(torch.stack([r + m, g + m, b + m], dim=-1), 0.0, 1.0)
|
||
|
||
#===================================================================================
|
||
def rgb_to_hsl(rgb: torch.Tensor):
|
||
maxc = rgb.max(dim=-1).values
|
||
minc = rgb.min(dim=-1).values
|
||
l = (maxc + minc) * 0.5
|
||
|
||
delta = maxc - minc
|
||
|
||
s = torch.zeros_like(l)
|
||
nz = delta > 1e-8
|
||
s[nz] = delta[nz] / (1.0 - torch.abs(2.0 * l[nz] - 1.0) + 1e-8)
|
||
|
||
r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
|
||
|
||
h = torch.zeros_like(l)
|
||
|
||
rmax = (maxc == r) & nz
|
||
gmax = (maxc == g) & nz
|
||
bmax = (maxc == b) & nz
|
||
|
||
h_r = ((g - b) / (delta + 1e-8)) % 6.0
|
||
h_g = ((b - r) / (delta + 1e-8)) + 2.0
|
||
h_b = ((r - g) / (delta + 1e-8)) + 4.0
|
||
|
||
h = torch.where(rmax, h_r, h)
|
||
h = torch.where(gmax, h_g, h)
|
||
h = torch.where(bmax, h_b, h)
|
||
|
||
h = (h * 60.0) % 360.0
|
||
|
||
return h, s, l
|
||
|
||
#===================================================================================
|
||
def rgb_to_hsv(rgb: torch.Tensor):
|
||
"""Convert RGB [0,1] to HSV. Returns (hue, saturation, value)."""
|
||
r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
|
||
|
||
max_val, max_idx = torch.max(rgb, dim=-1)
|
||
min_val, _ = torch.min(rgb, dim=-1)
|
||
diff = max_val - min_val
|
||
|
||
# Hue
|
||
hue = torch.zeros_like(max_val)
|
||
mask = diff != 0
|
||
h = torch.zeros_like(max_val)
|
||
h = torch.where((max_val == r) & mask, ((g - b) / diff) % 6, h)
|
||
h = torch.where((max_val == g) & mask, ((b - r) / diff) + 2, h)
|
||
h = torch.where((max_val == b) & mask, ((r - g) / diff) + 4, h)
|
||
hue = (h / 6) % 1.0
|
||
|
||
# Saturation
|
||
sat = torch.zeros_like(max_val)
|
||
sat = torch.where(max_val != 0, diff / max_val, sat)
|
||
|
||
# Value
|
||
val = max_val
|
||
|
||
return hue, sat, val
|
||
|
||
#===================================================================================
|
||
def hsv_to_rgb(hue: torch.Tensor, sat: torch.Tensor, val: torch.Tensor):
|
||
"""Convert HSV to RGB [0,1]."""
|
||
h = hue * 6.0
|
||
i = torch.floor(h).long()
|
||
f = h - i.float()
|
||
p = val * (1 - sat)
|
||
q = val * (1 - sat * f)
|
||
t = val * (1 - sat * (1 - f))
|
||
|
||
# Initialize output
|
||
rgb = torch.zeros(*hue.shape, 3, device=hue.device)
|
||
|
||
# Vectorized assignment using advanced indexing
|
||
# This is simplified - in practice you'd use torch.where for each case
|
||
mask0 = i == 0
|
||
mask1 = i == 1
|
||
mask2 = i == 2
|
||
mask3 = i == 3
|
||
mask4 = i == 4
|
||
mask5 = i == 5
|
||
|
||
rgb[..., 0] = torch.where(mask0, val, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask0, t, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask0, p, rgb[..., 2])
|
||
|
||
rgb[..., 0] = torch.where(mask1, q, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask1, val, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask1, p, rgb[..., 2])
|
||
|
||
rgb[..., 0] = torch.where(mask2, p, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask2, val, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask2, t, rgb[..., 2])
|
||
|
||
rgb[..., 0] = torch.where(mask3, p, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask3, q, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask3, val, rgb[..., 2])
|
||
|
||
rgb[..., 0] = torch.where(mask4, t, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask4, p, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask4, val, rgb[..., 2])
|
||
|
||
rgb[..., 0] = torch.where(mask5, val, rgb[..., 0])
|
||
rgb[..., 1] = torch.where(mask5, p, rgb[..., 1])
|
||
rgb[..., 2] = torch.where(mask5, q, rgb[..., 2])
|
||
|
||
return rgb
|
||
|
||
#===================================================================================
|
||
def rgb_to_lch(rgb: torch.Tensor):
|
||
"""Convert RGB to Lch (CIE L*a*b* with chroma)."""
|
||
# Convert RGB to XYZ (D65)
|
||
rgb_linear = rgb.clone()
|
||
mask = rgb_linear <= 0.04045
|
||
rgb_linear[mask] = rgb_linear[mask] / 12.92
|
||
rgb_linear[~mask] = ((rgb_linear[~mask] + 0.055) / 1.055) ** 2.4
|
||
|
||
# RGB to XYZ matrix
|
||
x = 0.4124564 * rgb_linear[..., 0] + 0.3575761 * rgb_linear[..., 1] + 0.1804375 * rgb_linear[..., 2]
|
||
y = 0.2126729 * rgb_linear[..., 0] + 0.7151522 * rgb_linear[..., 1] + 0.0721750 * rgb_linear[..., 2]
|
||
z = 0.0193339 * rgb_linear[..., 0] + 0.1191920 * rgb_linear[..., 1] + 0.9503041 * rgb_linear[..., 2]
|
||
|
||
# XYZ to Lab (D65 white point)
|
||
xn, yn, zn = 0.95047, 1.0, 1.08883
|
||
xr, yr, zr = x / xn, y / yn, z / zn
|
||
|
||
f = lambda t: torch.where(t > 0.008856, t ** (1/3), 7.787 * t + 16/116)
|
||
fx, fy, fz = f(xr), f(yr), f(zr)
|
||
|
||
L = 116 * fy - 16
|
||
a = 500 * (fx - fy)
|
||
b = 200 * (fy - fz)
|
||
|
||
# Chroma and hue
|
||
C = torch.sqrt(a**2 + b**2)
|
||
H = torch.atan2(b, a) % (2 * torch.pi)
|
||
|
||
return L, C, H
|
||
|
||
#===================================================================================
|
||
def lch_to_rgb(L: torch.Tensor, C: torch.Tensor, H: torch.Tensor):
|
||
"""Convert Lch to RGB."""
|
||
# Lch to Lab
|
||
a = C * torch.cos(H)
|
||
b = C * torch.sin(H)
|
||
|
||
# Lab to XYZ (D65)
|
||
fy = (L + 16) / 116
|
||
fx = a / 500 + fy
|
||
fz = fy - b / 200
|
||
|
||
f_inv = lambda t: torch.where(t > 0.206893, t**3, (t - 16/116) / 7.787)
|
||
xr, yr, zr = f_inv(fx), f_inv(fy), f_inv(fz)
|
||
|
||
xn, yn, zn = 0.95047, 1.0, 1.08883
|
||
x, y, z = xr * xn, yr * yn, zr * zn
|
||
|
||
# XYZ to RGB
|
||
r = 3.2404542 * x - 1.5371385 * y - 0.4985314 * z
|
||
g = -0.9692660 * x + 1.8760108 * y + 0.0415560 * z
|
||
b = 0.0556434 * x - 0.2040259 * y + 1.0572252 * z
|
||
|
||
rgb = torch.stack([r, g, b], dim=-1)
|
||
|
||
# Linear to sRGB
|
||
mask = rgb <= 0.0031308
|
||
rgb[mask] = 12.92 * rgb[mask]
|
||
rgb[~mask] = 1.055 * rgb[~mask] ** (1/2.4) - 0.055
|
||
|
||
return torch.clamp(rgb, 0, 1)
|
||
|
||
#===================================================================================
|
||
def rgb_to_yuv(rgb: torch.Tensor):
|
||
"""Convert RGB to Yuv (CIE Yuv)."""
|
||
# BT.709 coefficients
|
||
y = 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
|
||
u = rgb[..., 0] - y
|
||
v = rgb[..., 2] - y
|
||
return y, u, v
|
||
|
||
#===================================================================================
|
||
def yuv_to_rgb(y: torch.Tensor, u: torch.Tensor, v: torch.Tensor):
|
||
"""Convert Yuv to RGB."""
|
||
r = y + u
|
||
g = y - 0.21482 * u - 0.38129 * v # Simplified: approximate
|
||
b = y + v
|
||
rgb = torch.stack([r, g, b], dim=-1)
|
||
return torch.clamp(rgb, 0, 1)
|
||
|
||
#===================================================================================
|
||
def rgb_to_lch_torch(image: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Convert RGB [0,1] to CIE LCH(ab) using D65 white point.
|
||
Returns tensor with last dimension [L, C, H].
|
||
"""
|
||
# 1. sRGB → linear RGB
|
||
rgb_lin = srgb_to_linear_vec(image)
|
||
|
||
# 2. linear RGB → XYZ (D65)
|
||
M_rgb2xyz = torch.tensor([
|
||
[0.4124, 0.3576, 0.1805],
|
||
[0.2126, 0.7152, 0.0722],
|
||
[0.0193, 0.1192, 0.9505]
|
||
], device=image.device, dtype=image.dtype)
|
||
xyz = torch.einsum('ij,...j->...i', M_rgb2xyz, rgb_lin)
|
||
|
||
# 3. XYZ → CIE L*a*b* (D65 white point)
|
||
white = torch.tensor([0.95047, 1.0, 0.95883], device=image.device, dtype=image.dtype)
|
||
xyz_n = xyz / white
|
||
delta = 6.0 / 29.0
|
||
delta2 = delta * delta
|
||
delta3 = delta * delta * delta
|
||
|
||
f = torch.zeros_like(xyz_n)
|
||
mask_f = xyz_n > delta3
|
||
f[mask_f] = xyz_n[mask_f] ** (1.0/3.0)
|
||
f[~mask_f] = (xyz_n[~mask_f] / (3.0 * delta2)) + (4.0 / 29.0)
|
||
|
||
L = 116.0 * f[..., 1] - 16.0
|
||
a = 500.0 * (f[..., 0] - f[..., 1])
|
||
b = 200.0 * (f[..., 1] - f[..., 2])
|
||
|
||
# 4. Lab → LCH
|
||
C = torch.sqrt(a*a + b*b)
|
||
H = torch.atan2(b, a) * (180.0 / math.pi)
|
||
H = H % 360.0
|
||
H = torch.where(H >= 180.0, H - 360.0, H) # map to [-180, 180)
|
||
|
||
return torch.stack([L, C, H], dim=-1)
|
||
|
||
#===================================================================================
|
||
def lch_to_rgb_torch(lch: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Convert CIE LCH(ab) back to RGB [0,1].
|
||
"""
|
||
L, C, H = lch[..., 0], lch[..., 1], lch[..., 2]
|
||
H_rad = (H % 360.0) * (math.pi / 180.0)
|
||
|
||
# LCH → Lab
|
||
a = C * torch.cos(H_rad)
|
||
b = C * torch.sin(H_rad)
|
||
|
||
# Lab → XYZ
|
||
delta = 6.0 / 29.0
|
||
delta2 = delta * delta
|
||
delta3 = delta * delta * delta
|
||
|
||
fy = (L + 16.0) / 116.0
|
||
fx = fy + a / 500.0
|
||
fz = fy - b / 200.0
|
||
|
||
xyz = torch.zeros_like(lch)
|
||
# Inverse f(t)
|
||
mask = fx > delta
|
||
xyz[..., 0][mask] = fx[mask] ** 3
|
||
xyz[..., 0][~mask] = 3.0 * delta2 * (fx[~mask] - 4.0/29.0)
|
||
|
||
mask = fy > delta
|
||
xyz[..., 1][mask] = fy[mask] ** 3
|
||
xyz[..., 1][~mask] = 3.0 * delta2 * (fy[~mask] - 4.0/29.0)
|
||
|
||
mask = fz > delta
|
||
xyz[..., 2][mask] = fz[mask] ** 3
|
||
xyz[..., 2][~mask] = 3.0 * delta2 * (fz[~mask] - 4.0/29.0)
|
||
|
||
white = torch.tensor([0.95047, 1.0, 0.95883], device=lch.device, dtype=lch.dtype)
|
||
xyz = xyz * white
|
||
|
||
# XYZ → linear RGB
|
||
M_xyz2rgb = torch.tensor([
|
||
[ 3.2406, -1.5372, -0.4986],
|
||
[-0.9689, 1.8758, 0.0415],
|
||
[ 0.0557, -0.2040, 1.0570]
|
||
], device=lch.device, dtype=lch.dtype)
|
||
rgb_lin = torch.einsum('ij,...j->...i', M_xyz2rgb, xyz)
|
||
|
||
# linear → sRGB
|
||
rgb_srgb = linear_to_srgb_vec(rgb_lin)
|
||
return rgb_srgb
|
||
|
||
#===================================================================================
|
||
def pack_rgb(r, g, b):
|
||
"""r,g,b in 0..1, returns (packed_int, '#RRGGBB')"""
|
||
r = max(0.0, min(1.0, r))
|
||
g = max(0.0, min(1.0, g))
|
||
b = max(0.0, min(1.0, b))
|
||
ri = int(round(r * 255))
|
||
gi = int(round(g * 255))
|
||
bi = int(round(b * 255))
|
||
packed = (ri << 16) | (gi << 8) | bi
|
||
hex_str = f"#{ri:02X}{gi:02X}{bi:02X}"
|
||
return packed, hex_str
|
||
|