Add OKLab and Oklch image utilities and nodes

Author: JPPhoto

invokeai/app/invocations/image.py

@@ -25,6 +25,14 @@
 from invokeai.app.services.image_records.image_records_common import ImageCategory
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import SEED_MAX
+from invokeai.backend.image_util.color_conversion import (
+    linear_srgb_from_oklab,
+    linear_srgb_from_oklch,
+    linear_srgb_from_srgb,
+    oklab_from_linear_srgb,
+    oklch_from_oklab,
+    srgb_from_linear_srgb,
+)
 from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
 from invokeai.backend.image_util.safety_checker import SafetyChecker
 
@@ -397,6 +405,51 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         )
 
 
+@invocation(
+    "unsharp_mask_oklab",
+    title="Unsharp Mask (Oklab)",
+    tags=["image", "unsharp_mask", "oklab"],
+    category="image",
+    version="1.0.0",
+)
+class OklabUnsharpMaskInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Applies an unsharp mask filter to an image in the Oklab color space"""
+
+    image: ImageField = InputField(description="The image to use")
+    radius: float = InputField(gt=0, description="Unsharp mask radius", default=2)
+    strength: float = InputField(ge=0, description="Unsharp mask strength", default=50)
+
+    def pil_from_array(self, arr: numpy.ndarray) -> Image.Image:
+        return Image.fromarray((numpy.clip(arr, 0.0, 1.0) * 255).astype("uint8"))
+
+    def array_from_pil(self, img: Image.Image) -> numpy.ndarray:
+        return numpy.array(img, dtype=numpy.float32) / 255.0
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+        mode = image.mode
+
+        alpha_channel = image.getchannel("A") if mode == "RGBA" else None
+        image = image.convert("RGB")
+
+        image_blurred = self.array_from_pil(image.filter(ImageFilter.GaussianBlur(radius=self.radius)))
+        image_arr = self.array_from_pil(image)
+
+        image_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(image_arr))
+        image_blurred_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(image_blurred))
+
+        image_oklab += (image_oklab - image_blurred_oklab) * (self.strength / 100.0)
+        image_oklab = numpy.clip(image_oklab, -1.0, 1.0)
+
+        image = self.pil_from_array(srgb_from_linear_srgb(linear_srgb_from_oklab(image_oklab))).convert(mode)
+
+        if alpha_channel is not None:
+            image.putalpha(alpha_channel)
+
+        image_dto = context.images.save(image=image)
+        return ImageOutput.build(image_dto)
+
+
 PIL_RESAMPLING_MODES = Literal[
     "nearest",
     "box",
@@ -802,6 +855,40 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         return ImageOutput.build(image_dto)
 
 
+@invocation(
+    "img_hue_adjust_oklch",
+    title="Adjust Image Hue (Oklch)",
+    tags=["image", "hue", "oklch"],
+    category="image",
+    version="1.0.0",
+)
+class OklchImageHueAdjustmentInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Adjusts the hue of an image in Oklch space."""
+
+    image: ImageField = InputField(description="The image to adjust")
+    hue: int = InputField(default=0, description="The degrees by which to rotate the hue, 0-360")
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+        mode = image.mode
+        alpha_channel = image.getchannel("A") if mode == "RGBA" else None
+
+        rgb = numpy.asarray(image.convert("RGB"), dtype=numpy.float32) / 255.0
+        oklch = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(rgb)))
+        oklch[..., 2] = (oklch[..., 2] + self.hue) % 360.0
+
+        image = Image.fromarray(
+            numpy.clip(srgb_from_linear_srgb(linear_srgb_from_oklch(oklch)) * 255.0, 0.0, 255.0).astype(numpy.uint8),
+            mode="RGB",
+        ).convert(mode)
+
+        if alpha_channel is not None:
+            image.putalpha(alpha_channel)
+
+        image_dto = context.images.save(image=image)
+        return ImageOutput.build(image_dto)
+
+
 COLOR_CHANNELS = Literal[
     "Red (RGBA)",
     "Green (RGBA)",

invokeai/backend/image_util/color_conversion.py

@@ -0,0 +1,70 @@
+import numpy
+
+
+def linear_srgb_from_srgb(srgb: numpy.ndarray) -> numpy.ndarray:
+    return numpy.where(srgb <= 0.0404482362771082, srgb / 12.92, ((srgb + 0.055) / 1.055) ** 2.4)
+
+
+def srgb_from_linear_srgb(linear_srgb: numpy.ndarray) -> numpy.ndarray:
+    linear_srgb = numpy.clip(linear_srgb, 0.0, 1.0)
+    return numpy.where(
+        linear_srgb <= 0.0031308,
+        linear_srgb * 12.92,
+        1.055 * numpy.power(linear_srgb, 1.0 / 2.4) - 0.055,
+    )
+
+
+def oklab_from_linear_srgb(linear_srgb: numpy.ndarray) -> numpy.ndarray:
+    lms_l = 0.4122214708 * linear_srgb[..., 0] + 0.5363325363 * linear_srgb[..., 1] + 0.0514459929 * linear_srgb[..., 2]
+    lms_m = 0.2119034982 * linear_srgb[..., 0] + 0.6806995451 * linear_srgb[..., 1] + 0.1073969566 * linear_srgb[..., 2]
+    lms_s = 0.0883024619 * linear_srgb[..., 0] + 0.2817188376 * linear_srgb[..., 1] + 0.6299787005 * linear_srgb[..., 2]
+
+    lms_l_cbrt = numpy.cbrt(lms_l)
+    lms_m_cbrt = numpy.cbrt(lms_m)
+    lms_s_cbrt = numpy.cbrt(lms_s)
+
+    return numpy.stack(
+        [
+            0.2104542553 * lms_l_cbrt + 0.7936177850 * lms_m_cbrt - 0.0040720468 * lms_s_cbrt,
+            1.9779984951 * lms_l_cbrt - 2.4285922050 * lms_m_cbrt + 0.4505937099 * lms_s_cbrt,
+            0.0259040371 * lms_l_cbrt + 0.7827717662 * lms_m_cbrt - 0.8086757660 * lms_s_cbrt,
+        ],
+        axis=-1,
+    )
+
+
+def linear_srgb_from_oklab(oklab: numpy.ndarray) -> numpy.ndarray:
+    lms_l_cbrt = oklab[..., 0] + 0.3963377774 * oklab[..., 1] + 0.2158037573 * oklab[..., 2]
+    lms_m_cbrt = oklab[..., 0] - 0.1055613458 * oklab[..., 1] - 0.0638541728 * oklab[..., 2]
+    lms_s_cbrt = oklab[..., 0] - 0.0894841775 * oklab[..., 1] - 1.2914855480 * oklab[..., 2]
+
+    lms_l = lms_l_cbrt**3
+    lms_m = lms_m_cbrt**3
+    lms_s = lms_s_cbrt**3
+
+    return numpy.stack(
+        [
+            4.0767416621 * lms_l - 3.3077115913 * lms_m + 0.2309699292 * lms_s,
+            -1.2684380046 * lms_l + 2.6097574011 * lms_m - 0.3413193965 * lms_s,
+            -0.0041960863 * lms_l - 0.7034186147 * lms_m + 1.7076147010 * lms_s,
+        ],
+        axis=-1,
+    )
+
+
+def oklch_from_oklab(oklab: numpy.ndarray) -> numpy.ndarray:
+    lightness = oklab[..., 0]
+    chroma = numpy.sqrt(oklab[..., 1] ** 2 + oklab[..., 2] ** 2)
+    hue = numpy.degrees(numpy.arctan2(oklab[..., 2], oklab[..., 1])) % 360.0
+    return numpy.stack([lightness, chroma, hue], axis=-1)
+
+
+def oklab_from_oklch(oklch: numpy.ndarray) -> numpy.ndarray:
+    hue_radians = numpy.radians(oklch[..., 2])
+    a_channel = oklch[..., 1] * numpy.cos(hue_radians)
+    b_channel = oklch[..., 1] * numpy.sin(hue_radians)
+    return numpy.stack([oklch[..., 0], a_channel, b_channel], axis=-1)
+
+
+def linear_srgb_from_oklch(oklch: numpy.ndarray) -> numpy.ndarray:
+    return linear_srgb_from_oklab(oklab_from_oklch(oklch))

invokeai/invocation_api/__init__.py

@@ -80,6 +80,15 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 from invokeai.app.util.misc import SEED_MAX, get_random_seed
+from invokeai.backend.image_util.color_conversion import (
+    linear_srgb_from_oklab,
+    linear_srgb_from_oklch,
+    linear_srgb_from_srgb,
+    oklab_from_linear_srgb,
+    oklab_from_oklch,
+    oklch_from_oklab,
+    srgb_from_linear_srgb,
+)
 from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.model_manager.taxonomy import (
@@ -207,6 +216,14 @@
     # invokeai.app.util.misc
     "SEED_MAX",
     "get_random_seed",
+    # invokeai.backend.image_util.color_conversion
+    "linear_srgb_from_srgb",
+    "srgb_from_linear_srgb",
+    "oklab_from_linear_srgb",
+    "linear_srgb_from_oklab",
+    "oklch_from_oklab",
+    "oklab_from_oklch",
+    "linear_srgb_from_oklch",
     # invokeai.backend.model_manager.taxonomy
     "BaseModelType",
     "ModelType",

tests/app/invocations/test_image.py

@@ -0,0 +1,102 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy
+from PIL import Image, ImageFilter
+
+from invokeai.app.invocations.image import ImageField, OklabUnsharpMaskInvocation, OklchImageHueAdjustmentInvocation
+from invokeai.backend.image_util.color_conversion import (
+    linear_srgb_from_oklab,
+    linear_srgb_from_oklch,
+    linear_srgb_from_srgb,
+    oklab_from_linear_srgb,
+    oklch_from_oklab,
+    srgb_from_linear_srgb,
+)
+
+
+def test_oklab_unsharp_mask_invocation_preserves_alpha_and_sharpens_in_oklab() -> None:
+    input_image = Image.new("RGBA", (3, 1))
+    input_image.putdata(
+        [
+            (255, 0, 0, 32),
+            (0, 255, 0, 128),
+            (0, 0, 255, 224),
+        ]
+    )
+
+    context = MagicMock()
+    context.images.get_pil.return_value = input_image
+    context.images.save.side_effect = lambda image: SimpleNamespace(
+        image_name="out", width=image.width, height=image.height
+    )
+
+    invocation = OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=1.0, strength=50.0)
+    output = invocation.invoke(context)
+    saved_image = context.images.save.call_args.kwargs["image"]
+
+    assert output.image.image_name == "out"
+    assert output.width == 3
+    assert output.height == 1
+    assert numpy.asarray(saved_image.getchannel("A")).reshape(-1).tolist() == [32, 128, 224]
+
+    rgb = numpy.asarray(input_image.convert("RGB"), dtype=numpy.float32) / 255.0
+    blurred_rgb = (
+        numpy.asarray(input_image.convert("RGB").filter(ImageFilter.GaussianBlur(radius=1.0)), dtype=numpy.float32)
+        / 255.0
+    )
+
+    rgb_unsharp = numpy.clip(rgb + (rgb - blurred_rgb) * 0.5, 0.0, 1.0)
+    oklab_unsharp = srgb_from_linear_srgb(
+        linear_srgb_from_oklab(
+            numpy.clip(
+                oklab_from_linear_srgb(linear_srgb_from_srgb(rgb))
+                + (
+                    oklab_from_linear_srgb(linear_srgb_from_srgb(rgb))
+                    - oklab_from_linear_srgb(linear_srgb_from_srgb(blurred_rgb))
+                )
+                * 0.5,
+                -1.0,
+                1.0,
+            )
+        )
+    )
+
+    assert not numpy.allclose(oklab_unsharp, rgb_unsharp, atol=1e-3)
+    assert numpy.allclose(
+        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0, oklab_unsharp, atol=1 / 255.0
+    )
+
+
+def test_oklch_hue_adjustment_invocation_preserves_alpha_and_rotates_hue_in_oklch() -> None:
+    input_image = Image.new("RGBA", (2, 1))
+    input_image.putdata(
+        [
+            (210, 80, 30, 64),
+            (40, 160, 220, 192),
+        ]
+    )
+
+    context = MagicMock()
+    context.images.get_pil.return_value = input_image
+    context.images.save.side_effect = lambda image: SimpleNamespace(
+        image_name="out", width=image.width, height=image.height
+    )
+
+    invocation = OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=180)
+    output = invocation.invoke(context)
+    saved_image = context.images.save.call_args.kwargs["image"]
+
+    rgb = numpy.asarray(input_image.convert("RGB"), dtype=numpy.float32) / 255.0
+    oklch = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(rgb)))
+    rotated_oklch = oklch.copy()
+    rotated_oklch[..., 2] = (rotated_oklch[..., 2] + 180.0) % 360.0
+    expected_rgb = srgb_from_linear_srgb(linear_srgb_from_oklch(rotated_oklch))
+
+    assert output.image.image_name == "out"
+    assert output.width == 2
+    assert output.height == 1
+    assert numpy.asarray(saved_image.getchannel("A")).reshape(-1).tolist() == [64, 192]
+    assert numpy.allclose(
+        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0, expected_rgb, atol=1 / 255.0
+    )

tests/backend/image_util/test_color_conversion.py

@@ -0,0 +1,61 @@
+import numpy
+
+from invokeai.backend.image_util.color_conversion import (
+    linear_srgb_from_oklab,
+    linear_srgb_from_oklch,
+    linear_srgb_from_srgb,
+    oklab_from_linear_srgb,
+    oklab_from_oklch,
+    oklch_from_oklab,
+    srgb_from_linear_srgb,
+)
+
+
+def test_srgb_oklab_round_trip() -> None:
+    srgb = numpy.array(
+        [
+            [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]],
+            [[1.0, 0.0, 0.0], [0.1, 0.6, 0.9]],
+        ],
+        dtype=numpy.float32,
+    )
+
+    round_tripped = srgb_from_linear_srgb(linear_srgb_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(srgb))))
+
+    assert numpy.allclose(round_tripped, srgb, atol=1e-5)
+
+
+def test_oklab_from_pure_srgb_red_matches_reference_value() -> None:
+    srgb_red = numpy.array([[[1.0, 0.0, 0.0]]], dtype=numpy.float32)
+
+    oklab_red = oklab_from_linear_srgb(linear_srgb_from_srgb(srgb_red))
+
+    assert numpy.allclose(oklab_red[0, 0], [0.62795536, 0.22486306, 0.1258463], atol=1e-6)
+
+
+def test_oklab_oklch_round_trip() -> None:
+    oklab = numpy.array(
+        [
+            [[0.6, 0.2, 0.1], [0.4, -0.1, 0.05]],
+        ],
+        dtype=numpy.float32,
+    )
+
+    round_tripped = oklab_from_oklch(oklch_from_oklab(oklab))
+
+    assert numpy.allclose(round_tripped, oklab, atol=1e-6)
+
+
+def test_srgb_oklch_round_trip() -> None:
+    srgb = numpy.array(
+        [
+            [[0.2, 0.4, 0.8], [0.9, 0.3, 0.1]],
+        ],
+        dtype=numpy.float32,
+    )
+
+    round_tripped = srgb_from_linear_srgb(
+        linear_srgb_from_oklch(oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(srgb))))
+    )
+
+    assert numpy.allclose(round_tripped, srgb, atol=1e-5)