from u2net.u2net_inference import get_u2net_model, get_saliency_mask

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr

print('Loading model...')
model = get_u2net_model()
print('Successfully loaded model...')
examples = ['examples/1.jpg', 'examples/2.jpg', 'examples/3.jpg', 'examples/4.jpg','examples/5.jpg','examples/6.jpg']


def infer(image):
    image_out = get_saliency_mask(model, image)
    return image_out


iface = gr.Interface(
    fn=infer,
    title="U^2Net Based Saliency Estimatiion",
    description = "U^2Net Saliency Estimation",
    inputs=[gr.Image(label="image", type="numpy", shape=(640, 480))],
    outputs="image",
    cache_examples=True,
    examples=examples).launch(debug=True)