app.py

import gradio as gr
from PIL import Image
import numpy as np
import cv2
from lang_sam import LangSAM
from color_matcher import ColorMatcher
from color_matcher.normalizer import Normalizer
import torch
import warnings

# Suppress specific warnings if desired
warnings.filterwarnings("ignore", category=UserWarning)

# Device configuration: Use CUDA if available, else CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load the LangSAM model
model = LangSAM()  # Use the default model or specify custom checkpoint if necessary
# Note: Removed model.to(device) since LangSAM does not support it

def extract_masks(image_pil, prompts):
    """
    Extracts masks for each prompt using the LangSAM model.

    Args:
        image_pil (PIL.Image): The input image.
        prompts (str): Comma-separated prompts for segmentation.

    Returns:
        dict: A dictionary mapping each prompt to its corresponding binary mask.
    """
    prompts_list = [p.strip() for p in prompts.split(',') if p.strip()]
    masks_dict = {}
    with torch.no_grad():  # Disable gradient computation for inference
        for prompt in prompts_list:
            # Ensure the model uses the correct device internally
            masks, boxes, phrases, logits = model.predict(image_pil, prompt)
            if masks is not None and len(masks) > 0:
                # Move masks to CPU and convert to numpy
                masks_np = masks[0].cpu().numpy()
                mask = (masks_np > 0).astype(np.uint8) * 255  # Binary mask
                masks_dict[prompt] = mask
    return masks_dict

def apply_color_matching(source_img_np, ref_img_np):
    """
    Applies color matching from the reference image to the source image.

    Args:
        source_img_np (numpy.ndarray): Source image in NumPy array format.
        ref_img_np (numpy.ndarray): Reference image in NumPy array format.

    Returns:
        numpy.ndarray: Color-matched image.
    """
    # Initialize ColorMatcher
    cm = ColorMatcher()

    # Apply color matching
    img_res = cm.transfer(src=source_img_np, ref=ref_img_np, method='mkl')

    # Normalize the result
    img_res = Normalizer(img_res).uint8_norm()

    return img_res

def process_image(current_image_pil, selected_prompt, masks_dict, replacement_image_pil, color_ref_image_pil, apply_replacement, apply_color_grading, apply_color_to_full_image, blending_amount, image_history):
    """
    Processes the image by applying replacement and/or color grading based on user input.

    Args:
        current_image_pil (PIL.Image): The current image to be edited.
        selected_prompt (str): The selected segment prompt.
        masks_dict (dict): Dictionary of masks for each prompt.
        replacement_image_pil (PIL.Image): Replacement image (optional).
        color_ref_image_pil (PIL.Image): Color reference image (optional).
        apply_replacement (bool): Flag to apply replacement.
        apply_color_grading (bool): Flag to apply color grading.
        apply_color_to_full_image (bool): Flag to apply color grading to the full image.
        blending_amount (int): Amount for blending the mask.
        image_history (list): History of images for undo functionality.

    Returns:
        tuple: Updated image, status message, updated history, and image display.
    """
    # Check if current_image_pil is None
    if current_image_pil is None:
        return None, "No current image to edit.", image_history, None

    if not apply_replacement and not apply_color_grading:
        return current_image_pil, "No changes applied. Please select at least one operation.", image_history, current_image_pil

    if apply_replacement and replacement_image_pil is None:
        return current_image_pil, "Replacement image not provided.", image_history, current_image_pil

    if apply_color_grading and color_ref_image_pil is None:
        return current_image_pil, "Color reference image not provided.", image_history, current_image_pil

    # Get the mask from masks_dict
    if selected_prompt not in masks_dict:
        return current_image_pil, f"No mask available for selected segment: {selected_prompt}", image_history, current_image_pil

    mask = masks_dict[selected_prompt]

    # Save current image to history for undo
    if image_history is None:
        image_history = []
    image_history.append(current_image_pil.copy())

    # Proceed with replacement or color matching
    current_image_np = np.array(current_image_pil)
    result_image_np = current_image_np.copy()

    # Create mask with blending
    # First, normalize mask to range [0,1]
    mask_normalized = mask.astype(np.float32) / 255.0

    # Apply blending by blurring the mask
    if blending_amount > 0:
        # The kernel size for blurring; larger blending_amount means more blur
        kernel_size = int(blending_amount)
        if kernel_size % 2 == 0:
            kernel_size += 1  # Kernel size must be odd
        mask_blurred = cv2.GaussianBlur(mask_normalized, (kernel_size, kernel_size), 0)
    else:
        mask_blurred = mask_normalized

    # Convert mask to 3 channels
    mask_blurred_3ch = cv2.merge([mask_blurred, mask_blurred, mask_blurred])

    # If apply replacement
    if apply_replacement:
        # Resize replacement image to match current image
        replacement_image_resized = replacement_image_pil.resize(current_image_pil.size)
        replacement_image_np = np.array(replacement_image_resized)

        # Blend the replacement image with the current image using the mask
        result_image_np = (replacement_image_np.astype(np.float32) * mask_blurred_3ch + result_image_np.astype(np.float32) * (1 - mask_blurred_3ch)).astype(np.uint8)

    # If apply color grading
    if apply_color_grading:
        # Convert color reference image to numpy
        color_ref_image_np = np.array(color_ref_image_pil)

        if apply_color_to_full_image:
            # Apply color matching to the full image
            color_matched_image = apply_color_matching(result_image_np, color_ref_image_np)
            result_image_np = color_matched_image
        else:
            # Apply color matching only to the masked area
            # Extract the masked area
            masked_region = (result_image_np.astype(np.float32) * mask_blurred_3ch).astype(np.uint8)
            # Apply color matching
            color_matched_region = apply_color_matching(masked_region, color_ref_image_np)
            # Blend the color matched region back into the result image
            result_image_np = (color_matched_region.astype(np.float32) * mask_blurred_3ch + result_image_np.astype(np.float32) * (1 - mask_blurred_3ch)).astype(np.uint8)

    # Convert result back to PIL Image
    result_image_pil = Image.fromarray(result_image_np)

    # Update current_image_pil
    current_image_pil = result_image_pil

    return current_image_pil, f"Applied changes to '{selected_prompt}'", image_history, current_image_pil

def undo(image_history):
    """
    Undoes the last image edit by reverting to the previous image in the history.

    Args:
        image_history (list): History of images.

    Returns:
        tuple: Reverted image, updated history, and image display.
    """
    if image_history and len(image_history) > 1:
        # Pop the last image
        image_history.pop()
        # Return the previous image
        current_image_pil = image_history[-1]
        return current_image_pil, image_history, current_image_pil
    elif image_history and len(image_history) == 1:
        current_image_pil = image_history[0]
        return current_image_pil, image_history, current_image_pil
    else:
        # Cannot undo
        return None, [], None

def gradio_interface():
    """
    Defines and launches the Gradio interface for continuous image editing.
    """
    with gr.Blocks() as demo:
        # Define the state variables
        image_history = gr.State([])
        current_image_pil = gr.State(None)
        masks_dict = gr.State({})  # Store masks for each prompt

        gr.Markdown("## Continuous Image Editing with LangSAM")

        with gr.Row():
            with gr.Column():
                initial_image = gr.Image(type="pil", label="Upload Image")
                prompts = gr.Textbox(lines=1, placeholder="Enter prompts separated by commas (e.g., sky, grass)", label="Prompts")
                segment_button = gr.Button("Segment Image")
                segment_dropdown = gr.Dropdown(label="Select Segment", choices=[], allow_custom_value=True)
                replacement_image = gr.Image(type="pil", label="Replacement Image (optional)")
                color_ref_image = gr.Image(type="pil", label="Color Reference Image (optional)")
                apply_replacement = gr.Checkbox(label="Apply Replacement", value=False)
                apply_color_grading = gr.Checkbox(label="Apply Color Grading", value=False)
                apply_color_to_full_image = gr.Checkbox(label="Apply Color Correction to Full Image", value=False)
                blending_amount = gr.Slider(minimum=0, maximum=500, step=1, label="Blending Amount", value=150)
                apply_button = gr.Button("Apply Changes")
                undo_button = gr.Button("Undo")
            with gr.Column():
                current_image_display = gr.Image(type="pil", label="Edited Image", interactive=False)
                status = gr.Textbox(lines=2, interactive=False, label="Status")

        def initialize_image(initial_image_pil):
            """
            Initializes the image history and sets up the initial image.

            Args:
                initial_image_pil (PIL.Image): The uploaded initial image.

            Returns:
                tuple: Updated states and status message.
            """
            if initial_image_pil is not None:
                image_history = [initial_image_pil]
                current_image_pil = initial_image_pil
                return current_image_pil, image_history, initial_image_pil, {}, gr.update(choices=[], value=None), "Image loaded."
            else:
                return None, [], None, {}, gr.update(choices=[], value=None), "No image loaded."

        # When the initial image is uploaded, initialize the image history
        initial_image.upload(
            fn=initialize_image,
            inputs=initial_image,
            outputs=[current_image_pil, image_history, current_image_display, masks_dict, segment_dropdown, status]
        )

        # Segment button click
        def segment_image_wrapper(current_image_pil, prompts):
            """
            Handles the segmentation of the image based on user prompts.

            Args:
                current_image_pil (PIL.Image): The current image.
                prompts (str): Comma-separated prompts.

            Returns:
                tuple: Status message, updated masks, and dropdown updates.
            """
            if current_image_pil is None:
                return "No image uploaded.", {}, gr.update(choices=[], value=None)
            masks = extract_masks(current_image_pil, prompts)
            if not masks:
                return "No masks detected for the given prompts.", {}, gr.update(choices=[], value=None)
            dropdown_choices = list(masks.keys())
            return "Segmentation completed.", masks, gr.update(choices=dropdown_choices, value=dropdown_choices[0])

        segment_button.click(
            fn=segment_image_wrapper,
            inputs=[current_image_pil, prompts],
            outputs=[status, masks_dict, segment_dropdown]
        )

        # Apply button click
        apply_button.click(
            fn=process_image,
            inputs=[
                current_image_pil,
                segment_dropdown,
                masks_dict,
                replacement_image,
                color_ref_image,
                apply_replacement,
                apply_color_grading,
                apply_color_to_full_image,
                blending_amount,
                image_history
            ],
            outputs=[current_image_pil, status, image_history, current_image_display]
        )

        # Undo button click
        undo_button.click(
            fn=undo,
            inputs=image_history,
            outputs=[current_image_pil, image_history, current_image_display]
        )

    demo.launch(share=True)

# Run the Gradio Interface
if __name__ == "__main__":
    gradio_interface()