Update app.py

app.py

@@ -16,9 +16,9 @@ def predict(inp):
     print("***********Inference****************")
     mask = inference_img(model, inp) 
     inp_np = np.array(inp)
-    fg = np.uint8((mask*inp_np).permute(1,2,0).numpy())
+    fg = np.uint8((mask[:,:,None]*inp_np).numpy())
     print("***********Inference finish****************")
-    print("***********FG****************", fg.shape, mask.shape)
+    print("***********MASK****************", inp_np.max(), mask.max())
     
     return [mask, fg]