import gradio as gr import numpy as np from PIL import Image, ImageDraw import torch from transformers import SamModel, SamProcessor from diffusers import StableDiffusionInpaintPipeline # Constants IMG_SIZE = 512 # Global variables to store points and the original image input_points = [] input_image = None def generate_mask(image, points): """ Generates a mask using SAM based on input points. """ if not points: return None image = image.convert("RGB") points = [tuple(point) for point in points] # Initialize SAM model and processor on CPU sam_model = SamModel.from_pretrained("facebook/sam-vit-huge", torch_dtype=torch.float32).to("cpu") sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge") inputs = sam_processor(image, points=points, return_tensors="pt").to("cpu") with torch.no_grad(): outputs = sam_model(**inputs) masks = sam_processor.image_processor.post_process_masks( outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() ) if len(masks) == 0: return None best_mask = masks[0][0][outputs.iou_scores.argmax()] binary_mask = ~best_mask.numpy().astype(bool).astype(int) return binary_mask def replace_object(image, mask, prompt, negative_prompt, seed, guidance_scale): """ Replaces the object in the image based on the mask and prompt. """ if mask is None: return image # Initialize Inpainting pipeline on CPU with a compatible model inpaint_pipeline = StableDiffusionInpaintPipeline.from_pretrained( "stabilityai/stable-diffusion-2-inpainting", torch_dtype=torch.float32 ).to("cpu") mask_image = Image.fromarray((mask * 255).astype(np.uint8)) generator = torch.Generator("cpu").manual_seed(seed) try: result = inpaint_pipeline( prompt=prompt, image=image, mask_image=mask_image, negative_prompt=negative_prompt if negative_prompt else None, generator=generator, guidance_scale=guidance_scale ).images[0] return result except Exception as e: print(f"Inpainting error: {e}") return image def visualize_mask(image, mask): """ Overlays the mask on the image for visualization. """ if mask is None: return image bg_transparent = np.zeros(mask.shape + (4,), dtype=np.uint8) bg_transparent[mask == 1] = [0, 255, 0, 127] # Green with transparency mask_rgba = Image.fromarray(bg_transparent) overlay = Image.alpha_composite(image.convert("RGBA"), mask_rgba) return overlay.convert("RGB") def get_points(img, evt: gr.SelectData): """ Captures points selected by the user on the image. """ global input_points global input_image if len(input_points) == 0: input_image = img.copy() x = evt.index[0] y = evt.index[1] input_points.append([x, y]) # Generate mask based on selected points mask = generate_mask(input_image, input_points) # Mark selected points with a green crossmark draw = ImageDraw.Draw(img) size = 10 for point in input_points: px, py = point draw.line((px - size, py, px + size, py), fill="green", width=5) draw.line((px, py - size, px, py + size), fill="green", width=5) # Visualize the mask overlay masked_image = visualize_mask(input_image, mask) return masked_image, img def run_inpaint(prompt, negative_prompt, cfg, seed, invert): """ Runs the inpainting process based on user inputs. """ global input_image global input_points if input_image is None or len(input_points) == 0: raise gr.Error("No points provided. Click on the image to select the object to segment with SAM.") mask = generate_mask(input_image, input_points) if invert: what = 'subject' mask = ~mask else: what = 'background' try: inpainted = replace_object(input_image, mask, prompt, negative_prompt, seed, cfg) except Exception as e: raise gr.Error(str(e)) return inpainted.resize((IMG_SIZE, IMG_SIZE)) def reset_points_func(): """ Resets the selected points and the input image. """ global input_points global input_image input_points = [] input_image = None return None, None, None def preprocess(input_img): """ Preprocesses the uploaded image to ensure it is square and resized. """ if input_img is None: return None width, height = input_img.size if width != height: # Add white padding to make the image square new_size = max(width, height) new_image = Image.new("RGB", (new_size, new_size), 'white') left = (new_size - width) // 2 top = (new_size - height) // 2 new_image.paste(input_img, (left, top)) input_img = new_image return input_img.resize((IMG_SIZE, IMG_SIZE)) with gr.Blocks() as demo: gr.Markdown( """ # Object Replacement App Upload an image, select points on the object you want to replace, provide a text prompt for the replacement, and view the augmented image. **Instructions:** 1. **Upload Image:** Click on the first image box to upload your image. 2. **Select Points:** Click on the image to select points on the object you wish to replace. Use multiple points for better mask accuracy. 3. **Enter Prompts:** Provide a replacement prompt and optionally a negative prompt to refine the output. 4. **Adjust Settings:** Set the seed for reproducibility and adjust the guidance scale as needed. 5. **Replace Object:** Click the "Replace Object" button to generate the augmented image. 6. **Reset:** Click the "Reset" button to clear selections and start over. """) with gr.Row(): with gr.Column(): # Image upload and point selection upload_image = gr.Image( label="Upload Image", type="pil", interactive=True, height=IMG_SIZE, width=IMG_SIZE ) mask_visualization = gr.Image( label="Selected Object Mask Overlay", interactive=False, height=IMG_SIZE, width=IMG_SIZE ) selected_image = gr.Image( label="Image with Selected Points", type="pil", interactive=False, height=IMG_SIZE, width=IMG_SIZE, ) # Capture points using the select event upload_image.select(get_points, inputs=[upload_image], outputs=[mask_visualization, selected_image]) # Preprocess image on change upload_image.change(preprocess, inputs=[upload_image], outputs=[upload_image]) # Text inputs and settings prompt = gr.Textbox( label="Replacement Prompt", placeholder="e.g., a red sports car", lines=2 ) negative_prompt = gr.Textbox( label="Negative Prompt", placeholder="e.g., blurry, low quality", lines=2 ) cfg = gr.Slider( label="Classifier-Free Guidance Scale", minimum=1.0, maximum=20.0, value=7.5, step=0.5 ) seed = gr.Number( label="Seed", value=42, precision=0 ) invert = gr.Checkbox( label="Infill subject instead of background" ) # Buttons replace_button = gr.Button("Replace Object") reset_button = gr.Button("Reset") with gr.Column(): # Output images augmented_image = gr.Image( label="Augmented Image", type="pil", interactive=False, height=IMG_SIZE, width=IMG_SIZE, ) # Define button actions replace_button.click( fn=run_inpaint, inputs=[prompt, negative_prompt, cfg, seed, invert], outputs=[augmented_image] ) reset_button.click( fn=reset_points_func, inputs=[], outputs=[mask_visualization, selected_image, augmented_image] ) # Examples (optional) gr.Markdown( """ ## EXAMPLES Click on an example to load it. Then, follow the instructions above. """) with gr.Row(): examples = gr.Examples( examples=[ [ "car.png", "a red sports car", "blurry, low quality", 42 ], [ "monalisa.png", "a rockstar", "dark, overexposed", 123 ], ], inputs=[ upload_image, prompt, negative_prompt, seed ], label="Click to load examples", cache_examples=False # Set to False to avoid the error ) demo.queue(max_size=10).launch()