Add fast-foreground-estimation in masking image.

app.py

@@ -141,12 +141,10 @@ def predict(images, resolution, weights_file):
  image_proc = image_preprocessor.proc(image_pil)
  image_proc = image_proc.unsqueeze(0)
 
- # Perform the prediction
+ # Prediction
  with torch.no_grad():
- scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid()
-
- if device == 'cuda':
- scaled_pred_tensor = scaled_pred_tensor.cpu()
+ preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
+ pred = preds[0].squeeze()
 
  # Show Results
  pred_pil = transforms.ToPILImage()(pred)