import gradio as gr
import inpaint.infer_model as inpaint
import sod.infer_model as sod
import numpy as np
import torch
# import os
# cmd = 'sh download.sh'
# os.system(cmd)

device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
inpaint_model = inpaint.IVModel(device=device)
sod_model = sod.IVModel(device=torch.device("cpu"))
def sod_inpaint(img):
    img = img[:,:,::-1]
    res = sod_model.forward(img,None)
    res = np.uint8(res)
    res = inpaint_model.forward(res,None)
    res = np.uint8(res)
    return res[:,:,::-1]

iface = gr.Interface(fn=sod_inpaint, inputs="image", outputs="image", examples='examples', title='显著物体消除', description='这是一个图像API，功能是自动把画面中的显著物体消除', theme='huggingface')
iface.launch()