Fix(nodes): color correct invocation

invokeai/app/invocations/image.py

@@ -649,102 +649,104 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
     title="Color Correct",
     tags=["image", "color"],
     category="image",
-    version="1.2.2",
+    version="2.0.0",
 )
 class ColorCorrectInvocation(BaseInvocation, WithMetadata, WithBoard):
     """
-    Shifts the colors of a target image to match the reference image, optionally
-    using a mask to only color-correct certain regions of the target image.
+    Matches the color histogram of a base image to a reference image, optionally
+    using a mask to only color-correct certain regions of the base image.
     """
 
-    image: ImageField = InputField(description="The image to color-correct")
-    reference: ImageField = InputField(description="Reference image for color-correction")
-    mask: Optional[ImageField] = InputField(default=None, description="Mask to use when applying color-correction")
-    mask_blur_radius: float = InputField(default=8, description="Mask blur radius")
+    base_image: ImageField = InputField(description="The image to color-correct")
+    color_reference: ImageField = InputField(description="Reference image for color-correction")
+    mask: Optional[ImageField] = InputField(default=None, description="Optional mask to limit color correction area")
+    colorspace: Literal["RGB", "YCbCr", "YCbCr-Chroma", "YCbCr-Luma"] = InputField(
+        default="RGB", description="Colorspace in which to apply histogram matching", title="Color Space"
+    )
+
+    def _match_histogram_channel(self, source: numpy.ndarray, reference: numpy.ndarray) -> numpy.ndarray:
+        """Match histogram of source channel to reference channel using cumulative distribution functions."""
+        # Compute histograms
+        source_hist, _ = numpy.histogram(source.flatten(), bins=256, range=(0, 256))
+        reference_hist, _ = numpy.histogram(reference.flatten(), bins=256, range=(0, 256))
+
+        # Compute cumulative distribution functions
+        source_cdf = source_hist.cumsum()
+        reference_cdf = reference_hist.cumsum()
+
+        # Normalize CDFs (avoid division by zero)
+        if source_cdf[-1] > 0:
+            source_cdf = source_cdf / source_cdf[-1]
+        if reference_cdf[-1] > 0:
+            reference_cdf = reference_cdf / reference_cdf[-1]
+
+        # Create lookup table using linear interpolation
+        lookup_table = numpy.interp(source_cdf, reference_cdf, numpy.arange(256))
+
+        # Apply lookup table to source image
+        return lookup_table[source].astype(numpy.uint8)
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        pil_init_mask = None
-        if self.mask is not None:
-            pil_init_mask = context.images.get_pil(self.mask.image_name).convert("L")
-
-        init_image = context.images.get_pil(self.reference.image_name)
-
-        result = context.images.get_pil(self.image.image_name).convert("RGBA")
-
-        # if init_image is None or init_mask is None:
-        #    return result
-
-        # Get the original alpha channel of the mask if there is one.
-        # Otherwise it is some other black/white image format ('1', 'L' or 'RGB')
-        # pil_init_mask = (
-        #    init_mask.getchannel("A")
-        #    if init_mask.mode == "RGBA"
-        #    else init_mask.convert("L")
-        # )
-        pil_init_image = init_image.convert("RGBA")  # Add an alpha channel if one doesn't exist
-
-        # Build an image with only visible pixels from source to use as reference for color-matching.
-        init_rgb_pixels = numpy.asarray(init_image.convert("RGB"), dtype=numpy.uint8)
-        init_a_pixels = numpy.asarray(pil_init_image.getchannel("A"), dtype=numpy.uint8)
-        init_mask_pixels = numpy.asarray(pil_init_mask, dtype=numpy.uint8)
-
-        # Get numpy version of result
-        np_image = numpy.asarray(result.convert("RGB"), dtype=numpy.uint8)
-
-        # Mask and calculate mean and standard deviation
-        mask_pixels = init_a_pixels * init_mask_pixels > 0
-        np_init_rgb_pixels_masked = init_rgb_pixels[mask_pixels, :]
-        np_image_masked = np_image[mask_pixels, :]
-
-        if np_init_rgb_pixels_masked.size > 0:
-            init_means = np_init_rgb_pixels_masked.mean(axis=0)
-            init_std = np_init_rgb_pixels_masked.std(axis=0)
-            gen_means = np_image_masked.mean(axis=0)
-            gen_std = np_image_masked.std(axis=0)
-
-            # Color correct
-            np_matched_result = np_image.copy()
-            np_matched_result[:, :, :] = (
-                (
-                    (
-                        (np_matched_result[:, :, :].astype(numpy.float32) - gen_means[None, None, :])
-                        / gen_std[None, None, :]
-                    )
-                    * init_std[None, None, :]
-                    + init_means[None, None, :]
-                )
-                .clip(0, 255)
-                .astype(numpy.uint8)
-            )
-            matched_result = Image.fromarray(np_matched_result, mode="RGB")
+        # Load images as RGBA
+        base_image = context.images.get_pil(self.base_image.image_name, "RGBA")
+
+        # Store original alpha channel
+        original_alpha = base_image.getchannel("A")
+
+        # Convert to working colorspace
+        if self.colorspace == "RGB":
+            base_array = numpy.asarray(base_image.convert("RGB"), dtype=numpy.uint8)
+            ref_rgb = context.images.get_pil(self.color_reference.image_name, "RGB")
+            ref_array = numpy.asarray(ref_rgb, dtype=numpy.uint8)
+            channels_to_match = [0, 1, 2]  # R, G, B
         else:
-            matched_result = Image.fromarray(np_image, mode="RGB")
-
-        # Blur the mask out (into init image) by specified amount
-        if self.mask_blur_radius > 0:
-            nm = numpy.asarray(pil_init_mask, dtype=numpy.uint8)
-            inverted_nm = 255 - nm
-            dilation_size = int(round(self.mask_blur_radius) + 20)
-            dilating_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_size, dilation_size))
-            inverted_dilated_nm = cv2.dilate(inverted_nm, dilating_kernel)
-            dilated_nm = 255 - inverted_dilated_nm
-            nmd = cv2.erode(
-                dilated_nm,
-                kernel=numpy.ones((3, 3), dtype=numpy.uint8),
-                iterations=int(self.mask_blur_radius / 2),
+            # Convert to YCbCr colorspace
+            base_ycbcr = base_image.convert("YCbCr")
+            ref_ycbcr = context.images.get_pil(self.color_reference.image_name, "YCbCr")
+
+            base_array = numpy.asarray(base_ycbcr, dtype=numpy.uint8)
+            ref_array = numpy.asarray(ref_ycbcr, dtype=numpy.uint8)
+
+            # Determine which channels to match based on mode
+            if self.colorspace == "YCbCr":
+                channels_to_match = [0, 1, 2]  # Y, Cb, Cr
+            elif self.colorspace == "YCbCr-Chroma":
+                channels_to_match = [1, 2]  # Cb, Cr only
+            else:  # YCbCr-Luma
+                channels_to_match = [0]  # Y only
+
+        # Apply histogram matching to selected channels
+        corrected_array = base_array.copy()
+        for channel_idx in channels_to_match:
+            corrected_array[:, :, channel_idx] = self._match_histogram_channel(
+                base_array[:, :, channel_idx], ref_array[:, :, channel_idx]
             )
-            pmd = Image.fromarray(nmd, mode="L")
-            blurred_init_mask = pmd.filter(ImageFilter.BoxBlur(self.mask_blur_radius))
-        else:
-            blurred_init_mask = pil_init_mask
 
-        multiplied_blurred_init_mask = ImageChops.multiply(blurred_init_mask, result.split()[-1])
+        # Convert back to RGB if we were in YCbCr
+        if self.colorspace != "RGB":
+            corrected_image = Image.fromarray(corrected_array, mode="YCbCr").convert("RGB")
+        else:
+            corrected_image = Image.fromarray(corrected_array, mode="RGB")
 
-        # Paste original on color-corrected generation (using blurred mask)
-        matched_result.paste(init_image, (0, 0), mask=multiplied_blurred_init_mask)
+        # Apply mask if provided (white = original, black = result)
+        if self.mask is not None:
+            # Load mask as grayscale
+            mask_image = context.images.get_pil(self.mask.image_name, "L")
+            # Start with corrected image, paste base image where mask is white
+            result = corrected_image.copy()
+            if mask_image.size != result.size:
+                raise ValueError("Mask size must match base image size.")
+            else:
+                result.paste(base_image.convert("RGB"), mask=mask_image)
+        else:
+            result = corrected_image
 
-        image_dto = context.images.save(image=matched_result)
+        # Convert to RGBA and restore original alpha
+        result = result.convert("RGBA")
+        result.putalpha(original_alpha)
 
+        # Save and return
+        image_dto = context.images.save(image=result)
         return ImageOutput.build(image_dto)
 
 

invokeai/frontend/web/src/services/api/schema.ts

@@ -4776,8 +4776,8 @@ export type components = {
         };
         /**
          * Color Correct
-         * @description Shifts the colors of a target image to match the reference image, optionally
-         *     using a mask to only color-correct certain regions of the target image.
+         * @description Matches the color histogram of a base image to a reference image, optionally
+         *     using a mask to only color-correct certain regions of the base image.
          */
         ColorCorrectInvocation: {
             /**
@@ -4811,23 +4811,24 @@ export type components = {
              * @description The image to color-correct
              * @default null
              */
-            image?: components["schemas"]["ImageField"] | null;
+            base_image?: components["schemas"]["ImageField"] | null;
             /**
              * @description Reference image for color-correction
              * @default null
              */
-            reference?: components["schemas"]["ImageField"] | null;
+            color_reference?: components["schemas"]["ImageField"] | null;
             /**
-             * @description Mask to use when applying color-correction
+             * @description Optional mask to limit color correction area
              * @default null
              */
             mask?: components["schemas"]["ImageField"] | null;
             /**
-             * Mask Blur Radius
-             * @description Mask blur radius
-             * @default 8
+             * Color Space
+             * @description Colorspace in which to apply histogram matching
+             * @default RGB
+             * @enum {string}
              */
-            mask_blur_radius?: number;
+            colorspace?: "RGB" | "YCbCr" | "YCbCr-Chroma" | "YCbCr-Luma";
             /**
              * type
              * @default color_correct