| import os | |
| import cv2 | |
| import paddlehub as hub | |
| import gradio as gr | |
| import torch | |
| torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2016/10/21/14/46/fox-1758183_1280.jpg', 'fox.jpg') | |
| model = hub.Module(name='U2Net') | |
| def infer(webcam, img,option): | |
| if option == "webcam": | |
| webcam.save('temp.jpg') | |
| result = model.Segmentation( | |
| images=[cv2.imread("temp.jpg")], | |
| paths=None, | |
| batch_size=1, | |
| input_size=320, | |
| output_dir='output', | |
| visualization=True) | |
| else: | |
| img.save('temp.jpg') | |
| result = model.Segmentation( | |
| images=[cv2.imread("temp.jpg")], | |
| paths=None, | |
| batch_size=1, | |
| input_size=320, | |
| output_dir='output', | |
| visualization=True) | |
| return result[0]['front'][:,:,::-1], result[0]['mask'] | |
| inputs = [gr.inputs.Image(source="webcam", label="Webcam", type="pil",optional=True),gr.inputs.Image(source="upload", label="Input Image", type="pil",optional=True),gr.inputs.Radio(choices=["webcam","Image"], type="value", default="Image", label="Input Type")] | |
| outputs = [ | |
| gr.outputs.Image(type="numpy",label="Front"), | |
| gr.outputs.Image(type="numpy",label="Mask") | |
| ] | |
| title = "U^2-Net" | |
| description = "demo for U^2-Net. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2005.09007'>U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection</a> | <a href='https://github.com/xuebinqin/U-2-Net'>Github Repo</a></p>" | |
| examples = [ | |
| [None,'fox.jpg','Image'], | |
| ['parrot.jpg',None,'webcam'] | |
| ] | |
| gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch() |