CVPR/lama-example

1import os2os.system("gdown https://drive.google.com/uc?id=1-95IOJ-2y9BtmABiffIwndPqNZD_gLnV")3os.system("unzip big-lama.zip")4import cv25import paddlehub as hub6import gradio as gr7import torch8from PIL import Image, ImageOps9import numpy as np10os.mkdir("data")11os.mkdir("dataout")12model = hub.Module(name='U2Net')13def infer(img,mask,option):14  img = ImageOps.contain(img, (700,700))15  width, height = img.size16  img.save("./data/data.png")17  if option == "automatic (U2net)":18      result = model.Segmentation(19          images=[cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)],20          paths=None,21          batch_size=1,22          input_size=320,23          output_dir='output',24          visualization=True)25      im = Image.fromarray(result[0]['mask'])26  else:27      mask = mask.resize((width,height))28      im = mask29  im.save("./data/data_mask.png")30  os.system('python predict.py model.path=/home/user/app/big-lama/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')31  return "./dataout/data_mask.png",im32  33inputs = [gr.inputs.Image(type='pil', label="Original Image"),gr.inputs.Image(type='pil',source="canvas", label="Mask",invert_colors=True),gr.inputs.Radio(choices=["automatic (U2net)","manual"], type="value", default="manual", label="Masking option")]34outputs = [gr.outputs.Image(type="file",label="output"),gr.outputs.Image(type="pil",label="Mask")]35title = "LaMa Image Inpainting Example"36description = "Gradio demo for LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Masks are generated by U^2net"37article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.07161' target='_blank'>Resolution-robust Large Mask Inpainting with Fourier Convolutions</a> | <a href='https://github.com/saic-mdal/lama' target='_blank'>Github Repo</a></p><center><img src='https://visitor-badge.glitch.me/badge?page_id=cvpr_lama' alt='visitor badge'></center>"38examples = [39  ['person512.png',"canvas.png","automatic (U2net)"],40  ['person512.png',"maskexam.png","manual"]41]42gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch(enable_queue=True,cache_examples=True)