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, mask_image=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)