Add the rgba2rgb preprocessing for RGBA inputs.

app.py

@@ -29,6 +29,44 @@ torch.jit.script = lambda f: f
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+def rgba2rgb(img):
+    """
+    Convert RGBA image to RGB with white background.
+    Supports both PIL.Image and numpy.ndarray.
+    """
+
+    # 1. Handle PIL Image
+    if isinstance(img, Image.Image):
+        img = img.convert("RGBA")
+        bg = Image.new("RGBA", img.size, (255, 255, 255))
+        return Image.alpha_composite(bg, img).convert("RGB")
+
+    # 2. Handle Numpy Array (OpenCV)
+    elif isinstance(img, np.ndarray):
+        # Grayscale to RGB
+        if img.ndim == 2:
+            return np.stack([img] * 3, axis=-1)
+
+        # Already 3 channels
+        if img.shape[2] == 3:
+            return img
+
+        # RGBA to RGB (blending with white)
+        elif img.shape[2] == 4:
+            # Normalize alpha to 0-1 and keep shape (H, W, 1)
+            alpha = img[..., 3:4].astype(float) / 255.0
+            foreground = img[..., :3].astype(float)
+            background = 255.0
+
+            # Blend formula: source * alpha + bg * (1 - alpha)
+            out = foreground * alpha + background * (1.0 - alpha)
+
+            return out.clip(0, 255).astype(np.uint8)
+
+    else:
+        raise TypeError(f"Unsupported type: {type(img)}")
+
+
 ## CPU version refinement
 def FB_blur_fusion_foreground_estimator_cpu(image, FG, B, alpha, r=90):
     if isinstance(image, Image.Image):
@@ -112,7 +150,7 @@ def refine_foreground(image, mask, r=90, device='cuda'):
         mask = mask.unsqueeze(0)
 
         estimated_foreground = FB_blur_fusion_foreground_estimator_gpu_2(image, mask, r=r)
-        
+
         estimated_foreground = estimated_foreground.squeeze()
         estimated_foreground = (estimated_foreground.mul(255.0)).to(torch.uint8)
         estimated_foreground = estimated_foreground.permute(1, 2, 0).contiguous().cpu().numpy().astype(np.uint8)
@@ -215,6 +253,8 @@ def predict(images, resolution, weights_file):
                 image_ori = Image.open(image_data)
         else:
             image_ori = Image.fromarray(image_src)
+        if image_ori.mode == 'RGBA':
+            image_ori = rgba2rgb(image_ori)
 
         image = image_ori.convert('RGB')
         # Preprocess the image

app_local.py

@@ -25,6 +25,44 @@ torch.set_float32_matmul_precision('high')
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+def rgba2rgb(img):
+    """
+    Convert RGBA image to RGB with white background.
+    Supports both PIL.Image and numpy.ndarray.
+    """
+
+    # 1. Handle PIL Image
+    if isinstance(img, Image.Image):
+        img = img.convert("RGBA")
+        bg = Image.new("RGBA", img.size, (255, 255, 255))
+        return Image.alpha_composite(bg, img).convert("RGB")
+
+    # 2. Handle Numpy Array (OpenCV)
+    elif isinstance(img, np.ndarray):
+        # Grayscale to RGB
+        if img.ndim == 2:
+            return np.stack([img] * 3, axis=-1)
+
+        # Already 3 channels
+        if img.shape[2] == 3:
+            return img
+
+        # RGBA to RGB (blending with white)
+        elif img.shape[2] == 4:
+            # Normalize alpha to 0-1 and keep shape (H, W, 1)
+            alpha = img[..., 3:4].astype(float) / 255.0
+            foreground = img[..., :3].astype(float)
+            background = 255.0
+
+            # Blend formula: source * alpha + bg * (1 - alpha)
+            out = foreground * alpha + background * (1.0 - alpha)
+
+            return out.clip(0, 255).astype(np.uint8)
+
+    else:
+        raise TypeError(f"Unsupported type: {type(img)}")
+
+
 ## CPU version refinement
 def FB_blur_fusion_foreground_estimator_cpu(image, FG, B, alpha, r=90):
     if isinstance(image, Image.Image):
@@ -108,7 +146,7 @@ def refine_foreground(image, mask, r=90, device='cuda'):
         mask = mask.unsqueeze(0)
 
         estimated_foreground = FB_blur_fusion_foreground_estimator_gpu_2(image, mask, r=r)
-        
+
         estimated_foreground = estimated_foreground.squeeze()
         estimated_foreground = (estimated_foreground.mul(255.0)).to(torch.uint8)
         estimated_foreground = estimated_foreground.permute(1, 2, 0).contiguous().cpu().numpy().astype(np.uint8)
@@ -210,6 +248,8 @@ def predict(images, resolution, weights_file):
                 image_ori = Image.open(image_data)
         else:
             image_ori = Image.fromarray(image_src)
+        if image_ori.mode == 'RGBA':
+            image_ori = rgba2rgb(image_ori)
 
         image = image_ori.convert('RGB')
         # Preprocess the image