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_mask(image_pil, text_prompt):
    masks, boxes, phrases, logits = model.predict(image_pil, text_prompt)
    masks_np = masks[0].cpu().numpy()
    mask = (masks_np > 0).astype(np.uint8) * 255  # Binary mask
    return mask

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, prompt, replacement_image_pil, color_ref_image_pil, apply_replacement, apply_color_grading, 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
    
    # Save current image to history for undo
    if image_history is None:
        image_history = []
    image_history.append(current_image_pil.copy())
    
    # Extract mask
    mask = extract_mask(current_image_pil, prompt)
    
    # Check if mask is valid
    if mask.sum() == 0:
        return current_image_pil, f"No mask detected for prompt: {prompt}", image_history, current_image_pil
    
    # 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 fit the mask area
        # Get bounding box of the mask
        y_indices, x_indices = np.where(mask > 0)
        if y_indices.size == 0 or x_indices.size == 0:
            # No mask detected
            return current_image_pil, f"No mask detected for prompt: {prompt}", image_history, current_image_pil
        y_min, y_max = y_indices.min(), y_indices.max()
        x_min, x_max = x_indices.min(), x_indices.max()
        
        # Extract the region of interest
        mask_height = y_max - y_min + 1
        mask_width = x_max - x_min + 1
        
        # Resize replacement image to fit mask area
        replacement_image_resized = replacement_image_pil.resize((mask_width, mask_height))
        replacement_image_np = np.array(replacement_image_resized)
        
        # Create a mask for the ROI
        mask_roi = mask_blurred[y_min:y_max+1, x_min:x_max+1]
        mask_roi_3ch = cv2.merge([mask_roi, mask_roi, mask_roi])
        
        # Replace the masked area with the replacement image using blending
        region_to_replace = result_image_np[y_min:y_max+1, x_min:x_max+1]
        blended_region = (replacement_image_np.astype(np.float32) * mask_roi_3ch + region_to_replace.astype(np.float32) * (1 - mask_roi_3ch)).astype(np.uint8)
        result_image_np[y_min:y_max+1, x_min:x_max+1] = blended_region
        
    # If apply color grading
    if apply_color_grading:
        # Extract the masked area
        masked_region = (result_image_np.astype(np.float32) * mask_blurred_3ch).astype(np.uint8)
        # Convert color reference image to numpy
        color_ref_image_np = np.array(color_ref_image_pil)
        # 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 for prompt: {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)
        
        gr.Markdown("## Continuous Image Editing with LangSAM")
        
        with gr.Row():
            with gr.Column():
                initial_image = gr.Image(type="pil", label="Upload Image")
                prompt = gr.Textbox(lines=1, placeholder="Enter prompt for object detection", label="Prompt")
                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)
                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
            else:
                return None, [], None
        
        # 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])
        
        # Apply button click
        apply_button.click(fn=process_image, 
                           inputs=[current_image_pil, prompt, replacement_image, color_ref_image, apply_replacement, apply_color_grading, 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()