Spaces:
Runtime error
Runtime error
| 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() | |