Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from test import inference_img
 from models import *
+import numpy as np 
 
 device='cpu'
 model = StyleMatte()
@@ -13,17 +14,19 @@ model.eval()
 
 def predict(inp):
     print("***********Inference****************")
-    res = inference_img(model, inp) 
+    mask = inference_img(model, inp) 
     print("***********Inference finish****************")
+    inp_np = np.array(inp)
+    fg = np.uint8((mask*inp_np).permute(1,2,0).numpy())
     
-    return res
+    return [mask, fg]
 
 print("MODEL LOADED")
 print("************************************")
 
 iface = gr.Interface(fn=predict, 
              inputs=gr.Image(type="numpy"),
-             outputs=gr.Image(type="numpy"),
+             outputs=[gr.Image(type="numpy"),gr.Image(type="numpy")],
              examples=["./logo.jpeg"])
 print("****************Interface created******************")