creative-demo / app.py
omidreza's picture
Update app.py
f1db483
Raw
History Blame Contribute Delete
1.73 kB
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()