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, 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 # 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) mask_3ch = cv2.merge([mask, mask, mask]) result_image_np = current_image_np.copy() # If replacement image is provided if replacement_image_pil is not None: # 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[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 result_image_np[y_min:y_max+1, x_min:x_max+1] = np.where(mask_roi_3ch > 0, replacement_image_np, result_image_np[y_min:y_max+1, x_min:x_max+1]) # If color reference image is provided if color_ref_image_pil is not None: # Extract the masked area masked_region = cv2.bitwise_and(result_image_np, mask_3ch) # 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) # Combine the color matched region back into the result image result_image_np = np.where(mask_3ch > 0, color_matched_region, result_image_np) # 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_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, 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()