| 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) | |