Add fast-foreground-estimation in masking image.

app.py

@@ -23,6 +23,40 @@ torch.jit.script = lambda f: f
 
 device = "cuda" if torch.cuda.is_available() else "CPU"
 
+### image_proc.py
+def refine_foreground(image, mask, r=90):
+    if mask.size != image.size:
+        mask = mask.resize(image.size)
+    image = np.array(image) / 255.0
+    mask = np.array(mask) / 255.0
+    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
+    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
+    return image_masked
+
+
+def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
+    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
+    alpha = alpha[:, :, None]
+    F, blur_B = FB_blur_fusion_foreground_estimator(
+        image, image, image, alpha, r)
+    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
+
+
+def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
+    if isinstance(image, Image.Image):
+        image = np.array(image) / 255.0
+    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
+
+    blurred_FA = cv2.blur(F * alpha, (r, r))
+    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
+
+    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
+    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
+    F = blurred_F + alpha * \
+        (image - alpha * blurred_F - (1 - alpha) * blurred_B)
+    F = np.clip(F, 0, 1)
+    return F, blurred_B
+
 
 def array_to_pil_image(image: np.ndarray, size: Tuple[int, int] = (1024, 1024)) -> Image.Image:
     image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
@@ -114,19 +148,16 @@ def predict(images, resolution, weights_file):
         if device == 'cuda':
             scaled_pred_tensor = scaled_pred_tensor.cpu()
 
-        # Resize the prediction to match the original image shape
-        pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
-
-        # Apply the prediction mask to the original image
-        image_pil = image_pil.resize(pred.shape[::-1])
-        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
-        image_masked = (pred * np.array(image_pil)).astype(np.uint8)
+        # Show Results
+        pred_pil = transforms.ToPILImage()(pred)
+        image_masked = refine_foreground(image, pred_pil)
+        image_masked.putalpha(pred_pil.resize(image.size))
 
         torch.cuda.empty_cache()
 
         if tab_is_batch:
             save_file_path = os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0]))
-            cv2.imwrite(save_file_path, cv2.cvtColor(image_masked, cv2.COLOR_RGB2BGR))
+            image_masked.save(save_file_path)
             save_paths.append(save_file_path)
 
     if tab_is_batch: