1

app.py

@@ -27,6 +27,7 @@ import matplotlib.pyplot as pl
 from transformers import pipeline
 from dust3r.utils.viz_demo import convert_scene_output_to_glb
 import depth_pro
+import spaces
 pl.ion()
 
 # for gpu >= Ampere and pytorch >= 1.12

dust3r/cloud_opt_flow/optimizer.py

@@ -11,9 +11,9 @@ from dust3r.utils.geometry import xy_grid, geotrf, depthmap_to_pts3d
 from dust3r.utils.device import to_cpu, to_numpy
 from dust3r.utils.goem_opt import DepthBasedWarping, OccMask, WarpImage, depth_regularization_si_weighted, tum_to_pose_matrix
 from third_party.raft import load_RAFT
-# from sam2.build_sam import build_sam2_video_predictor
-# sam2_checkpoint = "third_party/sam2/checkpoints/sam2.1_hiera_large.pt"
-# model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+from sam2.build_sam import build_sam2_video_predictor
+sam2_checkpoint = "third_party/sam2/checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
 
 def smooth_L1_loss_fn(estimate, gt, mask, beta=1.0, per_pixel_thre=50.):
     loss_raw_shape = F.smooth_l1_loss(estimate*mask, gt*mask, beta=beta, reduction='none')
@@ -109,7 +109,6 @@ class PointCloudOptimizer(BasePCOptimizer):
             self.flow_ji.requires_grad_(False)
             self.flow_valid_mask_i.requires_grad_(False)
             self.flow_valid_mask_j.requires_grad_(False)
-            sam2_mask_refine = False
             if sam2_mask_refine: 
                 with torch.no_grad():
                     self.refine_motion_mask_w_sam2()

requirements.txt

@@ -18,4 +18,5 @@ imageio
 seaborn
 evo
 transformers
-git+https://github.com/apple/ml-depth-pro.git
+git+https://github.com/apple/ml-depth-pro.git
+git+https://github.com/facebookresearch/sam2.git