move sam to main

app.py

@@ -72,6 +72,8 @@ from diffusers import DDIMScheduler
 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
 
 pretrained_model_name_or_path = "kxic/EscherNet_demo"
 resolution = 256
@@ -115,6 +117,22 @@ pipeline.enable_vae_slicing()
 # pipeline.enable_xformers_memory_efficient_attention()
 
 
+#### object segmentation
+def sam_init():
+    sam_checkpoint = os.path.join("./sam_pt/sam_vit_h_4b8939.pth")
+    if os.path.exists(sam_checkpoint) is False:
+        os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P ./sam_pt/")
+    model_type = "vit_h"
+
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=device)
+    predictor = SamPredictor(sam)
+    return predictor
+
+rembg_session = rembg.new_session()
+predictor = sam_init()
+
+
+
 @spaces.GPU(duration=120)
 def run_eschernet(eschernet_input_dict, sample_steps, sample_seed, nvs_num, nvs_mode):
     # set the random seed
@@ -330,7 +348,7 @@ def get_reconstructed_scene(filelist, schedule, niter, min_conf_thr,
     if os.path.exists(outdir):
         shutil.rmtree(outdir)
     os.makedirs(outdir, exist_ok=True)
-    imgs, imgs_rgba = load_images(filelist, size=image_size, verbose=not silent, do_remove_background=True)
+    imgs, imgs_rgba = load_images(filelist, size=image_size, verbose=not silent, do_remove_background=True, rembg_session=rembg_session, predictor=predictor)
     if len(imgs) == 1:
         imgs = [imgs[0], copy.deepcopy(imgs[0])]
         imgs[1]['idx'] = 1
@@ -544,7 +562,7 @@ Image views are treated as tokens and the camera pose is encoded by <b>CaPE (Cam
 
 <h4><b>Tips:</b></h4>
 
-- Our model can take <b>any number input images</b>. The more images you provide, the better the results.
+- Our model can take <b>any number input images</b>. The more images you provide (>=3 for this demo), the better the results.
 
 - Our model can generate <b>any number and any pose</b> novel views. You can specify the number of views you want to generate. In this demo, we set novel views on an <b>archemedian spiral</b> for simplicity.
 

dust3r/utils/image.py

@@ -12,23 +12,9 @@ from PIL.ImageOps import exif_transpose
 import torchvision.transforms as tvf
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
 import cv2  # noqa
-
-import rembg
-rembg_session = rembg.new_session()
-
 import time
 from PIL import Image
 from rembg import remove
-from segment_anything import sam_model_registry, SamPredictor
-def sam_init():
-    sam_checkpoint = os.path.join("./sam_pt/sam_vit_h_4b8939.pth")
-    if os.path.exists(sam_checkpoint) is False:
-        os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P ./sam_pt/")
-    model_type = "vit_h"
-
-    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{0 if torch.cuda.is_available() else 'cpu'}")
-    predictor = SamPredictor(sam)
-    return predictor
 
 def sam_segment(predictor, input_image, *bbox_coords):
     bbox = np.array(bbox_coords)
@@ -50,7 +36,6 @@ def sam_segment(predictor, input_image, *bbox_coords):
     torch.cuda.empty_cache()
     return Image.fromarray(out_image_bbox, mode='RGBA')
 
-predictor = sam_init()
 
 try:
     from pillow_heif import register_heif_opener  # noqa
@@ -103,7 +88,7 @@ def _resize_pil_image(img, long_edge_size):
     new_size = tuple(int(round(x*long_edge_size/S)) for x in img.size)
     return img.resize(new_size, interp)
 
-def load_images(folder_or_list, size, square_ok=False, verbose=True, do_remove_background=True, rembg_session=None):
+def load_images(folder_or_list, size, square_ok=False, verbose=True, do_remove_background=True, rembg_session=None, predictor=None):
     """ open and convert all images in a list or folder to proper input format for DUSt3R
     """
     if isinstance(folder_or_list, str):