Refactor remove_background function in app.py to enhance mask handling. Added checks for mask dimensions and data type, ensuring proper conversion to RGBA format for images. This improves the output quality when masks are applied.

app.py

@@ -68,17 +68,24 @@ def remove_background(
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
-    # Ensure the ONNX model exists (download if needed)
     ensure_dis_onnx_model()
-    # Save PIL image to a temporary file
     with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp:
         image.save(temp.name)
         extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name)
-    # Convert to PIL Image if needed
+    # If extracted_img is a mask (single channel), use it as alpha for the original image
     if isinstance(extracted_img, np.ndarray):
-        if extracted_img.dtype != np.uint8:
-            extracted_img = (np.clip(extracted_img, 0, 1) * 255).astype(np.uint8)
-        extracted_img = Image.fromarray(extracted_img)
+        # If mask is float, convert to uint8
+        if mask.dtype != np.uint8:
+            mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
+        # Ensure mask is 2D
+        if mask.ndim == 3:
+            mask = mask[..., 0]
+        # Convert original image to RGBA
+        image = image.convert("RGBA")
+        image_np = np.array(image)
+        image_np[..., 3] = mask
+        return Image.fromarray(image_np)
+    # If extracted_img is already a color image, just return it
     return extracted_img
 
 def do_resize_content(original_image: Image, scale_rate):