Refactor background removal and image preprocessing in app.py

- Enhanced the remove_background function to handle mask processing and integrate it with the original image.
- Updated preprocess_image to support RGBA images and added functionality for alpha masking.
- Removed unused hex_to_rgb function and streamlined the image processing flow for improved clarity and performance.

app.py

@@ -72,9 +72,16 @@ def remove_background(
     with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp:
         # Save the PIL image to the temporary file
         image.save(temp.name)
-        # Call the background removal function with the file path
-        extracted_img, _ = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name)
-    # The function should return the PIL image directly
+        extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name)
+    if isinstance(extracted_img, np.ndarray):
+        if mask.dtype != np.uint8:
+            mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
+        if mask.ndim == 3:
+            mask = mask[..., 0]
+        image = image.convert("RGBA")
+        image_np = np.array(image)
+        image_np[..., 3] = mask
+        return Image.fromarray(image_np)
     return extracted_img
 
 def do_resize_content(original_image: Image, scale_rate):
@@ -98,37 +105,19 @@ def add_background(image, bg_color=(255, 255, 255)):
     return Image.alpha_composite(background, image)
 
 
-def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
-    """Converts a hex color string to an RGB tuple."""
-    hex_color = hex_color.lstrip('#')
-    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
-
-
 def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
     """
-    Preprocesses the input image by optionally removing the background, resizing,
-    and adding a new solid background.
-    Returns an RGB PIL Image.
+    input image is a pil image in RGBA, return RGB image
     """
-    if image.mode != 'RGBA':
-        image = image.convert('RGBA')
-
-    if background_choice == "Auto Remove background":
+    print(background_choice)
+    if background_choice == "Alpha as mask":
+        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
+        image = Image.alpha_composite(background, image)
+    else:
         image = remove_background(image)
         if image is None:
             raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
 
-    # Resize the content of the image
-    image = do_resize_content(image, foreground_ratio)
-    
-    # Add a solid background color
-    rgb_background = hex_to_rgb(backgroud_color)
-    image_with_bg = add_background(image, rgb_background)
-    
-    # Convert to RGB and return
-    return image_with_bg.convert("RGB")
-
-
 @spaces.GPU
 def gen_image(input_image, seed, scale, step):
     global pipeline, model, args