update

asset3d_gen/models/segment_model.py

@@ -6,6 +6,7 @@ import cv2
 import numpy as np
 import rembg
 import torch
+from transformers import pipeline
 from huggingface_hub import snapshot_download
 from PIL import Image
 from segment_anything import (
@@ -292,6 +293,30 @@ class RembgRemover(object):
         return output_image
 
 
+class BMGG14Remover(object):
+    def __init__(self) -> None:
+        self.model = pipeline(
+            "image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True
+        )
+        
+    def __call__(
+        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
+    ):
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        image = resize_pil(image)
+        output_image = self.model(image)
+        
+        if save_path is not None:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            output_image.save(save_path)
+
+        return output_image
+
+
 def invert_rgba_pil(
     image: Image.Image, mask: Image.Image, save_path: str = None
 ) -> Image.Image:
@@ -368,8 +393,11 @@ if __name__ == "__main__":
         model_type="vit_h",
     )
     remover = RembgRemover()
-    # clean_image = remover(input_image)
-    # clean_image.save(output_image)
+    clean_image = remover(input_image)
+    clean_image.save(output_image)
     get_segmented_image(
         Image.open(input_image), remover, remover, None, "./test_seg.png"
     )
+
+    remover = BMGG14Remover()
+    remover("asset3d_gen/models/test_seg.jpg", "./seg.png")

common.py

@@ -24,6 +24,7 @@ from asset3d_gen.models.segment_model import (
     RembgRemover,
     SAMPredictor,
     trellis_preprocess,
+    BMGG14Remover,
 )
 from asset3d_gen.models.sr_model import ImageRealESRGAN
 from asset3d_gen.scripts.render_gs import entrypoint as render_gs_api
@@ -149,7 +150,8 @@ def download_kolors_weights() -> None:
 
 
 if os.getenv("GRADIO_APP") == "imageto3d":
-    RBG_REMOVER = RembgRemover()
+    # RBG_REMOVER = RembgRemover()
+    RBG_REMOVER = BMGG14Remover()
     SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
     PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
         "JeffreyXiang/TRELLIS-image-large"
@@ -163,7 +165,8 @@ if os.getenv("GRADIO_APP") == "imageto3d":
         os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
     )
 elif os.getenv("GRADIO_APP") == "textto3d":
-    RBG_REMOVER = RembgRemover()
+    # RBG_REMOVER = RembgRemover()
+    RBG_REMOVER = BMGG14Remover()
     PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
         "JeffreyXiang/TRELLIS-image-large"
     )