From b6a2d1dd799f36feb8dc8e43988ddeedebf7bf16 Mon Sep 17 00:00:00 2001 From: MdONeilsl Date: Mon, 6 Jul 2026 08:17:05 -0400 Subject: [PATCH] Add MD Color Mod nodes, operators, and web UI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .gitignore | 5 +- README.md | 15 +- __init__.py | 32 + nodes/__init__.py | 0 nodes/color_nodes.py | 302 +++++ nodes/img_color_nodes.py | 1174 ++++++++++++++++++++ nodes/layer_nodes.py | 50 + py/color_auto_color_enhance.py | 137 +++ py/color_auto_equalize.py | 113 ++ py/color_auto_stretch_contrast.py | 280 +++++ py/color_auto_white_balance.py | 94 ++ py/color_balance.py | 129 +++ py/color_brightness_contrast.py | 131 +++ py/color_colorize.py | 145 +++ py/color_comp_channel_mixer.py | 177 +++ py/color_comp_extract.py | 237 ++++ py/color_comp_mono_mixer.py | 142 +++ py/color_desat_color_2_gray.py | 151 +++ py/color_desat_desturation.py | 138 +++ py/color_desat_sepia.py | 142 +++ py/color_exposure.py | 136 +++ py/color_hue_chroma.py | 154 +++ py/color_hue_saturation.py | 99 ++ py/color_invert.py | 182 +++ py/color_levels.py | 356 ++++++ py/color_map_alien.py | 246 ++++ py/color_map_color_exchange.py | 166 +++ py/color_map_color_rotate.py | 252 +++++ py/color_posterize.py | 168 +++ py/color_saturation.py | 185 +++ py/color_shadows_highlights.py | 220 ++++ py/color_temperature.py | 209 ++++ py/color_threshold.py | 163 +++ py/color_tone_mapping_fattal_et_al_2002.py | 198 ++++ py/layers_blending.py | 209 ++++ py/libs/__init__.py | 0 py/libs/blur.py | 33 + py/libs/color.py | 439 ++++++++ py/libs/process.py | 76 ++ pyproject.toml | 6 + requirements.txt | 0 web/extension_reg.js | 43 + web/libs/widgets_module.js | 9 + web/nodes/color_balance_widget.js | 103 ++ web/nodes/color_hue_saturation_widget.js | 104 ++ 45 files changed, 7347 insertions(+), 3 deletions(-) create mode 100644 __init__.py create mode 100644 nodes/__init__.py create mode 100644 nodes/color_nodes.py create mode 100644 nodes/img_color_nodes.py create mode 100644 nodes/layer_nodes.py create mode 100644 py/color_auto_color_enhance.py create mode 100644 py/color_auto_equalize.py create mode 100644 py/color_auto_stretch_contrast.py create mode 100644 py/color_auto_white_balance.py create mode 100644 py/color_balance.py create mode 100644 py/color_brightness_contrast.py create mode 100644 py/color_colorize.py create mode 100644 py/color_comp_channel_mixer.py create mode 100644 py/color_comp_extract.py create mode 100644 py/color_comp_mono_mixer.py create mode 100644 py/color_desat_color_2_gray.py create mode 100644 py/color_desat_desturation.py create mode 100644 py/color_desat_sepia.py create mode 100644 py/color_exposure.py create mode 100644 py/color_hue_chroma.py create mode 100644 py/color_hue_saturation.py create mode 100644 py/color_invert.py create mode 100644 py/color_levels.py create mode 100644 py/color_map_alien.py create mode 100644 py/color_map_color_exchange.py create mode 100644 py/color_map_color_rotate.py create mode 100644 py/color_posterize.py create mode 100644 py/color_saturation.py create mode 100644 py/color_shadows_highlights.py create mode 100644 py/color_temperature.py create mode 100644 py/color_threshold.py create mode 100644 py/color_tone_mapping_fattal_et_al_2002.py create mode 100644 py/layers_blending.py create mode 100644 py/libs/__init__.py create mode 100644 py/libs/blur.py create mode 100644 py/libs/color.py create mode 100644 py/libs/process.py create mode 100644 pyproject.toml create mode 100644 requirements.txt create mode 100644 web/extension_reg.js create mode 100644 web/libs/widgets_module.js create mode 100644 web/nodes/color_balance_widget.js create mode 100644 web/nodes/color_hue_saturation_widget.js diff --git a/.gitignore b/.gitignore index dd15ec4..fd68831 100644 --- a/.gitignore +++ b/.gitignore @@ -178,4 +178,7 @@ cython_debug/ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data # refer to https://docs.cursor.com/context/ignore-files .cursorignore -.cursorindexingignore \ No newline at end of file +.cursorindexingignore +note.txt +tmp +web/widgets diff --git a/README.md b/README.md index fa9ebdc..0e3021d 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,13 @@ -# ComfyUI-MdColorMod -GIMP color mod in confyUI +# GIMP/GEGL Color Mod for ComfyUI + +A standardized collection of color operation nodes designed for seamless integration with ComfyUI, leveraging GIMP’s powerful GEGL framework. + +--- + +## ✅ Features + +- **All Nodes Support Masks**: Each node accepts a mask input for precise control over affected areas. +- **Batch Processing**: Process multiple images and their corresponding masks simultaneously. +- **Full RGB/RGBA Support**: Input and output support both RGB and RGBA format. + +--- diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..20edede --- /dev/null +++ b/__init__.py @@ -0,0 +1,32 @@ +import importlib +import pkgutil +import os + +# Tells ComfyUI to load the JS files from the /web folder +WEB_DIRECTORY = "./web" + +# Dynamically collect all node mappings from the nodes/ folder +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +_current_dir = os.path.dirname(__file__) +_nodes_dir = os.path.join(_current_dir, "nodes") + +for module_info in pkgutil.iter_modules([_nodes_dir]): + module_name = module_info.name + module = importlib.import_module(f".nodes.{module_name}", package=__package__) + + if hasattr(module, "NODE_CLASS_MAPPINGS"): + for key, value in module.NODE_CLASS_MAPPINGS.items(): + if key in NODE_CLASS_MAPPINGS: + print(f"Warning: duplicate node key '{key}' from {module_name}") + NODE_CLASS_MAPPINGS[key] = value + + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + +# Clean up helper variables (optional, but keeps namespace tidy) +del _current_dir, _nodes_dir, module_info, module_name, module + +# Public API – yes, you should keep __all__ so users know what's available +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/color_nodes.py b/nodes/color_nodes.py new file mode 100644 index 0000000..d60c3d2 --- /dev/null +++ b/nodes/color_nodes.py @@ -0,0 +1,302 @@ + +from ..py.libs.color import pack_rgb + +module_cat = "md/color" + +MAX_RGB_INT = 0xFFFFFF + +# ============================================================================= +# RGB +# ============================================================================= +class mdPrimitiveColorRGB: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "Red": ("FLOAT", {"default": 0, "min": 0, "max": 255, "step": 1}), + "Green": ("FLOAT", {"default": 0, "min": 0, "max": 255, "step": 1}), + "Blue": ("FLOAT", {"default": 0, "min": 0, "max": 255, "step": 1}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, Red, Green, Blue): + # Round and clamp + r = max(0, min(255, int(round(Red)))) + g = max(0, min(255, int(round(Green)))) + b = max(0, min(255, int(round(Blue)))) + packed = (r << 16) | (g << 8) | b + hex_str = f"#{r:02X}{g:02X}{b:02X}" + return (packed, hex_str) + +# ============================================================================= +# HSV +# ============================================================================= +class mdPrimitiveColorHSV: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "Hue": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 0.1}), + "Saturation": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Value": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, Hue, Saturation, Value): + h = Hue / 360.0 + s = max(0.0, min(1.0, Saturation / 100.0)) + v = max(0.0, min(1.0, Value / 100.0)) + + if s == 0.0: + r = g = b = v + else: + h = h * 6.0 + i = int(h) + f = h - i + p = v * (1.0 - s) + q = v * (1.0 - s * f) + t = v * (1.0 - s * (1.0 - f)) + if i == 0: + r, g, b = v, t, p + elif i == 1: + r, g, b = q, v, p + elif i == 2: + r, g, b = p, v, t + elif i == 3: + r, g, b = p, q, v + elif i == 4: + r, g, b = t, p, v + else: + r, g, b = v, p, q + + return pack_rgb(r, g, b) + +# ============================================================================= +# HSL +# ============================================================================= +class mdPrimitiveColorHSL: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "Hue": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 0.1}), + "Saturation": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Lightness": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, Hue, Saturation, Lightness): + h = Hue / 360.0 + s = max(0.0, min(1.0, Saturation / 100.0)) + l = max(0.0, min(1.0, Lightness / 100.0)) + + if s == 0.0: + r = g = b = l + else: + def hue2rgb(p, q, t): + if t < 0.0: t += 1.0 + if t > 1.0: t -= 1.0 + if t < 1.0 / 6.0: return p + (q - p) * 6.0 * t + if t < 1.0 / 2.0: return q + if t < 2.0 / 3.0: return p + (q - p) * (2.0 / 3.0 - t) * 6.0 + return p + + q = l * (1.0 + s) if l < 0.5 else l + s - l * s + p = 2.0 * l - q + r = hue2rgb(p, q, h + 1.0 / 3.0) + g = hue2rgb(p, q, h) + b = hue2rgb(p, q, h - 1.0 / 3.0) + + return pack_rgb(r, g, b) + +# ============================================================================= +# CMYK +# ============================================================================= +class mdPrimitiveColorCMYK: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "Cyan": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Magenta": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Yellow": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Key": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, Cyan, Magenta, Yellow, Key): + c = max(0.0, min(1.0, Cyan / 100.0)) + m = max(0.0, min(1.0, Magenta / 100.0)) + y = max(0.0, min(1.0, Yellow / 100.0)) + k = max(0.0, min(1.0, Key / 100.0)) + + r = (1.0 - c) * (1.0 - k) + g = (1.0 - m) * (1.0 - k) + b = (1.0 - y) * (1.0 - k) + + return pack_rgb(r, g, b) + +# ============================================================================= +# CIE 1931 XYZ +# ============================================================================= +class mdPrimitiveColorXYZ: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "X": ("FLOAT", {"default": 0.95047, "min": 0.0, "max": 1.0, "step": 0.001}), + "Y": ("FLOAT", {"default": 1.00000, "min": 0.0, "max": 1.0, "step": 0.001}), + "Z": ("FLOAT", {"default": 1.08883, "min": 0.0, "max": 1.0, "step": 0.001}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, X, Y, Z): + # XYZ → linear sRGB (D65) + r_lin = 3.2404542 * X - 1.5371385 * Y - 0.4985314 * Z + g_lin = -0.9692660 * X + 1.8760108 * Y + 0.0415560 * Z + b_lin = 0.0556434 * X - 0.2040259 * Y + 1.0572252 * Z + + def gamma(c): + if c <= 0.0031308: + return 12.92 * c + else: + return 1.055 * (c ** (1.0 / 2.4)) - 0.055 + + r = gamma(r_lin) + g = gamma(g_lin) + b = gamma(b_lin) + + return pack_rgb(r, g, b) + +# ============================================================================= +# CIELAB +# ============================================================================= +class mdPrimitiveColorLAB: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "L": ("FLOAT", {"default": 50, "min": 0, "max": 100, "step": 0.1}), + "A": ("FLOAT", {"default": 0, "min": -128, "max": 128, "step": 0.1}), + "B": ("FLOAT", {"default": 0, "min": -128, "max": 128, "step": 0.1}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, L, A, B): + # D65 reference white + xn, yn, zn = 0.95047, 1.0, 1.08883 + + delta = 6.0 / 29.0 + fy = (L + 16.0) / 116.0 + fx = fy + (A / 500.0) + fz = fy - (B / 200.0) + + def f_inv(t): + if t > delta: + return t ** 3 + else: + return 3.0 * delta ** 2 * (t - 4.0 / 29.0) + + X = xn * f_inv(fx) + Y = yn * f_inv(fy) + Z = zn * f_inv(fz) + + r_lin = 3.2404542 * X - 1.5371385 * Y - 0.4985314 * Z + g_lin = -0.9692660 * X + 1.8760108 * Y + 0.0415560 * Z + b_lin = 0.0556434 * X - 0.2040259 * Y + 1.0572252 * Z + + def gamma(c): + if c <= 0.0031308: + return 12.92 * c + else: + return 1.055 * (c ** (1.0 / 2.4)) - 0.055 + + r = gamma(r_lin) + g = gamma(g_lin) + b = gamma(b_lin) + + return pack_rgb(r, g, b) + +# ============================================================================= +# Y′CbCr +# ============================================================================= +class mdPrimitiveColorYCbCr: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "Luma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "BlueDiff": ("FLOAT", {"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.001}), + "RedDiff": ("FLOAT", {"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.001}), + }, + } + + RETURN_TYPES = ("INT", "STRING") + RETURN_NAMES = ("int", "HEX") + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, Luma, BlueDiff, RedDiff): + Y = Luma + Cb = BlueDiff + Cr = RedDiff + + r = Y + 1.402 * Cr + g = Y - 0.344136 * Cb - 0.714136 * Cr + b = Y + 1.772 * Cb + + return pack_rgb(r, g, b) + +# ============================================================================= +# Node registry +# ============================================================================= + +NODE_CLASS_MAPPINGS = { + "mdPrimitiveColorHSV": mdPrimitiveColorHSV, + "mdPrimitiveColorHSL": mdPrimitiveColorHSL, + "mdPrimitiveColorCMYK": mdPrimitiveColorCMYK, + "mdPrimitiveColorXYZ": mdPrimitiveColorXYZ, + "mdPrimitiveColorLAB": mdPrimitiveColorLAB, + "mdPrimitiveColorYCbCr": mdPrimitiveColorYCbCr, + "mdPrimitiveColorRGB": mdPrimitiveColorRGB, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "mdPrimitiveColorRGB": "Primitive Color RGB (MD)", + "mdPrimitiveColorHSV": "Primitive Color HSV (MD)", + "mdPrimitiveColorHSL": "Primitive Color HSL (MD)", + "mdPrimitiveColorCMYK": "Primitive Color CMYK (MD)", + "mdPrimitiveColorXYZ": "Primitive Color CIE XYZ (MD)", + "mdPrimitiveColorLAB": "Primitive Color CIELAB (MD)", + "mdPrimitiveColorYCbCr": "Primitive Color YCbCr (MD)", +} diff --git a/nodes/img_color_nodes.py b/nodes/img_color_nodes.py new file mode 100644 index 0000000..456a8a5 --- /dev/null +++ b/nodes/img_color_nodes.py @@ -0,0 +1,1174 @@ +import torch +import json + +from typing import Tuple, Optional + +from ..py.color_balance import exec_operation_color_balance +from ..py.color_temperature import exec_operator_color_temperature +from ..py.color_hue_chroma import exec_operator_color_hue_chroma +from ..py.color_hue_saturation import exec_operator_color_hue_saturation +from ..py.color_saturation import exec_operator_color_saturation +from ..py.color_exposure import exec_operator_color_exposure +from ..py.color_shadows_highlights import exec_operator_color_shadows_highlights +from ..py.color_brightness_contrast import exec_operator_color_brightness_contrast +from ..py.color_levels import exec_operator_color_levels +from ..py.color_invert import exec_operator_color_invert +from ..py.color_auto_equalize import exec_operator_color_auto_equalize +from ..py.color_auto_white_balance import exec_operator_color_auto_white_balance +from ..py.color_auto_stretch_contrast import exec_operator_color_auto_stretch_contrast +from ..py.color_auto_color_enhance import exec_operator_color_auto_color_enhance +from ..py.color_comp_channel_mixer import exec_operator_color_comp_channel_mixer +from ..py.color_comp_extract import exec_operator_color_comp_extract +from ..py.color_comp_mono_mixer import exec_operator_color_comp_mono_mixer +from ..py.color_desat_color_2_gray import exec_operator_color_desat_color_2_gray +from ..py.color_desat_desturation import exec_operator_color_desat_desturation +from ..py.color_desat_sepia import exec_operator_color_desat_sepia +from ..py.color_map_alien import exec_operator_color_map_alien +from ..py.color_map_color_exchange import exec_operator_color_map_color_exchange +from ..py.color_map_color_rotate import exec_operator_color_map_color_rotate +from ..py.color_tone_mapping_fattal_et_al_2002 import exec_operator_color_tone_mapping_fattal_et_al_2002 +from ..py.color_threshold import exec_operator_color_threshold +from ..py.color_colorize import exec_operator_color_colorize +from ..py.color_posterize import exec_operator_color_posterize + +module_cat = "md/image/color" + +MAX_RGB_INT = 0xFFFFFF + +#=================================================================================== +class mdColorBalance: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "range": (["Shadow", "Midtonnes", "Highlights"],), + "Cyan-Red": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Magenta-Green": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Yellow-Blue": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Preserve luminosity": ("BOOLEAN", {"default": True}), + "jdata": ("STRING", {"default": "[[0,0,0],[0,0,0],[0,0,0]]"}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + mask: torch.Tensor = None, + range: str = "Shadow", + Preserve_luminosity: bool = True, + jdata: str = "[[0,0,0],[0,0,0],[0,0,0]]", + **kwargs + ) -> Tuple[torch.Tensor]: + # Parse the JSON data + try: + data_list = json.loads(jdata) + if not (isinstance(data_list, list) and len(data_list) == 3 and + all(isinstance(row, list) and len(row) == 3 for row in data_list)): + data_list = [[0, 0, 0], [0, 0, 0], [0, 0, 0]] + except Exception: + data_list = [[0, 0, 0], [0, 0, 0], [0, 0, 0]] + + # Extract adjustments for all three ranges + shadow_adj = tuple(data_list[0]) # [cyan-red, magenta-green, yellow-blue] + midtones_adj = tuple(data_list[1]) + highlights_adj = tuple(data_list[2]) + + # Delegate all processing to the operator function + return (exec_operation_color_balance( + image=image, + mask=mask, + shadow_adj=shadow_adj, + midtones_adj=midtones_adj, + highlights_adj=highlights_adj, + preserve_luminosity=Preserve_luminosity + ),) + +#=================================================================================== +class mdColorTemperature: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "original_temp": ("FLOAT", { + "default": 6500, "min": 1000, "max": 40000, "step": 1, + "display": "Original Temperature" + }), + "intended_temp": ("FLOAT", { + "default": 6500, "min": 1000, "max": 40000, "step": 1, + "display": "Intended Temperature" + }), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec( + self, + image: torch.Tensor, + original_temp: float, + intended_temp: float, + mask: torch.Tensor = None + ) -> Tuple[torch.Tensor]: + """ + Execute the color temperature adjustment. + The exec method calls exec_operator_color_temperature with all necessary + arguments, letting it do all the work as requested. + """ + return (exec_operator_color_temperature(image, original_temp, intended_temp, mask),) + +#=================================================================================== +class mdColorHueChroma: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Hue": ("FLOAT", {"default": 0, "min": -180, "max": 180, "step": 0.1}), + "Chroma": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Lightness": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec( + self, + image: torch.Tensor, + Hue: float, + Chroma: float, + Lightness: float, + mask: torch.Tensor = None + ) -> Tuple[torch.Tensor]: + """ + Execute the Hue-Chroma operation. + Calls exec_operator_color_hue_chroma with all necessary arguments. + """ + return (exec_operator_color_hue_chroma( + image=image, + hue=Hue, + chroma=Chroma, + lightness=Lightness, + mask=mask + ),) + +#=================================================================================== +class mdColorHueSaturation: + """ + Adjusts master and per‑colour hue, lightness, and saturation of an image. + + The adjustment is driven by a 7×3 matrix (`jdata`) that provides master + and six per‑hue shifts (Red, Yellow, Green, Cyan, Blue, Magenta). An + `Overlap` parameter controls how much neighbouring hues are affected + (0 = hard edges, 100 = broad blending). A luminance‑based `mask` can + optionally limit the effect to certain areas. Works on both RGB and + RGBA images; when alpha is present it is left untouched. + + Inputs + ------ + image : IMAGE tensor (B, H, W, C) with C = 3 or 4. + color : (unused in the computation – kept for UI compatibility) + Overlap : overlap spread (0–100 %) + Hue / Lightness / Saturation : (unused; shifts come from `jdata`) + jdata : JSON string representing a 7×3 array (master + 6 hues) + mask : optional MASK tensor (B, H, W) or (B, 1, H, W) to limit effect + + Returns + ------- + IMAGE tensor with the same number of channels as the input. + """ + + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "color": ( + ["Master", "Red", "Yellow", "Green", "Cyan", "Blue", "Magenta"], + {"default": "Master"}, + ), + "Overlap": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + "Hue": ("FLOAT", {"default": 0, "min": -180, "max": 180, "step": 0.1}), + "Lightness": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Saturation": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "jdata": ( + "STRING", + {"default": "[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]]"}, + ), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat # provided by the surrounding framework + + def exec(self, image, color, Overlap, Hue, Lightness, Saturation, jdata, mask=None): + # ---------- Parse and validate jdata ---------- + try: + data = json.loads(jdata) + if not isinstance(data, list) or len(data) != 7: + raise ValueError("jdata must be a list of exactly 7 entries") + for i, entry in enumerate(data): + if not isinstance(entry, list) or len(entry) != 3: + raise ValueError(f"jdata[{i}] must be a list of 3 numbers") + if not all(isinstance(v, (int, float)) for v in entry): + raise ValueError(f"jdata[{i}] contains non‑numeric values") + except Exception as e: + # Fall back to zeros on parsing failure (original behaviour) + print(f"mdColorHueSaturation: invalid jdata, using defaults: {e}") + data = [[0, 0, 0] for _ in range(7)] + + # ---------- Input validation (improved error messages) ---------- + if image.dim() != 4: + raise ValueError(f"Image tensor must be 4D (B,H,W,C), got shape {image.shape}") + if image.shape[-1] not in (3, 4): + raise ValueError(f"Unsupported number of channels {image.shape[-1]}; only 3 (RGB) or 4 (RGBA) allowed") + + if mask is not None: + if mask.dim() not in (3, 4): + raise ValueError(f"Mask must be 3D (B,H,W) or 4D (B,1,H,W), got shape {mask.shape}") + if mask.shape[0] != image.shape[0] and mask.shape[0] != 1: + raise ValueError( + f"Mask batch size {mask.shape[0]} must be 1 or match image batch size {image.shape[0]}" + ) + + # ---------- Delegate to the core operator ---------- + out = exec_operator_color_hue_saturation(image, mask, Overlap, data) + return (out,) + +#=================================================================================== +class mdColorSaturation: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Scale": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.01}), + "Color_Space": (["Native", "CIE Lab/Lch", "CIE Yuv"], {"default": "Native"}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec(self, image: torch.Tensor, Scale: float, Color_Space: str, mask: torch.Tensor = None) -> Tuple[torch.Tensor]: + return exec_operator_color_saturation(image, Scale, Color_Space, mask) + +# =================================================================================== +class mdColorExposure: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Black_level": ("FLOAT", {"default": 0, "min": -0.1, "max": 0.1, "step": 0.001}), + "Exposure": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec( + self, + image: torch.Tensor, + Black_level: float, + Exposure: float, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Minimum code in exec() - delegate everything to the operator function + return (exec_operator_color_exposure(image, Black_level, Exposure, mask),) + +# =================================================================================== +class mdColorShadowsHighlights: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Shadows": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Shadows_color_adj": ("FLOAT", {"default": 100, "min": 0, "max": 100, "step": 0.1}), + "Highlights": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "Highlights_color_adj": ("FLOAT", {"default": 50, "min": 0, "max": 100, "step": 0.1}), + "White_point_adj": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.001}), + "Radius": ("FLOAT", {"default": 100, "min": 0.1, "max": 1500, "step": 0.1}), + "Compress": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec(self, image, Shadows, Shadows_color_adj, Highlights, Highlights_color_adj, White_point_adj, Radius, Compress, mask=None) -> Tuple[torch.Tensor]: + return exec_operator_color_shadows_highlights( + image=image, + Shadows=Shadows, + Shadows_color_adj=Shadows_color_adj, + Highlights=Highlights, + Highlights_color_adj=Highlights_color_adj, + White_point_adj=White_point_adj, + Radius=Radius, + Compress=Compress, + mask=mask, + ) + +# =================================================================================== +class mdColorBrightnessContrast: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Brightness": ("FLOAT", {"default": 0, "min": -127, "max": 127, "step": 0.1}), + "Contrast": ("FLOAT", {"default": 0, "min": -127, "max": 127, "step": 0.1}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Brightness: float, + Contrast: float, + mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor]: + # Delegate all work to the standalone function + return exec_operator_color_brightness_contrast(image, Brightness, Contrast, mask) + +# =================================================================================== +class mdColorLevels: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Channel": (["Value", "Red", "Green", "Blue", "Alpha"], {"default": "Value"}), + "Interpolation": (["Linear", "Non-Linear", "Perceptual"], {"default": "Non-Linear"}), + + "Input_low": ("FLOAT", {"default": 0, "min": 0, "max": 255, "step": 1}), + "Input_high": ("FLOAT", {"default": 255, "min": 0, "max": 255, "step": 1}), + "Input_gamma": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}), + "Input_clamp": ("BOOLEAN", {"default": False}), + + "Output_low": ("FLOAT", {"default": 0, "min": 0, "max": 255, "step": 1}), + "Output_high": ("FLOAT", {"default": 255, "min": 0, "max": 255, "step": 1}), + "Output_clamp": ("BOOLEAN", {"default": False}), + + "Auto_input": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec( + self, + image: torch.Tensor, + Channel: str, + Interpolation: str, + 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: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + """ + Execute the Levels adjustment with minimal code - delegate all work to + exec_operator_color_levels. + """ + return (exec_operator_color_levels( + image=image, + channel=Channel, + interpolation=Interpolation, + input_low=Input_low, + input_high=Input_high, + input_gamma=Input_gamma, + input_clamp=Input_clamp, + output_low=Output_low, + output_high=Output_high, + output_clamp=Output_clamp, + auto_input=Auto_input, + mask=mask, + ),) + +#=================================================================================== +#=================================================================================== +class mdColorInvert: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "type": (["RGB", "Linear", "Value"], {"default": "RGB"}), + }, + "optional": {"mask": ("MASK",)}, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec(self, image: torch.Tensor, type: str, mask: torch.Tensor = None) -> Tuple[torch.Tensor]: + return exec_operator_color_invert(image, type, mask) + +#=================================================================================== +#=================================================================================== +class mdColorAutoEqualize: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": {"image": ("IMAGE",),}, + "optional": {"mask": ("MASK",),}, + } + + RETURN_TYPES: Tuple[torch.Tensor] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec(self, image: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]: + return exec_operator_color_auto_equalize(image, mask) + +#=================================================================================== +class mdColorAutoWhiteBalance: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": {"image": ("IMAGE",),}, + "optional": {"mask": ("MASK",),}, + } + + RETURN_TYPES: Tuple[str] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec(self, image: torch.Tensor, mask: torch.Tensor = None) -> Tuple[torch.Tensor]: + return exec_operator_color_auto_white_balance(image, mask) + +#=================================================================================== +class mdColorAutoStretchContrast: + """ + ComfyUI node that performs GIMP-like "Auto Stretch Contrast" + """ + + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "keep_colors": ("BOOLEAN", {"default": True}), + "non_linear": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES: Tuple[str] = ("IMAGE",) + RETURN_NAMES: Tuple[str] = ("image",) + FUNCTION: str = "exec" + CATEGORY: str = module_cat + + def exec( + self, + image: torch.Tensor, + keep_colors: bool = True, + non_linear: bool = False, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + + return exec_operator_color_auto_stretch_contrast( + image=image, + keep_colors=keep_colors, + non_linear=non_linear, + mask=mask, + ) + +#=================================================================================== +class mdColorAutoColorEnhance: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": {"image": ("IMAGE",),}, + "optional": {"mask": ("MASK",),}, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, image: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]: + # Delegate everything – this is the "minimum code" required in the method + return exec_operator_color_auto_color_enhance(image, mask) + +#=================================================================================== +class mdColorCompChannelMixer: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Red_in_red": ("FLOAT", {"default": 1, "min": -2, "max": 2, "step": 0.01}), + "Green_in_red": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Blue_in_red": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Red_in_green": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Green_in_green": ("FLOAT", {"default": 1, "min": -2, "max": 2, "step": 0.01}), + "Blue_in_green": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Red_in_blue": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Green_in_blue": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.01}), + "Blue_in_blue": ("FLOAT", {"default": 1, "min": -2, "max": 2, "step": 0.01}), + "Preserve_luminosity": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Red_in_red: float, + Green_in_red: float, + Blue_in_red: float, + Red_in_green: float, + Green_in_green: float, + Blue_in_green: float, + Red_in_blue: float, + Green_in_blue: float, + Blue_in_blue: float, + Preserve_luminosity: bool, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_comp_channel_mixer( + image, + Red_in_red, + Green_in_red, + Blue_in_red, + Red_in_green, + Green_in_green, + Blue_in_green, + Red_in_blue, + Green_in_blue, + Blue_in_blue, + Preserve_luminosity, + mask, + ) + +#=================================================================================== +class mdColorCompExtract: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Component": ([ + "RGB red", "RGB green", "RGB blue", + "Hue", + "HSV saturation", "HSV value", + "HSL saturation", "HSL lightness", + "CMYK cyan", "CMYK magenta", "CMYK yellow", "CMYK key", + "Y'CbCr Y'", "Y'CbCr Cb", "Y'CbCr Cr", + "LAB L", "LAB A", "LAB B", + "LCH C(ab)", "LCH H(ab)", + "Alpha" + ], {"default": "RGB red"}), + "Invert": ("BOOLEAN", {"default": False}), + "Linear": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, image: torch.Tensor, Component: str, Invert: bool, Linear: bool, mask: torch.Tensor = None) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_comp_extract( + image=image, + component=Component, + invert=Invert, + linear=Linear, + mask=mask + ) + +#=================================================================================== +class mdColorCompMonoMixer: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "preserve_luminosity": ("BOOLEAN", {"default": False}), + "red_multiplier": ("FLOAT", {"default": 0.333, "min": -5.0, "max": 5.0, "step": 0.01}), + "green_multiplier": ("FLOAT", {"default": 0.333, "min": -5.0, "max": 5.0, "step": 0.01}), + "blue_multiplier": ("FLOAT", {"default": 0.333, "min": -5.0, "max": 5.0, "step": 0.01}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + preserve_luminosity: bool, + red_multiplier: float, + green_multiplier: float, + blue_multiplier: float, + mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_comp_mono_mixer( + image, + preserve_luminosity, + red_multiplier, + green_multiplier, + blue_multiplier, + mask, + ) + +#=================================================================================== +#=================================================================================== +# Too slow +class mdColorDesaturateColor2Gray: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Radius": ("FLOAT", {"default": 300, "min": 2, "max": 6000, "step": 1}), + "Samples": ("FLOAT", {"default": 4, "min": 1, "max": 1000, "step": 1}), + "Iterations": ("FLOAT", {"default": 10, "min": 1, "max": 1000, "step": 1}), + "Enhance_Shadows": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Radius: float, + Samples: float, + Iterations: float, + Enhance_Shadows: bool, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_desat_color_2_gray( + image=image, + radius=int(Radius), + samples=int(Samples), + iterations=int(Iterations), + enhance_shadows=Enhance_Shadows, + mask=mask, + ) + +#=================================================================================== +class mdColorDesaturateDesaturate: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Mode": ( + ["Luminance", "Luma", "Lightness (HSL)", "Avg (HSI Intensity)", "Value (HSV)"], + {"default": "Luminance"} + ), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Mode: str, + mask: torch.Tensor = None + ) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_desat_desturation(image, Mode, mask) + +#=================================================================================== +class mdColorDesaturateSepia: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "sRGB": ("BOOLEAN", {"default": True}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + strength: float, + sRGB: bool, + mask: torch.Tensor | None = None, + ) -> Tuple[torch.Tensor]: + # Delegate all work to the core function + return exec_operator_color_desat_sepia(image, strength, sRGB, mask) + +#=================================================================================== +class mdColorMapAlien: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Model": (["RGB", "HSL"], {"default": "RGB"}), + + "Red_frequency": ("FLOAT", {"default": 1, "min": 0, "max": 20, "step": 0.025}), + "Green_frequency": ("FLOAT", {"default": 1, "min": 0, "max": 20, "step": 0.025}), + "Blue_frequency": ("FLOAT", {"default": 1, "min": 0, "max": 20, "step": 0.025}), + + "Red_shift": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 1}), + "Green_shift": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 1}), + "Blue_shift": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 1}), + + "Keep_red": ("BOOLEAN", {"default": False}), + "Keep_green": ("BOOLEAN", {"default": False}), + "Keep_blue": ("BOOLEAN", {"default": False}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Model: str, + Red_frequency: float, + Green_frequency: float, + Blue_frequency: float, + Red_shift: float, + Green_shift: float, + Blue_shift: float, + Keep_red: bool, + Keep_green: bool, + Keep_blue: bool, + mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor]: + """ + Minimal exec method – all work is delegated to exec_operator_color_map_alien. + """ + return exec_operator_color_map_alien( + image=image, + model=Model, + freq1=Red_frequency, + freq2=Green_frequency, + freq3=Blue_frequency, + shift1=Red_shift, + shift2=Green_shift, + shift3=Blue_shift, + keep1=Keep_red, + keep2=Keep_green, + keep3=Keep_blue, + mask=mask, + ) + +#=================================================================================== +class mdColorMapColorExchange: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "from_color": ("INT", {"default": 0x7A95C7, "min": 0, "max": MAX_RGB_INT}), + "to_color": ("INT", {"default": 0xFF0000, "min": 0, "max": MAX_RGB_INT}), + "Red_Threshold": ("FLOAT", {"default": 0.106, "min": 0, "max": 1, "step": 0.001}), + "Green_Threshold": ("FLOAT", {"default": 0.100, "min": 0, "max": 1, "step": 0.001}), + "Blue_Threshold": ("FLOAT", {"default": 0.190, "min": 0, "max": 1, "step": 0.001}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + from_color: int, + to_color: int, + Red_Threshold: float, + Green_Threshold: float, + Blue_Threshold: float, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Minimal exec – delegate all work to the helper + return exec_operator_color_map_color_exchange( + image=image, + from_color=from_color, + to_color=to_color, + red_threshold=Red_Threshold, + green_threshold=Green_Threshold, + blue_threshold=Blue_Threshold, + mask=mask, + ) + +#=================================================================================== +class mdColorMapColorRotate: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + + "Source_from": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 0.1}), + "Source_to": ("FLOAT", {"default": 90, "min": 0, "max": 360, "step": 0.1}), + "Source_clockwise": ("BOOLEAN", {"default": False}), + + "Destination_from": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 0.1}), + "Destination_to": ("FLOAT", {"default": 90, "min": 0, "max": 360, "step": 0.1}), + "Destination_clockwise": ("BOOLEAN", {"default": False}), + + "Gray_threshold": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), + "Mode": (["Change to this", "Treat as this"], {"default": "Change to this"}), + "Hue": ("FLOAT", {"default": 0, "min": 0, "max": 360, "step": 0.025}), + "Saturation": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Source_from: float, + Source_to: float, + Source_clockwise: bool, + Destination_from: float, + Destination_to: float, + Destination_clockwise: bool, + Gray_threshold: float, + Mode: str, + Hue: float, + Saturation: float, + mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor]: + # Minimal exec – delegate all work to the helper + return exec_operator_color_map_color_rotate( + image=image, + Source_from=Source_from, + Source_to=Source_to, + Source_clockwise=Source_clockwise, + Destination_from=Destination_from, + Destination_to=Destination_to, + Destination_clockwise=Destination_clockwise, + Gray_threshold=Gray_threshold, + Mode=Mode, + Hue=Hue, + Saturation=Saturation, + mask=mask + ) + +#=================================================================================== +#=================================================================================== +class mdColorToneMappingFattalEtAl2002: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.001}), + "Beta": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 2.0, "step": 0.001}), + "Saturation": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.001}), + "Noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Alpha: float, + Beta: float, + Saturation: float, + Noise: float, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Delegate all the work to the core function immediately + return exec_operator_color_tone_mapping_fattal_et_al_2002( + image, Alpha, Beta, Saturation, Noise, mask + ) + +#=================================================================================== +#=================================================================================== +class mdColorThreshold: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Channel": ( + ["Value", "Red", "Green", "Blue", "Alpha", "Luminance", "RGB"], + {"default": "Value"}, + ), + "Low": ("FLOAT", {"default": 127, "min": 0, "max": 255, "step": 1}), + "High": ("FLOAT", {"default": 255, "min": 0, "max": 255, "step": 1}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Channel: str, + Low: float, + High: float, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Immediately delegate all work to the core function + return exec_operator_color_threshold(image, Channel, Low, High, mask) + +#=================================================================================== +class mdColorColorize: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Hue": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}), + "Saturation": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}), + "Lightness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}), + "Color": ("INT", {"default": 0x7A95C7, "min": 0, "max": MAX_RGB_INT}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Hue: float, + Saturation: float, + Lightness: float, + Color: int, + mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor]: + # Immediately delegate all work to the core function + return exec_operator_color_colorize(image, Hue, Saturation, Lightness, Color, mask) + +#=================================================================================== +class mdColorPosterize: + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "image": ("IMAGE",), + "Color": ("INT", {"default": 3, "min": 2, "max": 256}), + }, + "optional": { + "mask": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec( + self, + image: torch.Tensor, + Color: int, + mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor]: + # Immediately delegate all processing to the core operator + return exec_operator_color_posterize(image, Color, mask) + +#=================================================================================== + +NODE_CLASS_MAPPINGS = { + "mdColorPosterize": mdColorPosterize, + "mdColorColorize": mdColorColorize, + "mdColorThreshold": mdColorThreshold, + #"mdColorToneMappingFattalEtAl2002": mdColorToneMappingFattalEtAl2002, + "mdColorMapColorRotate": mdColorMapColorRotate, + "mdColorMapColorExchange": mdColorMapColorExchange, + "mdColorMapAlien": mdColorMapAlien, + "mdColorDesaturateSepia": mdColorDesaturateSepia, + "mdColorDesaturateDesaturate": mdColorDesaturateDesaturate, + #"mdColorDesaturateColor2Gray": mdColorDesaturateColor2Gray, + "mdColorCompMonoMixer": mdColorCompMonoMixer, + "mdColorCompExtract": mdColorCompExtract, + "mdColorCompChannelMixer": mdColorCompChannelMixer, + "mdColorAutoColorEnhance": mdColorAutoColorEnhance, + "mdColorAutoStretchContrast": mdColorAutoStretchContrast, + "mdColorAutoWhiteBalance": mdColorAutoWhiteBalance, + "mdColorAutoEqualize": mdColorAutoEqualize, + "mdColorInvert": mdColorInvert, + "mdColorLevels": mdColorLevels, + "mdColorBrightnessContrast": mdColorBrightnessContrast, + "mdColorShadowsHighlights": mdColorShadowsHighlights, + "mdColorExposure": mdColorExposure, + "mdColorSaturation": mdColorSaturation, + "mdColorHueSaturation": mdColorHueSaturation, + "mdColorHueChroma": mdColorHueChroma, + "mdColorTemperature": mdColorTemperature, + "mdColorBalance": mdColorBalance, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "mdColorPosterize": "Color Posterize (MD)", + "mdColorColorize": "Color Colorize (MD)", + "mdColorThreshold": "Color Threshold (MD)", + #"mdColorToneMappingFattalEtAl2002": "Color Tone Mapping Fattal Et Al 2002 (MD)", + "mdColorMapColorRotate": "Color Map Color Rotate (MD)", + "mdColorMapColorExchange": "Color Map Color Exchange (MD)", + "mdColorMapAlien": "Color Map Alien (MD)", + "mdColorDesaturateSepia": "Color Desaturate Sepia (MD)", + "mdColorDesaturateDesaturate": "Color Desaturate (MD)", + #"mdColorDesaturateColor2Gray": "Color Desaturate Color to Gray (MD)", + "mdColorCompMonoMixer": "Color Components Mono Mixer (MD)", + "mdColorCompExtract": "Color Components Extract (MD)", + "mdColorCompChannelMixer": "Color Components Channel Mixer (MD)", + "mdColorAutoColorEnhance": "Color Auto Color Enhance (MD)", + "mdColorAutoStretchContrast": "Color Auto Stretch Contrast (MD)", + "mdColorAutoWhiteBalance": "Color Auto White Balance (MD)", + "mdColorAutoEqualize": "Color Auto Equalize (MD)", + "mdColorInvert": "Color Invert (MD)", + "mdColorLevels": "Color Levels (MD)", + "mdColorBrightnessContrast": "Color Brightness-Contrast (MD)", + "mdColorShadowsHighlights": "Color Shadows-Highlights (MD)", + "mdColorExposure":"Color Exposure (MD)", + "mdColorSaturation": "Color Saturation (MD)", + "mdColorHueSaturation": "Color Hue-Saturation (MD)", + "mdColorHueChroma": "Color Hue-Chroma (MD)", + "mdColorTemperature": "Color Temperature (MD)", + "mdColorBalance": "Color Balance (MD)", +} + diff --git a/nodes/layer_nodes.py b/nodes/layer_nodes.py new file mode 100644 index 0000000..cb05ec6 --- /dev/null +++ b/nodes/layer_nodes.py @@ -0,0 +1,50 @@ + +module_cat = "md/layers" + +from ..py.layers_blending import composite_layer_stack + +#=================================================================================== +# Node class – declaration only, execution delegated to the function below +#=================================================================================== +class mdLayerStackNode: + """ + Composites a layer over a background using a blend mode, a mask, and a global transparency. + Handles RGBA images (4 channels) and respects the layer's own alpha channel. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "background": ("IMAGE",), + "layer": ("IMAGE",), + "blending_mode": (["Normal", "Screen", "Multiply", "Overlay", "Difference", + "Dissolve", "Clear", "Add", "Subtract", "Replace", + "Darken Only", "Lighten Only", "Color Dodge", "Color Burn", + "Hard Light", "Soft Light", "Pin Light", "Linear Dodge (Add)", + "Linear Burn (Multiply)", "Vivid Light", "Halo", "Negate", + "Merge RGB", "Merge Cyan", "Merge Magenta", "Merge Yellow", + "Merge Black", "Fill with FG Color", "Fill with BG Color", + "Average", "Composite", "Behind", "Template"],), + "transparency": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "exec" + CATEGORY = module_cat + + def exec(self, background, layer, blending_mode, transparency, mask=None): + return composite_layer_stack(background, layer, blending_mode, transparency, mask) + +#=================================================================================== +NODE_CLASS_MAPPINGS = { + "mdLayerStackNode": mdLayerStackNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "mdLayerStackNode": "Layer Stack Node (MD)", +} diff --git a/py/color_auto_color_enhance.py b/py/color_auto_color_enhance.py new file mode 100644 index 0000000..1b7c52b --- /dev/null +++ b/py/color_auto_color_enhance.py @@ -0,0 +1,137 @@ +import torch +from typing import Tuple, Optional +from .libs.color import rgb_to_lch_torch, lch_to_rgb_torch +from .libs.process import compute_chunk_size + + +def exec_operator_color_auto_color_enhance( + image: torch.Tensor, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Apply GIMP‑style “Color Enhance” (chroma stretching) to an RGB or RGBA image. + + The enhancement linearly stretches the chroma (C) channel of each image in the + batch so that its minimum mapped chroma becomes 0 and its maximum becomes 100. + When a mask is supplied, only the masked region is used to determine the per‑ + image chroma range, and the result is blended with the original according to + the per‑pixel mask strength. + + Processing is performed in chunks to manage GPU memory; the chunk size is + determined by ``compute_chunk_size``. + + Args: + image: Image tensor of shape ``[B, H, W, C]`` with ``C = 3`` (RGB) or + ``C = 4`` (RGBA). Values are expected in [0, 1]. + mask: Optional mask tensor of shape ``[B, H, W]`` or ``[B, H, W, 1]``, + values in [0, 1]. Pixels with value > 0.5 are considered for + chroma range calculation; the mask also controls blending strength + when non‑zero. + + Returns: + A single‑element tuple ``(enhanced_image,)`` with the same shape and + dtype as the input, clamped to [0, 1]. Alpha channels (if present) are + passed through unchanged. + """ + # --- validation ------------------------------------------------------------ + if image.dim() != 4: + raise ValueError(f"Expected image tensor of shape [B, H, W, C], got {image.shape}") + B, H, W, C = image.shape + if C not in (3, 4): + raise ValueError(f"Image must have 3 (RGB) or 4 (RGBA) channels, got {C}") + + if mask is not None: + if mask.dim() not in (3, 4): + raise ValueError(f"Mask must be 3‑ or 4‑dimensional, got {mask.dim()}D") + if mask.shape[0] != B: + raise ValueError(f"Batch size mismatch: image {B} vs mask {mask.shape[0]}") + # Squeeze trailing dimension if needed → [B, H, W, 1] + if mask.dim() == 3: + mask = mask.unsqueeze(-1) + elif mask.shape[-1] != 1: + mask = mask[..., :1] # keep first channel only + mask = mask.to(device=image.device, dtype=image.dtype) + + # --- RGBA support: split alpha, process RGB only ------------------------- + has_alpha = (C == 4) + if has_alpha: + rgb = image[..., :3] + alpha = image[..., 3:] # [B, H, W, 1] + else: + rgb = image + + # --- chunked processing --------------------------------------------------- + chunk_size = compute_chunk_size(image, mask) + enhanced_chunks = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + rgb_chunk = rgb[start:end] # [chunk, H, W, 3] + m_chunk = mask[start:end] if mask is not None else None + + # ---------- per‑image chroma range with optional mask ----------------- + lch = rgb_to_lch_torch(rgb_chunk) # [chunk, H, W, 3] + + if m_chunk is not None: + # Boolean mask: True where mask > 0.5 + mask_bool = m_chunk > 0.5 # [chunk, H, W, 1] + + # Prepare C channel for masked min/max (vectorised) + C_chan = lch[..., 1] # [chunk, H, W] + # fill non‑masked with +inf for min, -inf for max + C_min = C_chan.masked_fill(~mask_bool.squeeze(-1), float('inf')) + C_max = C_chan.masked_fill(~mask_bool.squeeze(-1), float('-inf')) + + minC = C_min.amin(dim=(1, 2)) # [chunk] + maxC = C_max.amax(dim=(1, 2)) + + # Handle images with no valid mask pixels → set range to 0 + any_mask = mask_bool.any(dim=(1, 2, 3)).squeeze(-1) # [chunk] + minC = torch.where(any_mask, minC, torch.zeros_like(minC)) + maxC = torch.where(any_mask, maxC, torch.zeros_like(maxC)) + else: + C_chan = lch[..., 1] + minC = C_chan.amin(dim=(1, 2)) + maxC = C_chan.amax(dim=(1, 2)) + + # ---------- early exit if no chroma variation in the chunk ------------ + delta = maxC - minC # [chunk] + if not delta.any(): # all deltas ≈ 0 + # No change – just forward the original chunk + enhanced_chunks.append(rgb_chunk) + continue + + # ---------- stretch chroma -------------------------------------------- + valid = delta > 1e-12 # [chunk] + scale = torch.zeros_like(delta) + offset = torch.zeros_like(delta) + scale[valid] = 100.0 / delta[valid] + offset[valid] = -minC[valid] * scale[valid] + + # Broadcast to [chunk, 1, 1] + scale = scale.view(-1, 1, 1) + offset = offset.view(-1, 1, 1) + valid = valid.view(-1, 1, 1) + + C_new = torch.where(valid, lch[..., 1] * scale + offset, lch[..., 1]) + lch_new = torch.stack([lch[..., 0], C_new, lch[..., 2]], dim=-1) + + # ---------- convert back to RGB --------------------------------------- + rgb_enh = lch_to_rgb_torch(lch_new) + + # ---------- blend with mask if provided ------------------------------- + if m_chunk is not None: + # mask acts as alpha blending factor (per pixel) + rgb_enh = m_chunk * rgb_enh + (1.0 - m_chunk) * rgb_chunk + + enhanced_chunks.append(rgb_enh) + + # --- reassemble batch and handle alpha ------------------------------------ + enhanced_rgb = torch.cat(enhanced_chunks, dim=0) # [B, H, W, 3] + + if has_alpha: + output = torch.cat([enhanced_rgb, alpha], dim=-1) + else: + output = enhanced_rgb + + return (output.to(device=image.device, dtype=image.dtype),) diff --git a/py/color_auto_equalize.py b/py/color_auto_equalize.py new file mode 100644 index 0000000..92e38ae --- /dev/null +++ b/py/color_auto_equalize.py @@ -0,0 +1,113 @@ +import torch +from typing import Tuple, Optional + +from .libs.color import rgb_to_hsv, hsv_to_rgb +from .libs.process import compute_chunk_size + + +def exec_operator_color_auto_equalize( + image: torch.Tensor, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Auto‑equalise the V (value) channel of an RGB or RGBA image batch. + + The equalisation curve (LUT) is computed from the histogram of all pixels + in the batch. If a mask is provided, it acts as a per‑pixel strength factor: + - 1.0 → fully equalised + - 0.0 → original V channel preserved + - Values between 0 and 1 blend linearly between the two. + + Alpha channels are left untouched. + + Parameters + ---------- + image : torch.Tensor + (B, H, W, C) with C=3 (RGB) or C=4 (RGBA). + Dtype: torch.uint8 (0–255) or torch.float32 (0–1). + mask : torch.Tensor, optional + (B, H, W) tensor with values in [0,1] representing application strength. + If omitted, all pixels are fully equalised (mask = 1). + + Returns + ------- + tuple of torch.Tensor + A 1‑tuple containing the equalised image tensor, same shape and dtype + as the input. + """ + # ---------- validation -------------------------------------------------- + if image.dim() != 4: + raise ValueError(f"Expected 4D tensor (B,H,W,C), got {image.dim()}D") + B, H, W, C = image.shape + if C not in (3, 4): + raise ValueError(f"Only 3 (RGB) or 4 (RGBA) channels supported, got {C}") + + if mask is not None: + if mask.dim() != 3: + raise ValueError(f"Mask must be 3D (B,H,W), got {mask.dim()}D") + if mask.shape != (B, H, W): + raise ValueError( + f"Mask shape {mask.shape} incompatible with image shape {image.shape}" + ) + + # ---------- separate alpha if present ----------------------------------- + if C == 4: + alpha = image[..., 3:4] + img_rgb = image[..., :3] + else: + alpha = None + img_rgb = image + + # ---------- normalise to float [0,1] ------------------------------------ + if img_rgb.dtype == torch.uint8: + img_float = img_rgb.float() / 255.0 + else: + img_float = img_rgb.float() + + # ---------- RGB → HSV --------------------------------------------------- + h, s, v = rgb_to_hsv(img_float) # all shapes (B,H,W) + + # Early exit on truly empty input + if v.numel() == 0: + return (image,) + + # ---------- chunked histogram over ALL pixels (mask ignored) ------------ + bins = 256 + hist = torch.zeros(bins, device=v.device, dtype=torch.float32) + + chunk_size = compute_chunk_size(image, mask) + v_flat = v.reshape(-1) + n_vals = v_flat.numel() + for start in range(0, n_vals, chunk_size): + end = min(start + chunk_size, n_vals) + hist += torch.histc(v_flat[start:end], bins=bins, min=0.0, max=1.0) + + # Build normalised CDF (LUT) + cdf = torch.cumsum(hist, dim=0) + lut = cdf / cdf[-1].clamp_min(1e-8) + + # Apply LUT → equalised V + idx = (v * (bins - 1)).long().clamp(0, bins - 1) + v_eq = lut[idx] # (B,H,W) + + # ---------- blend with original V using mask strength ------------------- + if mask is not None: + # mask values in [0,1] directly control blend + v_final = v_eq * mask + v * (1.0 - mask) + else: + v_final = v_eq + + # ---------- HSV → RGB --------------------------------------------------- + out_rgb = hsv_to_rgb(h, s, v_final).clamp(0, 1) + + # ---------- restore original dtype and alpha ---------------------------- + if img_rgb.dtype == torch.uint8: + out_rgb = (out_rgb * 255).round().to(torch.uint8) + + if alpha is not None: + out = torch.cat([out_rgb, alpha], dim=-1) + else: + out = out_rgb + + return (out,) + diff --git a/py/color_auto_stretch_contrast.py b/py/color_auto_stretch_contrast.py new file mode 100644 index 0000000..27fa7bc --- /dev/null +++ b/py/color_auto_stretch_contrast.py @@ -0,0 +1,280 @@ +import torch +from typing import Optional, Tuple +from .libs.color import hsv_to_rgb, rgb_to_hsv +from .libs.process import compute_chunk_size + + +EPS = 1e-8 + + +def _normalize_mask(mask: torch.Tensor, device, B, H, W): + """ + Returns: + mask_f: float32 per-pixel strength mask in [0,1] + mask_bool: boolean mask for statistics + """ + if mask is None: + return None, None + + mask = mask.to(device) + + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(-1) + + if mask.shape[-1] != 1: + raise ValueError("Mask must have a single channel") + + if mask.shape[0] == 1 and B > 1: + mask = mask.expand(B, -1, -1, -1) + elif mask.shape[0] != B: + raise ValueError("Mask batch size mismatch") + + mask_f = mask.clamp(0.0, 1.0).float() + mask_bool = mask_f > 0.0 + return mask_f, mask_bool + + +def _check_input(image: torch.Tensor): + if image is None: + raise ValueError("image cannot be None") + + image = image.float() + + if image.ndim == 3: + image = image.unsqueeze(0) + + if image.shape[-1] not in (3, 4): + raise ValueError(f"Expected RGB/RGBA, got {image.shape}") + + return image + + +def compute_vram_chunk(image: torch.Tensor, mask: Optional[torch.Tensor]): + """ + Wrapper around user-provided compute_chunk_size with safety fallback. + """ + try: + cs = compute_chunk_size(image, mask) + return max(1, int(cs)) + except Exception: + return max(1, image.shape[0]) + + +def _early_exit_check(image, mask_bool): + """ + Fast-path detection: + - empty batch + - no usable mask + """ + if image.shape[0] == 0: + return True + + if mask_bool is not None and not mask_bool.any(): + return True + + return False + +def _stretch_chunk_rgb( + rgb: torch.Tensor, # (B,H,W,3) + mask_f: Optional[torch.Tensor], # (B,H,W,1) float or None + mask_bool: Optional[torch.Tensor], + keep_colors: bool, + non_linear: bool, +): + """ + Apply auto-stretch contrast in RGB space (GIMP-equivalent). + All operations are batch-vectorized and GPU-safe. + """ + + device = rgb.device + B, H, W, C = rgb.shape + + work = rgb + + # ── linear light conversion (optional) ───────────────────────── + if non_linear: + work = work.clamp(0.0, 1.0).pow(1.0 / 2.2) + + # ── reshape for reduction (B, N, C) ──────────────────────────── + flat = work.view(B, H * W, C) + + if mask_bool is not None: + mask_flat = mask_bool.view(B, H * W) + + # avoid invalid pixels affecting reduction + masked = torch.where( + mask_flat.unsqueeze(-1), + flat, + torch.tensor(float("inf"), device=device) + ) + masked_neg = torch.where( + mask_flat.unsqueeze(-1), + flat, + torch.tensor(float("-inf"), device=device) + ) + + mins = masked.amin(dim=1) + maxs = masked_neg.amax(dim=1) + + # fallback safety (empty mask per channel) + mins = torch.where(torch.isinf(mins), torch.zeros_like(mins), mins) + maxs = torch.where(torch.isinf(maxs), torch.ones_like(maxs), maxs) + + else: + mins = flat.amin(dim=1) + maxs = flat.amax(dim=1) + + # ── shared channel behavior (GIMP "preserve colors") ─────────── + if keep_colors: + lo = mins.amin(dim=1, keepdim=True) # (B,1) + hi = maxs.amax(dim=1, keepdim=True) # (B,1) + mins = lo.expand_as(mins) + maxs = hi.expand_as(maxs) + + # ── avoid division by zero ───────────────────────────────────── + denom = (maxs - mins).clamp(min=EPS) + + # ── reshape for broadcast ─────────────────────────────────────── + mins = mins.view(B, 1, 1, 3) + maxs = maxs.view(B, 1, 1, 3) + denom = denom.view(B, 1, 1, 3) + + # ── apply stretch (GIMP-equivalent formula) ───────────────────── + out = (work - mins) / denom + out = out.clamp(0.0, 1.0) + + # ── revert gamma if needed ────────────────────────────────────── + if non_linear: + out = out.pow(2.2) + + return out + +def exec_operator_color_auto_stretch_contrast( + image: torch.Tensor, + stretch_type: str = "RGB", + keep_colors: bool = True, + non_linear: bool = False, + mask: torch.Tensor = None, + low_percentile: float = 0.0, + high_percentile: float = 1.0, +) -> Tuple[torch.Tensor]: + """ + Applies auto-stretch contrast to an image batch. + + The operation expands the dynamic range of the image so that its + minimum and maximum values span the full [0,1] interval. + + When a mask is provided, it: + - defines the region used for statistical computation + - and controls per-pixel strength of the effect + + Alpha channels are preserved unchanged. + """ + + # ── input validation ──────────────────────────────────────────── + image = _check_input(image) + device = image.device + + if stretch_type.upper() not in ("RGB", "HSV"): + raise ValueError("stretch_type must be 'RGB' or 'HSV'") + + # ── split alpha ──────────────────────────────────────────────── + has_alpha = image.shape[-1] == 4 + rgb = image[..., :3].clamp(0.0, 1.0) + alpha = image[..., 3:] if has_alpha else None + + B, H, W, _ = rgb.shape + + # ── normalize mask ───────────────────────────────────────────── + mask_f, mask_bool = _normalize_mask(mask, device, B, H, W) + + # ── early exit ──────────────────────────────────────────────── + if _early_exit_check(rgb, mask_bool): + return (image.clone(),) + + # ── VRAM chunking ───────────────────────────────────────────── + chunk_size = compute_vram_chunk(rgb, mask_f) + + rgb_out = torch.empty_like(rgb) + + # ── process chunks ───────────────────────────────────────────── + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + + sub_rgb = rgb[start:end] + sub_mask_f = mask_f[start:end] if mask_f is not None else None + sub_mask_b = mask_bool[start:end] if mask_bool is not None else None + + if stretch_type.upper() == "RGB": + stretched = _stretch_chunk_rgb( + sub_rgb, + sub_mask_f, + sub_mask_b, + keep_colors, + non_linear, + ) + + else: + # ── HSV path (value channel stretch) ─────────────────── + hsv = rgb_to_hsv(sub_rgb) + + # ── normalize HSV format ─────────────────────────────────────── + if isinstance(hsv, tuple): + if len(hsv) == 3: + h, s, v = hsv + else: + hsv = hsv[0] + h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2] + else: + h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2] + + # ── compute min/max on V channel ─────────────────────────────── + flat_v = v.view(v.shape[0], -1) + + if sub_mask_b is not None: + mask_flat = sub_mask_b.view(v.shape[0], -1) + + v_masked = torch.where( + mask_flat.unsqueeze(-1), + v.view(v.shape[0], -1, 1), + torch.tensor(float("inf"), device=device) + ) + + v_min = v_masked.amin(dim=1) + v_max = torch.where( + mask_flat.unsqueeze(-1), + v.view(v.shape[0], -1, 1), + torch.tensor(float("-inf"), device=device) + ).amax(dim=1) + + v_min = torch.where(torch.isinf(v_min), torch.zeros_like(v_min), v_min) + v_max = torch.where(torch.isinf(v_max), torch.ones_like(v_max), v_max) + else: + v_min = flat_v.amin(dim=1) + v_max = flat_v.amax(dim=1) + + denom = (v_max - v_min).clamp(min=EPS).view(-1, 1, 1) + + v = (v - v_min.view(-1, 1, 1)) / denom + v = v.clamp(0.0, 1.0) + + # ── rebuild HSV → RGB correctly ──────────────────────────────── + stretched = hsv_to_rgb(h, s, v) + + # ── PER-PIXEL MASK STRENGTH BLENDING (REQUIRED) ───────────── + if mask_f is not None: + mf = sub_mask_f + + stretched = rgb[start:end] + (stretched - rgb[start:end]) * mf + + rgb_out[start:end] = stretched + + # ── reattach alpha ───────────────────────────────────────────── + if has_alpha: + rgb_out = torch.cat((rgb_out, alpha), dim=-1) + + return (rgb_out,) + + diff --git a/py/color_auto_white_balance.py b/py/color_auto_white_balance.py new file mode 100644 index 0000000..7af9b85 --- /dev/null +++ b/py/color_auto_white_balance.py @@ -0,0 +1,94 @@ +import torch +from typing import Tuple + +def exec_operator_color_auto_white_balance( + image: torch.Tensor, + mask: torch.Tensor = None +) -> Tuple[torch.Tensor]: + + """ + Auto White Balance (GIMP-style percentile stretch per channel). + + Args: + image: [B, H, W, C] float tensor in [0,1] + mask: [B, 1, H, W] optional weighting mask (0–1) + + Returns: + Tuple containing processed image [B, H, W, C] + """ + + if image.dtype != torch.float32: + image = image.float() + + B, H, W, C = image.shape + eps = 1e-6 + + # ----------------------------- + # Split RGB / Alpha + # ----------------------------- + rgb = image[..., :min(3, C)] # [B,H,W,3] + has_alpha = C > 3 + alpha = image[..., 3:] if has_alpha else None + + # Convert to NCHW + rgb = rgb.permute(0, 3, 1, 2).contiguous() # [B,3,H,W] + + B, Cc, H, W = rgb.shape + + # ----------------------------- + # Apply mask if provided + # ----------------------------- + if mask is not None: + if mask.dtype != torch.float32: + mask = mask.float() + if mask.ndim == 4: + mask = mask[:, 0:1, :, :] + mask = torch.nn.functional.interpolate(mask, size=(H, W), mode="bilinear", align_corners=False) + mask = mask.clamp(0.0, 1.0) + else: + mask = torch.ones((B, 1, H, W), device=rgb.device, dtype=rgb.dtype) + + # ----------------------------- + # Compute weighted percentiles + # ----------------------------- + flat_rgb = rgb.reshape(B, Cc, -1) # [B,3,N] + flat_mask = mask.reshape(B, 1, -1) # [B,1,N] + + # Sort values per channel + sorted_vals, indices = torch.sort(flat_rgb, dim=2) + sorted_weights = torch.gather(flat_mask.expand_as(flat_rgb), 2, indices) + + cum_weights = torch.cumsum(sorted_weights, dim=2) + total_weights = cum_weights[:, :, -1:] + eps + + cum_norm = cum_weights / total_weights + + lower_p = 0.0005 # 0.05% + upper_p = 0.9995 # 99.95% + + def percentile_value(p): + # first index where cum_norm >= p + idx = torch.searchsorted(cum_norm, torch.full((B, Cc, 1), p, device=rgb.device)) + idx = idx.clamp(0, sorted_vals.shape[2] - 1) + return torch.gather(sorted_vals, 2, idx) + + lower = percentile_value(lower_p) + upper = percentile_value(upper_p) + + # ----------------------------- + # Stretch + # ----------------------------- + denom = (upper - lower).clamp(min=eps) + stretched = (flat_rgb - lower) / denom + stretched = stretched.clamp(0.0, 1.0) + + rgb_out = stretched.reshape(B, Cc, H, W).permute(0, 2, 3, 1).contiguous() + + # ----------------------------- + # Reattach alpha + # ----------------------------- + if has_alpha: + rgb_out = torch.cat([rgb_out, alpha], dim=-1) + + return (rgb_out,) + diff --git a/py/color_balance.py b/py/color_balance.py new file mode 100644 index 0000000..5a35b34 --- /dev/null +++ b/py/color_balance.py @@ -0,0 +1,129 @@ +import torch +from typing import Tuple, Optional + +@torch.no_grad() # 1) no gradient tracking → saves VRAM & time +def exec_operation_color_balance( + image: torch.Tensor, # shape: (B, H, W, C) with C = 3 or 4 + mask: Optional[torch.Tensor], # can be (H,W), (B,H,W), (B,H,W,1), ... + shadow_adj: Tuple[float, float, float], + midtones_adj: Tuple[float, float, float], + highlights_adj: Tuple[float, float, float], + preserve_luminosity: bool +) -> torch.Tensor: + """ + Apply GIMP‑style color balance. + Each adjustment tuple = (Cyan‑Red, Magenta‑Green, Yellow‑Blue) in [-100, 100]. + + Improvements over original: + - Robust batch‑mask handling (2D mask → broadcast to whole batch) + - Full RGB / RGBA support (alpha channel is never color‑shifted) + - Input validation with clear error messages + - `torch.no_grad()` + single memory allocation → VRAM optimised + - Works on any device (CPU / CUDA / MPS) + """ + + # ---------------------------------------------------------------------- + # 1. Input validation & preparation + # ---------------------------------------------------------------------- + if image.dim() != 4: + raise ValueError(f"image must be 4D (B,H,W,C), got {image.dim()}D") + B, H, W, C = image.shape + if C not in (3, 4): + raise ValueError(f"only 3 (RGB) or 4 (RGBA) channels allowed, got {C}") + for adj in (shadow_adj, midtones_adj, highlights_adj): + if len(adj) != 3: + raise ValueError("each adjustment must be a sequence of 3 floats") + device = image.device + + # Work on a float copy (same device) + img = image.to(dtype=torch.float32, copy=True) # single allocation + rgb = img[..., :3] # view – no copy + + # ---------------------------------------------------------------------- + # 2. Prepare mask → broadcastable (B, H, W, 1) + # ---------------------------------------------------------------------- + mask_tensor: Optional[torch.Tensor] = None + if mask is not None: + m = mask.to(device=device, dtype=torch.float32) + # unify shapes + if m.dim() == 2: # (H, W) → apply to all images + m = m[None, ..., None] # (1, H, W, 1) + m = m.expand(B, -1, -1, -1) # (B, H, W, 1) + elif m.dim() == 3: + if m.shape[-1] == 1: # (H, W, 1) or (B, H, W, 1)? + if m.shape[0] == B: + m = m # (B, H, W, 1) – already fine + else: + m = m[None] # (1, H, W, 1) + m = m.expand(B, -1, -1, -1) + else: # (B, H, W) + m = m.unsqueeze(-1) # (B, H, W, 1) + elif m.dim() == 4: + if m.shape[0] == 1 and B > 1: + m = m.expand(B, -1, -1, -1) # broadcast batch dimension + else: + raise ValueError(f"mask shape {mask.shape} not supported") + # safety check + if m.shape[:3] != (B, H, W): + raise ValueError( + f"mask spatial dimensions must match image. " + f"Expected (B,H,W)=({B},{H},{W}), got {m.shape[:3]}" + ) + mask_tensor = m + + # ---------------------------------------------------------------------- + # 3. Luma & transfer‑function weights (Rec. 709) + # ---------------------------------------------------------------------- + luma_w = torch.tensor([0.299, 0.587, 0.114], device=device, dtype=torch.float32) + luma = (rgb * luma_w).sum(dim=-1, keepdim=True) # (B, H, W, 1) + + def smoothstep(edge0, edge1, x): + t = torch.clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0) + return t * t * (3.0 - 2.0 * t) + + shadow_w = smoothstep(0.5, 0.0, luma) # 1 at 0, 0 at 0.5 + midtones_w = 4.0 * luma * (1.0 - luma) # peak 1 at 0.5 + highlights_w = smoothstep(0.5, 1.0, luma) # 0 at 0.5, 1 at 1.0 + + # ---------------------------------------------------------------------- + # 4. Normalised adjustments → (1,1,1,3) + # ---------------------------------------------------------------------- + def to_tensor(adj): + return torch.tensor(adj, device=device, dtype=torch.float32).view(1,1,1,3) / 100.0 + + delta = ( + shadow_w * to_tensor(shadow_adj) + + midtones_w * to_tensor(midtones_adj) + + highlights_w * to_tensor(highlights_adj) + ) # (B, H, W, 3) + + # ---------------------------------------------------------------------- + # 5. Apply delta (with or without mask) + # ---------------------------------------------------------------------- + if mask_tensor is not None: + blended = rgb + delta * mask_tensor # only RGB + else: + blended = rgb + delta + + # Build output (keep alpha untouched) + if C == 4: + alpha = img[..., 3:4] + out = torch.cat([blended, alpha], dim=-1) + else: + out = blended + + # ---------------------------------------------------------------------- + # 6. Optional luminance preservation (RGB only) + # ---------------------------------------------------------------------- + if preserve_luminosity: + new_luma = (out[..., :3] * luma_w).sum(dim=-1, keepdim=True) + eps = 1e-8 + scale = luma / (new_luma + eps) + out[..., :3] = out[..., :3] * scale + + # ---------------------------------------------------------------------- + # 7. Final clamp (entire image, including alpha, exactly like original) + # ---------------------------------------------------------------------- + out.clamp_(0.0, 1.0) + + return out diff --git a/py/color_brightness_contrast.py b/py/color_brightness_contrast.py new file mode 100644 index 0000000..49360dd --- /dev/null +++ b/py/color_brightness_contrast.py @@ -0,0 +1,131 @@ +import torch +import math +from typing import Tuple, Optional + +from .libs.process import compute_chunk_size + + +def exec_operator_color_brightness_contrast( + image: torch.Tensor, + brightness: float, + contrast: float, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Apply GIMP‑style brightness and contrast adjustment to a batch of images. + + This function processes an entire batch in GPU‑friendly chunks to manage + VRAM consumption. Chunk size is determined by :func:`compute_chunk_size`. + + **Input**: + image (torch.Tensor): shape ``(B, H, W, C)``, values in [0, 1]. + Supported channel counts: ``C=3`` (RGB) or ``C=4`` (RGBA). + brightness (float): slider value in [-127, 127] (GIMP UI range). + contrast (float): slider value in [-127, 127]. + mask (Optional[torch.Tensor]): + Optional mask of shape ``(B, H, W)`` or ``(B, H, W, 1)``, + values in [0, 1]. When provided, the adjustment is blended + per‑pixel using ``mask``. + + **Output**: + Tuple containing a single tensor of shape ``(B, H, W, C)`` where ``C`` + matches the input’s channel count. Alpha channel (if present) is + preserved unchanged. + + **Processing logic**: + 1. Convert slider values to GIMP internal ranges: + ``b = (brightness/127) / 2``, ``slant = tan((contrast/127 + 1) * pi/4)``. + 2. Apply brightness (GIMP formula) and contrast to the RGB channels. + 3. Clamp results to [0, 1]. + 4. If a mask is given, blend the adjusted RGB with the original RGB. + 5. Re‑attach the alpha channel if the input was RGBA. + + **GPU memory management**: + The batch is split into chunks whose size is calculated by + ``compute_chunk_size(image, mask)``. Each chunk is processed + independently, and the final result is concatenated along the + batch dimension. No spatial tiling is performed. + + Returns: + ``(result,)`` where *result* is the adjusted batch tensor. + """ + # ----------------------- Early exit ----------------------- + if brightness == 0 and contrast == 0: + return (image,) + + # ----------------------- Input validation ----------------- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() != 4: + raise ValueError(f"image must be 4D (B,H,W,C), got {image.dim()}D") + if image.shape[-1] not in (3, 4): + raise ValueError(f"image last dim must be 3 or 4, got {image.shape[-1]}") + if not image.is_floating_point(): + raise TypeError("image tensor must be floating point") + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor") + if mask.dim() not in (3, 4): + raise ValueError(f"mask must be 3D (B,H,W) or 4D (B,H,W,1), got {mask.dim()}D") + if mask.shape[0] != image.shape[0]: + raise ValueError("mask batch size must equal image batch size") + if mask.shape[1:3] != image.shape[1:3]: + raise ValueError("mask spatial size must match image") + if mask.dim() == 4 and mask.shape[-1] != 1: + raise ValueError("4D mask must have 1 channel") + if mask.dtype != image.dtype: + mask = mask.to(dtype=image.dtype) # auto‑cast to avoid errors + + # ----------------------- Chunked processing ----------------- + chunk_size = compute_chunk_size(image, mask) + B = image.shape[0] + # If chunk_size is invalid or larger than batch, process all at once. + if chunk_size is None or chunk_size <= 0 or chunk_size >= B: + chunk_size = B + + # Precompute constants shared across all chunks + brightness_internal = brightness / 127.0 + contrast_internal = contrast / 127.0 + b = brightness_internal / 2.0 # [-0.5, 0.5] + slant = math.tan((contrast_internal + 1) * math.pi / 4) + + result_chunks = [] + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + img_chunk = image[start:end] # (cs, H, W, C) + m_chunk = mask[start:end] if mask is not None else None + + # ---- RGB channels ---- + rgb = img_chunk[..., :3] # (cs, H, W, 3) + + # Brightness (GIMP) + if b < 0: + rgb = rgb * (1.0 + b) + else: + rgb = rgb + (1.0 - rgb) * b + + # Contrast + rgb = (rgb - 0.5) * slant + 0.5 + rgb = torch.clamp(rgb, 0.0, 1.0) + + # Mask blending + if m_chunk is not None: + if m_chunk.dim() == 3: + m_chunk = m_chunk.unsqueeze(-1) # (cs, H, W, 1) + original = img_chunk[..., :3] + rgb = original * (1.0 - m_chunk) + rgb * m_chunk + + # Alpha preservation + if img_chunk.shape[-1] == 4: + result_chunk = torch.cat([rgb, img_chunk[..., 3:4]], dim=-1) + else: + result_chunk = rgb + + result_chunks.append(result_chunk) + + # Concatenate all chunks back along batch dimension + result = torch.cat(result_chunks, dim=0) + return (result,) + + diff --git a/py/color_colorize.py b/py/color_colorize.py new file mode 100644 index 0000000..304768e --- /dev/null +++ b/py/color_colorize.py @@ -0,0 +1,145 @@ +import torch +from typing import Tuple, Optional + +from .libs.process import compute_chunk_size + +# =================================================================================== +# GEGL-perceptual Colorize operator (faithful approximation) +# =================================================================================== +def exec_operator_color_colorize( + image: torch.Tensor, + Hue: float, + Saturation: float, + Lightness: float, + Color: int, + mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + """ + Apply a GEGL‑perceptual colourisation to a batch of images. + + The algorithm works in the following steps: + 1. Compute BT.709 perceptual luminance of the original image. + 2. Decode the target colour and build a chroma vector. + 3. Scale the chroma vector by **Hue × Saturation** and add it + to the luminance, effectively tinting the image. + 4. Apply a luminance bias based on the **Lightness** parameter. + 5. If a **mask** is provided, blend the result with the original + using per‑pixel mask values as strength. + + The function supports RGB and RGBA input: alpha is preserved + unchanged and only the RGB channels are colourised. Batch + processing is handled internally with GPU‑friendly chunking. + + Args: + image: [B, H, W, C] float tensor (C=3 or 4), values 0‑1. + Hue: float 0‑1, controls tint strength (together with Saturation). + Saturation: float 0‑1, scales the chroma direction. + Lightness: float 0‑1, brightness bias (0.5 = no change). + Color: 24‑bit integer (0xRRGGBB) of the target tint. + mask: Optional mask [H, W] or [B, H, W], values 0‑1. + Returns: + A 1‑tuple containing the colourised image tensor [B, H, W, C]. + """ + # --- input validation ---------------------------------------------------- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() != 4: + raise ValueError(f"image must be 4D [B,H,W,C], got shape {image.shape}") + if image.shape[-1] not in (3, 4): + raise ValueError(f"image must have 3 or 4 channels, got {image.shape[-1]}") + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor") + if mask.dim() not in (2, 3): + raise ValueError(f"mask must be 2D [H,W] or 3D [B,H,W], got shape {mask.shape}") + if mask.dim() == 3 and mask.shape[0] != image.shape[0]: + raise ValueError( + f"batch size mismatch: image {image.shape[0]}, mask {mask.shape[0]}" + ) + # device consistency check + if mask.device != image.device: + raise ValueError("image and mask must be on the same device") + + # --- early exit conditions ----------------------------------------------- + B = image.shape[0] + if B == 0: + return (image,) + if mask is not None and not mask.isnan().any() and mask.max() == 0.0: + # mask is completely black → original image is unchanged + return (image,) + + # --- prepare alpha handling ---------------------------------------------- + has_alpha = image.shape[-1] == 4 + if has_alpha: + rgb = image[..., :3] + alpha = image[..., 3:4] + else: + rgb = image + alpha = None + + # --- determine chunk size for GPU processing ----------------------------- + chunk_size = compute_chunk_size(image, mask) + results = [] + + # --- chunked batch loop -------------------------------------------------- + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + chunk_rgb = rgb[start:end] # [C, H, W, 3] + chunk_mask = None + if mask is not None: + if mask.dim() == 3: # batched mask + chunk_mask = mask[start:end] + else: # shared 2D mask + chunk_mask = mask + + # --- colourisation pipeline (identical to original) ------------------- + # 1. Ensure linear RGB (clamp) + chunk_rgb = chunk_rgb.clamp(0, 1) + + r, g, b = chunk_rgb[..., 0], chunk_rgb[..., 1], chunk_rgb[..., 2] + + # 2. BT.709 perceptual luminance + Y = 0.222488 * r + 0.716903 * g + 0.060609 * b + + # 3. Target colour + cr = ((Color >> 16) & 255) / 255.0 + cg = ((Color >> 8) & 255) / 255.0 + cb = (Color & 255) / 255.0 + Yc = 0.222488 * cr + 0.716903 * cg + 0.060609 * cb + + # 4. Chroma direction and tint scaling + dr, dg, db = cr - Yc, cg - Yc, cb - Yc + sat = Saturation * Hue + + rr = Y + dr * sat + gg = Y + dg * sat + bb = Y + db * sat + + out_rgb = torch.stack([rr, gg, bb], dim=-1) + + # 5. Lightness bias + light_bias = Lightness - 0.5 + out_rgb = out_rgb + light_bias + out_rgb = out_rgb.clamp(0, 1) + + # 6. Mask blending (per‑pixel strength) + if chunk_mask is not None: + # ensure mask has a trailing channel dimension for broadcasting + if chunk_mask.dim() == 2: + chunk_mask = chunk_mask.unsqueeze(0).unsqueeze(-1) # [1, H, W, 1] + elif chunk_mask.dim() == 3: + chunk_mask = chunk_mask.unsqueeze(-1) # [B, H, W, 1] + out_rgb = chunk_rgb * (1 - chunk_mask) + out_rgb * chunk_mask + + # 7. Reattach alpha if present + if has_alpha: + chunk_out = torch.cat([out_rgb, alpha[start:end]], dim=-1) + else: + chunk_out = out_rgb + + results.append(chunk_out) + + # concatenate all chunks back to full batch + final = torch.cat(results, dim=0) + return (final,) + diff --git a/py/color_comp_channel_mixer.py b/py/color_comp_channel_mixer.py new file mode 100644 index 0000000..ca1c3b8 --- /dev/null +++ b/py/color_comp_channel_mixer.py @@ -0,0 +1,177 @@ +import torch +from typing import Tuple, Optional +from .libs.process import compute_chunk_size + +_LUMA = torch.tensor([0.2126, 0.7152, 0.0722]) + + +def exec_operator_color_comp_channel_mixer( + image: torch.Tensor, + Red_in_red: float, + Green_in_red: float, + Blue_in_red: float, + Red_in_green: float, + Green_in_green: float, + Blue_in_green: float, + Red_in_blue: float, + Green_in_blue: float, + Blue_in_blue: float, + Preserve_luminosity: bool, + mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + """ + Applies a channel‑mixing matrix to an RGB or RGBA image, optionally + modulated by a per‑pixel mask and with luminosity preservation. + + Input: + image : torch.Tensor – batch of images, shape [B, H, W, C] with C=3 (RGB) or 4 (RGBA). + Red_in_red … Blue_in_blue : coefficients of the 3×3 mixing matrix. + Preserve_luminosity : if True, scales the mixed result so that the perceived + luminance matches the original before clamping. + mask : optional torch.Tensor – per‑pixel blend strength, broadcastable to + [B, H, W, 1]. Allowed shapes: [H, W], [B, H, W], [B, H, W, 1]. + A pixel value of 1.0 gives full mixing; 0.0 leaves the original unchanged. + + Output: + Tuple containing a single tensor of mixed images, same shape as input but + with RGB channels modified. If the input has an alpha channel, it is passed + through unchanged. + + """ + # ---------- input validation ---------- + if image.dim() != 4: + raise ValueError( + f"Input image must be a 4D tensor [B, H, W, C], got shape {image.shape}" + ) + if image.shape[-1] not in (3, 4): + raise ValueError( + f"Input image must have 3 (RGB) or 4 (RGBA) channels, got {image.shape[-1]}" + ) + if mask is not None: + if mask.dim() not in (2, 3, 4): + raise ValueError( + f"Mask must be 2D [H,W], 3D [B,H,W] or 4D [B,H,W,1], got shape {mask.shape}" + ) + # ensure mask can be broadcast to [B, H, W, 1] + if mask.dim() == 2: + mask = mask.unsqueeze(0).unsqueeze(-1) # [1, H, W, 1] + elif mask.dim() == 3: + if mask.shape[0] != image.shape[0]: + mask = mask.unsqueeze(0).unsqueeze(-1) # treat as [1, H, W, 1] + else: + mask = mask.unsqueeze(-1) # [B, H, W, 1] + # mask already 4D – verify last dim == 1 + if mask.shape[-1] != 1: + mask = mask.unsqueeze(-1) + if mask.shape[0] not in (1, image.shape[0]): + raise ValueError( + f"Mask batch dimension ({mask.shape[0]}) must be 1 or match image batch ({image.shape[0]})" + ) + if mask.shape[1:3] != image.shape[1:3]: + raise ValueError( + f"Mask spatial size {mask.shape[1:3]} does not match image size {image.shape[1:3]}" + ) + + # ---------- early exit conditions ---------- + identity = ( + Red_in_red == 1.0 + and Green_in_red == 0.0 + and Blue_in_red == 0.0 + and Red_in_green == 0.0 + and Green_in_green == 1.0 + and Blue_in_green == 0.0 + and Red_in_blue == 0.0 + and Green_in_blue == 0.0 + and Blue_in_blue == 1.0 + and not Preserve_luminosity + ) + if mask is None and identity: + return (image,) + if mask is not None and torch.all(mask == 0.0): + return (image,) + if mask is not None and identity: + # identity transform + mask all zeros already caught; + # if mask not all zero we still need blending, so continue + pass + + # ---------- chunked processing ---------- + B = image.shape[0] + chunk_size = compute_chunk_size(image, mask) + results = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + chunk_img = image[start:end] + + # slice mask if it has a batch dim, otherwise keep the single mask + if mask is not None: + if mask.shape[0] == B: + chunk_mask = mask[start:end] + else: + chunk_mask = mask + # move to same device/dtype as image + chunk_mask = chunk_mask.to(device=chunk_img.device, dtype=chunk_img.dtype) + else: + chunk_mask = None + + # -------- inner processing function (preserves original logic) -------- + def _process_chunk(img: torch.Tensor, msk: Optional[torch.Tensor]) -> torch.Tensor: + has_alpha = img.shape[-1] == 4 + if has_alpha: + alpha = img[..., 3:4] # preserve alpha + rgb = img[..., :3] + else: + rgb = img + + r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2] + + # matrix mixing (identical to original) + r_out = Red_in_red * r + Green_in_red * g + Blue_in_red * b + g_out = Red_in_green * r + Green_in_green * g + Blue_in_green * b + b_out = Red_in_blue * r + Green_in_blue * g + Blue_in_blue * b + + # luminosity preservation + if Preserve_luminosity: + device = rgb.device + dtype = rgb.dtype + luma = _LUMA.to(device=device, dtype=dtype) + + def channel_luma(a, b, c): + return a * luma[0] + b * luma[1] + c * luma[2] + + orig_luma = luma[0] * r + luma[1] * g + luma[2] * b + mix_luma = ( + channel_luma(Red_in_red, Green_in_red, Blue_in_red) * r + + channel_luma(Red_in_green, Green_in_green, Blue_in_green) * g + + channel_luma(Red_in_blue, Green_in_blue, Blue_in_blue) * b + ) + eps = 1e-8 + scale = orig_luma / (mix_luma + eps) + r_out = r_out * scale + g_out = g_out * scale + b_out = b_out * scale + + # stack and clamp (only RGB part) + mixed_rgb = torch.stack([r_out, g_out, b_out], dim=-1) + + # mask blending (applies only to RGB) + if msk is not None: + # broadcast mask from [B, H, W, 1] to [B, H, W, 3] + mixed_rgb = rgb * (1.0 - msk) + mixed_rgb * msk + + mixed_rgb = torch.clamp(mixed_rgb, 0.0, 1.0) + + if has_alpha: + return torch.cat([mixed_rgb, alpha], dim=-1) + return mixed_rgb + + # -------- process the chunk -------- + res_chunk = _process_chunk(chunk_img, chunk_mask) + results.append(res_chunk) + + # free memory after each chunk if on GPU + if image.is_cuda: + torch.cuda.empty_cache() + + return (torch.cat(results, dim=0),) + diff --git a/py/color_comp_extract.py b/py/color_comp_extract.py new file mode 100644 index 0000000..11ef94b --- /dev/null +++ b/py/color_comp_extract.py @@ -0,0 +1,237 @@ +import torch +import math +from typing import Tuple, Optional +from .libs.process import compute_chunk_size + +def exec_operator_color_comp_extract( + image: torch.Tensor, + component: str, + invert: bool = False, + linear: bool = False, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Extract a specific colour component from an image batch and return it as a mask tensor. + + The function supports the following colour spaces: + RGB (red, green, blue), HSV (hue, saturation, value), HSL (saturation, lightness), + CMYK (cyan, magenta, yellow, key), Y'CbCr (Y', Cb, Cr), CIE LAB (L, A, B), + CIE LCH (C(ab), H(ab)), and Alpha. + + Parameters + ---------- + image : torch.Tensor + Input image batch with shape ``(B, H, W, C)`` where ``C`` is 3 (RGB) or + 4 (RGBA). When a 3‑D tensor ``(H, W, C)`` is given it is treated as a + batch of size 1. + component : str + Name of the colour component to extract (case‑insensitive). Valid values + are those listed in the ``INPUT_TYPES`` of the calling node. + invert : bool + If ``True``, the output mask is inverted (``1 - value``) after the mask is + applied. + linear : bool + If ``True``, the extracted component is kept in linear light. When + ``False`` (default) a gamma encoding (sRGB transfer function) is applied + to the raw component values before the mask is multiplied. + mask : torch.Tensor or None + Optional per‑pixel strength mask with shape ``(B, H, W)`` or ``(B, H, W, 1)``. + Must match the batch size and spatial dimensions of ``image``. If ``None``, + the extracted component is used as is (subject to ``linear`` and ``invert``). + + Returns + ------- + Tuple[torch.Tensor] + A single‑element tuple containing the result mask of shape ``(B, H, W)``, + with values in the range ``[0, 1]``. + + Raises + ------ + TypeError + If ``image`` or ``mask`` is not a ``torch.Tensor``. + ValueError + If the shapes of ``image`` or ``mask`` are invalid or incompatible. + """ + # ---------- input validation ---------- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() not in (3, 4): + raise ValueError( + f"image must have shape (B,H,W,C) or (H,W,C), got {image.shape}" + ) + if image.dim() == 3: + image = image.unsqueeze(0) # treat as batch size 1 + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor or None") + # unify mask shape to (B, H, W) + if mask.dim() == 4: + mask = mask.squeeze(-1) + if mask.dim() != 3: + raise ValueError( + f"mask must be (B,H,W) or (B,H,W,1), got {mask.shape}" + ) + if mask.shape[0] != image.shape[0] or mask.shape[1:3] != image.shape[1:3]: + raise ValueError( + "mask batch and spatial dimensions must match those of image" + ) + + # ---------- early exit for empty input / all‑zero mask ---------- + if image.shape[0] == 0: + return (torch.empty(0, image.shape[1], image.shape[2], + device=image.device, dtype=torch.float32),) + + if mask is not None and mask.eq(0).all(): + # result will be all zeros or all ones depending on invert + result = torch.ones_like(mask) if invert else torch.zeros_like(mask) + return (result,) + + # ---------- determine GPU‑friendly chunk size ---------- + if image.is_cuda and torch.cuda.is_available(): + chunk_size = compute_chunk_size(image, mask) + if not isinstance(chunk_size, int) or chunk_size <= 0: + chunk_size = image.shape[0] + else: + chunk_size = image.shape[0] + chunk_size = max(1, min(chunk_size, image.shape[0])) + + comp = component.strip().lower() + results = [] + + # ---------- chunked processing loop ---------- + for start in range(0, image.shape[0], chunk_size): + end = min(start + chunk_size, image.shape[0]) + img_chunk = image[start:end] + m_chunk = mask[start:end] if mask is not None else None + + # --- colour space conversion (same as original) --- + rgb = img_chunk[..., :3] + r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2] + + if comp == "rgb red": + res = r + elif comp == "rgb green": + res = g + elif comp == "rgb blue": + res = b + + elif comp.startswith("hsv"): + maxc = torch.max(rgb, dim=-1).values + minc = torch.min(rgb, dim=-1).values + diff = maxc - minc + 1e-8 + + if comp == "hsv value": + res = maxc + elif comp == "hsv saturation": + res = diff / (maxc + 1e-8) + elif comp == "hsv hue": + hue = torch.zeros_like(maxc) + mask_nonzero = diff > 1e-8 + r_eq = (maxc == r) & mask_nonzero + g_eq = (maxc == g) & mask_nonzero + b_eq = (maxc == b) & mask_nonzero + hue_r = ((g - b) / diff) % 6.0 + hue_g = ((b - r) / diff) + 2.0 + hue_b = ((r - g) / diff) + 4.0 + hue = torch.where(r_eq, hue_r, hue) + hue = torch.where(g_eq, hue_g, hue) + hue = torch.where(b_eq, hue_b, hue) + res = (hue / 6.0) % 1.0 + else: + res = maxc + + elif comp.startswith("hsl"): + maxc = torch.max(rgb, dim=-1).values + minc = torch.min(rgb, dim=-1).values + l = (maxc + minc) * 0.5 + if comp == "hsl lightness": + res = l + elif comp == "hsl saturation": + denom = 1 - torch.abs(2 * l - 1) + 1e-8 + res = (maxc - minc) / denom + else: + res = l + + elif comp.startswith("cmyk"): + k = 1.0 - torch.max(rgb, dim=-1).values + denom = 1.0 - k + 1e-8 + if comp == "cmyk cyan": + res = (1.0 - r - k) / denom + elif comp == "cmyk magenta": + res = (1.0 - g - k) / denom + elif comp == "cmyk yellow": + res = (1.0 - b - k) / denom + elif comp == "cmyk key": + res = k + else: + res = k + + elif "ycbcr" in comp or "y'cbcr" in comp: + y = 0.299 * r + 0.587 * g + 0.114 * b + if "cb" in comp: + res = (b - y) * 0.564 + 0.5 + elif "cr" in comp: + res = (r - y) * 0.713 + 0.5 + else: + res = y + + elif comp.startswith("lab") or comp.startswith("lch"): + def f(t): + return torch.where(t > 0.008856, t ** (1/3), 7.787 * t + 16/116) + + rl = torch.where(r > 0.04045, ((r + 0.055) / 1.055) ** 2.4, r / 12.92) + gl = torch.where(g > 0.04045, ((g + 0.055) / 1.055) ** 2.4, g / 12.92) + bl = torch.where(b > 0.04045, ((b + 0.055) / 1.055) ** 2.4, b / 12.92) + + x = 0.4124564 * rl + 0.3575761 * gl + 0.1804375 * bl + y = 0.2126729 * rl + 0.7151522 * gl + 0.0721750 * bl + z = 0.0193339 * rl + 0.1191920 * gl + 0.9503041 * bl + + xn, yn, zn = 0.95047, 1.0, 1.08883 + x, y, z = x/xn, y/yn, z/zn + + fx, fy, fz = f(x), f(y), f(z) + a = 500 * (fx - fy) + bb = 200 * (fy - fz) + + if comp == "lab l": + res = (116 * fy - 16) / 100 + elif comp == "lab a": + res = a / 128 + 0.5 + elif comp == "lab b": + res = bb / 128 + 0.5 + elif comp == "lch c(ab)": + res = torch.sqrt(a*a + bb*bb) / 128 + elif comp == "lch h(ab)": + res = (torch.atan2(bb, a) / (2 * math.pi)) % 1.0 + else: + res = (116 * fy - 16) / 100 + + elif comp == "alpha": + res = img_chunk[..., 3] if img_chunk.shape[-1] == 4 else torch.ones_like(r) + + else: + res = r # fallback + + # --- post‑processing (identical to original) --- + res = res.clamp(0, 1) + + if not linear: + res = torch.where( + res <= 0.0031308, + 12.92 * res, + 1.055 * torch.pow(res, 1/2.4) - 0.055 + ).clamp(0, 1) + + if m_chunk is not None: + res = res * m_chunk + + if invert: + res = 1.0 - res + + results.append(res) + + final_result = torch.cat(results, dim=0) + return (final_result,) + diff --git a/py/color_comp_mono_mixer.py b/py/color_comp_mono_mixer.py new file mode 100644 index 0000000..69b3dcc --- /dev/null +++ b/py/color_comp_mono_mixer.py @@ -0,0 +1,142 @@ +import torch +from typing import Tuple, Optional + +# Assume this import is available in the target environment +from .libs.process import compute_chunk_size + + +def exec_operator_color_comp_mono_mixer( + image: torch.Tensor, + preserve_luminosity: bool, + red_multiplier: float, + green_multiplier: float, + blue_multiplier: float, + mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + """ + Apply a channel mixer that converts an RGB or RGBA image to grayscale using + per‑channel multipliers, optionally normalised to preserve perceptual luminance. + + The operation mimics GIMP's Colors → Components → Mono Mixer. If a mask is + supplied, it acts as a per‑pixel strength factor: the grayscale result is + blended with the original image according to the mask values. An existing + alpha channel is always kept unchanged. + + Processing runs on the same device as the input tensor. To avoid exhausting + GPU memory the batch is automatically split into chunks whose size is + determined by :func:`compute_chunk_size`. + + Args: + image: Input image tensor [B, H, W, C], C=3 (RGB) or C=4 (RGBA). + Values must lie in [0, 1]. + preserve_luminosity: If True, the three multipliers are normalised so that + their absolute sum equals 1. + red_multiplier: Coefficient for the red channel [-5.0, 5.0]. + green_multiplier: Coefficient for the green channel [-5.0, 5.0]. + blue_multiplier: Coefficient for the blue channel [-5.0, 5.0]. + mask: Optional mask tensor [B, H, W] or [B, H, W, 1]. + Values in [0, 1], 0 = no effect, 1 = full effect. + + Returns: + Tuple containing a single tensor [B, H, W, C] of the same channel count + as the input, with colour values in [0, 1]. + + Raises: + ValueError: If the input tensor does not have shape [B, H, W, 3] or + [B, H, W, 4], or if the mask dimensions / batch size are + incompatible. + """ + # ------------------------------------------------------------------ + # 1. Input validation + # ------------------------------------------------------------------ + if image.dim() != 4 or image.shape[-1] not in (3, 4): + raise ValueError( + f"image must be a 4D tensor [B,H,W,C] with C=3 or 4, got shape {tuple(image.shape)}" + ) + + B, H, W, C = image.shape + has_alpha = (C == 4) + + # ------------------------------------------------------------------ + # 2. Mask validation & early‑exit conditions + # ------------------------------------------------------------------ + if mask is not None: + if mask.dim() == 3: + # [B, H, W] -> [B, H, W, 1] + mask = mask.unsqueeze(-1) + elif mask.dim() == 4: + if mask.shape[-1] != 1: + raise ValueError( + f"mask must have shape [B,H,W] or [B,H,W,1], got {tuple(mask.shape)}" + ) + else: + raise ValueError( + f"mask must be 3D or 4D, got {mask.dim()}D with shape {tuple(mask.shape)}" + ) + + if mask.shape[0] != B: + raise ValueError( + f"Batch size mismatch: image has {B}, mask has {mask.shape[0]}" + ) + # Early exit: no pixels are affected + if mask.max() == 0: + return (image,) + + # Early exit: empty batch + if B == 0: + return (image,) + + # ------------------------------------------------------------------ + # 3. Chunk size for VRAM‑friendly processing + # ------------------------------------------------------------------ + chunk_size = compute_chunk_size(image, mask) + + # ------------------------------------------------------------------ + # 4. Per‑chunk processing + # ------------------------------------------------------------------ + results = [] + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + img_chunk = image[start:end] + msk_chunk = mask[start:end] if mask is not None else None + + # --- normalise multipliers (GEGL behaviour) --- + if preserve_luminosity: + m_sum = red_multiplier + green_multiplier + blue_multiplier + norm_factor = 1.0 if m_sum == 0.0 else abs(1.0 / m_sum) + else: + norm_factor = 1.0 + + r = red_multiplier * norm_factor + g = green_multiplier * norm_factor + b = blue_multiplier * norm_factor + + # --- grayscale via dot product on RGB channels --- + gray = img_chunk[..., 0] * r + img_chunk[..., 1] * g + img_chunk[..., 2] * b + gray = torch.clamp(gray, 0.0, 1.0) + + # --- build result RGB (all channels equal) --- + gray_rgb = torch.stack([gray, gray, gray], dim=-1) # [chunk, H, W, 3] + + # --- alpha preservation --- + if has_alpha: + alpha = img_chunk[..., 3:4] + res_chunk = torch.cat([gray_rgb, alpha], dim=-1) + else: + res_chunk = gray_rgb + + # --- mask blending (RGB channels only) --- + if msk_chunk is not None: + # msk_chunk is [chunk, H, W, 1], broadcast to RGB + res_chunk[..., :3] = ( + img_chunk[..., :3] * (1.0 - msk_chunk) + gray_rgb * msk_chunk + ) + + results.append(res_chunk) + + # ------------------------------------------------------------------ + # 5. Reassemble full batch + # ------------------------------------------------------------------ + final = torch.cat(results, dim=0) + return (final,) + diff --git a/py/color_desat_color_2_gray.py b/py/color_desat_color_2_gray.py new file mode 100644 index 0000000..73f3f28 --- /dev/null +++ b/py/color_desat_color_2_gray.py @@ -0,0 +1,151 @@ + + +import torch +import torch.nn.functional as F +from typing import Tuple + + +def exec_operator_color_desat_color_2_gray( + image: torch.Tensor, + radius: int, + samples: int, + iterations: int, + enhance_shadows: bool, + mask: torch.Tensor = None, +) -> Tuple[torch.Tensor]: + """ + GIMP-style Color to Gray using STRESS stochastic sampling. + Fully GPU-accelerated: all random offset sampling is vectorised, + no Python loops over pixels or samples. + """ + device = image.device + B, H, W, C = image.shape + + # Convert to [B, C, H, W] for grid_sample + img = image.permute(0, 3, 1, 2).to(device) + + # Initial grayscale: luminance (BT.709) + rgb_weights = torch.tensor([0.2126, 0.7152, 0.0722], device=device).view(1, 3, 1, 1) + gray = (img * rgb_weights).sum(dim=1, keepdim=True) # [B, 1, H, W] + + # Tiny radius → just return luminance + if radius < 2: + gray = gray.clamp(0, 1) + result = gray.expand(-1, 3, -1, -1).permute(0, 2, 3, 1) + if mask is not None: + # apply mask if provided (even in early exit) + if mask.dim() == 3: + mask = mask.unsqueeze(-1) + elif mask.dim() == 2: + mask = mask.unsqueeze(0).unsqueeze(-1) + mask = mask.to(device) + result = image * (1 - mask) + result * mask + return (result,) + + # Precompute normalisation factors for grid_sample + # (we'll use align_corners=False, so pixel centres at (i+0.5)/W-0.5 etc.) + # For grid_sample with align_corners=False, coordinates are in [-1, 1] + # where -1 maps to the left/top border of the first pixel, + # 1 maps to the right/bottom border of the last pixel. + # Pixel i's centre is at (i + 0.5) / W * 2 - 1 in x. + # But we need offsets in pixel units, so we convert pixel offsets to + # normalised grid offsets by dividing by (W-1) and multiplying by 2, + # then add to the base grid of pixel centres. + # A simpler way: generate a base grid of pixel centres, then add + # normalised random offsets. This is done below. + + # Create a base grid of pixel coordinates (x, y) for the whole image. + # Shape: [1, H, W, 2] (batch=1) + ys = torch.linspace(-1, 1, H, device=device) + xs = torch.linspace(-1, 1, W, device=device) + grid_y, grid_x = torch.meshgrid(ys, xs, indexing="ij") # both [H, W] + base_grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) # [1, H, W, 2] + # Expand to batch size + base_grid = base_grid.expand(B, H, W, 2).clone() # we'll add offsets + + # Normalise radius from pixels to [-1,1] scale + # 1 pixel in image space corresponds to (2 / (W-1)) in normalised coords (for x) + # but we need to handle both directions; use the smaller divisor to be safe + # or use isotropic scaling. We'll just use W-1 for x and H-1 for y. + rx = (radius / (W - 1)) * 2 if W > 1 else 0.0 + ry = (radius / (H - 1)) * 2 if H > 1 else 0.0 + + for _ in range(iterations): + # Generate random offsets for all pixels and samples at once + # Shape: [B, H, W, samples, 2] + offsets = torch.zeros(B, H, W, samples, 2, device=device) + offsets[..., 0].uniform_(-rx, rx) # x offset + offsets[..., 1].uniform_(-ry, ry) # y offset + + # Build the full grid for each sample by adding offsets to base_grid + # base_grid: [B, H, W, 1, 2] + offsets: [B, H, W, S, 2] + sample_grids = base_grid.unsqueeze(3) + offsets # [B, H, W, S, 2] + + # -------------------------- + # Sample the colour image at all offsets in one go + # We need to call grid_sample with a grid of shape [N, H_out, W_out, 2]. + # Here N = B * S, H_out = H, W_out = W. + # Reshape sample_grids to [B*S, H, W, 2] + grid_color = sample_grids.permute(0, 3, 1, 2, 4).reshape(B * samples, H, W, 2) + # Replicate the image S times along the batch dimension + img_rep = img.repeat(samples, 1, 1, 1) # [B*S, C, H, W] + sampled_col = F.grid_sample( + img_rep, grid_color, mode='bilinear', + padding_mode='border', align_corners=False + ) # [B*S, C, H, W] + # Reshape back to [B, S, C, H, W] then to [B, C, H, W, S] + sampled_col = sampled_col.reshape(B, samples, C, H, W).permute(0, 2, 3, 4, 1) # [B, C, H, W, S] + + # -------------------------- + # Sample the current grayscale at the same offsets + # grid_sample expects input with channels, so repeat gray to B*S batch + gray_rep = gray.repeat(samples, 1, 1, 1) # [B*S, 1, H, W] + sampled_gray = F.grid_sample( + gray_rep, grid_color, mode='bilinear', + padding_mode='border', align_corners=False + ) # [B*S, 1, H, W] + sampled_gray = sampled_gray.reshape(B, samples, 1, H, W).permute(0, 2, 3, 4, 1) # [B, 1, H, W, S] + + # -------------------------- + # Compute envelopes and contrast + min_col, _ = sampled_col.min(dim=-1) # [B, C, H, W] + max_col, _ = sampled_col.max(dim=-1) + color_range = max_col - min_col + color_contrast = torch.norm(color_range, dim=1, keepdim=True) # [B, 1, H, W] + + min_gray, _ = sampled_gray.min(dim=-1) # [B, 1, H, W] + max_gray, _ = sampled_gray.max(dim=-1) + gray_range = max_gray - min_gray + + gray_mean = sampled_gray.mean(dim=-1) # [B, 1, H, W] + + # STRESS scaling + eps = 1e-6 + scale = color_contrast / (gray_range + eps) + scale = scale.clamp(0.5, 2.0) + + # Update gray: local adjustment + gray_new = gray_mean + (gray - gray_mean) * scale + + if enhance_shadows: + shadow_boost = 1.0 + 0.3 * (1.0 - gray_new) + gray_new = gray_new * shadow_boost + + gray = gray_new.clamp(0, 1) + + # Convert to 3‑channel RGB + result = gray.expand(-1, 3, -1, -1).permute(0, 2, 3, 1) + + # Apply mask if provided + if mask is not None: + if mask.dim() == 3: + mask = mask.unsqueeze(-1) # [B, H, W] → [B, H, W, 1] + elif mask.dim() == 2: + mask = mask.unsqueeze(0).unsqueeze(-1) # [H, W] → [1, H, W, 1] + mask = mask.to(device) + result = image * (1 - mask) + result * mask + + return (result,) + + + diff --git a/py/color_desat_desturation.py b/py/color_desat_desturation.py new file mode 100644 index 0000000..e96cc2f --- /dev/null +++ b/py/color_desat_desturation.py @@ -0,0 +1,138 @@ + +import torch +from typing import Tuple, Optional +from .libs.process import compute_chunk_size + +def exec_operator_color_desat_desturation( + image: torch.Tensor, + mode: str, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Desaturate an RGB or RGBA image batch using one of five GIMP modes. + + Modes and formulas (based on GIMP source code): + - "Luminance": (0.22 * R) + (0.72 * G) + (0.06 * B) + - "Luma": (0.22 * R) + (0.72 * G) + (0.06 * B) (non‑linearized) + - "Lightness (HSL)": 0.5 * (max(R,G,B) + min(R,G,B)) + - "Avg (HSI Intensity)": (R + G + B) / 3 + - "Value (HSV)": max(R,G,B) + + Processing is performed in chunks along the batch dimension to limit + GPU memory usage. An optional mask can be supplied to blend between the + original and the desaturated image on a per‑pixel basis. If the mask is + entirely zero the input is returned unchanged; if it is entirely one the + blending step is skipped, providing an early‑exit optimisation. + + Alpha channels (RGBA input) are preserved and passed through unchanged. + + Args: + image: Input image tensor with shape [B, H, W, C], C = 3 or 4, + values in [0, 1]. + mode: Desaturation mode (case‑insensitive). + mask: Optional mask tensor with shape [B, H, W] or [B, H, W, 1], + values in [0, 1]. When provided, it controls the blend + strength per pixel. + + Returns: + A tuple containing the desaturated image tensor with the same + shape as the input. + + Raises: + ValueError: If the image does not have 3 or 4 channels, if the mask + shape is invalid, or if an unknown mode is given. + """ + # ---------- Input validation ---------- + if image.dim() != 4: + raise ValueError( + f"Expected image tensor of shape [B, H, W, C], got {image.shape}" + ) + C = image.shape[-1] + if C not in (3, 4): + raise ValueError( + f"Expected 3 or 4 colour channels, got {C}" + ) + + # Work in float32 throughout (original code always converts to float) + img = image.float() + B, H, W, _ = img.shape + + # ---------- Mask validation & normalisation ---------- + if mask is not None: + if mask.dim() == 3: + mask = mask.unsqueeze(-1) # [B,H,W] -> [B,H,W,1] + elif mask.dim() == 4 and mask.shape[-1] == 1: + pass # already [B,H,W,1] + else: + raise ValueError( + f"Mask must be [B, H, W] or [B, H, W, 1], got {mask.shape}" + ) + if mask.shape[0] != B or mask.shape[1] != H or mask.shape[2] != W: + raise ValueError( + f"Mask dimensions {mask.shape} do not match image {img.shape}" + ) + + # ---------- Early exits ---------- + if mask is not None: + if torch.all(mask == 0): + # No desaturation needed anywhere + return (img,) + # Flag: if every pixel is fully desaturated we can skip blending later + all_ones = torch.all(mask == 1) + else: + # No mask → fully desaturate (identical to all_ones == True) + all_ones = True + + # ---------- Chunked processing ---------- + mode_lower = mode.lower() + if mode_lower not in ("luminance", "luma", "lightness (hsl)", + "avg (hsi intensity)", "average", "value (hsv)"): + raise ValueError(f"Unknown desaturation mode: {mode}") + + chunk_size = compute_chunk_size(image, mask) # returns int + results = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + + # Slice batch + img_chunk = img[start:end] + mask_chunk = mask[start:end] if mask is not None else None + + # Extract RGB channels (drop alpha for calculation) + r = img_chunk[..., 0] + g = img_chunk[..., 1] + b = img_chunk[..., 2] + + # Compute greyscale according to the chosen mode + if mode_lower in ("luminance", "luma"): + gray = 0.22 * r + 0.72 * g + 0.06 * b + elif mode_lower == "lightness (hsl)": + max_rgb = torch.maximum(torch.maximum(r, g), b) + min_rgb = torch.minimum(torch.minimum(r, g), b) + gray = (max_rgb + min_rgb) * 0.5 + elif mode_lower in ("avg (hsi intensity)", "average"): + gray = (r + g + b) / 3.0 + elif mode_lower == "value (hsv)": + gray = torch.maximum(torch.maximum(r, g), b) + + # Expand greyscale to three channels + gray_rgb = gray.unsqueeze(-1).expand_as(img_chunk) + + # Blend with original using the mask (if needed) + if mask_chunk is not None and not all_ones: + # Per‑pixel blend + result_chunk = img_chunk * (1.0 - mask_chunk) + gray_rgb * mask_chunk + else: + # No blending required (no mask, or all_ones == True) + result_chunk = gray_rgb + + # Reattach the alpha channel if present in the input + if C == 4: + result_chunk = torch.cat([result_chunk, img_chunk[..., 3:4]], dim=-1) + + results.append(result_chunk) + + final_image = torch.cat(results, dim=0) + return (final_image,) + diff --git a/py/color_desat_sepia.py b/py/color_desat_sepia.py new file mode 100644 index 0000000..509ebc6 --- /dev/null +++ b/py/color_desat_sepia.py @@ -0,0 +1,142 @@ +import torch +from typing import Tuple +from .libs.process import compute_chunk_size + + +def exec_operator_color_desat_sepia( + image: torch.Tensor, + strength: float, + srgb: bool, + mask: torch.Tensor | None = None, +) -> Tuple[torch.Tensor]: + """ + Apply a GIMP‑style sepia desaturation to an image batch. + + The effect interpolates between the original image (strength=0) and full + sepia toning (strength=1). When a mask is supplied it controls the + per‑pixel blend strength: the sepia colour is applied only where the mask + is non‑zero, leaving the rest unchanged. The mask’s values must be in + [0,1]. + + **Supported formats** + * Input / output: RGB (3 channels) or RGBA (4 channels). + * Alpha (if present) is **never** modified by the sepia matrix or the + mask – it passes through unchanged. + + Args: + image: [B, H, W, C] float tensor in [0,1], C=3 or 4. + strength: 0.0 → no effect, 1.0 → full sepia. + srgb: Reserved for future colour‑space handling; currently unused. + mask: Optional float tensor in [0,1]. + Shape [B, H, W] or [B, H, W, 1]. Batch size must match + *image*, spatial dimensions must be H×W. + + Returns: + A tuple containing a single [B, H, W, C] float tensor with the + modified image. + """ + # ---------- 1. Input validation & early exit ---------- + if not isinstance(image, torch.Tensor): + raise TypeError("`image` must be a torch.Tensor") + if image.dim() != 4: + raise ValueError( + f"`image` must be 4‑dimensional [B,H,W,C], got {image.dim()}D" + ) + C = image.shape[-1] + if C not in (3, 4): + raise ValueError( + f"Only RGB (3) or RGBA (4) channels are supported, got {C}" + ) + + strength = max(0.0, min(1.0, float(strength))) + if strength == 0.0: + return (image,) # no change needed + + B, H, W, C = image.shape + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("`mask` must be a torch.Tensor") + if mask.dim() not in (3, 4): + raise ValueError( + f"`mask` must be 3‑ or 4‑dimensional, got {mask.dim()}D" + ) + if mask.shape[0] != B: + raise ValueError( + f"Batch size mismatch: image has {B}, mask has {mask.shape[0]}" + ) + if mask.shape[1:3] != (H, W): + raise ValueError( + f"Spatial size mismatch: image is {H}×{W}, " + f"mask is {mask.shape[1]}×{mask.shape[2]}" + ) + + # ---------- 2. Chunked processing ---------- + chunk_size = compute_chunk_size(image, mask) + if chunk_size <= 0: + chunk_size = B # fallback – whole batch at once + + chunks = [] + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + im_chunk = image[start:end] # [chunk, H, W, C] + mask_chunk = mask[start:end] if mask is not None else None + chunks.append(_sepia_chunk(im_chunk, strength, srgb, mask_chunk)) + + result = torch.cat(chunks, dim=0) + return (result,) + + +def _sepia_chunk( + im: torch.Tensor, + strength: float, + srgb: bool, # kept for signature compatibility + mask: torch.Tensor | None, +) -> torch.Tensor: + """Process a sub‑batch (already validated).""" + b, h, w, c = im.shape + has_alpha = c == 4 + + if has_alpha: + alpha = im[..., 3:4] # preserve original alpha + rgb = im[..., :3] + else: + alpha = None + rgb = im + + # ----- Build interpolated sepia matrix ----- + s = strength + full = torch.tensor( + [[0.393, 0.769, 0.189], + [0.349, 0.686, 0.168], + [0.272, 0.534, 0.131]], + dtype=rgb.dtype, device=rgb.device, + ) + eye = torch.eye(3, dtype=rgb.dtype, device=rgb.device) + mat = eye + (full - eye) * s # [3, 3] + + # ----- Apply matrix to RGB channels ----- + flat = rgb.reshape(-1, 3) # [b*h*w, 3] + sepia_flat = torch.mm(flat, mat.T) # [b*h*w, 3] + sepia_rgb = sepia_flat.reshape(b, h, w, 3) + + # ----- Blend with mask (only on RGB) ----- + if mask is not None: + # Bring mask to [b, h, w, 1] + if mask.dim() == 3: + mask_blend = mask.unsqueeze(-1) + else: + mask_blend = mask # already [b, h, w, 1] + mask_blend = mask_blend.to( + dtype=rgb.dtype, device=rgb.device + ).clamp_(0.0, 1.0) + # Broadcast across the 3 colour channels + rgb_result = rgb * (1 - mask_blend) + sepia_rgb * mask_blend + else: + rgb_result = sepia_rgb + + # ----- Reattach alpha if it was present ----- + if has_alpha: + return torch.cat([rgb_result, alpha], dim=-1) + return rgb_result + + diff --git a/py/color_exposure.py b/py/color_exposure.py new file mode 100644 index 0000000..cdedfe7 --- /dev/null +++ b/py/color_exposure.py @@ -0,0 +1,136 @@ +import torch + +from .libs.process import compute_chunk_size + +def exec_operator_color_exposure( + image: torch.Tensor, + black_level: float, + exposure: float, + mask: torch.Tensor = None, +) -> torch.Tensor: + """ + Apply a GEGL‑style exposure adjustment to an image batch, with optional + mask blending and full alpha‑channel preservation. + + Exposure is modeled as: + + white = 2^(-exposure) + gain = 1 / (white - black_level) (clamped for numerical stability) + result = (image - black_level) * gain + + When a mask is provided, the adjusted image is blended with the original + using the mask as the per‑pixel mixing weight: + + output = image * (1 - mask) + result * mask + + Args: + image (torch.Tensor): Input image tensor of shape ``(B, H, W, C)``, + with C ∈ {1, 3, 4}. + black_level (float): Black level offset (clamped to avoid division + by zero). + exposure (float): Exposure value (stops). 0 → no change, >0 brighter, + <0 darker. + mask (torch.Tensor, optional): Blend mask of shape ``(B, H, W)`` or + ``(1, H, W)``. If omitted, no blending. + + Returns: + torch.Tensor: Exposure‑adjusted image tensor with the same shape and + dtype as the input, but always float32 internally. + """ + # --- Input validation ---------------------------------------------------- + if not isinstance(image, torch.Tensor): + raise ValueError("`image` must be a torch.Tensor") + if image.dim() != 4: + raise ValueError( + f"`image` must have 4 dimensions (B,H,W,C), got {image.dim()}" + ) + if image.shape[-1] not in (1, 3, 4): + raise ValueError( + f"Unsupported number of channels: {image.shape[-1]}. " + "Expected 1, 3, or 4." + ) + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise ValueError("`mask` must be a torch.Tensor") + if mask.dim() not in (3, 4): + raise ValueError( + f"`mask` must have 3 or 4 dimensions, got {mask.dim()}" + ) + if mask.shape[0] not in (1, image.shape[0]): + raise ValueError( + f"Batch size of mask ({mask.shape[0]}) must be 1 or match " + f"image batch size ({image.shape[0]})" + ) + + # --- Early exit (identity transform) ------------------------------------- + if exposure == 0.0 and black_level == 0.0: + # Even with a mask, the blend would be identity → return original + return image + + # --- Device & precision setup -------------------------------------------- + image = image.to(dtype=torch.float32) + device = image.device + + # --- Alpha channel separation (if any) ----------------------------------- + has_alpha = image.shape[-1] == 4 + if has_alpha: + rgb = image[..., :3] # (B, H, W, 3) + alpha = image[..., 3:4] # (B, H, W, 1) + else: + rgb = image # operate on all channels + alpha = None + + B, H, W, C_rgb = rgb.shape + + # --- Mask preparation (batched or single) -------------------------------- + mask_tensor = None + if mask is not None: + mask_tensor = mask.to(device) + # Normalise mask shape to (B_mask, H, W, 1) + if mask_tensor.dim() == 3: + mask_tensor = mask_tensor.unsqueeze(-1) # (B, H, W, 1) + elif mask_tensor.dim() == 4 and mask_tensor.shape[1] == 1: + # ComfyUI sometimes uses (B, 1, H, W) + mask_tensor = mask_tensor.permute(0, 2, 3, 1) # (B, H, W, 1) + + # If a single mask is given, broadcast across the batch later + single_mask = mask_tensor.shape[0] == 1 + + # --- VRAM‑aware chunked processing --------------------------------------- + chunk_size = compute_chunk_size(image) # based on full image (includes alpha if present) + results = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + rgb_chunk = rgb[start:end] # (chunk_B, H, W, C_rgb) + + # --- Apply GEGL exposure transform (colour channels only) ------------ + white = torch.exp2(torch.tensor(-exposure, device=device)) + diff = torch.clamp(white - black_level, min=1e-6) + gain = 1.0 / diff + transformed = (rgb_chunk - black_level) * gain + + # --- Optional mask blending on colour channels ----------------------- + if mask_tensor is not None: + if single_mask: + # Use the single mask for every chunk + mask_chunk = mask_tensor.expand(end - start, -1, -1, -1) # (chunk_B, H, W, 1) + else: + mask_chunk = mask_tensor[start:end] # (chunk_B, H, W, 1) + + # Expand to match rgb_chunk shape (broadcast over channels) + mask_expanded = mask_chunk.expand_as(rgb_chunk) + result_chunk = rgb_chunk * (1.0 - mask_expanded) + transformed * mask_expanded + else: + result_chunk = transformed + + # --- Reattach alpha channel if present ------------------------------- + if has_alpha: + alpha_chunk = alpha[start:end] # (chunk_B, H, W, 1) + result_chunk = torch.cat([result_chunk, alpha_chunk], dim=-1) + + results.append(result_chunk) + + return torch.cat(results, dim=0) + + diff --git a/py/color_hue_chroma.py b/py/color_hue_chroma.py new file mode 100644 index 0000000..2874b3f --- /dev/null +++ b/py/color_hue_chroma.py @@ -0,0 +1,154 @@ +import torch +import torch.nn.functional as F + +from .libs.color import rgb_to_lch_torch, lch_to_rgb_torch + + +def exec_operator_color_hue_chroma( + image: torch.Tensor, + hue: float, + chroma: float, + lightness: float, + mask: torch.Tensor = None, +) -> torch.Tensor: + """ + Apply hue, chroma and lightness adjustments in LCH color space. + + The image is converted from RGB to LCH, the requested adjustments are + applied, then converted back to RGB. Hue and chroma adjustments are + skipped for achromatic pixels while lightness is always adjusted. + + If a mask is supplied, the adjusted image is blended with the original + image according to the mask values. + + Supports RGB and RGBA images and batched processing. + """ + + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + + if image.ndim != 4: + raise ValueError( + "image must have shape (B,H,W,C)" + ) + + if image.shape[-1] not in (3, 4): + raise ValueError( + "image must have 3 (RGB) or 4 (RGBA) channels" + ) + + if not image.is_floating_point(): + raise TypeError( + "image must be floating point" + ) + + rgb = image[..., :3] + alpha = image[..., 3:] if image.shape[-1] == 4 else None + + device = rgb.device + dtype = rgb.dtype + + if mask is not None: + + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor") + + mask = mask.to(device=device, dtype=dtype) + + if mask.ndim == 2: + mask = mask.unsqueeze(0).unsqueeze(-1) + + elif mask.ndim == 3: + mask = mask.unsqueeze(-1) + + elif mask.ndim != 4: + raise ValueError( + "mask must have shape (H,W), (B,H,W) or (B,H,W,1)" + ) + + if mask.shape[-1] != 1: + mask = mask[..., :1] + + # + # Pair image batch with corresponding mask. + # + if mask.shape[0] == 1 and rgb.shape[0] > 1: + mask = mask.expand(rgb.shape[0], -1, -1, -1) + + elif mask.shape[0] != rgb.shape[0]: + raise ValueError( + "mask batch size must be 1 or match image batch size" + ) + + if ( + mask.shape[1] != rgb.shape[1] + or mask.shape[2] != rgb.shape[2] + ): + mask = F.interpolate( + mask.permute(0, 3, 1, 2), + size=rgb.shape[1:3], + mode="bilinear", + align_corners=False, + ).permute(0, 2, 3, 1) + + EPSILON = 1e-4 + + # + # RGB -> LCH + # + lch = rgb_to_lch_torch(rgb) + + L = lch[..., 0] + C = lch[..., 1] + H = lch[..., 2] + + achromatic = torch.abs(C) < EPSILON + + L_adj = L + lightness + + # + # Only clone if modification is required. + # + if chroma != 0.0: + C_adj = C.clone() + C_adj.masked_scatter_( + ~achromatic, + C.masked_select(~achromatic) + chroma, + ) + else: + C_adj = C + + if hue != 0.0: + H_adj = H.clone() + H_adj.masked_scatter_( + ~achromatic, + H.masked_select(~achromatic) + hue, + ) + else: + H_adj = H + + C_adj = torch.clamp(C_adj, 0.0, 300.0) + + H_adj = H_adj.remainder(360.0) + H_adj = torch.where( + H_adj >= 180.0, + H_adj - 360.0, + H_adj, + ) + + result_rgb = lch_to_rgb_torch( + torch.stack( + (L_adj, C_adj, H_adj), + dim=-1, + ) + ) + + if mask is not None: + result_rgb = rgb * (1.0 - mask) + result_rgb * mask + + if alpha is not None: + return torch.cat((result_rgb, alpha), dim=-1) + + return result_rgb + + diff --git a/py/color_hue_saturation.py b/py/color_hue_saturation.py new file mode 100644 index 0000000..ffd30b3 --- /dev/null +++ b/py/color_hue_saturation.py @@ -0,0 +1,99 @@ + +import torch +from typing import Optional + +from .libs.color import rgb_to_hsl, hsl_to_rgb + +# ========================================================= +# MAIN OPERATOR (GIMP-style) +# ========================================================= +@torch.no_grad() +def exec_operator_color_hue_saturation( + image: torch.Tensor, + mask: Optional[torch.Tensor], + overlap: float, + data: list, +) -> torch.Tensor: + """ + Core hue/saturation/lightness adjustment with per‑colour ranges. + + Parameters + ---------- + image : torch.Tensor, shape (B, H, W, C) – C=3 (RGB) or C=4 (RGBA) + mask : optional torch.Tensor (B, H, W) or (B, 1, H, W) – luminance mask + overlap : float (0‑100) – hue blending spread + data : list of 7 lists of 3 floats – [master, red, yellow, green, cyan, blue, magenta] + + Returns + ------- + torch.Tensor of same shape as `image`, with alpha (if present) unchanged. + """ + + # ---------- Device and range ---------- + device = image.device + image = image.clamp(0.0, 1.0) + + # ---------- Handle RGBA ---------- + has_alpha = image.shape[-1] == 4 + if has_alpha: + rgb = image[..., :3] # (B, H, W, 3) + alpha = image[..., 3:] # (B, H, W, 1) + else: + rgb = image + + # ---------- Original algorithm on RGB channels ---------- + h, s, l = rgb_to_hsl(rgb) # all tensors (B, H, W) + + master_h, master_l, master_s = data[0] + + channels_h = torch.tensor([d[0] for d in data[1:]], device=device) # (6,) + channels_l = torch.tensor([d[1] for d in data[1:]], device=device) + channels_s = torch.tensor([d[2] for d in data[1:]], device=device) + + centers = torch.tensor([0, 60, 120, 180, 240, 300], device=device) # (6,) + + # ------- per‑hue weighting (computed once for all shifts) ------- + h_exp = h.unsqueeze(-1) # (B, H, W, 1) + diff = (h_exp - centers + 180.0) % 360.0 - 180.0 + absdiff = torch.abs(diff) + + spread = max(1.0, 180.0 * overlap / 100.0) + weight = torch.clamp(1.0 - absdiff / spread, 0.0, 1.0) # (B, H, W, 6) + + # ------- accumulate shifts ------- + h_shift = master_h + torch.sum(weight * channels_h, dim=-1) # (B, H, W) + l_shift = (master_l / 100.0) + torch.sum(weight * (channels_l / 100.0), dim=-1) + s_shift = (master_s / 100.0) + torch.sum(weight * (channels_s / 100.0), dim=-1) + + # Free the large weight tensor as soon as it is no longer needed + del weight + + # ------- apply shifts ------- + h = (h + h_shift) % 360.0 + l = torch.clamp(l + l_shift, 0.0, 1.0) + s = torch.clamp(s + s_shift, 0.0, 1.0) + + out_rgb = hsl_to_rgb(h, s, l) + + # ------- free HSL tensors ------- + del h, s, l, h_shift, l_shift, s_shift + + # ---------- Mask blend (on RGB part only) ---------- + if mask is not None: + # Normalise mask shape to (B, H, W, 1) for broadcasting + if mask.dim() == 3: + mask = mask.unsqueeze(-1) # (B, H, W, 1) + elif mask.dim() == 4 and mask.shape[1] == 1: + mask = mask.permute(0, 2, 3, 1) # (B, 1, H, W) → (B, H, W, 1) + + mask = mask.to(device).expand_as(out_rgb) # broadcast to 3 channels + out_rgb = rgb * (1.0 - mask) + out_rgb * mask + + # ---------- Reattach alpha if needed ---------- + if has_alpha: + out = torch.cat([out_rgb, alpha], dim=-1) + else: + out = out_rgb + + return out + diff --git a/py/color_invert.py b/py/color_invert.py new file mode 100644 index 0000000..fcd91d2 --- /dev/null +++ b/py/color_invert.py @@ -0,0 +1,182 @@ +import torch +from typing import Tuple, Optional + +from .libs.color import srgb_to_linear_vec, linear_to_srgb_vec, rgb_to_hsv, hsv_to_rgb +from .libs.process import compute_chunk_size + +# -------------------------------------------------------------------- +# Internal helpers for channel handling and chunked processing +# -------------------------------------------------------------------- + +def _ensure_channels(t: torch.Tensor, target_channels: int) -> torch.Tensor: + """Ensure tensor has exactly target_channels by adding/removing dummy dimensions. + Used only for internal shape adjustments; for actual channel handling see below. + """ + if t.ndim == 3 and t.shape[-1] == 1: + # single channel -> expand to target? Not used here. + pass + return t + + +def _split_rgb_alpha(image: torch.Tensor): + """Split image into RGB and alpha (if present).""" + if image.shape[-1] == 3: + return image, None + elif image.shape[-1] == 4: + return image[..., :3], image[..., 3:] + else: + raise ValueError(f"Unsupported number of channels: {image.shape[-1]}") + + +def _combine_rgb_alpha(rgb: torch.Tensor, alpha: Optional[torch.Tensor]) -> torch.Tensor: + """Recombine RGB and alpha, or return RGB if no alpha.""" + if alpha is None: + return rgb + return torch.cat([rgb, alpha], dim=-1) + + +def _process_chunk( + image_chunk: torch.Tensor, + inv_type: str, + mask_chunk: Optional[torch.Tensor] +) -> torch.Tensor: + """ + Process a single chunk (batch of images) with the chosen inversion type. + Handles RGB/RGBA, mask blending, and preserves alpha. + """ + # Separate RGB and alpha + rgb, alpha = _split_rgb_alpha(image_chunk) + + # Apply inversion to RGB channels only + if inv_type == "Value": + inverted_rgb = _exec_operator_color_value_invert_rgb(rgb) + elif inv_type == "Linear": + inverted_rgb = _exec_operator_color_linear_invert_rgb(rgb) + else: # "RGB" + inverted_rgb = 1.0 - rgb + + # Apply mask blending if mask is provided + if mask_chunk is not None: + # Ensure mask has proper shape: (B, H, W) or (B, H, W, 1) + if mask_chunk.ndim == 3: + mask_chunk = mask_chunk.unsqueeze(-1) # (B, H, W, 1) + elif mask_chunk.ndim == 4 and mask_chunk.shape[-1] == 1: + pass + else: + raise ValueError(f"Mask must have shape (B,H,W) or (B,H,W,1), got {mask_chunk.shape}") + + # Broadcast mask if needed (batch size 1 vs multiple) + if mask_chunk.shape[0] == 1 and image_chunk.shape[0] > 1: + mask_chunk = mask_chunk.expand(image_chunk.shape[0], -1, -1, -1) + elif mask_chunk.shape[0] != image_chunk.shape[0]: + raise ValueError( + f"Mask batch size {mask_chunk.shape[0]} does not match image batch size {image_chunk.shape[0]}" + ) + + # Blend: output = image * (1 - mask) + inverted * mask + rgb_blended = rgb * (1.0 - mask_chunk) + inverted_rgb * mask_chunk + # Clamp to [0,1] for safety (though operations should keep within bounds) + rgb_blended = torch.clamp(rgb_blended, 0.0, 1.0) + else: + rgb_blended = inverted_rgb + + # Recombine with original alpha + return _combine_rgb_alpha(rgb_blended, alpha) + + +# -------------------------------------------------------------------- +# Core inversion functions (unchanged logic, but work on RGB tensors) +# -------------------------------------------------------------------- + +def _exec_operator_color_linear_invert_rgb(rgb: torch.Tensor) -> torch.Tensor: + """Linear inversion in RGB space (sRGB → linear → invert → sRGB).""" + linear = srgb_to_linear_vec(rgb) + inv_linear = 1.0 - linear + return linear_to_srgb_vec(inv_linear) + + +def _exec_operator_color_value_invert_rgb(rgb: torch.Tensor) -> torch.Tensor: + """Value (luminosity) inversion in HSV space.""" + h, s, v = rgb_to_hsv(rgb) + v_inv = 1.0 - v + inverted = hsv_to_rgb(h, s, v_inv) + return torch.clamp(inverted, 0.0, 1.0) + + +# -------------------------------------------------------------------- +# Public entry point +# -------------------------------------------------------------------- + +def exec_operator_color_invert( + image: torch.Tensor, + inv_type: str, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Invert colours of an image batch, mimicking GIMP's Colors → Invert. + + Supports RGB and RGBA input/output; alpha channel is always passed through unchanged. + If a mask is given, the inverted image is blended with the original using the mask + as an alpha factor: output = image * (1 - mask) + inverted * mask. The mask must be + a single-channel tensor with shape (B, H, W) or (B, H, W, 1). Its batch size must + either match the image batch or be 1 (broadcast). + + Args: + image: (B, H, W, C) tensor, C = 3 or 4, values in [0, 1]. + inv_type: One of "RGB", "Linear", or "Value". + mask: Optional (B, H, W) or (B, H, W, 1) tensor, values in [0, 1]. + + Returns: + Tuple containing the inverted image tensor with same shape as input. + + Raises: + ValueError: for invalid input shapes, types, or values. + """ + # --- Input validation --- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.ndim != 4: + raise ValueError(f"image must be 4D (B,H,W,C), got {image.ndim}D") + if image.shape[-1] not in (3, 4): + raise ValueError(f"image must have 3 or 4 channels, got {image.shape[-1]}") + if image.min() < 0 or image.max() > 1: + raise ValueError("image values must be in [0, 1]") + + if inv_type not in ("RGB", "Linear", "Value"): + raise ValueError(f"inv_type must be 'RGB', 'Linear', or 'Value', got {inv_type}") + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor or None") + if mask.ndim not in (3, 4) or (mask.ndim == 4 and mask.shape[-1] != 1): + raise ValueError(f"mask must be (B,H,W) or (B,H,W,1), got {mask.shape}") + if mask.min() < 0 or mask.max() > 1: + raise ValueError("mask values must be in [0, 1]") + + # --- Early exit if mask is all zeros --- + if mask is not None and torch.all(mask == 0): + # No inversion applied – return original image (alpha preserved) + return (image,) + + # --- Chunked processing --- + batch_size = image.shape[0] + if batch_size == 0: + return (image,) + + # Compute optimal chunk size using the provided helper + chunk_size = compute_chunk_size(image, mask) + if chunk_size <= 0: + chunk_size = batch_size # fallback + + # Process in chunks + chunks = [] + for start in range(0, batch_size, chunk_size): + end = min(start + chunk_size, batch_size) + img_chunk = image[start:end] + mask_chunk = mask[start:end] if mask is not None else None + processed = _process_chunk(img_chunk, inv_type, mask_chunk) + chunks.append(processed) + + # Concatenate chunks along batch dimension + result = torch.cat(chunks, dim=0) + return (result,) diff --git a/py/color_levels.py b/py/color_levels.py new file mode 100644 index 0000000..31642d7 --- /dev/null +++ b/py/color_levels.py @@ -0,0 +1,356 @@ +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) + + diff --git a/py/color_map_alien.py b/py/color_map_alien.py new file mode 100644 index 0000000..d20f0a9 --- /dev/null +++ b/py/color_map_alien.py @@ -0,0 +1,246 @@ +import torch +import math +from typing import Tuple, Optional +from .libs.process import compute_chunk_size + +# =================================================================================== +def rgb_to_hsl_torch(image: torch.Tensor) -> torch.Tensor: + """ + Convert an RGB image to HSL. + Input: (B, H, W, 3) in [0, 1] + Output: (B, H, W, 3) with H, S, L in [0, 1] + """ + r, g, b = image[..., 0], image[..., 1], image[..., 2] + max_val, _ = torch.max(image, dim=-1) + min_val, _ = torch.min(image, dim=-1) + l = (max_val + min_val) / 2.0 + + delta = max_val - min_val + # Saturation: 0 if delta == 0 + s = torch.where(delta == 0, torch.zeros_like(delta), + delta / (1.0 - torch.abs(2.0 * l - 1.0))) + + h = torch.zeros_like(l) + mask = delta != 0 + r_eq_max = (max_val == r) & mask + g_eq_max = (max_val == g) & mask + # Otherwise b_eq_max (implicit) + h = torch.where(r_eq_max, ((g - b) / delta) % 6.0, h) + h = torch.where(g_eq_max, ((b - r) / delta) + 2.0, h) + h = torch.where(~r_eq_max & ~g_eq_max & mask, ((r - g) / delta) + 4.0, h) + h = h / 6.0 + h = torch.where(h < 0, h + 1.0, h) + h = torch.where(h >= 1.0, h - 1.0, h) + + return torch.stack([h, s, l], dim=-1) + + +def hsl_to_rgb_torch(image: torch.Tensor) -> torch.Tensor: + """ + Convert an HSL image to RGB. + Input: (B, H, W, 3) with H, S, L in [0, 1] + Output: (B, H, W, 3) in [0, 1] + """ + h, s, l = image[..., 0], image[..., 1], image[..., 2] + + c = (1.0 - torch.abs(2.0 * l - 1.0)) * s + x = c * (1.0 - torch.abs((h * 6.0) % 2.0 - 1.0)) + m = l - c / 2.0 + + sector = (h * 6.0).to(torch.int64) % 6 + zeros = torch.zeros_like(h) + + r = torch.where(sector == 0, c, + torch.where(sector == 1, x, + torch.where(sector == 2, zeros, + torch.where(sector == 3, zeros, + torch.where(sector == 4, x, c))))) + g = torch.where(sector == 0, x, + torch.where(sector == 1, c, + torch.where(sector == 2, c, + torch.where(sector == 3, x, + torch.where(sector == 4, zeros, zeros))))) + b = torch.where(sector == 0, zeros, + torch.where(sector == 1, zeros, + torch.where(sector == 2, x, + torch.where(sector == 3, c, + torch.where(sector == 4, c, x))))) + + return torch.stack([r + m, g + m, b + m], dim=-1) + + +def exec_operator_color_map_alien( + image: torch.Tensor, + model: str, + freq1: float, freq2: float, freq3: float, + shift1: float, shift2: float, shift3: float, + keep1: bool, keep2: bool, keep3: bool, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Apply the Alien Map colour transformation to a batch of images. + + The transformation mimics the GIMP “Alien Map” filter. For each colour + channel that is not “kept”, the operation + + out = 0.5 * (1 + sin( (2*in - 1) * freq + phase )) + + is applied, where *freq* = `frequency` × π and *phase* = π × `shift` / 180°. + The colour space can be either RGB or HSL; when HSL is chosen the image is + temporarily converted, transformed, and then converted back to RGB. + + **Batch processing** – if `mask` is a batch, the i‑th image is processed + with the i‑th mask. A single mask is broadcast across the whole batch. + + **RGB / RGBA** – 3‑channel (RGB) and 4‑channel (RGBA) images are accepted. + For RGBA the alpha channel is **preserved unchanged**. Mask blending + affects only the RGB channels. + + **Early exit** – when all three `keep` flags are `True` the transformation + is an identity; the function returns the original image immediately. + + **GPU acceleration & VRAM management** – all work is done on the same + device as the input tensor. To avoid out‑of‑memory errors the batch is + processed in chunks whose size is determined by + :func:`compute_chunk_size`. Spatial (tiled) splitting is **not** used. + + Parameters + ---------- + image : torch.Tensor + Input image(s), shape ``(B, H, W, 3)`` or ``(B, H, W, 4)``, values in [0, 1]. + model : str + ``"RGB"`` or ``"HSL"``. + freq1, freq2, freq3 : float + Frequency multipliers (0 – 20). + shift1, shift2, shift3 : float + Phase shifts in degrees (0 – 360). + keep1, keep2, keep3 : bool + If ``True`` the corresponding channel is left unchanged. + mask : torch.Tensor, optional + Per‑pixel strength of the effect, shape ``(B, H, W)`` or ``(B, H, W, 1)``, + values in [0, 1]. The output is blended linearly: + ``mask * transformed + (1 - mask) * original``. + + Returns + ------- + Tuple[torch.Tensor] + A single‑element tuple containing the transformed image tensor + ``(B, H, W, 3)`` or ``(B, H, W, 4)``, same shape and device as the input. + + Raises + ------ + ValueError + If the input image has an unsupported shape, or if the mask dimensions + are incompatible with the image. + """ + # ----------------------- input validation ---------------------------------- + if not isinstance(image, torch.Tensor): + raise ValueError("image must be a torch.Tensor") + if image.dim() != 4 or image.shape[-1] not in (3, 4): + raise ValueError( + f"image must be 4D (B, H, W, C) with C=3 or 4, got {image.shape}" + ) + + batch_size, height, width, channels = image.shape + has_alpha = (channels == 4) + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise ValueError("mask must be a torch.Tensor or None") + if mask.dim() == 3: + mask = mask.unsqueeze(-1) # (B, H, W, 1) + elif mask.dim() != 4 or mask.shape[-1] != 1: + raise ValueError( + f"mask must be 3D (B,H,W) or 4D (B,H,W,1), got {mask.shape}" + ) + if mask.shape[1:3] != (height, width): + raise ValueError( + f"mask spatial size {mask.shape[1:3]} must match image {(height, width)}" + ) + mask_batch = mask.shape[0] + if mask_batch != 1 and mask_batch != batch_size: + raise ValueError( + f"mask batch size {mask_batch} must be 1 or equal to image batch size {batch_size}" + ) + + # ----------------------- early exit ---------------------------------------- + if all([keep1, keep2, keep3]): + # The transformation is an identity; no need to compute anything. + return (image,) + + # ----------------------- prepare data -------------------------------------- + img = image.float() + img = torch.clamp(img, 0.0, 1.0) + + # Separate RGB and alpha (if present) + if has_alpha: + rgb = img[..., :3].clone() + alpha = img[..., 3:4].clone() + original_rgb = rgb.clone() + else: + rgb = img + alpha = None + original_rgb = img.clone() # used for mask blending later + + device = rgb.device + + # Colour‑space conversion (RGB → HSL) if required + if model == "HSL": + rgb = rgb_to_hsl_torch(rgb) + + # Transformation constants (on correct device) + freq = torch.tensor( + [freq1 * math.pi, freq2 * math.pi, freq3 * math.pi], + dtype=torch.float32, device=device + ) + phase = torch.tensor( + [math.pi * shift1 / 180.0, + math.pi * shift2 / 180.0, + math.pi * shift3 / 180.0], + dtype=torch.float32, device=device + ) + keep = torch.tensor([keep1, keep2, keep3], dtype=torch.bool, device=device) + + # ----------------------- chunked processing -------------------------------- + chunk_size = compute_chunk_size(image, mask) # returns int + if chunk_size <= 0 or chunk_size > batch_size: + chunk_size = batch_size + + out_rgb = torch.empty_like(rgb) + + for start in range(0, batch_size, chunk_size): + end = min(start + chunk_size, batch_size) + + rgb_chunk = rgb[start:end] + out_chunk = rgb_chunk.clone() + + # Apply the Alien Map formula per channel + for i in range(3): + if not keep[i]: + val = 2.0 * rgb_chunk[..., i] - 1.0 + out_chunk[..., i] = 0.5 * (1.0 + torch.sin(val * freq[i] + phase[i])) + + out_chunk = torch.clamp(out_chunk, 0.0, 1.0) + + # Back to RGB if we were in HSL + if model == "HSL": + out_chunk = hsl_to_rgb_torch(out_chunk) + + # Blend with mask (RGB channels only) + if mask is not None: + mask_chunk = mask[start:end] # will broadcast correctly if mask batch=1 + # original_rgb_chunk has the original clamped values + original_chunk = original_rgb[start:end] + out_chunk = mask_chunk * out_chunk + (1.0 - mask_chunk) * original_chunk + + out_rgb[start:end] = out_chunk + + # ----------------------- recombine alpha (if present) ----------------------- + if has_alpha: + out = torch.cat([out_rgb, alpha], dim=-1) + else: + out = out_rgb + + return (out,) + + diff --git a/py/color_map_color_exchange.py b/py/color_map_color_exchange.py new file mode 100644 index 0000000..349700e --- /dev/null +++ b/py/color_map_color_exchange.py @@ -0,0 +1,166 @@ +import torch +from typing import Optional, Tuple +from .libs.process import compute_chunk_size + +# =================================================================================== + +def exec_operator_color_map_color_exchange( + image: torch.Tensor, + from_color: int, + to_color: int, + red_threshold: float, + green_threshold: float, + blue_threshold: float, + mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + """ + Color Exchange (RGB/RGBA) with soft mask, batch chunking and alpha preservation. + + Pixels whose RGB components *all* lie within the given thresholds around a + source colour are recoloured to a target colour, preserving each channel's + original offset from the source:: + + out_channel = target_channel + (pixel_channel - source_channel) + + The result is clamped to [0,1]. When a soft mask is supplied, the + recolouring is blended per‑pixel using the mask value as the interpolation + strength. + + Parameters + ---------- + image : torch.Tensor + Batched image tensor of shape ``(B, H, W, 3)`` or ``(B, H, W, 4)``, + values in [0, 1], dtype float16/32/64. + from_color : int + Packed source colour (0xRRGGBB). + to_color : int + Packed target colour (0xRRGGBB). + red_threshold, green_threshold, blue_threshold : float + Per‑channel tolerance in [0, 1]. + mask : torch.Tensor, optional + Soft mask of shape ``(B, H, W)`` or ``(B, H, W, 1)``, or a single + mask broadcastable to the batch size. Values are clamped to [0, 1] + and used as per‑pixel blend strength. + + Returns + ------- + Tuple[torch.Tensor] + A single‑element tuple containing the recoloured image with the same + shape and number of channels as the input. + """ + # ---- Validate input -------------------------------------------------------- + if image.ndim != 4: + raise ValueError(f"image must be 4D (B,H,W,C), got shape {image.shape}") + C = image.shape[-1] + if C not in (3, 4): + raise ValueError(f"image must have 3 or 4 channels, got {C}") + if not torch.is_floating_point(image): + raise TypeError("image must be a floating point tensor") + + if not (0 <= from_color <= 0xFFFFFF): + raise ValueError(f"from_color must be in [0, 0xFFFFFF], got {from_color:#x}") + if not (0 <= to_color <= 0xFFFFFF): + raise ValueError(f"to_color must be in [0, 0xFFFFFF], got {to_color:#x}") + for th, name in zip( + (red_threshold, green_threshold, blue_threshold), + ("red_threshold", "green_threshold", "blue_threshold"), + ): + if not (0.0 <= th <= 1.0): + raise ValueError(f"{name} must be in [0, 1], got {th}") + + device = image.device + dtype = image.dtype + B = image.shape[0] + + # ---- Mask validation & early exit ------------------------------------------ + if mask is not None: + if mask.ndim == 3: + mask = mask.unsqueeze(-1) # (B,H,W) -> (B,H,W,1) + elif mask.ndim == 4: + if mask.shape[-1] != 1: + raise ValueError("mask with 4 dims must have last dim == 1") + else: + raise ValueError(f"mask must be 3D or 4D, got shape {mask.shape}") + + # Move mask to the correct device/dtype, clamp + mask = mask.to(device=device, dtype=dtype).clamp(0.0, 1.0) + + # Check batch compatibility + if mask.shape[0] != 1 and mask.shape[0] != B: + raise ValueError( + f"mask batch size {mask.shape[0]} must be 1 or match image batch size {B}" + ) + + # Early exit if mask is entirely zero + if mask.count_nonzero() == 0: + return (image,) + + # ---- Pre‑compute colour thresholds (RGB only) -------------------------------- + from_rgb = image.new_tensor([ + ((from_color >> 16) & 0xFF) / 255.0, + ((from_color >> 8) & 0xFF) / 255.0, + (from_color & 0xFF) / 255.0, + ]) + to_rgb = image.new_tensor([ + ((to_color >> 16) & 0xFF) / 255.0, + ((to_color >> 8) & 0xFF) / 255.0, + (to_color & 0xFF) / 255.0, + ]) + thresholds = image.new_tensor([red_threshold, green_threshold, blue_threshold]) + + # Broadcast dimensions: (1,1,1,3) + from_rgb = from_rgb.view(1, 1, 1, 3) + to_rgb = to_rgb.view(1, 1, 1, 3) + thresholds = thresholds.view(1, 1, 1, 3) + + # ---- RGBA support: separate RGB and alpha ---------------------------------- + if C == 4: + image_rgb = image[..., :3] + image_alpha = image[..., 3:4] # keep as (B,H,W,1) + else: + image_rgb = image + + # ---- Determine chunk size and process in sub‑batches ----------------------- + chunk_size = compute_chunk_size(image, mask) + result_chunks = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + img_chunk = image_rgb[start:end] # (chunk, H, W, 3) + + # Extract matching mask chunk (handle broadcast) + if mask is not None: + if mask.shape[0] == 1: + mask_chunk = mask # broadcast across chunk + else: + mask_chunk = mask[start:end] + else: + mask_chunk = None + + # ----- Original colour exchange logic (applied to RGB) ------------------ + diff = (img_chunk - from_rgb).abs() + match = (diff <= thresholds).all(dim=-1, keepdim=True) # (chunk,H,W,1) + + exchanged = to_rgb + (img_chunk - from_rgb) + exchanged = torch.clamp(exchanged, 0.0, 1.0) + result_rgb = torch.where(match, exchanged, img_chunk) + + # ----- Soft mask blending (applies only to RGB) ------------------------- + if mask_chunk is not None: + # mask_chunk is (1,H,W,1) or (chunk,H,W,1) – already prepared + result_rgb = img_chunk * (1.0 - mask_chunk) + result_rgb * mask_chunk + + # ----- Reassemble RGBA if needed ---------------------------------------- + if C == 4: + alpha_chunk = image_alpha[start:end] + result_chunk = torch.cat([result_rgb, alpha_chunk], dim=-1) + else: + result_chunk = result_rgb + + result_chunks.append(result_chunk) + + # Concatenate all chunks back to a single batch tensor + result = torch.cat(result_chunks, dim=0) + + return (result,) + diff --git a/py/color_map_color_rotate.py b/py/color_map_color_rotate.py new file mode 100644 index 0000000..1698913 --- /dev/null +++ b/py/color_map_color_rotate.py @@ -0,0 +1,252 @@ +import torch +from typing import Tuple, Optional + +from .libs.color import hsv_to_rgb, rgb_to_hsv +from .libs.process import compute_chunk_size + + +# ----------------------------------------------------------------------------- +# Angular helpers (unchanged) +# ----------------------------------------------------------------------------- + +def _cw_distance(start: float, angle: torch.Tensor) -> torch.Tensor: + """Clockwise angular distance [0,360).""" + return (angle - start) % 360.0 + + +def _ccw_distance(start: float, angle: torch.Tensor) -> torch.Tensor: + """Counter-clockwise angular distance [0,360).""" + return (start - angle) % 360.0 + + +# ----------------------------------------------------------------------------- +# Core color rotate logic (identical per-pixel behaviour, no mask/alpha) +# ----------------------------------------------------------------------------- + +def _color_rotate_core( + rgb_chunk: torch.Tensor, # (B, H, W, 3) in [0,1] + mask_chunk: Optional[torch.Tensor], # (B, H, W, 1) in [0,1] or None + Source_from: float, + Source_to: float, + Source_clockwise: bool, + Destination_from: float, + Destination_to: float, + Destination_clockwise: bool, + Gray_threshold: float, + Mode: str, + Hue: float, + Saturation: float, +) -> torch.Tensor: + """ + Applies GIMP-style color rotate to a chunk of RGB images, + optionally blending with a soft mask. + Returns an RGB tensor of the same shape as `rgb_chunk`. + """ + # RGB → HSV + hue, sat, val = rgb_to_hsv(rgb_chunk) + hue = hue * 360.0 + + # Gray handling + is_gray = sat <= Gray_threshold + replacement_hue = torch.full_like(hue, Hue % 360.0) + replacement_sat = torch.full_like(sat, float(max(0.0, min(1.0, Saturation)))) + + if Mode == "Change to this": + rotate_hue = hue + rotate_sat = sat + hue = torch.where(is_gray, replacement_hue, hue) + sat = torch.where(is_gray, replacement_sat, sat) + rotate_mask = ~is_gray + else: # "Treat as this" + rotate_hue = torch.where(is_gray, replacement_hue, hue) + rotate_sat = torch.where(is_gray, replacement_sat, sat) + hue = rotate_hue + sat = rotate_sat + rotate_mask = torch.ones_like(is_gray) + + # Source interval + if Source_clockwise: + span = (Source_to - Source_from) % 360.0 + dist = _cw_distance(Source_from, rotate_hue) + else: + span = (Source_from - Source_to) % 360.0 + dist = _ccw_distance(Source_from, rotate_hue) + + if span < 1e-6: + span = 360.0 + + inside = dist <= span + rotate_mask &= inside + + # Normalised position + t = (dist / span).clamp(0.0, 1.0) + + # Destination interval + if Destination_clockwise: + dst_span = (Destination_to - Destination_from) % 360.0 + new_hue = (Destination_from + t * dst_span) % 360.0 + else: + dst_span = (Destination_from - Destination_to) % 360.0 + new_hue = (Destination_from - t * dst_span) % 360.0 + + hue = torch.where(rotate_mask, new_hue, hue) + + # HSV → RGB + new_rgb = hsv_to_rgb(hue / 360.0, sat.clamp(0.0, 1.0), val) + + # Soft mask blending + if mask_chunk is not None: + rgb_out = rgb_chunk * (1.0 - mask_chunk) + new_rgb * mask_chunk + else: + rgb_out = new_rgb + + return rgb_out + + +# ----------------------------------------------------------------------------- +# Main entry point – batch processing, alpha preservation, chunking, early exit +# ----------------------------------------------------------------------------- + +def exec_operator_color_map_color_rotate( + image: torch.Tensor, + Source_from: float, + Source_to: float, + Source_clockwise: bool, + Destination_from: float, + Destination_to: float, + Destination_clockwise: bool, + Gray_threshold: float, + Mode: str, + Hue: float, + Saturation: float, + mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor]: + """ + Apply a GIMP‑style color‑rotate effect to a batch of RGB or RGBA images. + + The operation remaps a source hue interval to a destination hue interval. + An optional per‑pixel soft mask controls the blending strength. + Alpha channels (if present) are preserved unchanged. + Processing is performed in chunks along the batch dimension to manage + GPU memory usage. If the optional mask is entirely zero the function + returns the input immediately without any computation. + + Parameters + ---------- + image : torch.Tensor + Input batch of shape (B, H, W, 3) or (B, H, W, 4), values in [0, 1]. + Source_from, Source_to : float + Bounds of the source hue interval in degrees [0, 360]. + Source_clockwise : bool + Direction of the source interval (clockwise if True). + Destination_from, Destination_to : float + Bounds of the destination hue interval in degrees [0, 360]. + Destination_clockwise : bool + Direction of the destination interval. + Gray_threshold : float + Saturation below which a pixel is considered gray [0, 1]. + Mode : str + "Change to this" – replace gray pixels with the given Hue/Saturation; + "Treat as this" – treat gray pixels as if they had the given Hue/Saturation. + Hue : float + Replacement/treatment hue for gray pixels [0, 360]. + Saturation : float + Replacement/treatment saturation for gray pixels [0, 1]. + mask : torch.Tensor, optional + Soft mask of shape (B, H, W) or (B, H, W, 1), values in [0, 1]. + When provided, the effect is blended per‑pixel: output = image*(1‑mask) + effect*mask. + + Returns + ------- + Tuple[torch.Tensor] + A single‑element tuple containing the processed image batch of the same + shape and channel count as the input. + """ + # ----- Validation -------------------------------------------------------- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() != 4: + raise ValueError( + f"image must be 4‑dimensional (B,H,W,C), got shape {image.shape}" + ) + if image.shape[-1] not in (3, 4): + raise ValueError( + f"image must have 3 or 4 channels (RGB or RGBA), got {image.shape[-1]}" + ) + + # ----- Alpha preservation ------------------------------------------------ + has_alpha = image.shape[-1] == 4 + if has_alpha: + alpha = image[..., 3:4] # keep last dim for concatenation + rgb_image = image[..., :3] + else: + alpha = None + rgb_image = image + + # ----- Mask validation --------------------------------------------------- + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor") + # Unify shape to (B, H, W, 1) + if mask.ndim == 3: + mask = mask.unsqueeze(-1) + elif mask.ndim != 4 or mask.shape[-1] != 1: + raise ValueError( + f"mask must have shape (B,H,W) or (B,H,W,1), got {mask.shape}" + ) + if mask.shape[0] != image.shape[0]: + raise ValueError( + f"mask batch size ({mask.shape[0]}) does not match image ({image.shape[0]})" + ) + if mask.shape[1:3] != image.shape[1:3]: + raise ValueError( + f"mask spatial size {mask.shape[1:3]} does not match image {image.shape[1:3]}" + ) + # Move to image device/dtype + mask = mask.to(device=rgb_image.device, dtype=rgb_image.dtype) + + # ----- Early exit (mask is fully transparent) ---------------------------- + if mask is not None and mask.max() == 0: + return (image,) + + B = rgb_image.shape[0] + if B == 0: + return (image,) + + # ----- Chunked processing ------------------------------------------------ + chunk_size = compute_chunk_size(rgb_image, mask) + if chunk_size <= 0 or chunk_size > B: + chunk_size = B + + out_chunks = [] + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + rgb_chunk = rgb_image[start:end] + mask_chunk = mask[start:end] if mask is not None else None + + chunk_out = _color_rotate_core( + rgb_chunk=rgb_chunk, + mask_chunk=mask_chunk, + Source_from=Source_from, + Source_to=Source_to, + Source_clockwise=Source_clockwise, + Destination_from=Destination_from, + Destination_to=Destination_to, + Destination_clockwise=Destination_clockwise, + Gray_threshold=Gray_threshold, + Mode=Mode, + Hue=Hue, + Saturation=Saturation, + ) + out_chunks.append(chunk_out) + + rgb_result = torch.cat(out_chunks, dim=0) + + # ----- Recombine alpha --------------------------------------------------- + if has_alpha: + result = torch.cat([rgb_result, alpha], dim=-1) + else: + result = rgb_result + + return (result,) + diff --git a/py/color_posterize.py b/py/color_posterize.py new file mode 100644 index 0000000..91549a4 --- /dev/null +++ b/py/color_posterize.py @@ -0,0 +1,168 @@ +import torch +from typing import Tuple, Optional + +from .libs.color import srgb_to_linear_vec, linear_to_srgb_vec +from .libs.process import compute_chunk_size + +# ============================================================ +# GEGL-style posterization core +# ============================================================ + +def _posterize_channel_gegl(x: torch.Tensor, levels: int) -> torch.Tensor: + """ + GEGL-style uniform quantizer. + Equivalent to: + floor(x * (levels - 1) + 0.5) / (levels - 1) + """ + levels = max(2, int(levels)) + scale = levels - 1 + return torch.floor(x * scale + 0.5) / scale + + +# ============================================================ +# Main operator – improved with batch chunking, alpha preservation, +# error handling, and early exit. +# ============================================================ + +def exec_operator_color_posterize( + image: torch.Tensor, + levels: int, + mask: Optional[torch.Tensor] = None, + linear: bool = False, # GEGL toggle: operate in linear light +) -> Tuple[torch.Tensor]: + """ + Apply GEGL‑style posterisation to a batch of images. + + The input *image* is a float tensor of shape ``(B, H, W, C)`` with + ``C == 3`` (RGB) or ``C == 4`` (RGBA). *levels* controls the number + of discrete levels per colour channel (minimum 2). An optional *mask* + tensor of shape ``(H, W)`` (single mask) or ``(B, H, W)`` (per‑image + mask) defines the per‑pixel strength of the effect: ``0.0`` leaves the + pixel unchanged, ``1.0`` applies full posterisation. + + When the input contains an alpha channel (``C == 4``), the alpha + channel is **preserved unchanged**; only the RGB channels are + posterised and blended with the mask. + + The function validates input shapes, performs an early exit for empty + batches, and processes the batch in GPU‑friendly chunks as determined + by ``compute_chunk_size``. No spatial tiling is performed. + + Parameters + ---------- + image : torch.Tensor + Batch of images (B,H,W,C), float. + levels : int + Number of posterisation levels (2‑256). + mask : torch.Tensor or None + Optional strength mask (H,W) or (B,H,W). + linear : bool + If ``True``, the effect is applied in linear light + (sRGB ↔ linear conversion is used). + + Returns + ------- + Tuple[torch.Tensor] + A 1‑tuple containing the posterised image batch. + """ + + # ---------- input validation ---------- + if not isinstance(image, torch.Tensor): + raise TypeError(f"image must be a torch.Tensor, got {type(image)}") + if image.dim() != 4: + raise ValueError(f"image must be 4D (B,H,W,C), got shape {image.shape}") + C = image.shape[-1] + if C not in (1, 3, 4): + raise ValueError(f"Unsupported number of channels: {C}. Only 1, 3, or 4 are supported") + if image.dtype not in (torch.float32, torch.float64, torch.float16): + raise TypeError(f"image dtype must be float, got {image.dtype}") + + levels = max(2, int(levels)) # ensure >= 2 + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError(f"mask must be a torch.Tensor or None, got {type(mask)}") + if mask.dim() == 2: + pass # (H,W) will be broadcast + elif mask.dim() == 3: + if mask.shape[0] != image.shape[0]: + raise ValueError( + f"Batch dimension mismatch: image batch {image.shape[0]}, " + f"mask batch {mask.shape[0]}" + ) + else: + raise ValueError(f"mask must be 2D (H,W) or 3D (B,H,W), got shape {mask.shape}") + if mask.device != image.device: + raise ValueError("image and mask must reside on the same device") + + # ---------- early exit for empty batches ---------- + if image.shape[0] == 0: + return (image,) + + # ---------- alpha channel handling ---------- + has_alpha = (C == 4) + if has_alpha: + rgb = image[..., :3] + alpha = image[..., 3:4] # keep as (B,H,W,1) + else: + rgb = image + alpha = None + + B = rgb.shape[0] + + # ---------- determine chunk size (VRAM‑aware) ---------- + # The helper considers both image and mask sizes. + chunk_size = compute_chunk_size(image, mask) + chunk_size = max(1, min(chunk_size, B)) # clamp to valid range + + out_chunks = [] + + # ---------- chunked processing ---------- + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + + chunk_rgb = rgb[start:end].clone() # clone to avoid in‑place issues + # Prepare mask chunk + chunk_mask = None + if mask is not None: + if mask.ndim == 3: # per‑image masks + chunk_mask = mask[start:end] + else: # single (H,W) mask, reuse for all images + chunk_mask = mask + + # ---- posterisation (per‑channel) ---- + x = torch.clamp(chunk_rgb, 0.0, 1.0) + + if linear: + x = srgb_to_linear_vec(x) + + x_p = _posterize_channel_gegl(x, levels) + + if linear: + x_p = linear_to_srgb_vec(x_p) + + x_p = torch.clamp(x_p, 0.0, 1.0) + + # ---- mask blending (only on RGB) ---- + if chunk_mask is not None: + # Ensure mask has shape (B',H,W,1) for broadcasting + if chunk_mask.ndim == 2: + chunk_mask = chunk_mask.unsqueeze(0).unsqueeze(-1) # (1,H,W,1) + else: + chunk_mask = chunk_mask.unsqueeze(-1) # (B',H,W,1) + chunk_mask = chunk_mask.to(dtype=x.dtype, device=x.device) + x_p = chunk_rgb * (1.0 - chunk_mask) + x_p * chunk_mask + + # ---- reassemble alpha if present ---- + if has_alpha: + chunk_out = torch.cat([x_p, alpha[start:end]], dim=-1) + else: + chunk_out = x_p + + out_chunks.append(chunk_out) + + # ---------- concatenate chunks ---------- + final_image = torch.cat(out_chunks, dim=0) + + return (final_image,) + diff --git a/py/color_saturation.py b/py/color_saturation.py new file mode 100644 index 0000000..85e8ecb --- /dev/null +++ b/py/color_saturation.py @@ -0,0 +1,185 @@ +import torch +from typing import Tuple, Optional +import logging + +from .libs.color import rgb_to_hsv, hsv_to_rgb, rgb_to_lch, lch_to_rgb, rgb_to_yuv, yuv_to_rgb +from .libs.process import compute_chunk_size + +logger = logging.getLogger(__name__) + +# ---------------------------------------------------------------------- +# Core saturation algorithm – processes a single batch on one device +# ---------------------------------------------------------------------- +def _saturation_core( + image: torch.Tensor, # (B, H, W, C) C = 3 or 4 + scale: float, + color_space: str, + mask: Optional[torch.Tensor] # (B, H, W, C) or None +) -> torch.Tensor: + """ + Apply saturation scaling to an image batch (RGB or RGBA). + Handles exactly one device; no chunking, no fallback. + """ + + # Early exit when no change is requested + if scale == 1.0: + return image + + # Clamp input to valid colour range + img = image.clamp(0.0, 1.0) + + # Separate alpha if present + if img.shape[-1] == 4: + alpha = img[..., 3:4] + rgb = img[..., :3] + else: + alpha = None + rgb = img + + # ----- colour space conversion & scaling ----- + if color_space == "Native": + h, s, v = rgb_to_hsv(rgb) + s = torch.clamp(s * scale, 0.0, 1.0) + result_rgb = hsv_to_rgb(h, s, v) + elif color_space == "CIE Lab/Lch": + l, c, h = rgb_to_lch(rgb) + c = c * scale + result_rgb = lch_to_rgb(l, c, h) + else: # CIE Yuv + y, u, v = rgb_to_yuv(rgb) + u = u * scale + v = v * scale + result_rgb = yuv_to_rgb(y, u, v) + + result_rgb = torch.clamp(result_rgb, 0.0, 1.0) + + # Reassemble RGBA if needed + if alpha is not None: + result = torch.cat([result_rgb, alpha], dim=-1) + else: + result = result_rgb + + # Mask blending (applied to all channels, including alpha) + if mask is not None: + result = img * (1.0 - mask) + result * mask + + return result + + +# ---------------------------------------------------------------------- +# Public entry point – batch, device, VRAM management, error handling +# ---------------------------------------------------------------------- +def exec_operator_color_saturation( + image: torch.Tensor, + scale: float, + color_space: str, + mask: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor]: + """ + Adjust the saturation of an image or batch of images. + + Parameters + ---------- + image : torch.Tensor + Input image(s) in RGB or RGBA format. + Shape: (H, W, 3) or (H, W, 4) for a single image, + (B, H, W, 3) or (B, H, W, 4) for a batch. + scale : float + Saturation multiplier. 1.0 = no change, <1 reduces saturation, + >1 increases saturation. + color_space : str + Colour space in which saturation is scaled. + One of ``"Native"``, ``"CIE Lab/Lch"``, ``"CIE Yuv"``. + mask : torch.Tensor, optional + Optional mask defining where the adjustment is applied. + Shape must match the image batch size and spatial dimensions. + (B, H, W) or (B, 1, H, W) or (B, H, W, 1). + + Returns + ------- + Tuple[torch.Tensor] + A tuple containing the saturation‑adjusted image(s) in the same + format and shape as the input. + """ + # --- validate input shapes and types --- + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() not in (3, 4): + raise ValueError(f"image must be 3D (H,W,C) or 4D (B,H,W,C), got {image.dim()}D") + if image.shape[-1] not in (3, 4): + raise ValueError(f"image channels must be 3 or 4, got {image.shape[-1]}") + + if not isinstance(scale, (int, float)): + raise TypeError("scale must be a float") + if color_space not in ("Native", "CIE Lab/Lch", "CIE Yuv"): + raise ValueError(f"unknown color_space '{color_space}'") + + # Normalise to batch format (B, H, W, C) + single_input = image.dim() == 3 + if single_input: + image = image.unsqueeze(0) # (1, H, W, C) + if mask is not None: + mask = mask.unsqueeze(0) # (1, …) whatever dim + + batch_size, h, w, channels = image.shape + device = image.device + + # --- mask handling (ensure (B, H, W, C) compatible) --- + if mask is not None: + mask = mask.to(device) + # bring to (B, H, W, 1) + if mask.dim() == 2: # (H, W) + mask = mask.unsqueeze(0).unsqueeze(-1) + elif mask.dim() == 3: # (B, H, W) + mask = mask.unsqueeze(-1) + elif mask.dim() == 4: + if mask.shape[1] == 1: # (B, 1, H, W) + mask = mask.permute(0, 2, 3, 1) # (B, H, W, 1) + elif mask.shape[-1] != 1: + raise ValueError("mask must have a single channel") + else: + raise ValueError("mask must be 2D, 3D or 4D") + + if mask.shape[0] != batch_size: + raise ValueError(f"batch size mismatch: image {batch_size}, mask {mask.shape[0]}") + # Expand mask to full channel dimension + mask = mask.expand(-1, -1, -1, channels) + + # --- VRAM optimised chunked processing (GPU only) --- + VRAM_CHUNK_SIZE = compute_chunk_size(image) # process 4 images at once on GPU to limit peak memory + + def _process_chunked(img_batch, msk_batch): + """Process a whole batch, potentially in chunks.""" + if device.type == 'cuda' and batch_size > VRAM_CHUNK_SIZE: + results = [] + for i in range(0, batch_size, VRAM_CHUNK_SIZE): + end = min(i + VRAM_CHUNK_SIZE, batch_size) + sub_img = img_batch[i:end] + sub_msk = msk_batch[i:end] if msk_batch is not None else None + results.append(_saturation_core(sub_img, scale, color_space, sub_msk)) + # free GPU memory peaks + torch.cuda.empty_cache() + return torch.cat(results, dim=0) + else: + return _saturation_core(img_batch, scale, color_space, msk_batch) + + # --- execute with GPU OOM fallback --- + try: + result = _process_chunked(image, mask) + except RuntimeError as e: + if 'out of memory' in str(e) and device.type == 'cuda': + logger.warning("CUDA out of memory – falling back to CPU") + torch.cuda.empty_cache() + cpu_image = image.cpu() + cpu_mask = mask.cpu() if mask is not None else None + result_cpu = _saturation_core(cpu_image, scale, color_space, cpu_mask) + result = result_cpu.to(device) + else: + raise + + # Remove batch dimension if input was a single image + if single_input: + result = result.squeeze(0) + + return (result,) + diff --git a/py/color_shadows_highlights.py b/py/color_shadows_highlights.py new file mode 100644 index 0000000..955edf6 --- /dev/null +++ b/py/color_shadows_highlights.py @@ -0,0 +1,220 @@ +import torch +from typing import Tuple + +from .libs.blur import gaussian_blur_2d +from .libs.process import compute_chunk_size + + +def _process_core( + image: torch.Tensor, + Shadows: float, + Shadows_color_adj: float, + Highlights: float, + Highlights_color_adj: float, + White_point_adj: float, + Radius: float, + Compress: float, + mask: torch.Tensor = None, +) -> torch.Tensor: + """ + Apply the GEGL‑style shadows/highlights correction to a single batch + chunk. This function is identical to the original processing logic. + """ + img = image[..., :3].clamp(0.0, 1.0) + device = img.device + + # GEGL luminance (Rec.709) + lum = 0.2126 * img[..., 0] + 0.7152 * img[..., 1] + 0.0722 * img[..., 2] + + # White point correction + white = 1.0 - (White_point_adj / 100.0) + lum = torch.clamp(lum / max(white, 1e-6), 0.0, 1.0) + + # Local context (GEGL radius model) + local = gaussian_blur_2d(lum, Radius) + + # Normalise parameters + shadows = max(-1.0, min(1.0, float(Shadows) / 100.0)) + highlights = max(-1.0, min(1.0, float(Highlights) / 100.0)) + compress = max(0.0, min(0.99, float(Compress) / 100.0)) + + # GEGL‑like nonlinear weighting + shadow_w = (1.0 - local) ** 1.5 + highlight_w = local ** 1.5 + + lum_out = lum + shadows * shadow_w - highlights * highlight_w + + # Compression (midtone preservation) + mid = lum * (1.0 - lum) * 4.0 + lum_out = lum_out * (1.0 - compress * mid) + lum_out = torch.clamp(lum_out, 0.0, 1.0) + + # Saturation model (correct luminance) + gray = 0.2126 * img[..., 0:1] + 0.7152 * img[..., 1:2] + 0.0722 * img[..., 2:3] + + sat_shadow = 1.0 - (1.0 - Shadows_color_adj / 100.0) * (1.0 - local.unsqueeze(-1)) + sat_high = 1.0 - (1.0 - Highlights_color_adj / 100.0) * local.unsqueeze(-1) + sat = sat_shadow * sat_high + + rgb = gray + (img - gray) * sat + + # Luminance correction (safe division) + rgb = rgb * (lum_out.unsqueeze(-1) / (lum.unsqueeze(-1) + 1e-6)) + rgb = torch.clamp(rgb, 0.0, 1.0) + + # Mask blending + if mask is not None: + mask = mask.to(device) + if mask.dim() == 3: + mask = mask.unsqueeze(-1) + elif mask.dim() == 4 and mask.shape[1] == 1: + mask = mask.permute(0, 2, 3, 1) + mask = mask.expand_as(rgb) + rgb = img * (1.0 - mask) + rgb * mask + + # Preserve alpha channel + if image.shape[-1] == 4: + rgb = torch.cat([rgb, image[..., 3:4]], dim=-1) + + return rgb + + +def exec_operator_color_shadows_highlights( + image: torch.Tensor, + Shadows: float = 0.0, + Shadows_color_adj: float = 100.0, + Highlights: float = 0.0, + Highlights_color_adj: float = 50.0, + White_point_adj: float = 0.0, + Radius: float = 100.0, + Compress: float = 50.0, + mask: torch.Tensor = None, +) -> Tuple[torch.Tensor]: + """ + Apply GEGL‑inspired shadows/highlights correction to a batch of images. + + This operator adjusts the luminance and colour saturation of shadows and + highlights independently. It preserves the alpha channel when present and + supports optional soft masking. + + Parameters + ---------- + image : torch.Tensor + Input image tensor of shape (B, H, W, C) with C = 3 (RGB) or 4 (RGBA). + Values are expected in the [0, 1] range. + Shadows : float + Amount of shadow brightening. Range [-100, 100] → [-1, 1] internally. + Shadows_color_adj : float + Colour saturation adjustment for shadows, 0 = desaturate, 100 = keep. + Highlights : float + Amount of highlight darkening. Range [-100, 100] → [-1, 1] internally. + Highlights_color_adj : float + Colour saturation adjustment for highlights. + White_point_adj : float + White‑point correction, range [-10, 10]. + Radius : float + Spatial radius for the local luminance estimation (pixels). + Compress : float + Midtone compression strength, range [0, 100]. + mask : torch.Tensor, optional + Soft mask with shape (B, H, W) or (B, 1, H, W). 1 = full effect, + 0 = original image. + + Returns + ------- + Tuple[torch.Tensor] + A single‑element tuple containing the processed image tensor + with the same shape and number of channels as the input. + """ + # ------------------------------------------------------------ + # Early exit: empty input + # ------------------------------------------------------------ + if image is None or image.numel() == 0: + return (image,) + + # ------------------------------------------------------------ + # Early exit: identity operation (no changes) + # ------------------------------------------------------------ + if ( + Shadows == 0.0 + and Highlights == 0.0 + and White_point_adj == 0.0 + and Compress == 0.0 + and Shadows_color_adj == 100.0 + and Highlights_color_adj == 100.0 + and mask is None + ): + return (image,) + + # ------------------------------------------------------------ + # Input validation + # ------------------------------------------------------------ + if not isinstance(image, torch.Tensor): + raise TypeError("image must be a torch.Tensor") + if image.dim() != 4 or image.shape[-1] not in (3, 4): + raise ValueError( + f"image must have shape (B,H,W,C) with C=3 or 4, got {image.shape}" + ) + + batch_size = image.shape[0] + device = image.device + + if mask is not None: + if not isinstance(mask, torch.Tensor): + raise TypeError("mask must be a torch.Tensor") + if mask.shape[0] != batch_size: + raise ValueError( + f"mask batch size {mask.shape[0]} does not match image batch size {batch_size}" + ) + if mask.dim() == 4: + if mask.shape[1] != 1 or mask.shape[2:] != image.shape[1:3]: + raise ValueError( + f"mask shape {mask.shape} is incompatible with image shape {image.shape}" + ) + elif mask.dim() == 3: + if mask.shape[1:] != image.shape[1:3]: + raise ValueError( + f"mask shape {mask.shape} is incompatible with image shape {image.shape}" + ) + else: + raise ValueError("mask must be 3D (B,H,W) or 4D (B,1,H,W)") + + # ------------------------------------------------------------ + # Chunked processing for GPU VRAM management + # ------------------------------------------------------------ + try: + # Compute a safe chunk size (batch count) based on available VRAM + chunk_size = compute_chunk_size(image, mask) + chunk_size = max(1, int(chunk_size)) # ensure valid + + output_chunks = [] + for start in range(0, batch_size, chunk_size): + end = min(start + chunk_size, batch_size) + + img_chunk = image[start:end] + mask_chunk = mask[start:end] if mask is not None else None + + processed = _process_core( + img_chunk, + Shadows, + Shadows_color_adj, + Highlights, + Highlights_color_adj, + White_point_adj, + Radius, + Compress, + mask_chunk, + ) + output_chunks.append(processed) + + result = torch.cat(output_chunks, dim=0) + + except Exception as e: + # On any processing error, return the original image and print a warning + print(f"[Shadows/Highlights] Error: {e}. Returning original image.") + return (image,) + + return (result,) + + + diff --git a/py/color_temperature.py b/py/color_temperature.py new file mode 100644 index 0000000..4dbc149 --- /dev/null +++ b/py/color_temperature.py @@ -0,0 +1,209 @@ +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) + diff --git a/py/color_threshold.py b/py/color_threshold.py new file mode 100644 index 0000000..15ac3ae --- /dev/null +++ b/py/color_threshold.py @@ -0,0 +1,163 @@ +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 \ No newline at end of file diff --git a/py/color_tone_mapping_fattal_et_al_2002.py b/py/color_tone_mapping_fattal_et_al_2002.py new file mode 100644 index 0000000..e6322f5 --- /dev/null +++ b/py/color_tone_mapping_fattal_et_al_2002.py @@ -0,0 +1,198 @@ +import torch +import numpy as np +from typing import Tuple + +module_cat = "md/image/color" + +# ----------------------------------------------------------------------------------- +# Core tone mapping implementation (ported from GEGL's fattal02.c / pfstmo) +# ----------------------------------------------------------------------------------- +def exec_operator_color_tone_mapping_fattal_et_al_2002( + image: torch.Tensor, + Alpha: float, + Beta: float, + Saturation: float, + Noise: float, + mask: torch.Tensor = None, +) -> Tuple[torch.Tensor]: + """ + Applies Fattal et al. 2002 gradient-domain tone mapping. + Based on GIMP's implementation (GEGL operation 'fattal02'). + """ + # Convert to numpy (B, H, W, C) float32 in [0,1] + img_np = image.detach().cpu().numpy().astype(np.float32) + batch_size, H, W, C = img_np.shape + + # If mask is provided, use it; otherwise process full image + if mask is not None: + mask_np = mask.detach().cpu().numpy().astype(np.float32) + # mask is (B, H, W) or (B, 1, H, W) – ensure it's (B, H, W, 1) + if mask_np.ndim == 3: + mask_np = mask_np[..., np.newaxis] + elif mask_np.ndim == 4 and mask_np.shape[1] == 1: + mask_np = np.transpose(mask_np, (0, 2, 3, 1)) + else: + mask_np = np.ones((batch_size, H, W, 1), dtype=np.float32) + + out = [] + for b in range(batch_size): + img = img_np[b] # (H, W, C) + msk = mask_np[b] # (H, W, 1) + + # ---- 1. Convert to linear luminance (ITU-R BT.709) ---- + # GIMP uses: L = 0.2126*R + 0.7152*G + 0.0722*B + luminance = 0.2126 * img[..., 0] + 0.7152 * img[..., 1] + 0.0722 * img[..., 2] + luminance = np.clip(luminance, 1e-6, None) # avoid log(0) + + # ---- 2. Tone map the luminance channel ---- + tonemapped_lum = _fattal_tone_map(luminance, Alpha, Beta, Noise) + + # ---- 3. Apply saturation and recombine ---- + # Scale RGB channels by the ratio of new/old luminance + ratio = tonemapped_lum / (luminance + 1e-8) + # Saturation: blend between fully desaturated and original ratio + ratio_sat = 1.0 + Saturation * (ratio - 1.0) + result = img * ratio_sat[..., np.newaxis] + result = np.clip(result, 0.0, 1.0) + + # Apply mask + result = result * msk + img * (1.0 - msk) + out.append(result) + + # Stack batch and convert back to torch + out_np = np.stack(out, axis=0) + return (torch.from_numpy(out_np).to(image.device),) + + +# ----------------------------------------------------------------------------------- +# Internal Fattal tone mapping core (luminance-only) +# ----------------------------------------------------------------------------------- +def _fattal_tone_map(lum: np.ndarray, alpha: float, beta: float, noise: float) -> np.ndarray: + """ + Tone maps a single luminance channel using the Fattal et al. 2002 algorithm. + Based on GEGL's fattal02.c and the pfstmo reference implementation. + """ + H, W = lum.shape + # Work in log domain + log_lum = np.log(lum) + + # ---- Build Gaussian pyramid ---- + pyramid = [log_lum] + while pyramid[-1].shape[0] > 32 and pyramid[-1].shape[1] > 32: + blurred = _gaussian_blur(pyramid[-1], sigma=1.0) + pyramid.append(blurred[::2, ::2]) + + n_levels = len(pyramid) + + # ---- Compute gradient magnitudes at each level ---- + grad_mags = [] + for level, layer in enumerate(pyramid): + gy, gx = np.gradient(layer) + # Scale gradients by the level's resolution (as in GEGL) + scale = 2 ** level + mag = np.sqrt((gx / scale) ** 2 + (gy / scale) ** 2) + grad_mags.append(mag) + + # ---- Attenuate gradients using the Fattal formula ---- + # GEGL uses: attenuated = mag * ( (alpha / mag)^beta ) for mag > alpha + # and for mag <= alpha, attenuated = mag * (alpha / mag)^(beta*noise)? + # More precisely (from GEGL): + # if (mag > alpha) mag2 = mag * pow(alpha / mag, beta) + # else mag2 = mag * pow(alpha / mag, beta * noise) + # then clamp to [0, 1] and scale. + attenuated_mags = [] + for mag in grad_mags: + # Avoid division by zero + safe_mag = np.maximum(mag, 1e-8) + ratio = alpha / safe_mag + if noise > 0: + # For small gradients, use noise to control attenuation + mask = (mag > alpha) + exp = np.where(mask, beta, beta * noise) + else: + exp = beta + att = mag * np.power(ratio, exp) + # GEGL clamps attenuated gradients to [0,1] and then scales by original mag? + # Actually it computes: phi = att / mag (where att is the attenuated magnitude) + # Then phi = clamp(phi, 0, 1) + # Then new_grad = phi * original_gradient + phi = np.clip(att / safe_mag, 0.0, 1.0) + attenuated_mags.append(phi) + + # ---- Reconstruct the image from the modified gradients (solve Poisson equation) ---- + # We solve ∇²u = div( G' ) where G' is the attenuated gradient field. + # Use a simple multi-grid / conjugate gradient approach (GEGL uses linbcg). + # For simplicity, we use a fast Poisson solver via FFT (as in pfstmo). + # Since the pyramid has multiple levels, we propagate from coarsest to finest. + reconstructed = _solve_poisson_pyramid(pyramid, attenuated_mags) + + # ---- Convert back to linear domain and normalize ---- + result = np.exp(reconstructed) + # Normalize to [0, 1] (GIMP scales by max) + max_val = np.max(result) + if max_val > 0: + result = result / max_val + return np.clip(result, 0.0, 1.0) + + +def _gaussian_blur(arr: np.ndarray, sigma: float = 1.0) -> np.ndarray: + """Simple Gaussian blur using separable convolution.""" + from scipy.ndimage import gaussian_filter + return gaussian_filter(arr, sigma=sigma, mode='reflect') + + +def _solve_poisson_pyramid(pyramid: list, attenuated_mags: list) -> np.ndarray: + """ + Reconstructs the log-luminance image from the attenuated gradient magnitudes. + Uses a coarse-to-fine propagation with FFT-based Poisson solver at each level. + """ + # Start from the coarsest level + n_levels = len(pyramid) + # Coarsest level: just use the original log-lum (or solve directly) + # For simplicity, we use the coarsest level as-is (no gradient modification) + # In GEGL, they solve a linear system at each level and upsample. + # Here we implement a simplified version: solve at the finest level using FFT. + # (This is acceptable for a Python port and matches the pfstmo approach.) + + # Get the finest level gradients + # We need the original gradients and the attenuation factors + log_lum = pyramid[0] + gy, gx = np.gradient(log_lum) + # Attenuate gradients + phi = attenuated_mags[0] # shape (H, W) + # Apply phi to gradients + gx_att = gx * phi + gy_att = gy * phi + + # Solve Poisson equation: ∇²u = div(g_att) + # Using FFT-based solver (periodic boundary conditions) + from scipy.fft import fft2, ifft2, fftfreq + H, W = log_lum.shape + # Compute divergence of attenuated gradient + div_g = np.zeros_like(log_lum) + # Divergence: d(gx)/dx + d(gy)/dy + div_g[1:-1, :] = (gx_att[2:, :] - gx_att[:-2, :]) / 2.0 + div_g[:, 1:-1] += (gy_att[:, 2:] - gy_att[:, :-2]) / 2.0 + # Handle boundaries (simple forward/backward differences) + div_g[0, :] = gx_att[1, :] - gx_att[0, :] + div_g[-1, :] = gx_att[-1, :] - gx_att[-2, :] + div_g[:, 0] = gy_att[:, 1] - gy_att[:, 0] + div_g[:, -1] = gy_att[:, -1] - gy_att[:, -2] + + # Solve using FFT + fx = fftfreq(W, 1.0) + fy = fftfreq(H, 1.0) + FX, FY = np.meshgrid(fx, fy) + denom = -4.0 * (np.pi**2) * (FX**2 + FY**2) + denom[0, 0] = 1.0 # avoid division by zero + + F_div = fft2(div_g) + F_u = F_div / denom + F_u[0, 0] = 0.0 # set DC component to zero + u = np.real(ifft2(F_u)) + + # Add back the mean of the original log-lum to preserve brightness + u = u - np.mean(u) + np.mean(log_lum) + return u + + diff --git a/py/layers_blending.py b/py/layers_blending.py new file mode 100644 index 0000000..62c7820 --- /dev/null +++ b/py/layers_blending.py @@ -0,0 +1,209 @@ +import torch + +# -------------------------------------------------------------------- +# Blend mode helpers – kept as module-level functions for clean reuse +# -------------------------------------------------------------------- +def blend_normal(bg_rgb, fg_rgb): + return fg_rgb + +def blend_screen(bg_rgb, fg_rgb): + return 1.0 - (1.0 - bg_rgb) * (1.0 - fg_rgb) + +def blend_multiply(bg_rgb, fg_rgb): + return bg_rgb * fg_rgb + +def blend_overlay(bg_rgb, fg_rgb): + result = torch.where(bg_rgb <= 0.5, 2.0 * bg_rgb * fg_rgb, + 1.0 - 2.0 * (1.0 - bg_rgb) * (1.0 - fg_rgb)) + return result.clamp(0.0, 1.0) + +def blend_difference(bg_rgb, fg_rgb): + return torch.abs(bg_rgb - fg_rgb) + +def blend_dissolve(bg_rgb, fg_rgb): + noise = torch.rand_like(bg_rgb) + return torch.where(noise < 0.5, fg_rgb, bg_rgb) + +def blend_clear(bg_rgb, fg_rgb): + return bg_rgb + +def blend_add(bg_rgb, fg_rgb): + return torch.clamp(bg_rgb + fg_rgb, 0.0, 1.0) + +def blend_subtract(bg_rgb, fg_rgb): + return torch.clamp(bg_rgb - fg_rgb, 0.0, 1.0) + +def blend_replace(bg_rgb, fg_rgb): + return fg_rgb + +def blend_darken_only(bg_rgb, fg_rgb): + return torch.min(bg_rgb, fg_rgb) + +def blend_lighten_only(bg_rgb, fg_rgb): + return torch.max(bg_rgb, fg_rgb) + +def blend_color_dodge(bg_rgb, fg_rgb): + result = bg_rgb / (1.0 - fg_rgb + 1e-8) + return result.clamp(0.0, 1.0) + +def blend_color_burn(bg_rgb, fg_rgb): + result = 1.0 - (1.0 - bg_rgb) / (fg_rgb + 1e-8) + return result.clamp(0.0, 1.0) + +def blend_hard_light(bg_rgb, fg_rgb): + result = torch.where(fg_rgb <= 0.5, 2.0 * bg_rgb * fg_rgb, + 1.0 - 2.0 * (1.0 - bg_rgb) * (1.0 - fg_rgb)) + return result.clamp(0.0, 1.0) + +def blend_soft_light(bg_rgb, fg_rgb): + result = torch.where(fg_rgb <= 0.5, + (2.0 * bg_rgb * fg_rgb) + (bg_rgb * bg_rgb * (1.0 - 2.0 * fg_rgb)), + 2.0 * bg_rgb * (1.0 - fg_rgb) + torch.sqrt(bg_rgb) * (2.0 * fg_rgb - 1.0)) + return result.clamp(0.0, 1.0) + +def blend_pin_light(bg_rgb, fg_rgb): + result = torch.where(fg_rgb <= 0.5, + torch.min(bg_rgb, 2.0 * fg_rgb), + torch.max(bg_rgb, 2.0 * fg_rgb - 1.0)) + return result.clamp(0.0, 1.0) + +def blend_linear_dodge_add(bg_rgb, fg_rgb): + return torch.clamp(bg_rgb + fg_rgb, 0.0, 1.0) + +def blend_linear_burn_multiply(bg_rgb, fg_rgb): + return torch.clamp(bg_rgb + fg_rgb - 1.0, 0.0, 1.0) + +def blend_vivid_light(bg_rgb, fg_rgb): + result = torch.where(fg_rgb <= 0.5, + 1.0 - (1.0 - bg_rgb) / (2.0 * fg_rgb + 1e-8), + bg_rgb / (2.0 * (1.0 - fg_rgb) + 1e-8)) + return result.clamp(0.0, 1.0) + +def blend_halo(bg_rgb, fg_rgb): + edge = torch.abs(bg_rgb - fg_rgb) + glow = (bg_rgb + fg_rgb) * edge + return glow.clamp(0.0, 1.0) + +def blend_negate(bg_rgb, fg_rgb): + return 1.0 - torch.abs(bg_rgb - fg_rgb) + +def blend_merge_rgb(bg_rgb, fg_rgb): + return fg_rgb + +def blend_merge_cyan(bg_rgb, fg_rgb): + cyan = 1.0 - fg_rgb[..., 0:1] + return cyan.repeat(1, 1, 1, 3) + +def blend_merge_magenta(bg_rgb, fg_rgb): + magenta = 1.0 - fg_rgb[..., 1:2] + return magenta.repeat(1, 1, 1, 3) + +def blend_merge_yellow(bg_rgb, fg_rgb): + yellow = 1.0 - fg_rgb[..., 2:3] + return yellow.repeat(1, 1, 1, 3) + +def blend_merge_black(bg_rgb, fg_rgb): + intensity = fg_rgb.mean(dim=-1, keepdim=True) + return intensity.repeat(1, 1, 1, 3) + +def blend_fill_fg_color(bg_rgb, fg_rgb): + return fg_rgb + +def blend_fill_bg_color(bg_rgb, fg_rgb): + return bg_rgb + +def blend_average(bg_rgb, fg_rgb): + return (bg_rgb + fg_rgb) * 0.5 + +def blend_composite(bg_rgb, fg_rgb): + return fg_rgb + +def blend_behind(bg_rgb, fg_rgb): + return fg_rgb + +def blend_template(bg_rgb, fg_rgb): + return fg_rgb + +def composite_layer_stack(background, layer, blending_mode, transparency, mask=None): + """ + background, layer: tensors of shape [B, H, W, C] where C=3 or 4 + mask: optional tensor [B, H, W] or [H, W]; values in [0,1] + transparency: global opacity multiplier + """ + device = background.device + layer = layer.to(device) + bg_h, bg_w = background.shape[1:3] + + # ---------- 1. Ensure background and layer have 4 channels ---------- + bg = background.clone() + if bg.shape[-1] == 3: + alpha_bg = torch.ones((bg.shape[0], bg_h, bg_w, 1), device=device) + bg = torch.cat([bg, alpha_bg], dim=-1) + else: + alpha_bg = bg[..., 3:4] + + fg = layer.clone() + if fg.shape[-1] == 3: + alpha_fg = torch.ones((fg.shape[0], fg.shape[1], fg.shape[2], 1), device=device) + fg = torch.cat([fg, alpha_fg], dim=-1) + else: + alpha_fg = fg[..., 3:4] + + bg_rgb = bg[..., :3] + fg_rgb = fg[..., :3] + + # ---------- 2. Resize layer to match background dimensions ---------- + if fg.shape[1:3] != (bg_h, bg_w): + fg_rgb = torch.nn.functional.interpolate( + fg_rgb.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False + ).permute(0, 2, 3, 1) + alpha_fg = torch.nn.functional.interpolate( + alpha_fg.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False + ).permute(0, 2, 3, 1) + + # ---------- 3. Prepare mask ---------- + if mask is None: + mask = torch.ones((bg.shape[0], bg_h, bg_w, 1), device=device) + else: + mask = mask.to(device) + if mask.dim() == 2: # (H, W) + mask = mask.unsqueeze(0).unsqueeze(-1) + elif mask.dim() == 3: # (B, H, W) or (B, H, W, 1) + if mask.shape[-1] != 1: + mask = mask.unsqueeze(-1) + if mask.shape[1:3] != (bg_h, bg_w): + mask = torch.nn.functional.interpolate( + mask.permute(0, 3, 1, 2), size=(bg_h, bg_w), mode='bilinear', align_corners=False + ).permute(0, 2, 3, 1) + if mask.shape[0] < bg.shape[0]: + mask = mask.expand(bg.shape[0], -1, -1, -1) + + # ---------- 4. Compute final opacity factor ---------- + opacity = alpha_fg * mask * transparency # [B, H, W, 1] + + # Special handling for "Clear" and "Behind" + if blending_mode.lower() == "clear": + opacity = torch.zeros_like(opacity) + elif blending_mode.lower() == "behind": + opacity = opacity * (1.0 - alpha_bg) + + # ---------- 5. Apply RGB blend mode ---------- + # Build a mapping from the existing module-level blend functions + blend_func_name = blending_mode.lower().replace(' ', '_').replace('(', '').replace(')', '').replace('-', '_') + # Use the functions defined at module level (they are in the same scope) + blend_func = globals().get(f"blend_{blend_func_name}", blend_normal) + + blended_rgb = blend_func(bg_rgb, fg_rgb) + + # ---------- 6. Composite RGB using opacity ---------- + result_rgb = bg_rgb * (1.0 - opacity) + blended_rgb * opacity + + # ---------- 7. Compute result alpha (standard over compositing) ---------- + result_alpha = alpha_bg + (1.0 - alpha_bg) * opacity + + # ---------- 8. Concatenate and clamp ---------- + result = torch.cat([result_rgb, result_alpha], dim=-1).clamp(0.0, 1.0) + + return (result,) + + diff --git a/py/libs/__init__.py b/py/libs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/py/libs/blur.py b/py/libs/blur.py new file mode 100644 index 0000000..2ac118e --- /dev/null +++ b/py/libs/blur.py @@ -0,0 +1,33 @@ + +import torch +import torch.nn.functional as F + +def gaussian_blur_2d(x, radius: float): + if radius <= 0: + return x + + sigma = radius * 0.5 + k = int(max(3, radius * 2 + 1)) + if k % 2 == 0: + k += 1 + + device = x.device + + coords = torch.arange(k, device=device) - k // 2 + g = torch.exp(-(coords ** 2) / (2 * sigma ** 2)) + g = g / g.sum() + + # reshape for separable conv + g_h = g.view(1, 1, 1, k) + g_v = g.view(1, 1, k, 1) + + x = x.unsqueeze(1) # [B,1,H,W] + + # IMPORTANT: symmetric SAME padding + pad = k // 2 + + x = F.conv2d(x, g_h, padding=(0, pad)) + x = F.conv2d(x, g_v, padding=(pad, 0)) + + return x.squeeze(1) + diff --git a/py/libs/color.py b/py/libs/color.py new file mode 100644 index 0000000..7919bac --- /dev/null +++ b/py/libs/color.py @@ -0,0 +1,439 @@ + +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 + diff --git a/py/libs/process.py b/py/libs/process.py new file mode 100644 index 0000000..68d8b36 --- /dev/null +++ b/py/libs/process.py @@ -0,0 +1,76 @@ +import torch +from typing import Optional + +def compute_chunk_size( + image: torch.Tensor, + mask: Optional[torch.Tensor] = None, +) -> int: + """ + Determine a safe batch-size chunk for GPU processing based on the + currently available CUDA memory. + + The estimation accounts for the memory required by the input image, + an optional mask tensor, the output image, and temporary working + tensors created during color conversion. + + Args: + image (torch.Tensor): + Input image tensor of shape (B, H, W, C). + + mask (torch.Tensor, optional): + Optional mask tensor associated with the image batch. + + Returns: + int: + Maximum number of images that can be processed safely in a + single chunk. If the image resides on the CPU, the complete + batch size is returned. + """ + if image.device.type != "cuda": + return image.shape[0] + + free_mem, _ = torch.cuda.mem_get_info(image.device) + + image_batch = max(image.shape[0], 1) + + # Bytes required by one image sample + image_sample_bytes = ( + image[0].numel() * image.element_size() + ) + + # Bytes required by one mask sample + mask_sample_bytes = 0 + + if mask is not None: + if mask.device != image.device: + mask = mask.to(image.device) + + mask_batch = max(mask.shape[0], 1) + + if mask_batch == image_batch: + mask_sample_bytes = mask[0].numel() * mask.element_size() + elif mask_batch == 1: + # Single mask broadcast across the batch + mask_sample_bytes = mask[0].numel() * mask.element_size() + else: + raise ValueError( + f"Mask batch size ({mask_batch}) does not match image batch size ({image_batch})." + ) + + # Working memory: + # input image + # output image + # temporary color conversion tensors + # optional mask + working_factor = 4 + + required_per_sample = ( + image_sample_bytes * working_factor + + mask_sample_bytes + ) + + # Keep a safety margin so other CUDA allocations have room. + usable_memory = int(free_mem * 0.90) + + return max(1, min(image_batch, usable_memory // max(required_per_sample, 1))) + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..cb9a9ae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[project] +name = "MdColorMod" +version = "1.0.0" +description = "GIMP color mod nodes in comfyUI" +license = { file = "LICENSE" } +dependencies = [] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e69de29 diff --git a/web/extension_reg.js b/web/extension_reg.js new file mode 100644 index 0000000..3d9a8ca --- /dev/null +++ b/web/extension_reg.js @@ -0,0 +1,43 @@ +import { app } from "../../../scripts/app.js"; + +import { color_balance_widget } from "./nodes/color_balance_widget.js"; +import { color_hue_saturation_widget } from "./nodes/color_hue_saturation_widget.js"; + +//import { addRGBIntegerWidget } from "./widgets/rgb_int_widget.js"; + + +app.registerExtension({ + name: "md.color_mod.frontend_extension", + + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + + const onNodeCreated = nodeType.prototype.onNodeCreated; + + if (nodeData.name === "mdColorBalance") { + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + this._inst = new color_balance_widget(this); + }; + } + + if (nodeData.name === "mdColorHueSaturation") { + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + this._inst = new color_hue_saturation_widget(this); + }; + } + + /* + if (nodeData.name === "mdPrimitiveColorRGB") { + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated?.apply(this, arguments); + addRGBIntegerWidget(this, "rgb", 16711680); + return result; + }; + } + */ + }, +}); + + diff --git a/web/libs/widgets_module.js b/web/libs/widgets_module.js new file mode 100644 index 0000000..1dc8f26 --- /dev/null +++ b/web/libs/widgets_module.js @@ -0,0 +1,9 @@ + + +export const set_node_widget_value = (node, widget_id, value) => { + let widget = node.widgets?.find(w => w.name === widget_id); + if (!widget) return console.error(`Could not find ${widget_id} widget`); + //console.log(`set ${widget_id} to ${value}`); + widget.value = value; +}; + diff --git a/web/nodes/color_balance_widget.js b/web/nodes/color_balance_widget.js new file mode 100644 index 0000000..c6cf980 --- /dev/null +++ b/web/nodes/color_balance_widget.js @@ -0,0 +1,103 @@ +import { set_node_widget_value } from "../libs/widgets_module.js"; + +export class color_balance_widget { + constructor(node) { + this.node = node; + this.rangeNames = ["Shadow", "Midtonnes", "Highlights"]; + + // Find widgets + this.rangeWidget = node.widgets?.find(w => w.name === "range"); + this.cyanWidget = node.widgets?.find(w => w.name === "Cyan-Red"); + this.magentaWidget = node.widgets?.find(w => w.name === "Magenta-Green"); + this.yellowWidget = node.widgets?.find(w => w.name === "Yellow-Blue"); + this.dataWidget = node.widgets?.find(w => w.name === "jdata"); + + if (!this.rangeWidget || !this.cyanWidget || !this.magentaWidget || !this.yellowWidget || !this.dataWidget) { + console.warn("mdColorBalance: some widgets missing"); + return; + } + + this.dataWidget.hidden = true; + + this.setupCallbacks(); + // Initial load after the node is fully ready + setTimeout(() => this.loadFromData(), 0); + } + + getCurrentRangeIndex() { + return this.rangeNames.indexOf(this.rangeWidget.value); + } + + saveCurrentToData() { + let data; + try { + data = JSON.parse(this.dataWidget.value); + } catch { + data = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]; + } + const idx = this.getCurrentRangeIndex(); + if (idx === -1) return; + data[idx] = [ + parseFloat(this.cyanWidget.value), + parseFloat(this.magentaWidget.value), + parseFloat(this.yellowWidget.value) + ]; + this.dataWidget.value = JSON.stringify(data); + if (this.dataWidget.callback) this.dataWidget.callback(this.dataWidget.value); + } + + loadFromData() { + let data; + try { + data = JSON.parse(this.dataWidget.value); + } catch { + data = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]; + } + const idx = this.getCurrentRangeIndex(); + if (idx === -1) return; + const vals = data[idx] || [0, 0, 0]; + + // Use the shared utility to set widget values + set_node_widget_value(this.node, "Cyan-Red", vals[0]); + set_node_widget_value(this.node, "Magenta-Green", vals[1]); + set_node_widget_value(this.node, "Yellow-Blue", vals[2]); + + // Notify the widget callbacks so the UI updates + if (this.cyanWidget.callback) this.cyanWidget.callback(this.cyanWidget.value); + if (this.magentaWidget.callback) this.magentaWidget.callback(this.magentaWidget.value); + if (this.yellowWidget.callback) this.yellowWidget.callback(this.yellowWidget.value); + } + + setupCallbacks() { + // --- Range dropdown --- + const origRangeCb = this.rangeWidget.callback; + this.rangeWidget.callback = (value) => { + // Save current slider values for the old range + this.saveCurrentToData(); + // Load slider values for the new range + this.loadFromData(); + // Invoke any existing callback + if (origRangeCb) origRangeCb(value); + }; + + // --- Slider widgets --- + const onSliderChange = () => { + this.saveCurrentToData(); + }; + + // Preserve existing callbacks by wrapping them + const wrapSliderCallback = (widget, extra) => { + const orig = widget.callback; + widget.callback = function(value) { + extra(); + if (orig) orig.call(this, value); // ✅ Preserve `this` context + }; + }; + + wrapSliderCallback(this.cyanWidget, onSliderChange); + wrapSliderCallback(this.magentaWidget, onSliderChange); + wrapSliderCallback(this.yellowWidget, onSliderChange); + } +} + + diff --git a/web/nodes/color_hue_saturation_widget.js b/web/nodes/color_hue_saturation_widget.js new file mode 100644 index 0000000..d1ca805 --- /dev/null +++ b/web/nodes/color_hue_saturation_widget.js @@ -0,0 +1,104 @@ +import { set_node_widget_value } from "../libs/widgets_module.js"; + +export class color_hue_saturation_widget { + + constructor(node) { + this.node = node; + // Order must match the dropdown: Master, Red, Yellow, Green, Cyan, Blue, Magenta + this.rangeNames = ["Master", "Red", "Yellow", "Green", "Cyan", "Blue", "Magenta"]; + + // Find widgets + this.colorWidget = node.widgets?.find(w => w.name === "color"); + this.hueWidget = node.widgets?.find(w => w.name === "Hue"); + this.lightnessWidget = node.widgets?.find(w => w.name === "Lightness"); + this.saturationWidget = node.widgets?.find(w => w.name === "Saturation"); + this.dataWidget = node.widgets?.find(w => w.name === "jdata"); + + if (!this.colorWidget || !this.hueWidget || !this.lightnessWidget || + !this.saturationWidget || !this.dataWidget) { + console.warn("mdColorHueSaturation: some widgets missing"); + return; + } + + this.dataWidget.hidden = true; + + this.setupCallbacks(); + // Initial load after the node is fully ready + setTimeout(() => this.loadFromData(), 0); + } + + getCurrentRangeIndex() { + return this.rangeNames.indexOf(this.colorWidget.value); + } + + saveCurrentToData() { + let data; + try { + data = JSON.parse(this.dataWidget.value); + } catch { + data = this.rangeNames.map(() => [0, 0, 0]); + } + const idx = this.getCurrentRangeIndex(); + if (idx === -1) return; + data[idx] = [ + parseFloat(this.hueWidget.value), + parseFloat(this.lightnessWidget.value), + parseFloat(this.saturationWidget.value) + ]; + this.dataWidget.value = JSON.stringify(data); + if (this.dataWidget.callback) this.dataWidget.callback(this.dataWidget.value); + } + + loadFromData() { + let data; + try { + data = JSON.parse(this.dataWidget.value); + } catch { + data = this.rangeNames.map(() => [0, 0, 0]); + } + const idx = this.getCurrentRangeIndex(); + if (idx === -1) return; + const vals = data[idx] || [0, 0, 0]; + + // Update sliders using the shared utility + set_node_widget_value(this.node, "Hue", vals[0]); + set_node_widget_value(this.node, "Lightness", vals[1]); + set_node_widget_value(this.node, "Saturation", vals[2]); + + // Notify the widget callbacks so the UI updates + if (this.hueWidget.callback) this.hueWidget.callback(this.hueWidget.value); + if (this.lightnessWidget.callback) this.lightnessWidget.callback(this.lightnessWidget.value); + if (this.saturationWidget.callback) this.saturationWidget.callback(this.saturationWidget.value); + } + + setupCallbacks() { + // --- Color dropdown --- + const origColorCb = this.colorWidget.callback; + this.colorWidget.callback = (value) => { + // Save current slider values for the old range + this.saveCurrentToData(); + // Load slider values for the new range + this.loadFromData(); + // Invoke any existing callback + if (origColorCb) origColorCb(value); + }; + + // --- Slider widgets --- + const onSliderChange = () => { + this.saveCurrentToData(); + }; + + const wrapSliderCallback = (widget, extra) => { + const orig = widget.callback; + widget.callback = function(value) { + extra(); + if (orig) orig.call(this, value); + }; + }; + + wrapSliderCallback(this.hueWidget, onSliderChange); + wrapSliderCallback(this.lightnessWidget, onSliderChange); + wrapSliderCallback(this.saturationWidget, onSliderChange); + } +} +