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)