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.
210 lines
7.0 KiB
Python
210 lines
7.0 KiB
Python
import math
|
||
from typing import Optional
|
||
|
||
import torch
|
||
|
||
from .libs.color import srgb_to_linear
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Temperature limits
|
||
# -----------------------------------------------------------------------------
|
||
|
||
LOWEST_TEMPERATURE = 1000.0
|
||
HIGHEST_TEMPERATURE = 40000.0
|
||
|
||
|
||
def kelvin_to_linear_rgb(kelvin: float):
|
||
"""
|
||
Convert a colour temperature (in Kelvin) to a normalised linear‑RGB
|
||
white point using Tanner Helland's black‑body approximation.
|
||
|
||
The output is a 3‑tuple whose largest component is exactly 1.0.
|
||
Values are in linear light (not sRGB).
|
||
|
||
Parameters
|
||
----------
|
||
kelvin : float
|
||
Colour temperature in Kelvin. Clamped to [1000, 40000].
|
||
|
||
Returns
|
||
-------
|
||
tuple of float
|
||
Normalised linear‑RGB white point (r, g, b).
|
||
"""
|
||
T = max(LOWEST_TEMPERATURE, min(HIGHEST_TEMPERATURE, float(kelvin)))
|
||
t = T / 100.0
|
||
|
||
if t <= 66.0:
|
||
r = 255.0
|
||
g = 99.4708025861 * math.log(t) - 161.1195681661
|
||
if t <= 19.0:
|
||
b = 0.0
|
||
else:
|
||
b = 138.5177312231 * math.log(t - 10.0) - 305.0447927307
|
||
else:
|
||
r = 329.698727446 * ((t - 60.0) ** -0.1332047592)
|
||
g = 288.1221695283 * ((t - 60.0) ** -0.0755148492)
|
||
b = 255.0
|
||
|
||
r = min(255.0, max(0.0, r))
|
||
g = min(255.0, max(0.0, g))
|
||
b = min(255.0, max(0.0, b))
|
||
|
||
rgb = [
|
||
srgb_to_linear(r / 255.0),
|
||
srgb_to_linear(g / 255.0),
|
||
srgb_to_linear(b / 255.0),
|
||
]
|
||
|
||
m = max(rgb)
|
||
if m > 0:
|
||
rgb = [x / m for x in rgb]
|
||
|
||
return tuple(rgb)
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Colour‑temperature adjustment
|
||
# -----------------------------------------------------------------------------
|
||
|
||
@torch.no_grad()
|
||
def exec_operator_color_temperature(
|
||
image: torch.Tensor,
|
||
original_temp: float,
|
||
intended_temp: float,
|
||
mask: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
"""
|
||
Adjust the white balance of an RGB or RGBA image.
|
||
|
||
The image is assumed to be in a **linear** colour space. The
|
||
adjustment consists of mapping the white point from
|
||
``original_temp`` Kelvin to ``intended_temp`` Kelvin, preserving
|
||
approximate luminance. An optional mask allows a per‑pixel blend
|
||
between the original and the corrected image.
|
||
|
||
Parameters
|
||
----------
|
||
image : torch.Tensor
|
||
Linear‑RGB or linear‑RGBA image tensor.
|
||
Allowed shapes:
|
||
``(H, W, C)`` with ``C = 3`` or ``C = 4``
|
||
``(B, H, W, C)`` with ``C = 3`` or ``C = 4``
|
||
Dtype must be a floating‑point type (e.g. float32, float16).
|
||
original_temp : float
|
||
Colour temperature of the image's current white point (Kelvin).
|
||
intended_temp : float
|
||
Desired colour temperature for the white point (Kelvin).
|
||
mask : torch.Tensor, optional
|
||
Blend mask. Values are clamped to [0, 1] and used to linearly
|
||
interpolate between the original image (mask=0) and the
|
||
temperature‑adjusted image (mask=1).
|
||
Accepted shapes (the number of spatial dimensions must match
|
||
the image, and the batch dimension, if present, must be 1 or
|
||
equal to the image batch size):
|
||
``(H, W)``
|
||
``(H, W, 1)``
|
||
``(H, W, C)`` (averaged to a single channel)
|
||
``(B, H, W)``
|
||
``(B, H, W, 1)``
|
||
``(B, H, W, C)``
|
||
If the last dimension is not ``1``, the mask is reduced by
|
||
averaging its RGB channels (the alpha channel, if present, is
|
||
ignored). For a single‑channel mask no averaging takes place.
|
||
|
||
Returns
|
||
-------
|
||
torch.Tensor
|
||
Temperature‑adjusted image, with the same shape, dtype and
|
||
device as ``image``. Alpha channel (if present) is passed
|
||
through unchanged.
|
||
|
||
Raises
|
||
------
|
||
ValueError
|
||
If the image does not have 3 or 4 channels, or if the mask
|
||
shape is incompatible, or if tensors reside on different
|
||
devices.
|
||
TypeError
|
||
If ``image`` is not a floating‑point tensor.
|
||
"""
|
||
# ------------------------------------------------------------------
|
||
# Validation
|
||
# ------------------------------------------------------------------
|
||
if image.dtype not in (torch.float16, torch.float32, torch.float64):
|
||
raise TypeError("Image must be a floating‑point tensor.")
|
||
if image.ndim not in (3, 4):
|
||
raise ValueError(f"Image must have 3 or 4 dimensions, got {image.ndim}.")
|
||
if image.shape[-1] not in (3, 4):
|
||
raise ValueError(f"Expected 3 or 4 channels, got {image.shape[-1]}.")
|
||
|
||
if mask is not None:
|
||
if mask.device != image.device:
|
||
raise ValueError("Mask must be on the same device as the image.")
|
||
if mask.ndim not in (image.ndim, image.ndim - 1):
|
||
raise ValueError(
|
||
f"Mask must have {image.ndim - 1} or {image.ndim} dimensions, "
|
||
f"got {mask.ndim}."
|
||
)
|
||
# Check batch dimension compatibility
|
||
if image.ndim == 4 and mask.ndim == 4:
|
||
if mask.shape[0] not in (1, image.shape[0]):
|
||
raise ValueError(
|
||
f"Mask batch size must be 1 or {image.shape[0]}, "
|
||
f"got {mask.shape[0]}."
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# White‑point gains
|
||
# ------------------------------------------------------------------
|
||
src = torch.tensor(
|
||
kelvin_to_linear_rgb(original_temp),
|
||
dtype=image.dtype,
|
||
device=image.device,
|
||
)
|
||
dst = torch.tensor(
|
||
kelvin_to_linear_rgb(intended_temp),
|
||
dtype=image.dtype,
|
||
device=image.device,
|
||
)
|
||
|
||
eps = torch.finfo(image.dtype).eps
|
||
gains = dst / (src + eps)
|
||
# preserve approximate luminance
|
||
gains = gains / gains.mean()
|
||
|
||
# Broadcast gains to the image shape
|
||
shape = [1] * image.ndim
|
||
shape[-1] = 3 if image.shape[-1] == 4 else image.shape[-1]
|
||
gains = gains.view(*shape)
|
||
|
||
# Apply to RGB channels; alpha (if present) is left untouched
|
||
if image.shape[-1] == 4:
|
||
rgb_result = image[..., :3] * gains
|
||
result = torch.cat([rgb_result, image[..., 3:4]], dim=-1)
|
||
else:
|
||
result = image * gains
|
||
|
||
# ------------------------------------------------------------------
|
||
# Blend with mask
|
||
# ------------------------------------------------------------------
|
||
if mask is not None:
|
||
# Make mask broadcastable with image
|
||
if mask.ndim == image.ndim - 1:
|
||
mask = mask.unsqueeze(-1)
|
||
|
||
# If multi‑channel, reduce to a single luminance channel
|
||
if mask.shape[-1] != 1:
|
||
# Average only the colour channels (ignore alpha if present)
|
||
mask = mask[..., :3].mean(dim=-1, keepdim=True)
|
||
|
||
mask = mask.clamp(0.0, 1.0)
|
||
|
||
# Linear interpolation: mask=0 → original, mask=1 → adjusted
|
||
result = torch.lerp(image, result, mask)
|
||
|
||
# Ensure valid range (clamping happens in‑place for memory efficiency)
|
||
return result.clamp_(0.0, 1.0)
|
||
|