Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from PIL import Image, ImageEnhance
+from PIL import Image
 import numpy as np
 import cv2
 from lang_sam import LangSAM
@@ -8,9 +8,14 @@ from color_matcher.normalizer import Normalizer
 import torch
 
 # Load the LangSAM model
-model = LangSAM()  # Use the default model or specify custom checkpoint: LangSAM("<model_type>", "<path/to/checkpoint>")
+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
 
-# Function to apply color matching based on reference image
 def apply_color_matching(source_img_np, ref_img_np):
     # Initialize ColorMatcher
     cm = ColorMatcher()
@@ -23,60 +28,129 @@ def apply_color_matching(source_img_np, ref_img_np):
     
     return img_res
 
-# Function to extract sky and apply color matching using a reference image
-def extract_and_color_match_sky(image_pil, reference_image_pil, text_prompt="sky"):
-    # Use LangSAM to predict the mask for the sky
-    masks, boxes, phrases, logits = model.predict(image_pil, text_prompt)
-
-    # Ensure masks is converted from tensor to NumPy
-    masks_np = masks[0].cpu().numpy()  # Convert the tensor to NumPy array
-
-    # Convert the mask to a binary format and create a mask image
-    sky_mask = (masks_np > 0).astype(np.uint8) * 255  # Ensure it's a binary mask
+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
     
-    # Convert PIL image to numpy array for processing
-    img_np = np.array(image_pil)
+    # Save current image to history for undo
+    if image_history is None:
+        image_history = []
+    image_history.append(current_image_pil.copy())
     
-    # Convert sky mask to 3-channel format to blend with the original image
-    sky_mask_3ch = cv2.merge([sky_mask, sky_mask, sky_mask])
+    # Extract mask
+    mask = extract_mask(current_image_pil, prompt)
     
-    # Extract the sky region
-    sky_region = cv2.bitwise_and(img_np, sky_mask_3ch)
+    # 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
     
-    # Convert the reference image to a numpy array
-    ref_img_np = np.array(reference_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()
     
-    # Apply color matching using the reference image to the extracted sky region
-    sky_region_color_matched = apply_color_matching(sky_region, ref_img_np)
+    # 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])
     
-    # Combine the color-matched sky region back into the original image
-    result_img_np = np.where(sky_mask_3ch > 0, sky_region_color_matched, img_np)
+    # 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 the result back to PIL Image for final output
-    result_img_pil = Image.fromarray(result_img_np)
+    # Convert result back to PIL Image
+    result_image_pil = Image.fromarray(result_image_np)
     
-    return result_img_pil
-
-# Gradio Interface
-def gradio_interface():
-    # Gradio function to be called on input
-    def process_image(source_img, ref_img):
-        # Extract sky and apply color matching using reference image
-        result_img_pil = extract_and_color_match_sky(source_img, ref_img)
-        return result_img_pil
-
-    # Define Gradio input components
-    inputs = [
-        gr.Image(type="pil", label="Source Image"),
-        gr.Image(type="pil", label="Reference Image")  # Second input for reference image
-    ]
-
-    # Define Gradio output component
-    outputs = gr.Image(type="pil", label="Resulting Image")
+    # 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
 
-    # Launch Gradio app
-    gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="Sky Extraction and Color Matching").launch(share=True)
+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)", optional=True)
+                color_ref_image = gr.Image(type="pil", label="Color Reference Image (optional)", optional=True)
+                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()