Refactor background removal in app.py to utilize dis_bg_remover for improved image processing. Removed rembg dependency and updated related logic in remove_background function. Updated requirements.txt to reflect the change in background removal library.

app.py

@@ -9,7 +9,6 @@ import PIL
 from pipelines import TwoStagePipeline
 from huggingface_hub import hf_hub_download
 import os
-import rembg
 from typing import Any
 import json
 import os
@@ -18,9 +17,9 @@ import argparse
 
 from model import CRM
 from inference import generate3d
+from dis_bg_remover import remove_background as dis_remove_background
 
 pipeline = None
-rembg_session = rembg.new_session()
 
 
 def expand_to_square(image, bg_color=(0, 0, 0, 0)):
@@ -45,17 +44,9 @@ def remove_background(
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
-    do_remove = True
-    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
-        # explain why current do not rm bg
-        print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
-        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
-        image = Image.alpha_composite(background, image)
-        do_remove = False
-    do_remove = do_remove or force
-    if do_remove:
-        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
-    return image
+    model_path = "isnet_dis.onnx"  # Update with the correct path to the ONNX model
+    extracted_img, mask = dis_remove_background(model_path, image)
+    return extracted_img
 
 def do_resize_content(original_image: Image, scale_rate):
     # resize image content wile retain the original image size
@@ -87,7 +78,7 @@ def preprocess_image(image, background_choice, foreground_ratio, backgroud_color
         background = Image.new("RGBA", image.size, (0, 0, 0, 0))
         image = Image.alpha_composite(background, image)
     else:
-        image = remove_background(image, rembg_session, force=True)
+        image = remove_background(image, force=True)
     image = do_resize_content(image, foreground_ratio)
     image = expand_to_square(image)
     image = add_background(image, backgroud_color)

requirements.txt

@@ -9,7 +9,8 @@ opencv-contrib-python-headless==4.9.0.80
 opencv-python-headless==4.9.0.80
 xformers
 omegaconf
-rembg
+# rembg
+dis_bg_remover
 git+https://github.com/NVlabs/nvdiffrast
 pygltflib
 kiui