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)