from PIL import Image
import numpy as np
from rembg import remove
import cv2
import os
from torchvision.transforms import GaussianBlur
import gradio as gr
import requests
import time

def generate_image(input):
    input_path = 'input.png'
    bg_removed_path = 'bg_removed.png'
    mask_name = 'blured_mask.png'
    
    input.save(input_path)
    bg_removed = remove(input)
    bg_removed.save(bg_removed_path)

    img2_grayscale = bg_removed.convert('L')
    img2_a = np.array(img2_grayscale)

    mask = np.array(img2_grayscale)
    threshhold = 0
    mask[img2_a==threshhold] = 1  # this is white
    mask[img2_a>threshhold] = 0 # this is gray
    #The mask structure is white for inpainting and black for keeping as is

    strength = 1  # This controls the strength of our prompt relative to the init image.
    d = int(255 * (1-strength))
    mask *= 255-d # Converts our range from [0,1] to [0,255]
    mask += d

    mask = Image.fromarray(mask)

    blur = GaussianBlur(11,20)
    mask = blur(mask)
    mask.save(mask_name)

    return Image.open(bg_removed_path), Image.open(mask_name)

with gr.Blocks() as demo:
    gr.Markdown("Remove photo backgrounds with AI")
    gr.Markdown("you can use the generated mask in stable diffusion for inpainting")
    with gr.Row():

      with gr.Column():
        input_image = gr.Image(label = "Upload your product's photo", type = 'pil')
        image_button = gr.Button("Generate")
      
      with gr.Column():
        gallery = gr.Gallery(
                label="Generated images", show_label=False, elem_id="gallery"
            ).style(grid=[2], height="auto")
          
    image_button.click(generate_image, inputs=input_image, outputs=gallery)
    
    

demo.launch()