clahe.py

from typing import Literal

import cv2
import numpy as np
from PIL import Image

from invokeai.invocation_api import (
    BaseInvocation,
    InputField,
    InvocationContext,
    invocation,
    ImageField,
    ImageOutput,
    WithMetadata,
    WithBoard
)

@invocation(
    "clahe_node",
    title="CLAHE Contrast Enhancement",
    tags=["image", "enhancement", "cv2", "contrast"],
    category="image",
    version="1.3.0",
)
class CLAHEInvocation(BaseInvocation, WithMetadata, WithBoard):
    """
    Enhances image contrast using Contrast Limited Adaptive Histogram Equalization (CLAHE).
    This node processes the Luminance (L) channel in LAB color space to preserve colors.
    Includes a blend 'amount' to control the effect strength.
    """

    # --- Inputs ---
    image: ImageField = InputField(description="The input image for contrast optimization.")
    
    clip_limit: float = InputField(
        default=2.0, 
        description="Threshold for contrast limiting (higher values = more contrast). Minimum sensible value is 1.0."
    )
    
    grid_size: int = InputField(
        default=8, 
        description="Size of the grid for histogram equalization (e.g., 8 creates an 8x8 tile grid)."
    )
    
    amount: float = InputField(
        default=1.0,
        description="Strength of the effect (0.0 = original image, 1.0 = full CLAHE)."
    )

    def invoke(self, context: InvocationContext) -> ImageOutput:
        # 1. Load the original image from the Invoke system
        original_pil = context.images.get_pil(self.image.image_name)
        
        # Ensure we are working with standard RGB (removes alpha if present for OpenCV)
        original_pil = original_pil.convert("RGB")
        
        # 2. Convert PIL to OpenCV format (numpy array)
        image_np = np.array(original_pil)
        
        # 3. Perform CLAHE processing in LAB space
        # Convert RGB to LAB space. LAB separates luminance (L) from color (A, B).
        # We apply CLAHE only to the L-channel to avoid shifting colors.
        lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
        l_channel, a_channel, b_channel = cv2.split(lab)
        
        # Create the CLAHE object
        clahe = cv2.createCLAHE(
            clipLimit=self.clip_limit, 
            tileGridSize=(self.grid_size, self.grid_size)
        )
        
        # Apply CLAHE to the L-channel
        cl_processed = clahe.apply(l_channel)
        
        # Merge the processed L-channel back with the original color channels
        limg_merged = cv2.merge((cl_processed, a_channel, b_channel))
        
        # Convert back from LAB space to standard RGB
        processed_image_np = cv2.cvtColor(limg_merged, cv2.COLOR_LAB2RGB)
        
        # 4. Convert back to PIL Image
        processed_pil = Image.fromarray(processed_image_np)
        
        # 5. Blend original and processed image based on 'amount'
        # original_pil (0.0) <---> processed_pil (1.0)
        final_pil = processed_pil # Default to full effect
        
        if self.amount == 0.0:
            final_pil = original_pil
        elif self.amount < 1.0:
            # Linear interpolation: original_pil * (1.0 - amount) + processed_pil * amount
            final_pil = Image.blend(original_pil, processed_pil, self.amount)
        
        # 6. Save the final image and return the output
        # (This also handles board assignment and metadata preservation)
        image_dto = context.images.save(image=final_pil)
        
        return ImageOutput(
            image=ImageField(image_name=image_dto.image_name),
            width=image_dto.width,
            height=image_dto.height,
        )