Ahsen Khaliq commited on
Commit
ba1e91a
1 Parent(s): ac14ec4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +20 -10
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
- img.save('temp.jpg')
13
- result = model.Segmentation(
14
- images=[cv2.imread("temp.jpg")],
15
- paths=None,
16
- batch_size=1,
17
- input_size=320,
18
- output_dir='output',
19
- visualization=True)
 
 
 
 
 
 
 
 
 
 
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")