Update app.py

app.py

@@ -15,9 +15,9 @@ def predict(inp):
     print("***********Inference****************")
     res = inference_img(model, inp) 
     print(res.shape, res.dtype)
-    # main_subject = inp.copy()
-    # main_subject = main_subject * res[res > 0.1]
-    return res
+    main_subject = inp.copy()
+    main_subject = main_subject * (res > 0.1)
+    return main_subject
 
 print("MODEL LOADED")
 print("************************************")