Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionInpaintPipeline | |
| import cv2 | |
| from PIL import Image | |
| import numpy as np | |
| pipe = StableDiffusionInpaintPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-2-inpainting", | |
| torch_dtype=torch.float16, | |
| ) | |
| def zoom_out(image, n): | |
| # Original dimensions | |
| original_width, original_height = image.size | |
| # Calculate new dimensions after shrinking | |
| new_width = int(original_width // n) | |
| new_height = int(original_height // n) | |
| print(new_width,new_height) | |
| # Resize the image (shrink the object) | |
| resized_image = image.resize((new_width, new_height), Image.ANTIALIAS) | |
| # Create a new white canvas that is n times larger than the original | |
| new_canvas_width = original_width | |
| new_canvas_height = original_height | |
| new_image = Image.new("RGB", (new_canvas_width, new_canvas_height), "white") | |
| # Calculate the position to paste the resized object in the center | |
| x_offset = (new_canvas_width - new_width) // 2 | |
| y_offset = (new_canvas_height - new_height) // 2 | |
| # Paste the resized image onto the center of the new canvas | |
| new_image.paste(resized_image, (x_offset, y_offset)) | |
| return new_image | |
| def create_contour_mask(image, tolerance=10): | |
| # Convert the image to a numpy array | |
| image_array = np.array(image) | |
| # Define a tolerance range around white | |
| lower_bound = np.array([255 - tolerance] * 3) | |
| upper_bound = np.array([255] * 3) | |
| # Create a mask: black for background, white for object | |
| mask = np.ones((image_array.shape[0], image_array.shape[1]), dtype=np.uint8) * 255 | |
| non_white_pixels = np.any((image_array < lower_bound) | (image_array > upper_bound), axis=-1) | |
| mask[non_white_pixels] = 0 # Non-white pixels become black in the mask | |
| # Perform dilation to fill in any gaps (like white text within the object) | |
| kernel = np.ones((2, 2), np.uint8) # You can adjust the size of the kernel as needed | |
| dilated_mask = cv2.dilate(mask, kernel, iterations=2) | |
| # Invert the dilated mask to get the final object mask | |
| # object_mask = 255 - dilated_mask | |
| # Convert to PIL image for easier visualization and further processing | |
| mask_image = Image.fromarray(dilated_mask) | |
| return mask_image | |
| def process_image(image: Image.Image, prompt: str, slider_value: float) -> Image.Image: | |
| # Placeholder function for processing | |
| # Replace this with your actual processing logic | |
| # For example, modifying the image based on the slider value and prompt | |
| if slider_value != 1: | |
| image = zoom_out(image,slider_value) | |
| mask = create_contour_mask(image,10) | |
| # processed_image = image.copy() # Just returning a copy for now | |
| processed_image = pipe(prompt=prompt, image=image_big, mask_image=contour_mask).images[0] | |
| return processed_image | |
| with gr.Blocks() as demo: | |
| # Title at the top center | |
| gr.Markdown("<h1 style='text-align: center;'>Imagine Backgrounds</h1>") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Image upload on the left | |
| image_input = gr.Image(type='pil', label='Upload Image') | |
| # Slider below the image upload | |
| slider = gr.Slider(minimum=1, maximum=4, step=0.2, value=1, label='Select Zoom') | |
| # Textbox for prompt | |
| prompt_input = gr.Textbox(label='Enter Prompt') | |
| # Submit button | |
| submit_btn = gr.Button("Submit") | |
| with gr.Column(scale=1): | |
| # Output image on the right | |
| image_output = gr.Image(label='Output Image') | |
| # Event handler to process the image when the button is clicked | |
| submit_btn.click(fn=process_image, inputs=[image_input, prompt_input, slider], outputs=image_output) | |
| # Launch the Gradio app | |
| demo.launch(debug=True) | |