diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 10aafb3c52c..1d5ba44b24c 100644 --- a/invokeai/app/invocations/image.py +++ b/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) diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index a7964ce12ae..c5782374794 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/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