import os
os.system("gdown https://drive.google.com/uc?id=1-95IOJ-2y9BtmABiffIwndPqNZD_gLnV")
os.system("unzip big-lama.zip")
import cv2
import paddlehub as hub
import gradio as gr
import torch
from PIL import Image, ImageOps
import numpy as np
os.mkdir("data")
os.mkdir("dataout")
model = hub.Module(name='U2Net')
def infer(img,mask,option):
  img = ImageOps.contain(img, (700,700))
  width, height = img.size
  img.save("./data/data.png")
  if option == "automatic (U2net)":
      result = model.Segmentation(
          images=[cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)],
          paths=None,
          batch_size=1,
          input_size=320,
          output_dir='output',
          visualization=True)
      im = Image.fromarray(result[0]['mask'])
  else:
      mask = mask.resize((width,height))
      im = mask
  im.save("./data/data_mask.png")
  os.system('python predict.py model.path=/home/user/app/big-lama/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')
  return "./dataout/data_mask.png",im
  
inputs = [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")]
outputs = [gr.outputs.Image(type="file",label="output"),gr.outputs.Image(type="pil",label="Mask")]
title = "LaMa Image Inpainting Example"
description = "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"
article = "<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>"
examples = [
  ['person512.png',"canvas.png","automatic (U2net)"],
  ['person512.png',"maskexam.png","manual"]
]
gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch(enable_queue=True,cache_examples=True)