shir0ha commited on
Commit
23fa3d9
1 Parent(s): 81c7e37

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -16,7 +16,7 @@ def predict(inp):
16
  res = inference_img(model, inp)
17
  print(res.shape, res.dtype)
18
  main_subject = inp.copy()
19
- main_subject = main_subject * (res > 0.1)
20
  return main_subject
21
 
22
  print("MODEL LOADED")
 
16
  res = inference_img(model, inp)
17
  print(res.shape, res.dtype)
18
  main_subject = inp.copy()
19
+ main_subject = main_subject * (res > 0.1)[:,:,np.newaxis]
20
  return main_subject
21
 
22
  print("MODEL LOADED")