Update app.py

app.py

@@ -1,30 +1,86 @@
 import gradio as gr
 import torch
 from diffusers import StableDiffusionInpaintPipeline
-
-# pipe = StableDiffusionInpaintPipeline.from_pretrained(
-#     "stabilityai/stable-diffusion-2-inpainting",
-#     torch_dtype=torch.float16,
-# )
+import cv2
 from PIL import Image
 
-def process_image(image: Image.Image, prompt: str, slider_value: int) -> Image.Image:
+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
-    processed_image = image.copy()  # Just returning a copy for now
+    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;'>Image Inprinting</h1>")
+    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=1, value=1, label='Select Zoom')
+            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