carvekit

app.py

@@ -73,7 +73,24 @@ from dataset import get_pose
 from CN_encoder import CN_encoder
 from pipeline_zero1to3 import Zero1to3StableDiffusionPipeline
 from segment_anything import sam_model_registry, SamPredictor
+
 import rembg
+from carvekit.api.high import HiInterface
+
+def create_carvekit_interface():
+    # Check doc strings for more information
+    interface = HiInterface(object_type="object",  # Can be "object" or "hairs-like".
+                            batch_size_seg=6,
+                            batch_size_matting=1,
+                            device=device,
+                            seg_mask_size=640,  # Use 640 for Tracer B7 and 320 for U2Net
+                            matting_mask_size=2048,
+                            trimap_prob_threshold=231,
+                            trimap_dilation=30,
+                            trimap_erosion_iters=5,
+                            fp16=True)
+
+    return interface
 
 pretrained_model_name_or_path = "kxic/EscherNet_demo"
 resolution = 256
@@ -128,7 +145,8 @@ def sam_init():
     predictor = SamPredictor(sam)
     return predictor
 
-rembg_session = rembg.new_session()
+# rembg_session = rembg.new_session()
+rembg_session = create_carvekit_interface()
 predictor = sam_init()
 
 

dust3r/utils/image.py

@@ -118,12 +118,10 @@ def load_images(folder_or_list, size, square_ok=False, verbose=True, do_remove_b
         img = exif_transpose(PIL.Image.open(os.path.join(root, path))).convert('RGB')
         # remove background if needed
         if do_remove_background:
-            # if rembg_session is None:
-            #     rembg_session = rembg.new_session()
-            # image = rembg.remove(img, session=rembg_session)
-            # foreground = np.array(image)[..., -1] > 127
-
-            image_nobg = remove(img, alpha_matting=True, session=rembg_session)
+            # use rembg
+            # image_nobg = remove(img, alpha_matting=True, session=rembg_session)
+            # use carvekit
+            image_nobg = rembg_session([img])[0]
             arr = np.asarray(image_nobg)[:, :, -1]
             x_nonzero = np.nonzero(arr.sum(axis=0))
             y_nonzero = np.nonzero(arr.sum(axis=1))

requirements.txt

@@ -7,6 +7,7 @@ transformers
 gradio
 spaces
 rembg==2.0.56
+carvekit-colab==4.1.0
 open3d
 trimesh
 einops