Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import gradio as gr
2
  from test import inference_img
3
  from models import *
 
4
 
5
  device='cpu'
6
  model = StyleMatte()
@@ -13,17 +14,19 @@ model.eval()
13
 
14
  def predict(inp):
15
  print("***********Inference****************")
16
- res = inference_img(model, inp)
17
  print("***********Inference finish****************")
 
 
18
 
19
- return res
20
 
21
  print("MODEL LOADED")
22
  print("************************************")
23
 
24
  iface = gr.Interface(fn=predict,
25
  inputs=gr.Image(type="numpy"),
26
- outputs=gr.Image(type="numpy"),
27
  examples=["./logo.jpeg"])
28
  print("****************Interface created******************")
29
 
 
1
  import gradio as gr
2
  from test import inference_img
3
  from models import *
4
+ import numpy as np
5
 
6
  device='cpu'
7
  model = StyleMatte()
 
14
 
15
  def predict(inp):
16
  print("***********Inference****************")
17
+ mask = inference_img(model, inp)
18
  print("***********Inference finish****************")
19
+ inp_np = np.array(inp)
20
+ fg = np.uint8((mask*inp_np).permute(1,2,0).numpy())
21
 
22
+ return [mask, fg]
23
 
24
  print("MODEL LOADED")
25
  print("************************************")
26
 
27
  iface = gr.Interface(fn=predict,
28
  inputs=gr.Image(type="numpy"),
29
+ outputs=[gr.Image(type="numpy"),gr.Image(type="numpy")],
30
  examples=["./logo.jpeg"])
31
  print("****************Interface created******************")
32