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: