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

# Load the LangSAM model
model = LangSAM()  # Use the default model or specify custom checkpoint if necessary

def extract_masks(image_pil, prompts):
    prompts_list = [p.strip() for p in prompts.split(',') if p.strip()]
    masks_dict = {}
    for prompt in prompts_list:
        masks, boxes, phrases, logits = model.predict(image_pil, prompt)
        if masks:
            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):
    # 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):
    # 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):
    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():
    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=[])
                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=50, step=1, label="Blending Amount", value=0)
                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):
            # Initialize image history with the initial image
            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, {}, [], "Image loaded."
            else:
                return None, [], 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):
            if current_image_pil is None:
                return "No image uploaded.", {}, []
            masks = extract_masks(current_image_pil, prompts)
            if not masks:
                return "No masks detected for the given prompts.", {}, []
            dropdown_choices = list(masks.keys())
            return "Segmentation completed.", masks, gr.Dropdown.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()