Spaces:
Runtime error
Runtime error
| 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() |