diff --git a/datasets_preprocess/long_prepare_bonn.py b/datasets_preprocess/long_prepare_bonn.py index 5481083..3cd8de4 100644 --- a/datasets_preprocess/long_prepare_bonn.py +++ b/datasets_preprocess/long_prepare_bonn.py @@ -4,10 +4,10 @@ import numpy as np START_FRAME = 30 # inital frame -for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]: +for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]: END_FRAME = START_FRAME + TARGET_FRAMES # end frame - dirs = glob.glob("/home/xingyu/monst3r/data/bonn/rgbd_bonn_dataset/*/") + dirs = glob.glob("./data/bonn/rgbd_bonn_dataset/*/") dirs = sorted(dirs) # create new base directory diff --git a/datasets_preprocess/long_prepare_kitti.py b/datasets_preprocess/long_prepare_kitti.py index 2900817..4b9df77 100644 --- a/datasets_preprocess/long_prepare_kitti.py +++ b/datasets_preprocess/long_prepare_kitti.py @@ -21,7 +21,7 @@ def depth_read(filename): import shutil for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]: - depth_dirs = glob.glob("/home/xingyu/monst3r/data/kitti/val/*/proj_depth/groundtruth/image_02") + depth_dirs = glob.glob("./data/kitti/val/*/proj_depth/groundtruth/image_02") for dir in depth_dirs: # new depth dir new_depth_dir = f"./data/long_kitti_s1/depth_selection/val_selection_cropped/groundtruth_depth_gathered_{TARGET_FRAMES}/" + dir.split("/")[-4]+"_02" diff --git a/demo.py b/demo.py index 274f96d..862e33b 100644 --- a/demo.py +++ b/demo.py @@ -95,6 +95,13 @@ def parse_args(): default="cut3r", help="model update type: cut3r or ttt3r", ) + parser.add_argument( + "--beta_reduce_mode", + type=str, + default="mean", + choices=["mean", "max", "variance", "entropy"], + help="How TTT3R aggregates cross-attention into per-token update weights.", + ) parser.add_argument( "--frame_interval", type=int, @@ -463,6 +470,7 @@ def run_inference(args): print(f"Loading model from {args.model_path}...") model = ARCroco3DStereo.from_pretrained(args.model_path).to(device) model.config.model_update_type = args.model_update_type + model.config.beta_reduce_mode = args.beta_reduce_mode model.eval() diff --git a/eval/mv_recon/launch.py b/eval/mv_recon/launch.py index 6d8520d..971a847 100644 --- a/eval/mv_recon/launch.py +++ b/eval/mv_recon/launch.py @@ -40,6 +40,9 @@ def get_args_parser(): parser.add_argument("--freeze", action="store_true") parser.add_argument("--max_frames", type=int, default=None, help="max frames limit") parser.add_argument("--model_update_type", type=str, default="cut3r", help="model update type") + parser.add_argument("--beta_reduce_mode", type=str, default="mean", + choices=["mean", "max", "variance", "entropy"], + help="TTT3R beta aggregation mode") parser.add_argument("--voxel_size", type=float, default=0.0, help="voxel size for voxel grid downsampling, 0 means no downsampling") return parser @@ -58,7 +61,7 @@ def main(args): datasets_all = { "7scenes": SevenScenes( split="test", - ROOT="/home/share/Dataset/3D_scene/7scenes/", # "./data/7scenes", + ROOT="./data/7scenes", resolution=resolution, num_seq=1, full_video=True, @@ -103,6 +106,7 @@ def main(args): model = ARCroco3DStereo.from_pretrained(args.weights).to(device) model.config.model_update_type = args.model_update_type + model.config.beta_reduce_mode = args.beta_reduce_mode model.eval() # else: @@ -183,7 +187,7 @@ def main(args): with torch.cuda.amp.autocast(enabled=False): start = time.time() output = model(batch) - # preds, batch = model.forward_recurrent_light(batch) + # preds, batch = model.forward_recurrent_light(batch) end = time.time() preds, batch = output.ress, output.views valid_length = len(preds) // revisit diff --git a/eval/mv_recon/run.sh b/eval/mv_recon/run.sh index c19e9b2..ff42fb4 100644 --- a/eval/mv_recon/run.sh +++ b/eval/mv_recon/run.sh @@ -4,23 +4,28 @@ set -e workdir='.' model_names=('ttt3r') # ttt3r cut3r +beta_reduce_modes=('mean' 'entropy') # mean max variance entropy ckpt_name='cut3r_512_dpt_4_64' model_weights="${workdir}/src/${ckpt_name}.pth" for model_name in "${model_names[@]}"; do +for beta_reduce_mode in "${beta_reduce_modes[@]}"; do # for max_frames in 50 100 150 200 250 300 350 400 for max_frames in 200 do - output_dir="${workdir}/eval_results/video_recon/7scenes_${max_frames}/${model_name}" + + output_dir="${workdir}/eval_results/video_recon/7scenes_${max_frames}/${model_name}_${beta_reduce_mode}" echo "$output_dir" NCCL_TIMEOUT=360000 accelerate launch --num_processes 1 --main_process_port 29502 eval/mv_recon/launch.py \ --weights "$model_weights" \ --output_dir "$output_dir" \ --model_name "$model_name" \ --model_update_type "$model_name" \ - --max_frames $max_frames \ + --beta_reduce_mode "$beta_reduce_mode" \ + --max_frames $max_frames done done +done diff --git a/eval/relpose/launch.py b/eval/relpose/launch.py index 952ce52..f9c6fa7 100644 --- a/eval/relpose/launch.py +++ b/eval/relpose/launch.py @@ -54,6 +54,14 @@ def get_args_parser(): help="model type for state update strategy: cut3r or ttt3r", ) + parser.add_argument( + "--beta_reduce_mode", + type=str, + default="mean", + choices=["mean", "max", "variance", "entropy"], + help="How TTT3R aggregates cross-attention into per-token update weights.", + ) + parser.add_argument( "--pose_eval_stride", default=1, type=int, help="stride for pose evaluation" ) @@ -461,5 +469,6 @@ def prepare_output(outputs, revisit=1, solve_pose=False): # set model type model.config.model_update_type = args.model_update_type + model.config.beta_reduce_mode = args.beta_reduce_mode eval_pose_estimation(args, model, save_dir=args.output_dir) diff --git a/eval/relpose/run_tum.sh b/eval/relpose/run_tum.sh index ad2bc35..ec20685 100644 --- a/eval/relpose/run_tum.sh +++ b/eval/relpose/run_tum.sh @@ -4,6 +4,7 @@ set -e workdir='.' model_names=('ttt3r') # ttt3r cut3r +beta_reduce_modes=('mean' 'entropy') # mean max variance entropy ckpt_name='cut3r_512_dpt_4_64' model_weights="${workdir}/src/${ckpt_name}.pth" @@ -13,15 +14,18 @@ model_weights="${workdir}/src/${ckpt_name}.pth" datasets=('tum_s1_1000') for model_name in "${model_names[@]}"; do +for beta_reduce_mode in "${beta_reduce_modes[@]}"; do for data in "${datasets[@]}"; do - output_dir="${workdir}/eval_results/relpose/${data}/${model_name}" + output_dir="${workdir}/eval_results/relpose/${data}/${model_name}_${beta_reduce_mode}" echo "$output_dir" accelerate launch --num_processes 2 --main_process_port 29551 eval/relpose/launch.py \ --weights "$model_weights" \ --output_dir "$output_dir" \ --eval_dataset "$data" \ --size 512 \ - --model_update_type "$model_name" + --model_update_type "$model_name" \ + --beta_reduce_mode "$beta_reduce_mode" +done done done diff --git a/eval/video_depth/launch.py b/eval/video_depth/launch.py index 108efa4..1f5ff06 100644 --- a/eval/video_depth/launch.py +++ b/eval/video_depth/launch.py @@ -53,6 +53,14 @@ def get_args_parser(): help="model type for state update strategy: cut3r or ttt3r", ) + parser.add_argument( + "--beta_reduce_mode", + type=str, + default="mean", + choices=["mean", "max", "variance", "entropy"], + help="How TTT3R aggregates cross-attention into per-token update weights.", + ) + parser.add_argument( "--pose_eval_stride", default=1, type=int, help="stride for pose evaluation" @@ -339,5 +347,6 @@ def prepare_output(outputs, revisit=1): # set model type model.config.model_update_type = args.model_update_type + model.config.beta_reduce_mode = args.beta_reduce_mode eval_pose_estimation(args, model, save_dir=args.output_dir) diff --git a/eval/video_depth/run_bonn.sh b/eval/video_depth/run_bonn.sh index 1e6d63a..f40a720 100644 --- a/eval/video_depth/run_bonn.sh +++ b/eval/video_depth/run_bonn.sh @@ -4,6 +4,7 @@ set -e workdir='.' model_names=('ttt3r') # ttt3r cut3r +beta_reduce_modes=('mean' 'entropy') # mean max variance entropy ckpt_name='cut3r_512_dpt_4_64' model_weights="${workdir}/src/${ckpt_name}.pth" # datasets=('bonn_s1_50' 'bonn_s1_100' 'bonn_s1_110' 'bonn_s1_150' 'bonn_s1_200' 'bonn_s1_250' 'bonn_s1_300' 'bonn_s1_350' 'bonn_s1_400' 'bonn_s1_450' 'bonn_s1_500') @@ -11,8 +12,9 @@ datasets=('bonn_s1_500') for model_name in "${model_names[@]}"; do +for beta_reduce_mode in "${beta_reduce_modes[@]}"; do for data in "${datasets[@]}"; do - output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}" + output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}" echo "$output_dir" accelerate launch --num_processes 1 --main_process_port 29556 eval/video_depth/launch.py \ @@ -20,7 +22,8 @@ for data in "${datasets[@]}"; do --output_dir "$output_dir" \ --eval_dataset "$data" \ --size 512 \ - --model_update_type "$model_name" + --model_update_type "$model_name" \ + --beta_reduce_mode "$beta_reduce_mode" # scale&shift scale metric python eval/video_depth/eval_depth.py \ @@ -39,3 +42,4 @@ for data in "${datasets[@]}"; do --align "scale&shift" done done +done diff --git a/eval/video_depth/run_kitti.sh b/eval/video_depth/run_kitti.sh index 4d64388..2fa34aa 100644 --- a/eval/video_depth/run_kitti.sh +++ b/eval/video_depth/run_kitti.sh @@ -4,6 +4,7 @@ set -e workdir='.' model_names=('ttt3r') # ttt3r cut3r +beta_reduce_modes=('mean' 'entropy') # mean max variance entropy ckpt_name='cut3r_512_dpt_4_64' model_weights="${workdir}/src/${ckpt_name}.pth" # datasets=('kitti_s1_50' 'kitti_s1_100' 'kitti_s1_110' 'kitti_s1_150' 'kitti_s1_200' 'kitti_s1_250' 'kitti_s1_300' 'kitti_s1_350' 'kitti_s1_400' 'kitti_s1_450' 'kitti_s1_500') @@ -11,8 +12,9 @@ datasets=('kitti_s1_500') for model_name in "${model_names[@]}"; do +for beta_reduce_mode in "${beta_reduce_modes[@]}"; do for data in "${datasets[@]}"; do - output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}" + output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}" echo "$output_dir" accelerate launch --num_processes 1 --main_process_port 29555 eval/video_depth/launch.py \ @@ -20,7 +22,8 @@ for data in "${datasets[@]}"; do --output_dir "$output_dir" \ --eval_dataset "$data" \ --size 512 \ - --model_update_type "$model_name" + --model_update_type "$model_name" \ + --beta_reduce_mode "$beta_reduce_mode" # scale&shift scale metric python eval/video_depth/eval_depth.py \ @@ -39,3 +42,4 @@ for data in "${datasets[@]}"; do --align "scale&shift" done done +done diff --git a/eval/video_depth/run_sintel.sh b/eval/video_depth/run_sintel.sh index 0d2ac08..a20dd98 100644 --- a/eval/video_depth/run_sintel.sh +++ b/eval/video_depth/run_sintel.sh @@ -4,14 +4,16 @@ set -e workdir='.' model_names=('ttt3r') # ttt3r cut3r +beta_reduce_modes=('mean' 'entropy') # mean max variance entropy ckpt_name='cut3r_512_dpt_4_64' model_weights="${workdir}/src/${ckpt_name}.pth" datasets=('sintel') for model_name in "${model_names[@]}"; do +for beta_reduce_mode in "${beta_reduce_modes[@]}"; do for data in "${datasets[@]}"; do - output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}" + output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}" echo "$output_dir" accelerate launch --num_processes 1 --main_process_port 29555 eval/video_depth/launch.py \ @@ -19,7 +21,8 @@ for data in "${datasets[@]}"; do --output_dir "$output_dir" \ --eval_dataset "$data" \ --size 512 \ - --model_update_type "$model_name" + --model_update_type "$model_name" \ + --beta_reduce_mode "$beta_reduce_mode" # scale&shift scale metric python eval/video_depth/eval_depth.py \ @@ -38,3 +41,4 @@ for data in "${datasets[@]}"; do --align "scale&shift" done done +done diff --git a/eval/video_depth/utils.py b/eval/video_depth/utils.py index b34c421..aaef447 100644 --- a/eval/video_depth/utils.py +++ b/eval/video_depth/utils.py @@ -107,7 +107,7 @@ def get_vertical_colorbar(h, vmin, vmax, cmap_name="jet", label=None, cbar_preci # Do some plotting. ax = fig.add_subplot(111) - cmap = cm.get_cmap(cmap_name) + cmap = mpl.colormaps[cmap_name] norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) tick_cnt = 6 @@ -177,7 +177,7 @@ def colorize_np( x = (x - vmin) / (vmax - vmin) # x = np.clip(x, 0., 1.) - cmap = cm.get_cmap(cmap_name) + cmap = mpl.colormaps[cmap_name] x_new = cmap(x)[:, :, :3] if mask is not None: diff --git a/src/dust3r/model.py b/src/dust3r/model.py index 503c7bc..c60095e 100644 --- a/src/dust3r/model.py +++ b/src/dust3r/model.py @@ -74,7 +74,9 @@ def strip_module(state_dict): def load_model(model_path, device, verbose=True): if verbose: print("... loading model from", model_path) - ckpt = torch.load(model_path, map_location="cpu") + # PyTorch 2.6+ defaults to weights_only=True, which breaks older + # checkpoints storing OmegaConf/config objects alongside weights. + ckpt = torch.load(model_path, map_location="cpu", weights_only=False) args = ckpt["args"].model.replace( "ManyAR_PatchEmbed", "PatchEmbedDust3R" ) # ManyAR only for aspect ratio not consistent @@ -117,6 +119,7 @@ def __init__( pose_conf_head=False, pose_head=False, model_update_type="cut3r", + beta_reduce_mode="mean", **croco_kwargs, ): super().__init__() @@ -138,6 +141,7 @@ def __init__( self.pose_conf_head = pose_conf_head self.pose_head = pose_head self.model_update_type = model_update_type + self.beta_reduce_mode = beta_reduce_mode self.croco_kwargs = croco_kwargs @@ -326,6 +330,30 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kw): ) return model + def _compute_ttt3r_beta(self, cross_attn_state): + cross_attn_state = rearrange( + torch.cat(cross_attn_state, dim=0), + "l h nstate nimg -> 1 nstate nimg (l h)", + ) + reduce_mode = getattr(self.config, "beta_reduce_mode", "mean") + + if reduce_mode == "mean": + beta = torch.sigmoid(cross_attn_state.mean(dim=(-1, -2))) + elif reduce_mode == "max": + beta = torch.sigmoid(cross_attn_state.amax(dim=(-1, -2))) + elif reduce_mode == "variance": + beta = torch.sigmoid(cross_attn_state.var(dim=(-1, -2), unbiased=False)) + elif reduce_mode == "entropy": + attn_probs = torch.softmax(cross_attn_state, dim=2) + entropy = -(attn_probs * attn_probs.clamp_min(1e-8).log()).sum(dim=2) + norm = torch.log(torch.tensor(cross_attn_state.shape[2], device=entropy.device, dtype=entropy.dtype)) + entropy = entropy / norm.clamp_min(1e-8) + beta = 1.0 - entropy.mean(dim=-1) + else: + raise ValueError(f"Invalid beta reduction mode: {reduce_mode}") + + return beta[..., None] + def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768): self.patch_embed = get_patch_embed( self.patch_embed_cls, img_size, patch_size, enc_embed_dim, in_chans=3 @@ -901,9 +929,7 @@ def _forward_impl(self, views, ret_state=False): if self.config.model_update_type == "cut3r": update_mask1 = update_mask elif self.config.model_update_type == "ttt3r": - cross_attn_state = rearrange(torch.cat(cross_attn_state, dim=0), 'l h nstate nimg -> 1 nstate nimg (l h)') # [12, 16, 768, 1 + 576] -> [1, 768, 1 + 576, 12*16] - state_query_img_key = cross_attn_state.mean(dim=(-1, -2)) - update_mask1 = update_mask * torch.sigmoid(state_query_img_key)[..., None] * 1.0 + update_mask1 = update_mask * self._compute_ttt3r_beta(cross_attn_state) else: raise ValueError(f"Invalid model type: {self.config.model_update_type}") @@ -1271,9 +1297,7 @@ def forward_recurrent_lighter(self, views, device='cuda', ret_state=False): if self.config.model_update_type == "cut3r": update_mask1 = update_mask elif self.config.model_update_type == "ttt3r": - cross_attn_state = rearrange(torch.cat(cross_attn_state, dim=0), 'l h nstate nimg -> 1 nstate nimg (l h)') # [12, 16, 768, 1 + 576] -> [1, 768, 1 + 576, 12*16] - state_query_img_key = cross_attn_state.mean(dim=(-1, -2)) - update_mask1 = update_mask * torch.sigmoid(state_query_img_key)[..., None] * 1.0 + update_mask1 = update_mask * self._compute_ttt3r_beta(cross_attn_state) else: raise ValueError(f"Invalid model type: {self.config.model_update_type}")