Ahsen Khaliq
commited on
Commit
•
ba1e91a
1
Parent(s):
ac14ec4
Update app.py
Browse files
app.py
CHANGED
@@ -8,18 +8,28 @@ torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2016/10/21/14/46/f
|
|
8 |
|
9 |
model = hub.Module(name='U2Net')
|
10 |
|
11 |
-
def infer(img):
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
20 |
return result[0]['front'][:,:,::-1], result[0]['mask']
|
21 |
|
22 |
-
inputs = gr.inputs.Image(source="webcam", label="Webcam", type="pil")
|
23 |
outputs = [
|
24 |
gr.outputs.Image(type="numpy",label="Front"),
|
25 |
gr.outputs.Image(type="numpy",label="Mask")
|
|
|
8 |
|
9 |
model = hub.Module(name='U2Net')
|
10 |
|
11 |
+
def infer(webcam, img,option):
|
12 |
+
if option == "webcam":
|
13 |
+
webcam.save('temp.jpg')
|
14 |
+
result = model.Segmentation(
|
15 |
+
images=[cv2.imread("temp.jpg")],
|
16 |
+
paths=None,
|
17 |
+
batch_size=1,
|
18 |
+
input_size=320,
|
19 |
+
output_dir='output',
|
20 |
+
visualization=True)
|
21 |
+
else:
|
22 |
+
img.save('temp.jpg')
|
23 |
+
result = model.Segmentation(
|
24 |
+
images=[cv2.imread("temp.jpg")],
|
25 |
+
paths=None,
|
26 |
+
batch_size=1,
|
27 |
+
input_size=320,
|
28 |
+
output_dir='output',
|
29 |
+
visualization=True)
|
30 |
return result[0]['front'][:,:,::-1], result[0]['mask']
|
31 |
|
32 |
+
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")]
|
33 |
outputs = [
|
34 |
gr.outputs.Image(type="numpy",label="Front"),
|
35 |
gr.outputs.Image(type="numpy",label="Mask")
|