Upload 99 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+FastVGGT/assets/attn_map.png filter=lfs diff=lfs merge=lfs -text
+FastVGGT/assets/autolab_logo.png filter=lfs diff=lfs merge=lfs -text
+FastVGGT/assets/main.png filter=lfs diff=lfs merge=lfs -text
+FastVGGT/assets/vs.png filter=lfs diff=lfs merge=lfs -text

FastVGGT/.gitignore

@@ -0,0 +1,160 @@
+.hydra/
+output/
+ckpt/
+.vscode/
+dependency/
+# Byte-compiled / optimized / DLL files
+__pycache__/
+**/__pycache__/
+*.py[cod]
+*$py.class
+test_logs/
+quick_start_logs/
+logs/
+*.pth
+/data/
+*.png
+eval_results/
+.vscode/
+.curosr/
+
+# C extensions
+*.so
+LightGlue/
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Profiling data
+.prof
+
+# Folder specific to your needs
+**/tmp/
+**/outputs/skyseg.onnx
+skyseg.onnx
+
+# pixi environments
+.pixi
+*.egg-info

FastVGGT/.vscode/launch.json

@@ -0,0 +1,85 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+
+        {
+            "name": "launch",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "/home/sy/code/vggt_0625/training/launch.py",
+            "console": "integratedTerminal",
+            "args": "${command:pickArgs}",
+            "env": {
+                "CUDA_VISIBLE_DEVICES": "3",
+            },
+            "cwd": "/home/sy/code/vggt_0625/training",
+            "justMyCode": true,
+            "python": "/home/sy/anaconda3/envs/vggt/bin/python"
+        }
+        ,{
+            "name": "train_scannet",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "/home/sy/code/vggt_0625/training/launch_scannet.py",
+            "console": "integratedTerminal",
+            "args": [
+                // "--config_name", "scannet",
+                // "--exp_name", "scannet_exp001",
+                // "--resume_checkpoint_path", "/home/sy/code/vggt_0625/ckpt/model_tracker_fixed_e20.pt"
+            ],
+            "env": {
+                "CUDA_VISIBLE_DEVICES": "7",
+                "WORLD_SIZE": "1",
+                "RANK": "0",
+                "MASTER_ADDR": "localhost",
+                "MASTER_PORT": "12345"
+            },
+            "cwd": "/home/sy/code/vggt_0625/training",
+            "justMyCode": true,
+            "python": "/home/sy/anaconda3/envs/vggt/bin/python"
+        }
+        ,{
+            "name": "eval_scannet",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "/home/sy/code/FastVGGT/eval/eval_scannet.py",
+            "console": "integratedTerminal",
+            "args": [
+                "--data_dir","/data/sy/scannetv2/process_scannet",
+                "--gt_ply_dir","/data/sy/scannetv2/OpenDataLab___ScanNet_v2/raw/scans",
+                "--output_path", "/home/sy/code/FastVGGT/eval_results",
+                "--merging", "0",
+                "--ckpt_path","/home/sy/code/vggt_0625/ckpt/model_tracker_fixed_e20.pt",
+                "--vis_attn_map"
+            ],
+            "env": {
+                "CUDA_VISIBLE_DEVICES": "2"
+            },
+            "justMyCode": true,
+            "python": "/home/sy/anaconda3/envs/fastvggt/bin/python"
+        },
+        {
+            "name": "eval_cd",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "/home/sy/code/FastVGGT/eval/eval_custom.py",
+            "console": "integratedTerminal",
+            "args": [
+                "--merging", "0",
+                // "--kf","10",
+                // "--output_dir","/home/sy/code/vggt_0625/eval_results_cd",
+                "--data_path","/data/sy/segment-102751/",
+                "--vis_attn_map"
+            ],
+            "env": {
+                "CUDA_VISIBLE_DEVICES": "3"
+            },
+            "justMyCode": true,
+            // "python": "/home/sy/anaconda3/envs/fastvggt/bin/python"
+        }
+
+    ]
+}

FastVGGT/README.md

@@ -0,0 +1,163 @@
+<div align="center">
+<h2>⚡️ FastVGGT: Training-Free Acceleration of Visual Geometry Transformer</h2>
+  
+<p align="center">
+  <a href="https://arxiv.org/abs/2509.02560"><img src="https://img.shields.io/badge/arXiv-FastVGGT-red?logo=arxiv" alt="Paper PDF"></a>
+  <a href="https://mystorm16.github.io/fastvggt/"><img src="https://img.shields.io/badge/Project_Page-FastVGGT-yellow" alt="Project Page"></a>
+</p>
+  
+<img src="assets/maclab_logo.png" alt="Maclab Logo" width="110" style="margin-right: 40px;">
+<img src="assets/autolab_logo.png" alt="Autolab Logo" width="110">
+
+
+**[Media Analytics & Computing Laboratory](https://mac.xmu.edu.cn/)**; **[AUTOLAB](https://zhipengzhang.cn/)**
+
+
+[You Shen](https://mystorm16.github.io/), [Zhipeng Zhang](https://zhipengzhang.cn/), [Yansong Qu](https://quyans.github.io/), [Liujuan Cao](https://mac.xmu.edu.cn/ljcao/)
+</div>
+
+
+## 📰 News
+- [Sep 8, 2025] Added custom dataset evaluation.
+- [Sep 3, 2025] Paper release.
+- [Sep 2, 2025] Code release.
+
+## 🔭 Overview
+
+FastVGGT observes **strong similarity** in attention maps and leverages it to design a training-free acceleration method for long-sequence 3D reconstruction, **achieving up to 4× faster inference without sacrificing accuracy.**
+
+<img src="assets/main.png" alt="Autolab Logo" width="">
+
+
+## ⚙️ Environment Setup
+First, create a virtual environment using Conda, clone this repository to your local machine, and install the required dependencies.
+
+
+```bash
+conda create -n fastvggt python=3.10
+conda activate fastvggt
+git clone git@github.com:mystorm16/FastVGGT.git
+cd FastVGGT
+pip install -r requirements.txt
+```
+
+Next, prepare the ScanNet dataset: http://www.scan-net.org/ScanNet/
+
+Then, download the VGGT checkpoint (we use the checkpoint link provided in https://github.com/facebookresearch/vggt/tree/evaluation/evaluation):
+```bash
+wget https://huggingface.co/facebook/VGGT_tracker_fixed/resolve/main/model_tracker_fixed_e20.pt
+```
+
+Finally, configure the dataset path and VGGT checkpoint path. For example:
+```bash
+    parser.add_argument(
+        "--data_dir", type=Path, default="/data/scannetv2/process_scannet"
+    )
+    parser.add_argument(
+        "--gt_ply_dir",
+        type=Path,
+        default="/data/scannetv2/OpenDataLab___ScanNet_v2/raw/scans",
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="./ckpt/model_tracker_fixed_e20.pt",
+    )
+```
+
+
+## 💎 Observation
+
+Note: A large number of input_frames may significantly slow down saving the visualization results. Please try using a smaller number first.
+```bash
+python eval/eval_scannet.py --input_frame 30 --vis_attn_map --merging 0
+```
+
+We observe that many token-level attention maps are highly similar in each block, motivating our optimization of the Global Attention module.
+
+<img src="assets/attn_map.png" alt="Autolab Logo" width="">
+
+
+
+## 🏀 Evaluation
+### Custom Dataset
+Please organize the data according to the following directory:
+```
+<data_path>/
+├── images/       
+│   ├── 000000.jpg
+│   ├── 000001.jpg
+│   └── ...
+├── pose/                # Optional: Camera poses
+│   ├── 000000.txt 
+│   ├── 000001.txt
+│   └── ...
+└── gt_ply/              # Optional: GT point cloud
+    └── scene_xxx.ply   
+```
+- Required: `images/`
+- Additionally required when `--enable_evaluation` is enabled: `pose/` and `gt_ply/`
+
+Inference only:
+
+```bash
+python eval/eval_custom.py \
+  --data_path /path/to/your_dataset \
+  --output_path ./eval_results_custom \
+  --plot
+```
+
+Inference + Evaluation (requires `pose/` and `gt_ply/`):
+
+```bash
+python eval/eval_custom.py \
+  --data_path /path/to/your_dataset \
+  --enable_evaluation \
+  --output_path ./eval_results_custom \
+  --plot
+```
+
+### ScanNet
+Evaluate FastVGGT on the ScanNet dataset with 1,000 input images. The **--merging** parameter specifies the block index at which the merging strategy is applied:
+
+```bash
+python eval/eval_scannet.py --input_frame 1000 --merging 0
+```
+
+Evaluate Baseline VGGT on the ScanNet dataset with 1,000 input images:
+```bash
+python eval/eval_scannet.py --input_frame 1000
+```
+<img src="assets/vs.png" alt="Autolab Logo" width="">
+
+### 7 Scenes & NRGBD
+Evaluate across two datasets, sampling keyframes every 10 frames:
+```bash
+python eval/eval_7andN.py --kf 10
+```
+
+## 🍺 Acknowledgements
+
+- Thanks to these great repositories: [VGGT](https://github.com/facebookresearch/vggt), [Dust3r](https://github.com/naver/dust3r),  [Fast3R](https://github.com/facebookresearch/fast3r), [CUT3R](https://github.com/CUT3R/CUT3R), [MV-DUSt3R+](https://github.com/facebookresearch/mvdust3r), [StreamVGGT](https://github.com/wzzheng/StreamVGGT), [VGGT-Long](https://github.com/DengKaiCQ/VGGT-Long), [ToMeSD](https://github.com/dbolya/tomesd) and many other inspiring works in the community.
+
+- Special thanks to [Jianyuan Wang](https://jytime.github.io/) for his valuable discussions and suggestions on this work.
+
+<!-- ## ✍️ Checklist
+
+- [ ] Release the evaluation code on 7 Scenes / NRGBD -->
+
+
+## ⚖️ License
+See the [LICENSE](./LICENSE.txt) file for details about the license under which this code is made available.
+
+## Citation
+
+If you find this project helpful, please consider citing the following paper:
+```
+@article{shen2025fastvggt,
+  title={FastVGGT: Training-Free Acceleration of Visual Geometry Transformer},
+  author={Shen, You and Zhang, Zhipeng and Qu, Yansong and Cao, Liujuan},
+  journal={arXiv preprint arXiv:2509.02560},
+  year={2025}
+}
+```

FastVGGT/eval/base.py

@@ -0,0 +1,273 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# base class for implementing datasets
+# --------------------------------------------------------
+import PIL
+import numpy as np
+import torch
+
+from dataset_utils.transforms import ImgNorm
+import dataset_utils.cropping as cropping
+from utils import depthmap_to_absolute_camera_coordinates
+
+
+class BaseStereoViewDataset:
+    """Define all basic options.
+
+    Usage:
+        class MyDataset (BaseStereoViewDataset):
+            def _get_views(self, idx, rng):
+                # overload here
+                views = []
+                views.append(dict(img=, ...))
+                return views
+    """
+
+    def __init__(
+        self,
+        *,  # only keyword arguments
+        split=None,
+        resolution=None,  # square_size or (width, height) or list of [(width,height), ...]
+        transform=ImgNorm,
+        aug_crop=False,
+        seed=None,
+    ):
+        self.num_views = 2
+        self.split = split
+        self._set_resolutions(resolution)
+
+        self.transform = transform
+        if isinstance(transform, str):
+            transform = eval(transform)
+
+        self.aug_crop = aug_crop
+        self.seed = seed
+
+    def __len__(self):
+        return len(self.scenes)
+
+    def get_stats(self):
+        return f"{len(self)} pairs"
+
+    def __repr__(self):
+        resolutions_str = "[" + ";".join(f"{w}x{h}" for w, h in self._resolutions) + "]"
+        return (
+            f"""{type(self).__name__}({self.get_stats()},
+            {self.split=},
+            {self.seed=},
+            resolutions={resolutions_str},
+            {self.transform=})""".replace(
+                "self.", ""
+            )
+            .replace("\n", "")
+            .replace("   ", "")
+        )
+
+    def _get_views(self, idx, resolution, rng):
+        raise NotImplementedError()
+
+    def __getitem__(self, idx):
+        if isinstance(idx, tuple):
+            # the idx is specifying the aspect-ratio
+            idx, ar_idx = idx
+        else:
+            assert len(self._resolutions) == 1
+            ar_idx = 0
+
+        # set-up the rng
+        if self.seed:  # reseed for each __getitem__
+            self._rng = np.random.default_rng(seed=self.seed + idx)
+        elif not hasattr(self, "_rng"):
+            seed = torch.initial_seed()  # this is different for each dataloader process
+            self._rng = np.random.default_rng(seed=seed)
+
+        # over-loaded code
+        resolution = self._resolutions[
+            ar_idx
+        ]  # DO NOT CHANGE THIS (compatible with BatchedRandomSampler)
+        views = self._get_views(idx, resolution, self._rng)
+
+        # check data-types
+        for v, view in enumerate(views):
+            assert (
+                "pts3d" not in view
+            ), f"pts3d should not be there, they will be computed afterwards based on intrinsics+depthmap for view {view_name(view)}"
+            view["idx"] = v
+
+            # encode the image
+            width, height = view["img"].size
+            view["true_shape"] = np.int32((height, width))
+            view["img"] = self.transform(view["img"])
+
+            assert "camera_intrinsics" in view
+            if "camera_pose" not in view:
+                view["camera_pose"] = np.full((4, 4), np.nan, dtype=np.float32)
+            else:
+                assert np.isfinite(
+                    view["camera_pose"]
+                ).all(), f"NaN in camera pose for view {view_name(view)}"
+            assert "pts3d" not in view
+            assert "valid_mask" not in view
+            assert np.isfinite(
+                view["depthmap"]
+            ).all(), f"NaN in depthmap for view {view_name(view)}"
+            pts3d, valid_mask = depthmap_to_absolute_camera_coordinates(**view)
+
+            view["pts3d"] = pts3d
+            view["valid_mask"] = valid_mask & np.isfinite(pts3d).all(axis=-1)
+
+            # check all datatypes
+            for key, val in view.items():
+                res, err_msg = is_good_type(key, val)
+                assert res, f"{err_msg} with {key}={val} for view {view_name(view)}"
+            K = view["camera_intrinsics"]
+            view["img_mask"] = True
+            view["ray_mask"] = False
+            view["ray_map"] = torch.full(
+                (6, view["img"].shape[-2], view["img"].shape[-1]), torch.nan
+            )
+            view["update"] = True
+            view["reset"] = False
+
+        # last thing done!
+        for view in views:
+            # transpose to make sure all views are the same size
+            transpose_to_landscape(view)
+            # this allows to check whether the RNG is is the same state each time
+            view["rng"] = int.from_bytes(self._rng.bytes(4), "big")
+        return views
+
+    def _set_resolutions(self, resolutions):
+        """Set the resolution(s) of the dataset.
+        Params:
+            - resolutions: int or tuple or list of tuples
+        """
+        assert resolutions is not None, "undefined resolution"
+
+        if not isinstance(resolutions, list):
+            resolutions = [resolutions]
+
+        self._resolutions = []
+        for resolution in resolutions:
+            if isinstance(resolution, int):
+                width = height = resolution
+            else:
+                width, height = resolution
+            assert isinstance(
+                width, int
+            ), f"Bad type for {width=} {type(width)=}, should be int"
+            assert isinstance(
+                height, int
+            ), f"Bad type for {height=} {type(height)=}, should be int"
+            assert width >= height
+            self._resolutions.append((width, height))
+
+    def _crop_resize_if_necessary(
+        self, image, depthmap, intrinsics, resolution, rng=None, info=None
+    ):
+        """This function:
+        - first downsizes the image with LANCZOS inteprolation,
+          which is better than bilinear interpolation in
+        """
+        if not isinstance(image, PIL.Image.Image):
+            image = PIL.Image.fromarray(image)
+
+        # downscale with lanczos interpolation so that image.size == resolution
+        # cropping centered on the principal point
+        W, H = image.size
+        cx, cy = intrinsics[:2, 2].round().astype(int)
+
+        # calculate min distance to margin
+        min_margin_x = min(cx, W - cx)
+        min_margin_y = min(cy, H - cy)
+        assert min_margin_x > W / 5, f"Bad principal point in view={info}"
+        assert min_margin_y > H / 5, f"Bad principal point in view={info}"
+
+        ## Center crop
+        # Crop on the principal point, make it always centered
+        # the new window will be a rectangle of size (2*min_margin_x, 2*min_margin_y) centered on (cx,cy)
+        l, t = cx - min_margin_x, cy - min_margin_y
+        r, b = cx + min_margin_x, cy + min_margin_y
+        crop_bbox = (l, t, r, b)
+
+        image, depthmap, intrinsics = cropping.crop_image_depthmap(
+            image, depthmap, intrinsics, crop_bbox
+        )
+
+        # # transpose the resolution if necessary
+        W, H = image.size  # new size
+        assert resolution[0] >= resolution[1]
+        if H > 1.1 * W:
+            # image is portrait mode
+            resolution = resolution[::-1]
+        elif 0.9 < H / W < 1.1 and resolution[0] != resolution[1]:
+            # image is square, so we chose (portrait, landscape) randomly
+            if rng.integers(2):
+                resolution = resolution[::-1]
+
+        # high-quality Lanczos down-scaling
+        target_resolution = np.array(resolution)
+        # # if self.aug_crop > 1:
+        # #     target_resolution += rng.integers(0, self.aug_crop)
+        # if resolution != (224, 224):
+        #     halfw, halfh = ((2*(W//2))//16)*8, ((2*(H//2))//16)*8
+        #     ## Recale with max factor, so  one of width or height might be larger than target_resolution
+        #     image, depthmap, intrinsics = cropping.rescale_image_depthmap(image, depthmap, intrinsics, (2*halfw, 2*halfh))
+        # else:
+        image, depthmap, intrinsics = cropping.rescale_image_depthmap(
+            image, depthmap, intrinsics, target_resolution
+        )
+        # actual cropping (if necessary) with bilinear interpolation
+        # if resolution == (224, 224):
+        intrinsics2 = cropping.camera_matrix_of_crop(
+            intrinsics, image.size, resolution, offset_factor=0.5
+        )
+        crop_bbox = cropping.bbox_from_intrinsics_in_out(
+            intrinsics, intrinsics2, resolution
+        )
+        image, depthmap, intrinsics = cropping.crop_image_depthmap(
+            image, depthmap, intrinsics, crop_bbox
+        )
+        return image, depthmap, intrinsics
+
+
+def is_good_type(key, v):
+    """returns (is_good, err_msg)"""
+    if isinstance(v, (str, int, tuple)):
+        return True, None
+    if v.dtype not in (np.float32, torch.float32, bool, np.int32, np.int64, np.uint8):
+        return False, f"bad {v.dtype=}"
+    return True, None
+
+
+def view_name(view, batch_index=None):
+    def sel(x):
+        return x[batch_index] if batch_index not in (None, slice(None)) else x
+
+    db = sel(view["dataset"])
+    label = sel(view["label"])
+    instance = sel(view["instance"])
+    return f"{db}/{label}/{instance}"
+
+
+def transpose_to_landscape(view):
+    height, width = view["true_shape"]
+
+    if width < height:
+        # rectify portrait to landscape
+        assert view["img"].shape == (3, height, width)
+        view["img"] = view["img"].swapaxes(1, 2)
+
+        assert view["valid_mask"].shape == (height, width)
+        view["valid_mask"] = view["valid_mask"].swapaxes(0, 1)
+
+        assert view["depthmap"].shape == (height, width)
+        view["depthmap"] = view["depthmap"].swapaxes(0, 1)
+
+        assert view["pts3d"].shape == (height, width, 3)
+        view["pts3d"] = view["pts3d"].swapaxes(0, 1)
+
+        # transpose x and y pixels
+        view["camera_intrinsics"] = view["camera_intrinsics"][[1, 0, 2]]

FastVGGT/eval/criterion.py

@@ -0,0 +1,534 @@
+import torch
+import torch.nn as nn
+from copy import copy, deepcopy
+
+from eval.dataset_utils.corr import geotrf, inv
+
+
+def invalid_to_nans(arr, valid_mask, ndim=999):
+    if valid_mask is not None:
+        arr = arr.clone()
+        arr[~valid_mask] = float("nan")
+    if arr.ndim > ndim:
+        arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+    return arr
+
+
+def invalid_to_zeros(arr, valid_mask, ndim=999):
+    if valid_mask is not None:
+        arr = arr.clone()
+        arr[~valid_mask] = 0
+        nnz = valid_mask.view(len(valid_mask), -1).sum(1)
+    else:
+        nnz = arr.numel() // len(arr) if len(arr) else 0  # number of point per image
+    if arr.ndim > ndim:
+        arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+    return arr, nnz
+
+
+class BaseCriterion(nn.Module):
+    def __init__(self, reduction="mean"):
+        super().__init__()
+        self.reduction = reduction
+
+
+class Criterion(nn.Module):
+    def __init__(self, criterion=None):
+        super().__init__()
+        assert isinstance(
+            criterion, BaseCriterion
+        ), f"{criterion} is not a proper criterion!"
+        self.criterion = copy(criterion)
+
+    def get_name(self):
+        return f"{type(self).__name__}({self.criterion})"
+
+    def with_reduction(self, mode="none"):
+        res = loss = deepcopy(self)
+        while loss is not None:
+            assert isinstance(loss, Criterion)
+            loss.criterion.reduction = mode  # make it return the loss for each sample
+            loss = loss._loss2  # we assume loss is a Multiloss
+        return res
+
+
+class MultiLoss(nn.Module):
+    """Easily combinable losses (also keep track of individual loss values):
+        loss = MyLoss1() + 0.1*MyLoss2()
+    Usage:
+        Inherit from this class and override get_name() and compute_loss()
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._alpha = 1
+        self._loss2 = None
+
+    def compute_loss(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def get_name(self):
+        raise NotImplementedError()
+
+    def __mul__(self, alpha):
+        assert isinstance(alpha, (int, float))
+        res = copy(self)
+        res._alpha = alpha
+        return res
+
+    __rmul__ = __mul__  # same
+
+    def __add__(self, loss2):
+        assert isinstance(loss2, MultiLoss)
+        res = cur = copy(self)
+
+        while cur._loss2 is not None:
+            cur = cur._loss2
+        cur._loss2 = loss2
+        return res
+
+    def __repr__(self):
+        name = self.get_name()
+        if self._alpha != 1:
+            name = f"{self._alpha:g}*{name}"
+        if self._loss2:
+            name = f"{name} + {self._loss2}"
+        return name
+
+    def forward(self, *args, **kwargs):
+        loss = self.compute_loss(*args, **kwargs)
+        if isinstance(loss, tuple):
+            loss, details = loss
+        elif loss.ndim == 0:
+            details = {self.get_name(): float(loss)}
+        else:
+            details = {}
+        loss = loss * self._alpha
+
+        if self._loss2:
+            loss2, details2 = self._loss2(*args, **kwargs)
+            loss = loss + loss2
+            details |= details2
+
+        return loss, details
+
+
+class LLoss(BaseCriterion):
+    """L-norm loss"""
+
+    def forward(self, a, b):
+        assert (
+            a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
+        ), f"Bad shape = {a.shape}"
+        dist = self.distance(a, b)
+
+        if self.reduction == "none":
+            return dist
+        if self.reduction == "sum":
+            return dist.sum()
+        if self.reduction == "mean":
+            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
+        raise ValueError(f"bad {self.reduction=} mode")
+
+    def distance(self, a, b):
+        raise NotImplementedError()
+
+
+class L21Loss(LLoss):
+    """Euclidean distance between 3d points"""
+
+    def distance(self, a, b):
+        return torch.norm(a - b, dim=-1)  # normalized L2 distance
+
+
+L21 = L21Loss()
+
+
+def get_pred_pts3d(gt, pred, use_pose=False):
+    assert use_pose is True
+    return pred["pts3d_in_other_view"]  # return!
+
+
+def Sum(losses, masks, conf=None):
+    loss, mask = losses[0], masks[0]
+    if loss.ndim > 0:
+        # we are actually returning the loss for every pixels
+        if conf is not None:
+            return losses, masks, conf
+        return losses, masks
+    else:
+        # we are returning the global loss
+        for loss2 in losses[1:]:
+            loss = loss + loss2
+        return loss
+
+
+def get_norm_factor(pts, norm_mode="avg_dis", valids=None, fix_first=True):
+    assert pts[0].ndim >= 3 and pts[0].shape[-1] == 3
+    assert pts[1] is None or (pts[1].ndim >= 3 and pts[1].shape[-1] == 3)
+    norm_mode, dis_mode = norm_mode.split("_")
+
+    nan_pts = []
+    nnzs = []
+
+    if norm_mode == "avg":
+        # gather all points together (joint normalization)
+
+        for i, pt in enumerate(pts):
+            nan_pt, nnz = invalid_to_zeros(pt, valids[i], ndim=3)
+            nan_pts.append(nan_pt)
+            nnzs.append(nnz)
+
+            if fix_first:
+                break
+        all_pts = torch.cat(nan_pts, dim=1)
+
+        # compute distance to origin
+        all_dis = all_pts.norm(dim=-1)
+        if dis_mode == "dis":
+            pass  # do nothing
+        elif dis_mode == "log1p":
+            all_dis = torch.log1p(all_dis)
+        else:
+            raise ValueError(f"bad {dis_mode=}")
+
+        norm_factor = all_dis.sum(dim=1) / (torch.cat(nnzs).sum() + 1e-8)
+    else:
+        raise ValueError(f"Not implemented {norm_mode=}")
+
+    norm_factor = norm_factor.clip(min=1e-8)
+    while norm_factor.ndim < pts[0].ndim:
+        norm_factor.unsqueeze_(-1)
+
+    return norm_factor
+
+
+def normalize_pointcloud_t(
+    pts, norm_mode="avg_dis", valids=None, fix_first=True, gt=False
+):
+    if gt:
+        norm_factor = get_norm_factor(pts, norm_mode, valids, fix_first)
+        res = []
+
+        for i, pt in enumerate(pts):
+            res.append(pt / norm_factor)
+
+    else:
+        # pts_l, pts_r = pts
+        # use pts_l and pts_r[-1] as pts to normalize
+        norm_factor = get_norm_factor(pts, norm_mode, valids, fix_first)
+
+        res = []
+
+        for i in range(len(pts)):
+            res.append(pts[i] / norm_factor)
+            # res_r.append(pts_r[i] / norm_factor)
+
+        # res = [res_l, res_r]
+
+    return res, norm_factor
+
+
+@torch.no_grad()
+def get_joint_pointcloud_depth(zs, valid_masks=None, quantile=0.5):
+    # set invalid points to NaN
+    _zs = []
+    for i in range(len(zs)):
+        valid_mask = valid_masks[i] if valid_masks is not None else None
+        _z = invalid_to_nans(zs[i], valid_mask).reshape(len(zs[i]), -1)
+        _zs.append(_z)
+
+    _zs = torch.cat(_zs, dim=-1)
+
+    # compute median depth overall (ignoring nans)
+    if quantile == 0.5:
+        shift_z = torch.nanmedian(_zs, dim=-1).values
+    else:
+        shift_z = torch.nanquantile(_zs, quantile, dim=-1)
+    return shift_z  # (B,)
+
+
+@torch.no_grad()
+def get_joint_pointcloud_center_scale(pts, valid_masks=None, z_only=False, center=True):
+    # set invalid points to NaN
+
+    _pts = []
+    for i in range(len(pts)):
+        valid_mask = valid_masks[i] if valid_masks is not None else None
+        _pt = invalid_to_nans(pts[i], valid_mask).reshape(len(pts[i]), -1, 3)
+        _pts.append(_pt)
+
+    _pts = torch.cat(_pts, dim=1)
+
+    # compute median center
+    _center = torch.nanmedian(_pts, dim=1, keepdim=True).values  # (B,1,3)
+    if z_only:
+        _center[..., :2] = 0  # do not center X and Y
+
+    # compute median norm
+    _norm = ((_pts - _center) if center else _pts).norm(dim=-1)
+    scale = torch.nanmedian(_norm, dim=1).values
+    return _center[:, None, :, :], scale[:, None, None, None]
+
+
+class Regr3D_t(Criterion, MultiLoss):
+    def __init__(self, criterion, norm_mode="avg_dis", gt_scale=False, fix_first=True):
+        super().__init__(criterion)
+        self.norm_mode = norm_mode
+        self.gt_scale = gt_scale
+        self.fix_first = fix_first
+
+    def get_all_pts3d_t(self, gts, preds, dist_clip=None):
+        # everything is normalized w.r.t. camera of view1
+        in_camera1 = inv(gts[0]["camera_pose"])
+
+        gt_pts = []
+        valids = []
+        pr_pts = []
+
+        for i, gt in enumerate(gts):
+            # in_camera1: Bs, 4, 4 gt['pts3d']: Bs, H, W, 3
+            gt_pts.append(geotrf(in_camera1, gt["pts3d"]))
+            valid = gt["valid_mask"].clone()
+
+            if dist_clip is not None:
+                # points that are too far-away == invalid
+                dis = gt["pts3d"].norm(dim=-1)
+                valid = valid & (dis <= dist_clip)
+
+            valids.append(valid)
+            pr_pts.append(get_pred_pts3d(gt, preds[i], use_pose=True))
+            # if i != len(gts)-1:
+            #     pr_pts_l.append(get_pred_pts3d(gt, preds[i][0], use_pose=(i!=0)))
+
+            # if i != 0:
+            #     pr_pts_r.append(get_pred_pts3d(gt, preds[i-1][1], use_pose=(i!=0)))
+
+        # pr_pts = (pr_pts_l, pr_pts_r)
+
+        if self.norm_mode:
+            pr_pts, pr_factor = normalize_pointcloud_t(
+                pr_pts, self.norm_mode, valids, fix_first=self.fix_first, gt=False
+            )
+        else:
+            pr_factor = None
+
+        if self.norm_mode and not self.gt_scale:
+            gt_pts, gt_factor = normalize_pointcloud_t(
+                gt_pts, self.norm_mode, valids, fix_first=self.fix_first, gt=True
+            )
+        else:
+            gt_factor = None
+
+        return gt_pts, pr_pts, gt_factor, pr_factor, valids, {}
+
+    def compute_frame_loss(self, gts, preds, **kw):
+        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = (
+            self.get_all_pts3d_t(gts, preds, **kw)
+        )
+
+        pred_pts_l, pred_pts_r = pred_pts
+
+        loss_all = []
+        mask_all = []
+        conf_all = []
+
+        loss_left = 0
+        loss_right = 0
+        pred_conf_l = 0
+        pred_conf_r = 0
+
+        for i in range(len(gt_pts)):
+
+            # Left (Reference)
+            if i != len(gt_pts) - 1:
+                frame_loss = self.criterion(
+                    pred_pts_l[i][masks[i]], gt_pts[i][masks[i]]
+                )
+
+                loss_all.append(frame_loss)
+                mask_all.append(masks[i])
+                conf_all.append(preds[i][0]["conf"])
+
+                # To compare target/reference loss
+                if i != 0:
+                    loss_left += frame_loss.cpu().detach().numpy().mean()
+                    pred_conf_l += preds[i][0]["conf"].cpu().detach().numpy().mean()
+
+            # Right (Target)
+            if i != 0:
+                frame_loss = self.criterion(
+                    pred_pts_r[i - 1][masks[i]], gt_pts[i][masks[i]]
+                )
+
+                loss_all.append(frame_loss)
+                mask_all.append(masks[i])
+                conf_all.append(preds[i - 1][1]["conf"])
+
+                # To compare target/reference loss
+                if i != len(gt_pts) - 1:
+                    loss_right += frame_loss.cpu().detach().numpy().mean()
+                    pred_conf_r += preds[i - 1][1]["conf"].cpu().detach().numpy().mean()
+
+        if pr_factor is not None and gt_factor is not None:
+            filter_factor = pr_factor[pr_factor > gt_factor]
+        else:
+            filter_factor = []
+
+        if len(filter_factor) > 0:
+            factor_loss = (filter_factor - gt_factor).abs().mean()
+        else:
+            factor_loss = 0.0
+
+        self_name = type(self).__name__
+        details = {
+            self_name + "_pts3d_1": float(loss_all[0].mean()),
+            self_name + "_pts3d_2": float(loss_all[1].mean()),
+            self_name + "loss_left": float(loss_left),
+            self_name + "loss_right": float(loss_right),
+            self_name + "conf_left": float(pred_conf_l),
+            self_name + "conf_right": float(pred_conf_r),
+        }
+
+        return Sum(loss_all, mask_all, conf_all), (details | monitoring), factor_loss
+
+
+class ConfLoss_t(MultiLoss):
+    """Weighted regression by learned confidence.
+        Assuming the input pixel_loss is a pixel-level regression loss.
+
+    Principle:
+        high-confidence means high conf = 0.1 ==> conf_loss = x / 10 + alpha*log(10)
+        low  confidence means low  conf = 10  ==> conf_loss = x * 10 - alpha*log(10)
+
+        alpha: hyperparameter
+    """
+
+    def __init__(self, pixel_loss, alpha=1):
+        super().__init__()
+        assert alpha > 0
+        self.alpha = alpha
+        self.pixel_loss = pixel_loss.with_reduction("none")
+
+    def get_name(self):
+        return f"ConfLoss({self.pixel_loss})"
+
+    def get_conf_log(self, x):
+        return x, torch.log(x)
+
+    def compute_frame_loss(self, gts, preds, **kw):
+        # compute per-pixel loss
+        (losses, masks, confs), details, loss_factor = (
+            self.pixel_loss.compute_frame_loss(gts, preds, **kw)
+        )
+
+        # weight by confidence
+        conf_losses = []
+        conf_sum = 0
+        for i in range(len(losses)):
+            conf, log_conf = self.get_conf_log(confs[i][masks[i]])
+            conf_sum += conf.mean()
+            conf_loss = losses[i] * conf - self.alpha * log_conf
+            conf_loss = conf_loss.mean() if conf_loss.numel() > 0 else 0
+            conf_losses.append(conf_loss)
+
+        conf_losses = torch.stack(conf_losses) * 2.0
+        conf_loss_mean = conf_losses.mean()
+
+        return (
+            conf_loss_mean,
+            dict(
+                conf_loss_1=float(conf_losses[0]),
+                conf_loss2=float(conf_losses[1]),
+                conf_mean=conf_sum / len(losses),
+                **details,
+            ),
+            loss_factor,
+        )
+
+
+class Regr3D_t_ShiftInv(Regr3D_t):
+    """Same than Regr3D but invariant to depth shift."""
+
+    def get_all_pts3d_t(self, gts, preds):
+        # compute unnormalized points
+        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = (
+            super().get_all_pts3d_t(gts, preds)
+        )
+
+        # pred_pts_l, pred_pts_r = pred_pts
+        gt_zs = [gt_pt[..., 2] for gt_pt in gt_pts]
+
+        pred_zs = [pred_pt[..., 2] for pred_pt in pred_pts]
+        # pred_zs.append(pred_pts_r[-1][..., 2])
+
+        # compute median depth
+        gt_shift_z = get_joint_pointcloud_depth(gt_zs, masks)[:, None, None]
+        pred_shift_z = get_joint_pointcloud_depth(pred_zs, masks)[:, None, None]
+
+        # subtract the median depth
+        for i in range(len(gt_pts)):
+            gt_pts[i][..., 2] -= gt_shift_z
+
+        for i in range(len(pred_pts)):
+            # for j in range(len(pred_pts[i])):
+            pred_pts[i][..., 2] -= pred_shift_z
+
+        monitoring = dict(
+            monitoring,
+            gt_shift_z=gt_shift_z.mean().detach(),
+            pred_shift_z=pred_shift_z.mean().detach(),
+        )
+        return gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring
+
+
+class Regr3D_t_ScaleInv(Regr3D_t):
+    """Same than Regr3D but invariant to depth shift.
+    if gt_scale == True: enforce the prediction to take the same scale than GT
+    """
+
+    def get_all_pts3d_t(self, gts, preds):
+        # compute depth-normalized points
+        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = (
+            super().get_all_pts3d_t(gts, preds)
+        )
+
+        # measure scene scale
+
+        # pred_pts_l, pred_pts_r = pred_pts
+
+        pred_pts_all = [
+            x.clone() for x in pred_pts
+        ]  # [pred_pt for pred_pt in pred_pts_l]
+        # pred_pts_all.append(pred_pts_r[-1])
+
+        _, gt_scale = get_joint_pointcloud_center_scale(gt_pts, masks)
+        _, pred_scale = get_joint_pointcloud_center_scale(pred_pts_all, masks)
+
+        # prevent predictions to be in a ridiculous range
+        pred_scale = pred_scale.clip(min=1e-3, max=1e3)
+
+        # subtract the median depth
+        if self.gt_scale:
+            for i in range(len(pred_pts)):
+                # for j in range(len(pred_pts[i])):
+                pred_pts[i] *= gt_scale / pred_scale
+
+        else:
+            for i in range(len(pred_pts)):
+                # for j in range(len(pred_pts[i])):
+                pred_pts[i] *= pred_scale / gt_scale
+
+            for i in range(len(gt_pts)):
+                gt_pts[i] *= gt_scale / pred_scale
+
+        monitoring = dict(
+            monitoring, gt_scale=gt_scale.mean(), pred_scale=pred_scale.mean().detach()
+        )
+
+        return gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring
+
+
+class Regr3D_t_ScaleShiftInv(Regr3D_t_ScaleInv, Regr3D_t_ShiftInv):
+    # calls Regr3D_ShiftInv first, then Regr3D_ScaleInv
+    pass

FastVGGT/eval/data.py

@@ -0,0 +1,338 @@
+import os
+import cv2
+import numpy as np
+import os.path as osp
+from collections import deque
+from base import BaseStereoViewDataset
+import dataset_utils.cropping as cropping
+from vggt.utils.eval_utils import imread_cv2, shuffle_deque
+
+
+class SevenScenes(BaseStereoViewDataset):
+    def __init__(
+        self,
+        num_seq=1,
+        num_frames=5,
+        min_thresh=10,
+        max_thresh=100,
+        test_id=None,
+        full_video=False,
+        tuple_list=None,
+        seq_id=None,
+        rebuttal=False,
+        shuffle_seed=-1,
+        kf_every=1,
+        *args,
+        ROOT,
+        **kwargs,
+    ):
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+        self.num_seq = num_seq
+        self.num_frames = num_frames
+        self.max_thresh = max_thresh
+        self.min_thresh = min_thresh
+        self.test_id = test_id
+        self.full_video = full_video
+        self.kf_every = kf_every
+        self.seq_id = seq_id
+        self.rebuttal = rebuttal
+        self.shuffle_seed = shuffle_seed
+
+        # load all scenes
+        self.load_all_tuples(tuple_list)
+        self.load_all_scenes(ROOT)
+
+    def __len__(self):
+        if self.tuple_list is not None:
+            return len(self.tuple_list)
+        return len(self.scene_list) * self.num_seq
+
+    def load_all_tuples(self, tuple_list):
+        if tuple_list is not None:
+            self.tuple_list = tuple_list
+            # with open(tuple_path) as f:
+            #     self.tuple_list = f.read().splitlines()
+
+        else:
+            self.tuple_list = None
+
+    def load_all_scenes(self, base_dir):
+
+        if self.tuple_list is not None:
+            # Use pre-defined simplerecon scene_ids
+            self.scene_list = [
+                "stairs/seq-06",
+                "stairs/seq-02",
+                "pumpkin/seq-06",
+                "chess/seq-01",
+                "heads/seq-02",
+                "fire/seq-02",
+                "office/seq-03",
+                "pumpkin/seq-03",
+                "redkitchen/seq-07",
+                "chess/seq-02",
+                "office/seq-01",
+                "redkitchen/seq-01",
+                "fire/seq-01",
+            ]
+            print(f"Found {len(self.scene_list)} sequences in split {self.split}")
+            return
+
+        scenes = os.listdir(base_dir)
+
+        file_split = {"train": "TrainSplit.txt", "test": "TestSplit.txt"}[self.split]
+
+        self.scene_list = []
+        for scene in scenes:
+            if self.test_id is not None and scene != self.test_id:
+                continue
+            # read file split
+            with open(osp.join(base_dir, scene, file_split)) as f:
+                seq_ids = f.read().splitlines()
+
+                for seq_id in seq_ids:
+                    # seq is string, take the int part and make it 01, 02, 03
+                    # seq_id = 'seq-{:2d}'.format(int(seq_id))
+                    num_part = "".join(filter(str.isdigit, seq_id))
+                    seq_id = f"seq-{num_part.zfill(2)}"
+                    if self.seq_id is not None and seq_id != self.seq_id:
+                        continue
+                    self.scene_list.append(f"{scene}/{seq_id}")
+
+        print(f"Found {len(self.scene_list)} sequences in split {self.split}")
+
+    def _get_views(self, idx, resolution, rng):
+
+        if self.tuple_list is not None:
+            line = self.tuple_list[idx].split(" ")
+            scene_id = line[0]
+            img_idxs = line[1:]
+
+        else:
+            scene_id = self.scene_list[idx // self.num_seq]
+            seq_id = idx % self.num_seq
+
+            data_path = osp.join(self.ROOT, scene_id)
+            num_files = len([name for name in os.listdir(data_path) if "color" in name])
+            img_idxs = [f"{i:06d}" for i in range(num_files)]
+            img_idxs = img_idxs[:: self.kf_every]
+
+        # Intrinsics used in SimpleRecon
+        fx, fy, cx, cy = 525, 525, 320, 240
+        intrinsics_ = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+        views = []
+        imgs_idxs = deque(img_idxs)
+        if self.shuffle_seed >= 0:
+            imgs_idxs = shuffle_deque(imgs_idxs)
+
+        while len(imgs_idxs) > 0:
+            im_idx = imgs_idxs.popleft()
+            impath = osp.join(self.ROOT, scene_id, f"frame-{im_idx}.color.png")
+            depthpath = osp.join(self.ROOT, scene_id, f"frame-{im_idx}.depth.proj.png")
+            posepath = osp.join(self.ROOT, scene_id, f"frame-{im_idx}.pose.txt")
+
+            rgb_image = imread_cv2(impath)
+
+            depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED)
+            rgb_image = cv2.resize(rgb_image, (depthmap.shape[1], depthmap.shape[0]))
+
+            depthmap[depthmap == 65535] = 0
+            depthmap = np.nan_to_num(depthmap.astype(np.float32), 0.0) / 1000.0
+
+            depthmap[depthmap > 10] = 0
+            depthmap[depthmap < 1e-3] = 0
+
+            camera_pose = np.loadtxt(posepath).astype(np.float32)
+
+            if resolution != (224, 224) or self.rebuttal:
+                rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                    rgb_image, depthmap, intrinsics_, resolution, rng=rng, info=impath
+                )
+            else:
+                rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                    rgb_image, depthmap, intrinsics_, (512, 384), rng=rng, info=impath
+                )
+                W, H = rgb_image.size
+                cx = W // 2
+                cy = H // 2
+                l, t = cx - 112, cy - 112
+                r, b = cx + 112, cy + 112
+                crop_bbox = (l, t, r, b)
+                rgb_image, depthmap, intrinsics = cropping.crop_image_depthmap(
+                    rgb_image, depthmap, intrinsics, crop_bbox
+                )
+
+            views.append(
+                dict(
+                    img=rgb_image,
+                    depthmap=depthmap,
+                    camera_pose=camera_pose,
+                    camera_intrinsics=intrinsics,
+                    dataset="7scenes",
+                    label=osp.join(scene_id, im_idx),
+                    instance=impath,
+                )
+            )
+        return views
+
+
+class NRGBD(BaseStereoViewDataset):
+    def __init__(
+        self,
+        num_seq=1,
+        num_frames=5,
+        min_thresh=10,
+        max_thresh=100,
+        test_id=None,
+        full_video=False,
+        tuple_list=None,
+        seq_id=None,
+        rebuttal=False,
+        shuffle_seed=-1,
+        kf_every=1,
+        *args,
+        ROOT,
+        **kwargs,
+    ):
+
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+        self.num_seq = num_seq
+        self.num_frames = num_frames
+        self.max_thresh = max_thresh
+        self.min_thresh = min_thresh
+        self.test_id = test_id
+        self.full_video = full_video
+        self.kf_every = kf_every
+        self.seq_id = seq_id
+        self.rebuttal = rebuttal
+        self.shuffle_seed = shuffle_seed
+
+        # load all scenes
+        self.load_all_tuples(tuple_list)
+        self.load_all_scenes(ROOT)
+
+    def __len__(self):
+        if self.tuple_list is not None:
+            return len(self.tuple_list)
+        return len(self.scene_list) * self.num_seq
+
+    def load_all_tuples(self, tuple_list):
+        if tuple_list is not None:
+            self.tuple_list = tuple_list
+            # with open(tuple_path) as f:
+            #     self.tuple_list = f.read().splitlines()
+
+        else:
+            self.tuple_list = None
+
+    def load_all_scenes(self, base_dir):
+
+        scenes = [
+            d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))
+        ]
+
+        if self.test_id is not None:
+            self.scene_list = [self.test_id]
+
+        else:
+            self.scene_list = scenes
+
+        print(f"Found {len(self.scene_list)} sequences in split {self.split}")
+
+    def load_poses(self, path):
+        file = open(path, "r")
+        lines = file.readlines()
+        file.close()
+        poses = []
+        valid = []
+        lines_per_matrix = 4
+        for i in range(0, len(lines), lines_per_matrix):
+            if "nan" in lines[i]:
+                valid.append(False)
+                poses.append(np.eye(4, 4, dtype=np.float32).tolist())
+            else:
+                valid.append(True)
+                pose_floats = [
+                    [float(x) for x in line.split()]
+                    for line in lines[i : i + lines_per_matrix]
+                ]
+                poses.append(pose_floats)
+
+        return np.array(poses, dtype=np.float32), valid
+
+    def _get_views(self, idx, resolution, rng):
+
+        if self.tuple_list is not None:
+            line = self.tuple_list[idx].split(" ")
+            scene_id = line[0]
+            img_idxs = line[1:]
+
+        else:
+            scene_id = self.scene_list[idx // self.num_seq]
+
+            num_files = len(os.listdir(os.path.join(self.ROOT, scene_id, "images")))
+            img_idxs = [f"{i}" for i in range(num_files)]
+            img_idxs = img_idxs[:: min(self.kf_every, len(img_idxs) // 2)]
+
+        fx, fy, cx, cy = 554.2562584220408, 554.2562584220408, 320, 240
+        intrinsics_ = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+        posepath = osp.join(self.ROOT, scene_id, f"poses.txt")
+        camera_poses, valids = self.load_poses(posepath)
+
+        imgs_idxs = deque(img_idxs)
+        if self.shuffle_seed >= 0:
+            imgs_idxs = shuffle_deque(imgs_idxs)
+        views = []
+
+        while len(imgs_idxs) > 0:
+            im_idx = imgs_idxs.popleft()
+
+            impath = osp.join(self.ROOT, scene_id, "images", f"img{im_idx}.png")
+            depthpath = osp.join(self.ROOT, scene_id, "depth", f"depth{im_idx}.png")
+
+            rgb_image = imread_cv2(impath)
+            depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED)
+            depthmap = np.nan_to_num(depthmap.astype(np.float32), 0.0) / 1000.0
+            depthmap[depthmap > 10] = 0
+            depthmap[depthmap < 1e-3] = 0
+
+            rgb_image = cv2.resize(rgb_image, (depthmap.shape[1], depthmap.shape[0]))
+
+            camera_pose = camera_poses[int(im_idx)]
+            # gl to cv
+            camera_pose[:, 1:3] *= -1.0
+            if resolution != (224, 224) or self.rebuttal:
+                rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                    rgb_image, depthmap, intrinsics_, resolution, rng=rng, info=impath
+                )
+            else:
+                rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                    rgb_image, depthmap, intrinsics_, (512, 384), rng=rng, info=impath
+                )
+                W, H = rgb_image.size
+                cx = W // 2
+                cy = H // 2
+                l, t = cx - 112, cy - 112
+                r, b = cx + 112, cy + 112
+                crop_bbox = (l, t, r, b)
+                rgb_image, depthmap, intrinsics = cropping.crop_image_depthmap(
+                    rgb_image, depthmap, intrinsics, crop_bbox
+                )
+
+            views.append(
+                dict(
+                    img=rgb_image,
+                    depthmap=depthmap,
+                    camera_pose=camera_pose,
+                    camera_intrinsics=intrinsics,
+                    dataset="nrgbd",
+                    label=osp.join(scene_id, im_idx),
+                    instance=impath,
+                )
+            )
+
+        return views

FastVGGT/eval/dataset_utils/__init__.py

@@ -0,0 +1 @@
+

FastVGGT/eval/dataset_utils/corr.py

@@ -0,0 +1,234 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+
+import numpy as np
+import torch
+
+
+def todevice(batch, device, callback=None, non_blocking=False):
+    """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).
+
+    batch: list, tuple, dict of tensors or other things
+    device: pytorch device or 'numpy'
+    callback: function that would be called on every sub-elements.
+    """
+    if callback:
+        batch = callback(batch)
+
+    if isinstance(batch, dict):
+        return {k: todevice(v, device) for k, v in batch.items()}
+
+    if isinstance(batch, (tuple, list)):
+        return type(batch)(todevice(x, device) for x in batch)
+
+    x = batch
+    if device == "numpy":
+        if isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+    elif x is not None:
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x)
+        if torch.is_tensor(x):
+            x = x.to(device, non_blocking=non_blocking)
+    return x
+
+
+to_device = todevice  # alias
+
+
+def to_numpy(x):
+    return todevice(x, "numpy")
+
+
+def geotrf(Trf, pts, ncol=None, norm=False):
+    """Apply a geometric transformation to a list of 3-D points.
+
+    H: 3x3 or 4x4 projection matrix (typically a Homography)
+    p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3)
+
+    ncol: int. number of columns of the result (2 or 3)
+    norm: float. if != 0, the resut is projected on the z=norm plane.
+
+    Returns an array of projected 2d points.
+    """
+    assert Trf.ndim >= 2
+    if isinstance(Trf, np.ndarray):
+        pts = np.asarray(pts)
+    elif isinstance(Trf, torch.Tensor):
+        pts = torch.as_tensor(pts, dtype=Trf.dtype)
+
+    output_reshape = pts.shape[:-1]
+    ncol = ncol or pts.shape[-1]
+
+    if (
+        isinstance(Trf, torch.Tensor)
+        and isinstance(pts, torch.Tensor)
+        and Trf.ndim == 3
+        and pts.ndim == 4
+    ):
+        d = pts.shape[3]
+        if Trf.shape[-1] == d:
+            pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts)
+        elif Trf.shape[-1] == d + 1:
+            pts = (
+                torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts)
+                + Trf[:, None, None, :d, d]
+            )
+        else:
+            raise ValueError(f"bad shape, not ending with 3 or 4, for {pts.shape=}")
+    else:
+        if Trf.ndim >= 3:
+            n = Trf.ndim - 2
+            assert Trf.shape[:n] == pts.shape[:n], "batch size does not match"
+            Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1])
+
+            if pts.ndim > Trf.ndim:
+
+                pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1])
+            elif pts.ndim == 2:
+
+                pts = pts[:, None, :]
+
+        if pts.shape[-1] + 1 == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :]
+        elif pts.shape[-1] == Trf.shape[-1]:
+            Trf = Trf.swapaxes(-1, -2)  # transpose Trf
+            pts = pts @ Trf
+        else:
+            pts = Trf @ pts.T
+            if pts.ndim >= 2:
+                pts = pts.swapaxes(-1, -2)
+
+    if norm:
+        pts = pts / pts[..., -1:]  # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG
+        if norm != 1:
+            pts *= norm
+
+    res = pts[..., :ncol].reshape(*output_reshape, ncol)
+    return res
+
+
+def inv(mat):
+    """Invert a torch or numpy matrix"""
+    if isinstance(mat, torch.Tensor):
+        return torch.linalg.inv(mat)
+    if isinstance(mat, np.ndarray):
+        return np.linalg.inv(mat)
+    raise ValueError(f"bad matrix type = {type(mat)}")
+
+
+def reproject_view(pts3d, view2):
+    shape = view2["pts3d"].shape[:2]
+    return reproject(
+        pts3d, view2["camera_intrinsics"], inv(view2["camera_pose"]), shape
+    )
+
+
+def reproject(pts3d, K, world2cam, shape):
+    H, W, THREE = pts3d.shape
+    assert THREE == 3
+
+    with np.errstate(divide="ignore", invalid="ignore"):
+        pos = geotrf(K @ world2cam[:3], pts3d, norm=1, ncol=2)
+
+    return (H, W), ravel_xy(pos, shape)
+
+
+def ravel_xy(pos, shape):
+    H, W = shape
+    with np.errstate(invalid="ignore"):
+        qx, qy = pos.reshape(-1, 2).round().astype(np.int32).T
+    quantized_pos = qx.clip(min=0, max=W - 1, out=qx) + W * qy.clip(
+        min=0, max=H - 1, out=qy
+    )
+    return quantized_pos
+
+
+def unravel_xy(pos, shape):
+
+    return np.unravel_index(pos, shape)[0].base[:, ::-1].copy()
+
+
+def reciprocal_1d(corres_1_to_2, corres_2_to_1, ret_recip=False):
+    is_reciprocal1 = corres_2_to_1[corres_1_to_2] == np.arange(len(corres_1_to_2))
+    pos1 = is_reciprocal1.nonzero()[0]
+    pos2 = corres_1_to_2[pos1]
+    if ret_recip:
+        return is_reciprocal1, pos1, pos2
+    return pos1, pos2
+
+
+def extract_correspondences_from_pts3d(
+    view1, view2, target_n_corres, rng=np.random, ret_xy=True, nneg=0
+):
+    view1, view2 = to_numpy((view1, view2))
+
+    shape1, corres1_to_2 = reproject_view(view1["pts3d"], view2)
+    shape2, corres2_to_1 = reproject_view(view2["pts3d"], view1)
+
+    is_reciprocal1, pos1, pos2 = reciprocal_1d(
+        corres1_to_2, corres2_to_1, ret_recip=True
+    )
+    is_reciprocal2 = corres1_to_2[corres2_to_1] == np.arange(len(corres2_to_1))
+
+    if target_n_corres is None:
+        if ret_xy:
+            pos1 = unravel_xy(pos1, shape1)
+            pos2 = unravel_xy(pos2, shape2)
+        return pos1, pos2
+
+    available_negatives = min((~is_reciprocal1).sum(), (~is_reciprocal2).sum())
+    target_n_positives = int(target_n_corres * (1 - nneg))
+    n_positives = min(len(pos1), target_n_positives)
+    n_negatives = min(target_n_corres - n_positives, available_negatives)
+
+    if n_negatives + n_positives != target_n_corres:
+
+        n_positives = target_n_corres - n_negatives
+        assert n_positives <= len(pos1)
+
+    assert n_positives <= len(pos1)
+    assert n_positives <= len(pos2)
+    assert n_negatives <= (~is_reciprocal1).sum()
+    assert n_negatives <= (~is_reciprocal2).sum()
+    assert n_positives + n_negatives == target_n_corres
+
+    valid = np.ones(n_positives, dtype=bool)
+    if n_positives < len(pos1):
+
+        perm = rng.permutation(len(pos1))[:n_positives]
+        pos1 = pos1[perm]
+        pos2 = pos2[perm]
+
+    if n_negatives > 0:
+
+        def norm(p):
+            return p / p.sum()
+
+        pos1 = np.r_[
+            pos1,
+            rng.choice(
+                shape1[0] * shape1[1],
+                size=n_negatives,
+                replace=False,
+                p=norm(~is_reciprocal1),
+            ),
+        ]
+        pos2 = np.r_[
+            pos2,
+            rng.choice(
+                shape2[0] * shape2[1],
+                size=n_negatives,
+                replace=False,
+                p=norm(~is_reciprocal2),
+            ),
+        ]
+        valid = np.r_[valid, np.zeros(n_negatives, dtype=bool)]
+
+    if ret_xy:
+        pos1 = unravel_xy(pos1, shape1)
+        pos2 = unravel_xy(pos2, shape2)
+    return pos1, pos2, valid

FastVGGT/eval/dataset_utils/cropping.py

@@ -0,0 +1,140 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+
+import PIL.Image
+import os
+
+from utils import colmap_to_opencv_intrinsics, opencv_to_colmap_intrinsics
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2  # noqa
+import numpy as np  # noqa
+
+try:
+    lanczos = PIL.Image.Resampling.LANCZOS
+    bicubic = PIL.Image.Resampling.BICUBIC
+except AttributeError:
+    lanczos = PIL.Image.LANCZOS
+    bicubic = PIL.Image.BICUBIC
+
+
+class ImageList:
+    """Convenience class to aply the same operation to a whole set of images."""
+
+    def __init__(self, images):
+        if not isinstance(images, (tuple, list, set)):
+            images = [images]
+        self.images = []
+        for image in images:
+            if not isinstance(image, PIL.Image.Image):
+                image = PIL.Image.fromarray(image)
+            self.images.append(image)
+
+    def __len__(self):
+        return len(self.images)
+
+    def to_pil(self):
+        return tuple(self.images) if len(self.images) > 1 else self.images[0]
+
+    @property
+    def size(self):
+        sizes = [im.size for im in self.images]
+        assert all(sizes[0] == s for s in sizes)
+        return sizes[0]
+
+    def resize(self, *args, **kwargs):
+        return ImageList(self._dispatch("resize", *args, **kwargs))
+
+    def crop(self, *args, **kwargs):
+        return ImageList(self._dispatch("crop", *args, **kwargs))
+
+    def _dispatch(self, func, *args, **kwargs):
+        return [getattr(im, func)(*args, **kwargs) for im in self.images]
+
+
+def rescale_image_depthmap(
+    image, depthmap, camera_intrinsics, output_resolution, force=True
+):
+    """Jointly rescale a (image, depthmap)
+    so that (out_width, out_height) >= output_res
+    """
+    image = ImageList(image)
+    input_resolution = np.array(image.size)  # (W,H)
+    output_resolution = np.array(output_resolution)
+    if depthmap is not None:
+
+        assert tuple(depthmap.shape[:2]) == image.size[::-1]
+
+    assert output_resolution.shape == (2,)
+    scale_final = max(output_resolution / image.size) + 1e-8
+    if scale_final >= 1 and not force:  # image is already smaller than what is asked
+        return (image.to_pil(), depthmap, camera_intrinsics)
+    output_resolution = np.floor(input_resolution * scale_final).astype(int)
+
+    image = image.resize(
+        output_resolution, resample=lanczos if scale_final < 1 else bicubic
+    )
+    if depthmap is not None:
+        depthmap = cv2.resize(
+            depthmap,
+            output_resolution,
+            fx=scale_final,
+            fy=scale_final,
+            interpolation=cv2.INTER_NEAREST,
+        )
+
+    camera_intrinsics = camera_matrix_of_crop(
+        camera_intrinsics, input_resolution, output_resolution, scaling=scale_final
+    )
+
+    return image.to_pil(), depthmap, camera_intrinsics
+
+
+def camera_matrix_of_crop(
+    input_camera_matrix,
+    input_resolution,
+    output_resolution,
+    scaling=1,
+    offset_factor=0.5,
+    offset=None,
+):
+
+    margins = np.asarray(input_resolution) * scaling - output_resolution
+    assert np.all(margins >= 0.0)
+    if offset is None:
+        offset = offset_factor * margins
+
+    output_camera_matrix_colmap = opencv_to_colmap_intrinsics(input_camera_matrix)
+    output_camera_matrix_colmap[:2, :] *= scaling
+    output_camera_matrix_colmap[:2, 2] -= offset
+    output_camera_matrix = colmap_to_opencv_intrinsics(output_camera_matrix_colmap)
+
+    return output_camera_matrix
+
+
+def crop_image_depthmap(image, depthmap, camera_intrinsics, crop_bbox):
+    """
+    Return a crop of the input view.
+    """
+    image = ImageList(image)
+    l, t, r, b = crop_bbox
+
+    image = image.crop((l, t, r, b))
+    depthmap = depthmap[t:b, l:r]
+
+    camera_intrinsics = camera_intrinsics.copy()
+    camera_intrinsics[0, 2] -= l
+    camera_intrinsics[1, 2] -= t
+
+    return image.to_pil(), depthmap, camera_intrinsics
+
+
+def bbox_from_intrinsics_in_out(
+    input_camera_matrix, output_camera_matrix, output_resolution
+):
+    out_width, out_height = output_resolution
+    l, t = np.int32(np.round(input_camera_matrix[:2, 2] - output_camera_matrix[:2, 2]))
+    crop_bbox = (l, t, l + out_width, t + out_height)
+    return crop_bbox

FastVGGT/eval/dataset_utils/transforms.py

@@ -0,0 +1,78 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+
+import torchvision.transforms as tvf
+
+ImgNorm = tvf.Compose([tvf.ToTensor(), tvf.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+
+ColorJitter = tvf.Compose([tvf.ColorJitter(0.5, 0.5, 0.5, 0.1), ImgNorm])
+
+
+def _check_input(value, center=1, bound=(0, float("inf")), clip_first_on_zero=True):
+    if isinstance(value, (int, float)):
+        if value < 0:
+            raise ValueError(f"If  is a single number, it must be non negative.")
+        value = [center - float(value), center + float(value)]
+        if clip_first_on_zero:
+            value[0] = max(value[0], 0.0)
+    elif isinstance(value, (tuple, list)) and len(value) == 2:
+        value = [float(value[0]), float(value[1])]
+    else:
+        raise TypeError(f"should be a single number or a list/tuple with length 2.")
+
+    if not bound[0] <= value[0] <= value[1] <= bound[1]:
+        raise ValueError(f"values should be between {bound}, but got {value}.")
+
+    if value[0] == value[1] == center:
+        return None
+    else:
+        return tuple(value)
+
+
+import torch
+import torchvision.transforms.functional as F
+
+
+def SeqColorJitter():
+    """
+    Return a color jitter transform with same random parameters
+    """
+    brightness = _check_input(0.5)
+    contrast = _check_input(0.5)
+    saturation = _check_input(0.5)
+    hue = _check_input(0.1, center=0, bound=(-0.5, 0.5), clip_first_on_zero=False)
+
+    fn_idx = torch.randperm(4)
+    brightness_factor = (
+        None
+        if brightness is None
+        else float(torch.empty(1).uniform_(brightness[0], brightness[1]))
+    )
+    contrast_factor = (
+        None
+        if contrast is None
+        else float(torch.empty(1).uniform_(contrast[0], contrast[1]))
+    )
+    saturation_factor = (
+        None
+        if saturation is None
+        else float(torch.empty(1).uniform_(saturation[0], saturation[1]))
+    )
+    hue_factor = None if hue is None else float(torch.empty(1).uniform_(hue[0], hue[1]))
+
+    def _color_jitter(img):
+        for fn_id in fn_idx:
+            if fn_id == 0 and brightness_factor is not None:
+                img = F.adjust_brightness(img, brightness_factor)
+            elif fn_id == 1 and contrast_factor is not None:
+                img = F.adjust_contrast(img, contrast_factor)
+            elif fn_id == 2 and saturation_factor is not None:
+                img = F.adjust_saturation(img, saturation_factor)
+            elif fn_id == 3 and hue_factor is not None:
+                img = F.adjust_hue(img, hue_factor)
+        return ImgNorm(img)
+
+    return _color_jitter

FastVGGT/eval/eval_7andN.py

@@ -0,0 +1,497 @@
+import os
+import sys
+
+# Ensure project root is on sys.path for absolute imports like `vggt.*`
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+import time
+import torch
+import argparse
+import numpy as np
+import open3d as o3d
+import os.path as osp
+from torch.utils.data import DataLoader
+from torch.utils.data._utils.collate import default_collate
+from tqdm import tqdm
+from collections import defaultdict
+import torchvision.transforms as transforms
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser("3D Reconstruction evaluation", add_help=False)
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="/home/sy/code/FastVGGT/ckpt/model_tracker_fixed_e20.pt",
+        help="ckpt name",
+    )
+    parser.add_argument("--device", type=str, default="cuda:0", help="device")
+    parser.add_argument("--model_name", type=str, default="VGGT")
+    parser.add_argument(
+        "--conf_thresh", type=float, default=0.0, help="confidence threshold"
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="/home/sy/code/FastVGGT/eval_results",
+        help="value for outdir",
+    )
+    parser.add_argument("--size", type=int, default=518)
+    parser.add_argument("--revisit", type=int, default=1, help="revisit times")
+    parser.add_argument("--freeze", action="store_true")
+    parser.add_argument("--use_proj", action="store_true")
+    parser.add_argument(
+        "--merging", type=int, default=0, help="VGGT aggregator merging steps"
+    )
+    parser.add_argument("--kf", type=int, default=2, help="key frame")
+    return parser
+
+
+def main(args):
+    from data import SevenScenes, NRGBD
+    from utils import accuracy, completion
+
+    if args.size == 512:
+        resolution = (512, 384)
+    elif args.size == 224:
+        resolution = 224
+    elif args.size == 518:
+        resolution = (518, 392)
+    else:
+        raise NotImplementedError
+    datasets_all = {
+        "7scenes": SevenScenes(
+            split="test",
+            ROOT="/data/sy/7scenes",
+            resolution=resolution,
+            num_seq=1,
+            full_video=True,
+            kf_every=args.kf,
+        ),  # 20),
+        "NRGBD": NRGBD(
+            split="test",
+            ROOT="/data/sy/neural_rgbd_data",
+            resolution=resolution,
+            num_seq=1,
+            full_video=True,
+            kf_every=args.kf,
+        ),
+    }
+
+    device = args.device
+    model_name = args.model_name
+
+    from vggt.models.vggt import VGGT
+    from vggt.utils.pose_enc import pose_encoding_to_extri_intri
+    from vggt.utils.geometry import unproject_depth_map_to_point_map
+    from criterion import Regr3D_t_ScaleShiftInv, L21
+
+    # Force use of bf16 data type
+    dtype = torch.bfloat16
+    # Load VGGT model
+    model = VGGT(merging=args.merging, enable_point=True)
+    ckpt = torch.load(args.ckpt_path, map_location="cpu")
+
+    # ✅ Fix: load pre-trained weights
+    model.load_state_dict(
+        ckpt, strict=False
+    )  # Use strict=False due to enable_point=True difference
+
+    model = model.cuda().eval()
+    model = model.to(torch.bfloat16)
+
+    del ckpt
+    os.makedirs(osp.join(args.output_dir, f"{args.kf}"), exist_ok=True)
+
+    criterion = Regr3D_t_ScaleShiftInv(L21, norm_mode=False, gt_scale=True)
+
+    with torch.no_grad():
+        for name_data, dataset in datasets_all.items():
+            save_path = osp.join(osp.join(args.output_dir, f"{args.kf}"), name_data)
+            os.makedirs(save_path, exist_ok=True)
+            log_file = osp.join(save_path, "logs.txt")
+
+            acc_all = 0
+            acc_all_med = 0
+            comp_all = 0
+            comp_all_med = 0
+            nc1_all = 0
+            nc1_all_med = 0
+            nc2_all = 0
+            nc2_all_med = 0
+            scene_infer_times = defaultdict(list)
+
+            for data_idx in tqdm(range(len(dataset))):
+                batch = default_collate([dataset[data_idx]])
+                ignore_keys = set(
+                    [
+                        "depthmap",
+                        "dataset",
+                        "label",
+                        "instance",
+                        "idx",
+                        "true_shape",
+                        "rng",
+                    ]
+                )
+                for view in batch:
+                    for name in view.keys():  # pseudo_focal
+                        if name in ignore_keys:
+                            continue
+                        if isinstance(view[name], tuple) or isinstance(
+                            view[name], list
+                        ):
+                            view[name] = [
+                                x.to(device, non_blocking=True) for x in view[name]
+                            ]
+                        else:
+                            view[name] = view[name].to(device, non_blocking=True)
+
+                pts_all = []
+                pts_gt_all = []
+                images_all = []
+                masks_all = []
+                conf_all = []
+                in_camera1 = None
+
+                dtype = (
+                    torch.bfloat16
+                    if torch.cuda.get_device_capability()[0] >= 8
+                    else torch.float16
+                )
+                with torch.cuda.amp.autocast(dtype=dtype):
+                    if isinstance(batch, dict) and "img" in batch:
+                        batch["img"] = (batch["img"] + 1.0) / 2.0
+                    elif isinstance(batch, list) and all(
+                        isinstance(v, dict) and "img" in v for v in batch
+                    ):
+                        for view in batch:
+                            view["img"] = (view["img"] + 1.0) / 2.0
+                        # Gather all `img` tensors into a single tensor of shape [N, C, H, W]
+                        imgs_tensor = torch.cat([v["img"] for v in batch], dim=0)
+
+                with torch.cuda.amp.autocast(dtype=dtype):
+                    with torch.no_grad():
+                        torch.cuda.synchronize()
+                        start = time.time()
+                        preds = model(imgs_tensor)
+                        torch.cuda.synchronize()
+                        end = time.time()
+                        inference_time_ms = (end - start) * 1000
+                        print(f"Inference time: {inference_time_ms:.2f}ms")
+
+                    # Wrap model outputs per-view to align with batch later
+                    predictions = preds
+                    views = batch  # list[dict]
+                    if "pose_enc" in predictions:
+                        B, S = predictions["pose_enc"].shape[:2]
+                    elif "world_points" in predictions:
+                        B, S = predictions["world_points"].shape[:2]
+                    else:
+                        raise KeyError(
+                            "predictions is missing a key to infer sequence length"
+                        )
+
+                    ress = []
+                    for s in range(S):
+                        res = {
+                            "pts3d_in_other_view": predictions["world_points"][:, s],
+                            "conf": predictions["world_points_conf"][:, s],
+                            "depth": predictions["depth"][:, s],
+                            "depth_conf": predictions["depth_conf"][:, s],
+                            "camera_pose": predictions["pose_enc"][:, s, :],
+                        }
+                        if (
+                            isinstance(views, list)
+                            and s < len(views)
+                            and "valid_mask" in views[s]
+                        ):
+                            res["valid_mask"] = views[s]["valid_mask"]
+                        if "track" in predictions:
+                            res.update(
+                                {
+                                    "track": predictions["track"][:, s],
+                                    "vis": (
+                                        predictions.get("vis", None)[:, s]
+                                        if "vis" in predictions
+                                        else None
+                                    ),
+                                    "track_conf": (
+                                        predictions.get("conf", None)[:, s]
+                                        if "conf" in predictions
+                                        else None
+                                    ),
+                                }
+                            )
+                        ress.append(res)
+
+                    preds = ress
+
+                    valid_length = len(preds) // args.revisit
+                    if args.revisit > 1:
+                        preds = preds[-valid_length:]
+                        batch = batch[-valid_length:]
+
+                    # Evaluation
+                    print(f"Evaluation for {name_data} {data_idx+1}/{len(dataset)}")
+                    gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = (
+                        criterion.get_all_pts3d_t(batch, preds)
+                    )
+
+                    in_camera1 = None
+                    pts_all = []
+                    pts_gt_all = []
+                    images_all = []
+                    masks_all = []
+                    conf_all = []
+
+                    for j, view in enumerate(batch):
+                        if in_camera1 is None:
+                            in_camera1 = view["camera_pose"][0].cpu()
+
+                        image = view["img"].permute(0, 2, 3, 1).cpu().numpy()[0]
+                        mask = view["valid_mask"].cpu().numpy()[0]
+
+                        pts = pred_pts[j].cpu().numpy()[0]
+                        conf = preds[j]["conf"].cpu().data.numpy()[0]
+
+                        # mask = mask & (conf > 1.8)
+
+                        pts_gt = gt_pts[j].detach().cpu().numpy()[0]
+
+                        H, W = image.shape[:2]
+                        cx = W // 2
+                        cy = H // 2
+                        l, t = cx - 112, cy - 112
+                        r, b = cx + 112, cy + 112
+                        image = image[t:b, l:r]
+                        mask = mask[t:b, l:r]
+                        pts = pts[t:b, l:r]
+                        pts_gt = pts_gt[t:b, l:r]
+
+                        images_all.append(image[None, ...])
+                        pts_all.append(pts[None, ...])
+                        pts_gt_all.append(pts_gt[None, ...])
+                        masks_all.append(mask[None, ...])
+                        conf_all.append(conf[None, ...])
+
+                images_all = np.concatenate(images_all, axis=0)
+                pts_all = np.concatenate(pts_all, axis=0)
+                pts_gt_all = np.concatenate(pts_gt_all, axis=0)
+                masks_all = np.concatenate(masks_all, axis=0)
+
+                scene_id = view["label"][0].rsplit("/", 1)[0]
+                # Record average inference time per scene
+                try:
+                    scene_infer_times[scene_id].append(float(inference_time_ms))
+                except Exception:
+                    pass
+
+                save_params = {}
+
+                save_params["images_all"] = images_all
+                save_params["pts_all"] = pts_all
+                save_params["pts_gt_all"] = pts_gt_all
+                save_params["masks_all"] = masks_all
+
+                pts_all_masked = pts_all[masks_all > 0]
+                pts_gt_all_masked = pts_gt_all[masks_all > 0]
+                images_all_masked = images_all[masks_all > 0]
+
+                mask = np.isfinite(pts_all_masked)
+                pts_all_masked = pts_all_masked[mask]
+
+                mask_gt = np.isfinite(pts_gt_all_masked)
+                pts_gt_all_masked = pts_gt_all_masked[mask_gt]
+                images_all_masked = images_all_masked[mask]
+
+                # Reshape to point cloud (N, 3) before sampling
+                pts_all_masked = pts_all_masked.reshape(-1, 3)
+                pts_gt_all_masked = pts_gt_all_masked.reshape(-1, 3)
+                images_all_masked = images_all_masked.reshape(-1, 3)
+
+                # If number of points exceeds threshold, sample by points
+                if pts_all_masked.shape[0] > 999999:
+                    sample_indices = np.random.choice(
+                        pts_all_masked.shape[0], 999999, replace=False
+                    )
+                    pts_all_masked = pts_all_masked[sample_indices]
+                    images_all_masked = images_all_masked[sample_indices]
+
+                # Apply the same sampling to GT point cloud
+                if pts_gt_all_masked.shape[0] > 999999:
+                    sample_indices_gt = np.random.choice(
+                        pts_gt_all_masked.shape[0], 999999, replace=False
+                    )
+                    pts_gt_all_masked = pts_gt_all_masked[sample_indices_gt]
+
+                if args.use_proj:
+
+                    def umeyama_alignment(
+                        src: np.ndarray, dst: np.ndarray, with_scale: bool = True
+                    ):
+                        assert src.shape == dst.shape
+                        N, dim = src.shape
+
+                        mu_src = src.mean(axis=0)
+                        mu_dst = dst.mean(axis=0)
+                        src_c = src - mu_src
+                        dst_c = dst - mu_dst
+
+                        Sigma = dst_c.T @ src_c / N  # (3,3)
+
+                        U, D, Vt = np.linalg.svd(Sigma)
+
+                        S = np.eye(dim)
+                        if np.linalg.det(U) * np.linalg.det(Vt) < 0:
+                            S[-1, -1] = -1
+
+                        R = U @ S @ Vt
+
+                        if with_scale:
+                            var_src = (src_c**2).sum() / N
+                            s = (D * S.diagonal()).sum() / var_src
+                        else:
+                            s = 1.0
+
+                        t = mu_dst - s * R @ mu_src
+
+                        return s, R, t
+
+                    pts_all_masked = pts_all_masked.reshape(-1, 3)
+                    pts_gt_all_masked = pts_gt_all_masked.reshape(-1, 3)
+                    s, R, t = umeyama_alignment(
+                        pts_all_masked, pts_gt_all_masked, with_scale=True
+                    )
+                    pts_all_aligned = (s * (R @ pts_all_masked.T)).T + t  # (N,3)
+                    pts_all_masked = pts_all_aligned
+
+                pcd = o3d.geometry.PointCloud()
+                pcd.points = o3d.utility.Vector3dVector(pts_all_masked)
+                pcd.colors = o3d.utility.Vector3dVector(images_all_masked)
+
+                pcd_gt = o3d.geometry.PointCloud()
+                pcd_gt.points = o3d.utility.Vector3dVector(pts_gt_all_masked)
+                pcd_gt.colors = o3d.utility.Vector3dVector(images_all_masked)
+
+                trans_init = np.eye(4)
+
+                threshold = 0.1
+                reg_p2p = o3d.pipelines.registration.registration_icp(
+                    pcd,
+                    pcd_gt,
+                    threshold,
+                    trans_init,
+                    o3d.pipelines.registration.TransformationEstimationPointToPoint(),
+                )
+
+                transformation = reg_p2p.transformation
+
+                pcd = pcd.transform(transformation)
+                pcd.estimate_normals()
+                pcd_gt.estimate_normals()
+
+                gt_normal = np.asarray(pcd_gt.normals)
+                pred_normal = np.asarray(pcd.normals)
+
+                acc, acc_med, nc1, nc1_med = accuracy(
+                    pcd_gt.points, pcd.points, gt_normal, pred_normal
+                )
+                comp, comp_med, nc2, nc2_med = completion(
+                    pcd_gt.points, pcd.points, gt_normal, pred_normal
+                )
+                print(
+                    f"Idx: {scene_id}, Acc: {acc}, Comp: {comp}, NC1: {nc1}, NC2: {nc2} - Acc_med: {acc_med}, Compc_med: {comp_med}, NC1c_med: {nc1_med}, NC2c_med: {nc2_med}"
+                )
+                print(
+                    f"Idx: {scene_id}, Acc: {acc}, Comp: {comp}, NC1: {nc1}, NC2: {nc2} - Acc_med: {acc_med}, Compc_med: {comp_med}, NC1c_med: {nc1_med}, NC2c_med: {nc2_med}",
+                    file=open(log_file, "a"),
+                )
+
+                acc_all += acc
+                comp_all += comp
+                nc1_all += nc1
+                nc2_all += nc2
+
+                acc_all_med += acc_med
+                comp_all_med += comp_med
+                nc1_all_med += nc1_med
+                nc2_all_med += nc2_med
+
+                # release cuda memory
+                torch.cuda.empty_cache()
+
+            # Get depth from pcd and run TSDFusion
+            to_write = ""
+            # Read the log file
+            if os.path.exists(osp.join(save_path, "logs.txt")):
+                with open(osp.join(save_path, "logs.txt"), "r") as f_sub:
+                    to_write += f_sub.read()
+
+            with open(osp.join(save_path, f"logs_all.txt"), "w") as f:
+                log_data = to_write
+                metrics = defaultdict(list)
+                for line in log_data.strip().split("\n"):
+                    match = regex.match(line)
+                    if match:
+                        data = match.groupdict()
+                        # Exclude 'scene_id' from metrics as it's an identifier
+                        for key, value in data.items():
+                            if key != "scene_id":
+                                metrics[key].append(float(value))
+                        metrics["nc"].append(
+                            (float(data["nc1"]) + float(data["nc2"])) / 2
+                        )
+                        metrics["nc_med"].append(
+                            (float(data["nc1_med"]) + float(data["nc2_med"])) / 2
+                        )
+                mean_metrics = {
+                    metric: sum(values) / len(values)
+                    for metric, values in metrics.items()
+                }
+
+                c_name = "mean"
+                print_str = f"{c_name.ljust(20)}: "
+                for m_name in mean_metrics:
+                    print_num = np.mean(mean_metrics[m_name])
+                    print_str = print_str + f"{m_name}: {print_num:.3f} | "
+                print_str = print_str + "\n"
+                # Summarize per-scene average inference time
+                time_lines = []
+                for sid, times in scene_infer_times.items():
+                    if len(times) > 0:
+                        time_lines.append(
+                            f"Idx: {sid}, Time_avg_ms: {np.mean(times):.2f}"
+                        )
+                time_block = "\n".join(time_lines) + (
+                    "\n" if len(time_lines) > 0 else ""
+                )
+
+                f.write(to_write + time_block + print_str)
+
+
+from collections import defaultdict
+import re
+
+pattern = r"""
+    Idx:\s*(?P<scene_id>[^,]+),\s*
+    Acc:\s*(?P<acc>[^,]+),\s*
+    Comp:\s*(?P<comp>[^,]+),\s*
+    NC1:\s*(?P<nc1>[^,]+),\s*
+    NC2:\s*(?P<nc2>[^,]+)\s*-\s*
+    Acc_med:\s*(?P<acc_med>[^,]+),\s*
+    Compc_med:\s*(?P<comp_med>[^,]+),\s*
+    NC1c_med:\s*(?P<nc1_med>[^,]+),\s*
+    NC2c_med:\s*(?P<nc2_med>[^,]+)
+"""
+
+regex = re.compile(pattern, re.VERBOSE)
+
+
+if __name__ == "__main__":
+    parser = get_args_parser()
+    args = parser.parse_args()
+
+    main(args)

FastVGGT/eval/eval_custom.py

@@ -0,0 +1,467 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+import os
+import sys
+import matplotlib.pyplot as plt
+from scipy.spatial.transform import Rotation
+
+# Ensure project root is in sys.path for absolute imports like `vggt.*`
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+from vggt.models.vggt import VGGT
+from vggt.utils.eval_utils import (
+    load_poses,
+    get_vgg_input_imgs,
+    get_sorted_image_paths,
+    build_frame_selection,
+    load_images_rgb,
+    infer_vggt_and_reconstruct,
+    evaluate_scene_and_save,
+)
+
+# Import pose visualization libraries (optional EVO support)
+try:
+    from evo.core.trajectory import PoseTrajectory3D
+    import evo.tools.plot as plot
+
+    EVO_AVAILABLE = True
+except ImportError:
+    # EVO is optional; we have a matplotlib-based fallback
+    EVO_AVAILABLE = False
+
+
+def visualize_predicted_poses(
+    all_cam_to_world_mat, frame_ids, output_scene_dir, scene_name="custom_dataset"
+):
+    """
+    Visualize the predicted camera pose trajectory (no GT comparison required).
+
+    Args:
+        all_cam_to_world_mat: List of camera-to-world transform matrices
+        frame_ids: List of frame IDs
+        output_scene_dir: Output directory
+        scene_name: Scene name
+    """
+    # Provide basic pose visualization even without EVO
+    if not EVO_AVAILABLE:
+        print("⚠️  EVO not installed; using basic matplotlib visualization")
+
+    try:
+        # Convert to numpy array
+        poses_est = np.array(all_cam_to_world_mat)
+
+        if len(poses_est) < 2:
+            print("⚠️  Not enough poses to generate trajectory plot")
+            return
+
+        print(f"🎨 Generating pose trajectory visualization...")
+
+        # Extract translation part
+        positions = poses_est[:, :3, 3]  # shape: (N, 3)
+
+        # Create figure - show XZ-plane projection only
+        fig, ax = plt.subplots(1, 1, figsize=(10, 8))
+
+        # XZ-plane projection
+        ax.plot(
+            positions[:, 0],
+            positions[:, 2],
+            "b-",
+            linewidth=2,
+            label="Predicted Trajectory",
+        )
+        ax.scatter(
+            positions[0, 0], positions[0, 2], color="green", s=100, label="Start"
+        )
+        ax.scatter(positions[-1, 0], positions[-1, 2], color="red", s=100, label="End")
+        ax.set_xlabel("X (m)")
+        ax.set_ylabel("Z (m)")
+        ax.set_title(f"{scene_name} - XZ-plane projection")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+        # Save image
+        pose_plot_path = output_scene_dir / "predicted_trajectory.png"
+        plt.savefig(pose_plot_path, dpi=300, bbox_inches="tight")
+        plt.close()
+
+        print(f"📊 Trajectory visualization saved: {pose_plot_path}")
+
+    except Exception as e:
+        print(f"⚠️  Failed to generate pose visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+def main():
+    """
+    Evaluation script for a Custom Dataset.
+    Supports optional evaluation and custom dataset structure.
+    """
+    parser = argparse.ArgumentParser(
+        description="Run FastVGGT evaluation on a Custom Dataset"
+    )
+
+    # Required: dataset path
+    parser.add_argument(
+        "--data_path",
+        type=Path,
+        required=True,
+        help="Dataset path containing subfolders: color, depth, gt_ply, pose",
+    )
+
+    # Optional: enable evaluation
+    parser.add_argument(
+        "--enable_evaluation",
+        action="store_true",
+        help="Enable evaluation (requires pose and ply data)",
+    )
+
+    # Output path
+    parser.add_argument(
+        "--output_path",
+        type=Path,
+        default="./eval_results_custom",
+        help="Output path for evaluation results",
+    )
+
+    # Model parameters
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="/home/sy/code/FastVGGT/ckpt/model_tracker_fixed_e20.pt",
+        help="Model checkpoint file path",
+    )
+
+    parser.add_argument("--merging", type=int, default=0, help="Merging parameter")
+
+    # Processing parameters
+    parser.add_argument(
+        "--input_frame",
+        type=int,
+        default=200,
+        help="Maximum number of frames to process per scene",
+    )
+
+    parser.add_argument(
+        "--depth_conf_thresh",
+        type=float,
+        default=3.0,
+        help="Depth confidence threshold to filter low-confidence depth values",
+    )
+
+    # Evaluation parameters (only used when evaluation is enabled)
+    parser.add_argument(
+        "--chamfer_max_dist",
+        type=float,
+        default=0.5,
+        help="Maximum distance threshold used in Chamfer Distance computation",
+    )
+
+    parser.add_argument("--plot", action="store_true", help="Whether to generate plots")
+
+    parser.add_argument(
+        "--vis_attn_map",
+        action="store_true",
+        help="Visualize attention maps during inference",
+    )
+
+    args = parser.parse_args()
+    torch.manual_seed(33)
+
+    # Check data path exists
+    if not args.data_path.exists():
+        print(f"❌ Error: Data path does not exist: {args.data_path}")
+        return
+
+    # Check required subdirectories
+    color_dir = args.data_path / "images"
+    pose_dir = args.data_path / "pose"
+
+    if not color_dir.exists():
+        print(f"❌ Error: color directory does not exist: {color_dir}")
+        return
+
+    print(f"📁 Dataset path: {args.data_path}")
+    # print(f"🔧 Enable evaluation: {'Yes' if args.enable_evaluation else 'No'}")
+
+    # If evaluation is enabled, check pose and gt_ply directories
+    if args.enable_evaluation:
+        if not pose_dir.exists():
+            print(f"❌ Error: Evaluation requires pose directory: {pose_dir}")
+            return
+
+        gt_ply_dir = args.data_path / "gt_ply"
+        if not gt_ply_dir.exists():
+            print(f"❌ Error: Evaluation requires gt_ply directory: {gt_ply_dir}")
+            return
+        print(f"📊 Evaluation will use Ground Truth")
+    else:
+        print(f"🏃 Inference only, no evaluation")
+
+    # Create output directory
+    args.output_path.mkdir(parents=True, exist_ok=True)
+    output_scene_dir = args.output_path / "custom_dataset"
+
+    # Check if already processed
+    if (output_scene_dir / "metrics.json").exists() and args.enable_evaluation:
+        print(
+            f"⚠️  Results already exist, skipping: {output_scene_dir / 'metrics.json'}"
+        )
+        return
+
+    # Force use of bf16 dtype
+    dtype = torch.bfloat16
+
+    # Load VGGT model
+    print(f"🔄 Loading model: {args.ckpt_path}")
+    model = VGGT(merging=args.merging, vis_attn_map=args.vis_attn_map)
+    ckpt = torch.load(args.ckpt_path, map_location="cpu")
+    incompat = model.load_state_dict(ckpt, strict=False)
+    # if incompat.missing_keys or incompat.unexpected_keys:
+    #     print(f"⚠️  Partially incompatible keys when loading model: {incompat}")
+    model = model.cuda().eval()
+    model = model.to(torch.bfloat16)
+    print(f"✅ Model loaded")
+
+    # Load scene data
+    image_paths = get_sorted_image_paths(color_dir)
+    if len(image_paths) == 0:
+        print(f"❌ Error: No images found in {color_dir}")
+        return
+
+    print(f"🖼️  Found {len(image_paths)} images")
+
+    # Process pose data (if evaluation is enabled)
+    poses_gt = None
+    first_gt_pose = None
+    available_pose_frame_ids = None
+    c2ws = None
+
+    if args.enable_evaluation:
+        poses_gt, first_gt_pose, available_pose_frame_ids = load_poses(pose_dir)
+        if (
+            poses_gt is None
+            or first_gt_pose is None
+            or available_pose_frame_ids is None
+        ):
+            print(f"❌ Error: Failed to load pose data")
+            return
+        print(f"📐 Loaded {len(poses_gt)} poses")
+
+    # Frame selection
+    if args.enable_evaluation and available_pose_frame_ids is not None:
+        # Use pose data for frame selection
+        selected_frame_ids, selected_image_paths, selected_pose_indices = (
+            build_frame_selection(
+                image_paths, available_pose_frame_ids, args.input_frame
+            )
+        )
+        c2ws = poses_gt[selected_pose_indices]
+        image_paths = selected_image_paths
+    else:
+        # Simply take the first N frames
+        num_frames = min(len(image_paths), args.input_frame)
+        selected_frame_ids = list(range(num_frames))
+        image_paths = image_paths[:num_frames]
+
+    print(f"📋 Selected {len(image_paths)} frames for processing")
+
+    try:
+        # Load images
+        print(f"🔄 Loading images...")
+        images = load_images_rgb(image_paths)
+
+        if not images or len(images) < 3:
+            print(f"❌ Error: Not enough valid images (need at least 3)")
+            return
+
+        frame_ids = selected_frame_ids
+        images_array = np.stack(images)
+        vgg_input, patch_width, patch_height = get_vgg_input_imgs(images_array)
+        print(f"📐 Image patch dimensions: {patch_width}x{patch_height}")
+
+        # Update attention layer patch dimensions in the model
+        model.update_patch_dimensions(patch_width, patch_height)
+
+        # Inference + Reconstruction
+        print(f"🚀 Start inference and reconstruction...")
+        (
+            extrinsic_np,
+            intrinsic_np,
+            all_world_points,
+            all_point_colors,
+            all_cam_to_world_mat,
+            inference_time_ms,
+        ) = infer_vggt_and_reconstruct(
+            model, vgg_input, dtype, args.depth_conf_thresh, image_paths
+        )
+        print(f"⏱️  Inference time: {inference_time_ms:.2f}ms")
+
+        # Check results
+        if not all_cam_to_world_mat or not all_world_points:
+            print(f"❌ Error: Failed to obtain valid camera poses or point clouds")
+            return
+
+        # print(f"✅ Inference done, obtained {len(all_world_points)} point sets")
+
+        # Evaluation and saving
+        if args.enable_evaluation:
+            print(f"📊 Start evaluation...")
+            gt_ply_dir = args.data_path / "gt_ply"
+            metrics = evaluate_scene_and_save(
+                "custom_dataset",
+                c2ws,
+                first_gt_pose,
+                frame_ids,
+                all_cam_to_world_mat,
+                all_world_points,
+                output_scene_dir,
+                gt_ply_dir,
+                args.chamfer_max_dist,
+                inference_time_ms,
+                args.plot,
+            )
+            if metrics is not None:
+                print("📈 Evaluation results:")
+                for key, value in metrics.items():
+                    if key in [
+                        "chamfer_distance",
+                        "ate",
+                        "are",
+                        "rpe_rot",
+                        "rpe_trans",
+                        "inference_time_ms",
+                    ]:
+                        print(f"  {key}: {float(value):.4f}")
+
+            # Also visualize predicted poses in evaluation branch
+            if args.plot:
+                visualize_predicted_poses(
+                    all_cam_to_world_mat, frame_ids, output_scene_dir, "custom_dataset"
+                )
+        else:
+            # Save reconstruction only, no evaluation
+            print(f"💾 Saving reconstruction...")
+            output_scene_dir.mkdir(parents=True, exist_ok=True)
+
+            # Save camera poses
+            poses_output_path = output_scene_dir / "estimated_poses.txt"
+            with open(poses_output_path, "w") as f:
+                for i, pose in enumerate(all_cam_to_world_mat):
+                    f.write(f"# Frame {frame_ids[i]}\n")
+                    for row in pose:
+                        f.write(" ".join(map(str, row)) + "\n")
+                    f.write("\n")
+
+            # Save point cloud
+            if all_world_points:
+                points_output_path = output_scene_dir / "reconstructed_points.ply"
+
+                # Merge all frames' point clouds and colors
+                try:
+                    merged_point_cloud = np.vstack(all_world_points)
+                    merged_colors = (
+                        np.vstack(all_point_colors).astype(np.uint8)
+                        if all_point_colors is not None and len(all_point_colors) > 0
+                        else None
+                    )
+                    print(
+                        f"📊 Merged point clouds: {len(all_world_points)} frames, total {len(merged_point_cloud)} points"
+                    )
+
+                    # If too many points, randomly sample 100000 points
+                    max_points = 100000
+                    if len(merged_point_cloud) > max_points:
+                        print(
+                            f"🔽 Too many points, randomly sampling {max_points} points..."
+                        )
+                        # Randomly choose indices
+                        indices = np.random.choice(
+                            len(merged_point_cloud), size=max_points, replace=False
+                        )
+                        merged_point_cloud = merged_point_cloud[indices]
+                        if merged_colors is not None:
+                            merged_colors = merged_colors[indices]
+                        print(
+                            f"✅ Sampling done, kept {len(merged_point_cloud)} points"
+                        )
+
+                    # Save as PLY (with color)
+                    with open(points_output_path, "w") as f:
+                        f.write("ply\n")
+                        f.write("format ascii 1.0\n")
+                        f.write(f"element vertex {len(merged_point_cloud)}\n")
+                        f.write("property float x\n")
+                        f.write("property float y\n")
+                        f.write("property float z\n")
+                        if merged_colors is not None:
+                            f.write("property uchar red\n")
+                            f.write("property uchar green\n")
+                            f.write("property uchar blue\n")
+                        f.write("end_header\n")
+                        if merged_colors is None:
+                            for point in merged_point_cloud:
+                                if not (np.isnan(point).any() or np.isinf(point).any()):
+                                    f.write(
+                                        f"{point[0]:.6f} {point[1]:.6f} {point[2]:.6f}\n"
+                                    )
+                        else:
+                            for point, color in zip(merged_point_cloud, merged_colors):
+                                # Check point validity
+                                if not (np.isnan(point).any() or np.isinf(point).any()):
+                                    r = int(np.clip(color[0], 0, 255))
+                                    g = int(np.clip(color[1], 0, 255))
+                                    b = int(np.clip(color[2], 0, 255))
+                                    f.write(
+                                        f"{point[0]:.6f} {point[1]:.6f} {point[2]:.6f} {r} {g} {b}\n"
+                                    )
+
+                    print(f"💾 Point cloud saved to: {points_output_path}")
+
+                except Exception as e:
+                    print(f"⚠️  Error saving point cloud: {e}")
+                    # If merge fails, try to log per-frame info
+                    print(f"🔍 Point cloud debug info:")
+                    for i, frame_points in enumerate(all_world_points):
+                        print(
+                            f"  Frame {i}: {frame_points.shape if hasattr(frame_points, 'shape') else type(frame_points)}"
+                        )
+                        if (
+                            hasattr(frame_points, "shape")
+                            and len(frame_points.shape) >= 2
+                        ):
+                            print(
+                                f"    Shape: {frame_points.shape}, Dtype: {frame_points.dtype}"
+                            )
+                            if frame_points.shape[0] > 0:
+                                print(
+                                    f"    Range: x[{np.min(frame_points[:, 0]):.3f}, {np.max(frame_points[:, 0]):.3f}] "
+                                    f"y[{np.min(frame_points[:, 1]):.3f}, {np.max(frame_points[:, 1]):.3f}] "
+                                    f"z[{np.min(frame_points[:, 2]):.3f}, {np.max(frame_points[:, 2]):.3f}]"
+                                )
+
+            print(f"📁 Results saved to: {output_scene_dir}")
+
+            # Visualize predicted pose trajectory
+            if args.plot:
+                visualize_predicted_poses(
+                    all_cam_to_world_mat, frame_ids, output_scene_dir, "custom_dataset"
+                )
+
+        print(f"🎉 Done!")
+
+    except Exception as e:
+        print(f"❌ Error occurred during processing: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    main()

FastVGGT/eval/eval_scannet.py

@@ -0,0 +1,208 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+import os
+import sys
+
+# Ensure project root is in sys.path for absolute imports like `vggt.*`
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+from vggt.models.vggt import VGGT
+from vggt.utils.eval_utils import (
+    load_poses,
+    get_vgg_input_imgs,
+    get_sorted_image_paths,
+    get_all_scenes,
+    build_frame_selection,
+    load_images_rgb,
+    infer_vggt_and_reconstruct,
+    evaluate_scene_and_save,
+    compute_average_metrics_and_save,
+)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_dir", type=Path, default="/data/scannetv2/process_scannet"
+    )
+    parser.add_argument(
+        "--gt_ply_dir",
+        type=Path,
+        default="/data/scannetv2/OpenDataLab___ScanNet_v2/raw/scans",
+    )
+    parser.add_argument("--output_path", type=Path, default="./eval_results")
+    parser.add_argument("--merging", type=int, default=None)
+    parser.add_argument("--plot", type=bool, default=True)
+    parser.add_argument(
+        "--depth_conf_thresh",
+        type=float,
+        default=3.0,
+        help="Depth confidence threshold for filtering low confidence depth values",
+    )
+    parser.add_argument(
+        "--chamfer_max_dist",
+        type=float,
+        default=0.5,
+        help="Maximum distance threshold in Chamfer Distance computation, distances exceeding this value will be clipped",
+    )
+    parser.add_argument(
+        "--input_frame",
+        type=int,
+        default=200,
+        help="Maximum number of frames selected for processing per scene",
+    )
+    parser.add_argument(
+        "--num_scenes",
+        type=int,
+        default=50,
+        help="Maximum number of scenes to evaluate",
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="./ckpt/model_tracker_fixed_e20.pt",
+        help="Path to the model checkpoint file",
+    )
+    parser.add_argument(
+        "--vis_attn_map",
+        action="store_true",
+        help="Whether to visualize attention maps during inference",
+    )
+    args = parser.parse_args()
+    torch.manual_seed(33)
+
+    # Scene sampling
+    scannet_scenes = get_all_scenes(args.data_dir, args.num_scenes)
+    print(f"Evaluate {len(scannet_scenes)} scenes")
+
+    all_scenes_metrics = {"scenes": {}, "average": {}}
+    # Force use of bf16 data type
+    dtype = torch.bfloat16
+    # Load VGGT model
+    model = VGGT(merging=args.merging, vis_attn_map=args.vis_attn_map)
+    ckpt = torch.load(args.ckpt_path, map_location="cpu")
+    incompat = model.load_state_dict(ckpt, strict=False)
+    model = model.cuda().eval()
+    model = model.to(torch.bfloat16)
+
+    # Process each scene
+    for scene in scannet_scenes:
+        scene_dir = args.data_dir / f"{scene}"
+        output_scene_dir = args.output_path / f"input_frame_{args.input_frame}" / scene
+        if (output_scene_dir / "metrics.json").exists():
+            continue
+
+        # Load scene data
+        images_dir = scene_dir / "color"
+        pose_path = scene_dir / "pose"
+        image_paths = get_sorted_image_paths(images_dir)
+        poses_gt, first_gt_pose, available_pose_frame_ids = load_poses(pose_path)
+        if (
+            poses_gt is None
+            or first_gt_pose is None
+            or available_pose_frame_ids is None
+        ):
+            print(f"Skipping scene {scene}: no pose data")
+            continue
+
+        # Frame filtering
+        selected_frame_ids, selected_image_paths, selected_pose_indices = (
+            build_frame_selection(
+                image_paths, available_pose_frame_ids, args.input_frame
+            )
+        )
+
+        # Get corresponding poses
+        c2ws = poses_gt[selected_pose_indices]
+        image_paths = selected_image_paths
+
+        if len(image_paths) == 0:
+            print(f"No images found in {images_dir}")
+            continue
+
+        print("🚩Processing", scene, f"Found {len(image_paths)} images")
+        all_cam_to_world_mat = []
+        all_world_points = []
+
+        try:
+            # Load images
+            images = load_images_rgb(image_paths)
+
+            if not images or len(images) < 3:
+                print(f"Skipping {scene}: insufficient valid images")
+                continue
+
+            frame_ids = selected_frame_ids
+            images_array = np.stack(images)
+            vgg_input, patch_width, patch_height = get_vgg_input_imgs(images_array)
+            print(f"Patch dimensions: {patch_width}x{patch_height}")
+
+            # Update model attention layers with dynamic patch dimensions
+            model.update_patch_dimensions(patch_width, patch_height)
+
+            # Inference + Reconstruction
+            (
+                extrinsic_np,
+                intrinsic_np,
+                all_world_points,
+                all_point_colors,
+                all_cam_to_world_mat,
+                inference_time_ms,
+            ) = infer_vggt_and_reconstruct(
+                model, vgg_input, dtype, args.depth_conf_thresh, image_paths
+            )
+            print(f"Inference time: {inference_time_ms:.2f}ms")
+
+            # Process results
+            if not all_cam_to_world_mat or not all_world_points:
+                print(
+                    f"Skipping {scene}: failed to obtain valid camera poses or point clouds"
+                )
+                continue
+
+            # Evaluate and save
+            metrics = evaluate_scene_and_save(
+                scene,
+                c2ws,
+                first_gt_pose,
+                frame_ids,
+                all_cam_to_world_mat,
+                all_world_points,
+                output_scene_dir,
+                args.gt_ply_dir,
+                args.chamfer_max_dist,
+                inference_time_ms,
+                args.plot,
+            )
+            if metrics is not None:
+                all_scenes_metrics["scenes"][scene] = {
+                    key: float(value)
+                    for key, value in metrics.items()
+                    if key
+                    in [
+                        "chamfer_distance",
+                        "ate",
+                        "are",
+                        "rpe_rot",
+                        "rpe_trans",
+                        "inference_time_ms",
+                    ]
+                }
+                print("Complete metrics", all_scenes_metrics["scenes"][scene])
+
+        except Exception as e:
+            print(f"Error processing scene {scene}: {e}")
+            import traceback
+
+            traceback.print_exc()
+
+    # Summarize average metrics and save
+    compute_average_metrics_and_save(
+        all_scenes_metrics,
+        args.output_path,
+        args.input_frame,
+    )

FastVGGT/eval/utils.py

@@ -0,0 +1,142 @@
+import numpy as np
+from scipy.spatial import cKDTree as KDTree
+
+
+def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None):
+    """
+    Args:
+        - depthmap (HxW array):
+        - camera_intrinsics: a 3x3 matrix
+    Returns:
+        pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.
+    """
+    camera_intrinsics = np.float32(camera_intrinsics)
+    H, W = depthmap.shape
+
+    assert camera_intrinsics[0, 1] == 0.0
+    assert camera_intrinsics[1, 0] == 0.0
+    if pseudo_focal is None:
+        fu = camera_intrinsics[0, 0]
+        fv = camera_intrinsics[1, 1]
+    else:
+        assert pseudo_focal.shape == (H, W)
+        fu = fv = pseudo_focal
+    cu = camera_intrinsics[0, 2]
+    cv = camera_intrinsics[1, 2]
+
+    u, v = np.meshgrid(np.arange(W), np.arange(H))
+    z_cam = depthmap
+    x_cam = (u - cu) * z_cam / fu
+    y_cam = (v - cv) * z_cam / fv
+    X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32)
+
+    valid_mask = depthmap > 0.0
+    return X_cam, valid_mask
+
+
+def depthmap_to_absolute_camera_coordinates(
+    depthmap, camera_intrinsics, camera_pose, **kw
+):
+    """
+    Args:
+        - depthmap (HxW array):
+        - camera_intrinsics: a 3x3 matrix
+        - camera_pose: a 4x3 or 4x4 cam2world matrix
+    Returns:
+        pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.
+    """
+    X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics)
+
+    X_world = X_cam  # default
+    if camera_pose is not None:
+
+        R_cam2world = camera_pose[:3, :3]
+        t_cam2world = camera_pose[:3, 3]
+
+        X_world = (
+            np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :]
+        )
+
+    return X_world, valid_mask
+
+
+def completion_ratio(gt_points, rec_points, dist_th=0.05):
+    gen_points_kd_tree = KDTree(rec_points)
+    distances, _ = gen_points_kd_tree.query(gt_points)
+    comp_ratio = np.mean((distances < dist_th).astype(np.float32))
+    return comp_ratio
+
+
+def accuracy(gt_points, rec_points, gt_normals=None, rec_normals=None):
+    gt_points_kd_tree = KDTree(gt_points)
+    distances, idx = gt_points_kd_tree.query(rec_points, workers=-1)
+    acc = np.mean(distances)
+
+    acc_median = np.median(distances)
+
+    if gt_normals is not None and rec_normals is not None:
+        normal_dot = np.sum(gt_normals[idx] * rec_normals, axis=-1)
+        normal_dot = np.abs(normal_dot)
+
+        return acc, acc_median, np.mean(normal_dot), np.median(normal_dot)
+
+    return acc, acc_median
+
+
+def completion(gt_points, rec_points, gt_normals=None, rec_normals=None):
+    gt_points_kd_tree = KDTree(rec_points)
+    distances, idx = gt_points_kd_tree.query(gt_points, workers=-1)
+    comp = np.mean(distances)
+    comp_median = np.median(distances)
+
+    if gt_normals is not None and rec_normals is not None:
+        normal_dot = np.sum(gt_normals * rec_normals[idx], axis=-1)
+        normal_dot = np.abs(normal_dot)
+
+        return comp, comp_median, np.mean(normal_dot), np.median(normal_dot)
+
+    return comp, comp_median
+
+
+def compute_iou(pred_vox, target_vox):
+    # Get voxel indices
+    v_pred_indices = [voxel.grid_index for voxel in pred_vox.get_voxels()]
+    v_target_indices = [voxel.grid_index for voxel in target_vox.get_voxels()]
+
+    # Convert to sets for set operations
+    v_pred_filled = set(tuple(np.round(x, 4)) for x in v_pred_indices)
+    v_target_filled = set(tuple(np.round(x, 4)) for x in v_target_indices)
+
+    # Compute intersection and union
+    intersection = v_pred_filled & v_target_filled
+    union = v_pred_filled | v_target_filled
+
+    # Compute IoU
+    iou = len(intersection) / len(union)
+    return iou
+
+
+def colmap_to_opencv_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] -= 0.5
+    K[1, 2] -= 0.5
+    return K
+
+
+def opencv_to_colmap_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] += 0.5
+    K[1, 2] += 0.5
+    return K

FastVGGT/merging/__init__.py

@@ -0,0 +1,3 @@
+from . import merge
+
+__all__ = ["merge"]

FastVGGT/merging/merge.py

@@ -0,0 +1,370 @@
+import torch
+from typing import Tuple, Callable, Optional, Union
+
+
+@torch.jit.script
+def fast_similarity_chunks(
+    a: torch.Tensor, b_transposed: torch.Tensor, chunk_size: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    B, num_src, C = a.shape
+    original_dtype = a.dtype
+
+    # Convert to bf16 for computation to improve performance and reduce memory usage
+    a_bf16 = a.to(torch.bfloat16)
+    b_transposed_bf16 = b_transposed.to(torch.bfloat16)
+    node_max = torch.empty(B, num_src, device=a.device, dtype=original_dtype)
+    node_idx = torch.empty(B, num_src, device=a.device, dtype=torch.long)
+
+    # Process in chunks
+    for i in range(0, num_src, chunk_size):
+        end_i = min(i + chunk_size, num_src)
+        a_chunk = a_bf16[:, i:end_i, :]  # [B, chunk_size, C]
+        scores_chunk = torch.bmm(a_chunk, b_transposed_bf16)
+        chunk_max_bf16, chunk_idx = torch.max(scores_chunk, dim=2)
+        chunk_max = chunk_max_bf16.to(original_dtype)
+        node_max[:, i:end_i] = chunk_max
+        node_idx[:, i:end_i] = chunk_idx
+    return node_max, node_idx
+
+
+def do_nothing(
+    x: torch.Tensor,
+    extra_tensors=None,
+    extra_tensors_2=None,
+) -> Union[
+    torch.Tensor,
+    Tuple[torch.Tensor, torch.Tensor],
+    Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+]:
+    if extra_tensors is not None and extra_tensors_2 is not None:
+        return x, extra_tensors, extra_tensors_2
+    elif extra_tensors is not None:
+        return x, extra_tensors
+    else:
+        return x
+
+
+def token_merge_bipartite2d(
+    metric: torch.Tensor,
+    w: int,
+    h: int,
+    sx: int,
+    sy: int,
+    r: int,
+    no_rand: bool = False,
+    generator: Optional[torch.Generator] = None,
+    enable_protection: bool = False,
+) -> Tuple[Callable, Callable]:
+    """
+    Divide tokens into source (src) and destination (dst) groups, and merge r tokens from src to dst.
+    dst tokens are selected by randomly choosing one token from each (sx, sy) region.
+    Optionally protect the top 10% of tokens from merging based on importance scores.
+
+    Args:
+     - metric [B, N, C]: Tensor for similarity computation, B=batch size, N=token count, C=feature dimension
+     - w: Image width in tokens
+     - h: Image height in tokens
+     - sx: dst stride in x dimension, must divide w evenly
+     - sy: dst stride in y dimension, must divide h evenly
+     - r: Number of tokens to remove through merging
+     - no_rand: If True, disable randomness (use only top-left token)
+     - generator: Random number generator if no_rand is False and not None
+     - enable_protection: If True, enable importance protection feature
+
+    Returns:
+     - (merge, unmerge): Two functions for merging tokens and restoring pre-merge state
+    """
+    B, N, _ = metric.shape  # Batch size B, total tokens N
+    if r <= 0:
+        return do_nothing, do_nothing
+
+    gather = torch.gather
+
+    tokens_per_img = w * h + 5
+    num_imgs = N // tokens_per_img
+    assert tokens_per_img * num_imgs == N, "Token count doesn't match (w*h+5)*num_imgs"
+
+    with torch.no_grad():
+        # Determine whether to compute importance scores based on enable_protection
+        if enable_protection:
+            num_protected = int(N * 0.1)
+            step = max(1, N // num_protected)
+            protected_indices = torch.arange(0, N, step, device=metric.device)[
+                :num_protected
+            ]
+        else:
+            protected_indices = None
+            num_protected = 0
+
+        # Global idx_buffer_seq of length N; -1 indicates dst, 0 indicates src (maintain original logic)
+        idx_buffer_seq = torch.zeros(N, device=metric.device, dtype=torch.int64)
+        hsy, wsx = h // sy, w // sx  # Number of blocks within each image
+
+        # Mark first image entirely as dst
+        if num_imgs > 0:
+            idx_buffer_seq[:tokens_per_img] = -1
+
+        # Process other images - fully vectorized batch operations
+        if num_imgs > 1:
+            cls_indices = (
+                torch.arange(1, num_imgs, device=metric.device) * tokens_per_img
+            )
+            cls_indices = cls_indices[:, None] + torch.arange(5, device=metric.device)
+            idx_buffer_seq[cls_indices.flatten()] = -1
+            effective_h = min(hsy * sy, h)
+            effective_w = min(wsx * sx, w)
+            effective_grid_size = effective_h * effective_w
+
+            if no_rand:
+                base_pattern = torch.zeros(
+                    effective_grid_size, device=metric.device, dtype=torch.int64
+                )
+                grid_starts = (
+                    torch.arange(1, num_imgs, device=metric.device) * tokens_per_img + 5
+                )
+                grid_indices = grid_starts[:, None] + torch.arange(
+                    effective_grid_size, device=metric.device
+                )
+                idx_buffer_seq[grid_indices.flatten()] = base_pattern.repeat(
+                    num_imgs - 1
+                )
+            else:
+                total_other_imgs = num_imgs - 1
+                all_rand_idx = torch.randint(
+                    sy * sx,
+                    size=(total_other_imgs, hsy, wsx),
+                    device=metric.device,
+                    generator=generator,
+                )
+
+                scatter_src = -torch.ones(
+                    total_other_imgs, hsy, wsx, device=metric.device, dtype=torch.int64
+                )
+
+                idx_buffer_batch = torch.zeros(
+                    total_other_imgs,
+                    hsy,
+                    wsx,
+                    sy * sx,
+                    device=metric.device,
+                    dtype=torch.int64,
+                )
+                idx_buffer_batch.scatter_(
+                    dim=3,
+                    index=all_rand_idx.unsqueeze(-1),
+                    src=scatter_src.unsqueeze(-1),
+                )
+
+                idx_buffer_batch = (
+                    idx_buffer_batch.view(total_other_imgs, hsy, wsx, sy, sx)
+                    .transpose(2, 3)
+                    .reshape(total_other_imgs, hsy * sy, wsx * sx)
+                )
+
+                # Batch fill to target positions - still needs a small loop here, but operations are greatly reduced
+                for i in range(total_other_imgs):
+                    img_idx = i + 1
+                    grid_start = img_idx * tokens_per_img + 5
+                    flat_view = idx_buffer_batch[
+                        i, :effective_h, :effective_w
+                    ].flatten()
+                    idx_buffer_seq[grid_start : grid_start + effective_grid_size] = (
+                        flat_view
+                    )
+
+        rand_idx = idx_buffer_seq.reshape(1, -1, 1).argsort(dim=1)
+        num_dst_orig = int((idx_buffer_seq == -1).sum())
+
+        # Original src and dst indices
+        a_idx_orig = rand_idx[:, num_dst_orig:, :]
+        b_idx_orig = rand_idx[:, :num_dst_orig, :]
+        a_idx = a_idx_orig
+        b_idx = b_idx_orig
+
+        if enable_protection:
+            protected_idx = protected_indices.unsqueeze(0).unsqueeze(-1)
+            num_protected_actual = protected_idx.shape[1]
+        else:
+            protected_idx = None
+            num_protected_actual = 0
+
+        num_src = a_idx.shape[1]
+        num_dst = b_idx.shape[1]
+
+        # Define an internal function to separate src, dst, and protected tokens
+        def split(x):
+            C = x.shape[-1]
+
+            if enable_protection:
+                src = gather(x, dim=1, index=a_idx.expand(B, num_src, C))
+                dst = gather(x, dim=1, index=b_idx.expand(B, num_dst, C))
+                protected = gather(
+                    x, dim=1, index=protected_idx.expand(B, num_protected_actual, C)
+                )
+                return src, dst, protected
+            else:
+                src = gather(x, dim=1, index=a_idx.expand(B, num_src, C))
+                dst = gather(x, dim=1, index=b_idx.expand(B, num_dst, C))
+                return src, dst
+
+        # Compute cosine similarity (normalize first then dot product)
+        metric = metric / metric.norm(dim=-1, keepdim=True)
+        if enable_protection:
+            a, b, protected = split(metric)
+        else:
+            a, b = split(metric)
+
+        r = min(a.shape[1], r)
+        num_src_actual = a.shape[1]
+        chunk_size = min(5000, num_src_actual)
+
+        node_max = torch.empty(B, num_src_actual, device=a.device, dtype=a.dtype)
+        node_idx = torch.empty(B, num_src_actual, device=a.device, dtype=torch.long)
+
+        b_transposed = b.transpose(-1, -2)
+        node_max, node_idx = fast_similarity_chunks(a, b_transposed, chunk_size)
+        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
+
+        # If protection is enabled, filter out protected tokens to ensure they are not merged
+        if enable_protection:
+            src_indices = a_idx[0, :, 0]
+            protected_mask_src = torch.isin(src_indices, protected_indices)
+            edge_flat = edge_idx[0, :, 0]
+            valid_mask = ~protected_mask_src[edge_flat]
+            valid_edges = edge_flat[valid_mask]
+
+            valid_count = valid_edges.shape[0]
+            r_actual = min(r, valid_count)
+
+            unm_idx = valid_edges[r_actual:].unsqueeze(0).unsqueeze(-1)
+            src_idx = valid_edges[:r_actual].unsqueeze(0).unsqueeze(-1)
+        else:
+            unm_idx = edge_idx[..., r:, :]
+            src_idx = edge_idx[..., :r, :]
+            r_actual = r
+
+        # Get dst token indices corresponding to each src token to be merged
+        dst_idx = gather(node_idx[..., None], dim=-2, index=src_idx)
+        r = r_actual
+
+    # Define merge function to merge selected src tokens to corresponding dst tokens
+    def merge(
+        x: torch.Tensor,
+        mode: str = "mean",
+        extra_tensors=None,
+        extra_tensors_2=None,
+    ) -> Union[
+        torch.Tensor,
+        Tuple[torch.Tensor, torch.Tensor],
+        Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    ]:
+        if enable_protection:
+            src, dst, protected = split(x)
+        else:
+            src, dst = split(x)
+
+        n, t1, c = src.shape
+
+        # Extract unmerged src tokens - using actual unm_idx size
+        unm_len = unm_idx.shape[1]
+        unm = gather(src, dim=-2, index=unm_idx.expand(n, unm_len, c))
+        src_len = src_idx.shape[1]
+        src = gather(src, dim=-2, index=src_idx.expand(n, src_len, c))
+        dst = dst.scatter_reduce(-2, dst_idx.expand(n, src_len, c), src, reduce=mode)
+
+        # ---------------- Extra tensor processing ----------------
+        merged_extra_1 = None
+        merged_extra_2 = None
+        if extra_tensors is not None:
+            E_dim = extra_tensors.shape[-1]
+            if enable_protection:
+                src_e, dst_e, protected_e = split(extra_tensors)
+            else:
+                src_e, dst_e = split(extra_tensors)
+
+            # Consistent with main tensor, only select r src tokens to be merged
+            src_e_r = gather(src_e, dim=-2, index=src_idx.expand(n, src_len, E_dim))
+            unm_e = gather(src_e, dim=-2, index=unm_idx.expand(n, unm_len, E_dim))
+
+            dst_e = dst_e.scatter_reduce(
+                -2, dst_idx.expand(n, src_len, E_dim), src_e_r, reduce=mode
+            )
+            if enable_protection:
+                merged_extra_1 = torch.cat([unm_e, dst_e, protected_e], dim=1)
+            else:
+                merged_extra_1 = torch.cat([unm_e, dst_e], dim=1)
+
+        if extra_tensors_2 is not None:
+            E_dim_2 = extra_tensors_2.shape[-1]
+            if enable_protection:
+                src_e2, dst_e2, protected_e2 = split(extra_tensors_2)
+            else:
+                src_e2, dst_e2 = split(extra_tensors_2)
+
+            src_e2_r = gather(src_e2, dim=-2, index=src_idx.expand(n, src_len, E_dim_2))
+            unm_e2 = gather(src_e2, dim=-2, index=unm_idx.expand(n, unm_len, E_dim_2))
+
+            dst_e2 = dst_e2.scatter_reduce(
+                -2, dst_idx.expand(n, src_len, E_dim_2), src_e2_r, reduce=mode
+            )
+            if enable_protection:
+                merged_extra_2 = torch.cat([unm_e2, dst_e2, protected_e2], dim=1)
+            else:
+                merged_extra_2 = torch.cat([unm_e2, dst_e2], dim=1)
+
+        if enable_protection:
+            main_result = torch.cat([unm, dst, protected], dim=1)
+        else:
+            main_result = torch.cat([unm, dst], dim=1)
+
+        if merged_extra_1 is not None and merged_extra_2 is not None:
+            return main_result, merged_extra_1, merged_extra_2
+        elif merged_extra_1 is not None:
+            return main_result, merged_extra_1
+        else:
+            return main_result
+
+    # Define unmerge function to restore pre-merge state (for decoder)
+    def unmerge(x: torch.Tensor) -> torch.Tensor:
+        unm_len = unm_idx.shape[1]
+        dst_len = num_dst
+        src_len = src_idx.shape[1]
+        unm = x[..., :unm_len, :]
+        dst = x[..., unm_len : unm_len + dst_len, :]
+
+        if enable_protection:
+            protected = x[
+                ..., unm_len + dst_len : unm_len + dst_len + num_protected_actual, :
+            ]
+
+        _, _, c = unm.shape
+        src = gather(dst, dim=-2, index=dst_idx.expand(B, src_len, c))
+        out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype)
+        out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst)
+        out.scatter_(
+            dim=-2,
+            index=gather(
+                a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=unm_idx
+            ).expand(B, unm_len, c),
+            src=unm,
+        )
+
+        out.scatter_(
+            dim=-2,
+            index=gather(
+                a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=src_idx
+            ).expand(B, src_len, c),
+            src=src,
+        )
+
+        if enable_protection:
+            out.scatter_(
+                dim=-2,
+                index=protected_idx.expand(B, num_protected_actual, c),
+                src=protected,
+            )
+
+        return out
+
+    return merge, unmerge

FastVGGT/requirements.txt

@@ -0,0 +1,15 @@
+torch==2.3.1
+torchvision==0.18.1
+numpy==1.26.1
+Pillow
+huggingface_hub
+einops
+safetensors
+evo
+open3d
+matplotlib
+scipy
+opencv-python
+scikit-image
+tqdm
+

FastVGGT/vggt/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree. 

FastVGGT/vggt/dependency/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree. 

FastVGGT/vggt/dependency/distortion.py

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+
+
+def apply_distortion(points, distortion_params):
+    """
+    Apply distortion to normalized camera coordinates.
+    
+    Args:
+        points: Array of normalized camera coordinates
+        distortion_params: Distortion parameters
+        
+    Returns:
+        Distorted coordinates
+    """
+    # Simple passthrough for now - implement actual distortion if needed
+    return points
+
+
+def iterative_undistortion(points, distortion_params, max_iter=10):
+    """
+    Remove distortion from normalized camera coordinates using iterative method.
+    
+    Args:
+        points: Array of distorted normalized camera coordinates
+        distortion_params: Distortion parameters
+        max_iter: Maximum number of iterations
+        
+    Returns:
+        Undistorted coordinates
+    """
+    # Simple passthrough for now - implement actual undistortion if needed
+    return points
+
+
+def single_undistortion(points, distortion_params):
+    """
+    Remove distortion from normalized camera coordinates using single step.
+    
+    Args:
+        points: Array of distorted normalized camera coordinates
+        distortion_params: Distortion parameters
+        
+    Returns:
+        Undistorted coordinates
+    """
+    # Simple passthrough for now - implement actual undistortion if needed
+    return points 

FastVGGT/vggt/heads/camera_head.py

@@ -0,0 +1,149 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vggt.layers import Mlp
+from vggt.layers.block import Block
+from vggt.heads.head_act import activate_pose
+
+
+class CameraHead(nn.Module):
+    """
+    CameraHead predicts camera parameters from token representations using iterative refinement.
+
+    It applies a series of transformer blocks (the "trunk") to dedicated camera tokens.
+    """
+
+    def __init__(
+        self,
+        dim_in: int = 2048,
+        trunk_depth: int = 4,
+        pose_encoding_type: str = "absT_quaR_FoV",
+        num_heads: int = 16,
+        mlp_ratio: int = 4,
+        init_values: float = 0.01,
+        trans_act: str = "linear",
+        quat_act: str = "linear",
+        fl_act: str = "relu",  # Field of view activations: ensures FOV values are positive.
+    ):
+        super().__init__()
+
+        if pose_encoding_type == "absT_quaR_FoV":
+            self.target_dim = 9
+        else:
+            raise ValueError(f"Unsupported camera encoding type: {pose_encoding_type}")
+
+        self.trans_act = trans_act
+        self.quat_act = quat_act
+        self.fl_act = fl_act
+        self.trunk_depth = trunk_depth
+
+        # Build the trunk using a sequence of transformer blocks.
+        self.trunk = nn.Sequential(
+            *[
+                Block(dim=dim_in, num_heads=num_heads, mlp_ratio=mlp_ratio, init_values=init_values)
+                for _ in range(trunk_depth)
+            ]
+        )
+
+        # Normalizations for camera token and trunk output.
+        self.token_norm = nn.LayerNorm(dim_in)
+        self.trunk_norm = nn.LayerNorm(dim_in)
+
+        # Learnable empty camera pose token.
+        self.empty_pose_tokens = nn.Parameter(torch.zeros(1, 1, self.target_dim))
+        self.embed_pose = nn.Linear(self.target_dim, dim_in)
+
+        # Module for producing modulation parameters: shift, scale, and a gate.
+        self.poseLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True))
+
+        # Adaptive layer normalization without affine parameters.
+        self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6)
+        self.pose_branch = Mlp(in_features=dim_in, hidden_features=dim_in // 2, out_features=self.target_dim, drop=0)
+
+    def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list:
+        """
+        Forward pass to predict camera parameters.
+
+        Args:
+            aggregated_tokens_list (list): List of token tensors from the network;
+                the last tensor is used for prediction.
+            num_iterations (int, optional): Number of iterative refinement steps. Defaults to 4.
+
+        Returns:
+            list: A list of predicted camera encodings (post-activation) from each iteration.
+        """
+        # Use tokens from the last block for camera prediction.
+        tokens = aggregated_tokens_list[-1]
+
+        # Extract the camera tokens
+        pose_tokens = tokens[:, :, 0]
+        pose_tokens = self.token_norm(pose_tokens)
+
+        pred_pose_enc_list = self.trunk_fn(pose_tokens, num_iterations)
+        return pred_pose_enc_list
+
+    def trunk_fn(self, pose_tokens: torch.Tensor, num_iterations: int) -> list:
+        """
+        Iteratively refine camera pose predictions.
+
+        Args:
+            pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1, C].
+            num_iterations (int): Number of refinement iterations.
+
+        Returns:
+            list: List of activated camera encodings from each iteration.
+        """
+        B, S, C = pose_tokens.shape  # S is expected to be 1.
+        pred_pose_enc = None
+        pred_pose_enc_list = []
+
+        for _ in range(num_iterations):
+            # Use a learned empty pose for the first iteration.
+            if pred_pose_enc is None:
+                module_input = self.embed_pose(self.empty_pose_tokens.expand(B, S, -1))
+            else:
+                # Detach the previous prediction to avoid backprop through time.
+                pred_pose_enc = pred_pose_enc.detach()
+                module_input = self.embed_pose(pred_pose_enc)
+
+            # Generate modulation parameters and split them into shift, scale, and gate components.
+            shift_msa, scale_msa, gate_msa = self.poseLN_modulation(module_input).chunk(3, dim=-1)
+
+            # Adaptive layer normalization and modulation.
+            pose_tokens_modulated = gate_msa * modulate(self.adaln_norm(pose_tokens), shift_msa, scale_msa)
+            pose_tokens_modulated = pose_tokens_modulated + pose_tokens
+
+            pose_tokens_modulated = self.trunk(pose_tokens_modulated)
+            # Compute the delta update for the pose encoding.
+            pred_pose_enc_delta = self.pose_branch(self.trunk_norm(pose_tokens_modulated))
+
+            if pred_pose_enc is None:
+                pred_pose_enc = pred_pose_enc_delta
+            else:
+                pred_pose_enc = pred_pose_enc + pred_pose_enc_delta
+
+            # Apply final activation functions for translation, quaternion, and field-of-view.
+            activated_pose = activate_pose(
+                pred_pose_enc, trans_act=self.trans_act, quat_act=self.quat_act, fl_act=self.fl_act
+            )
+            pred_pose_enc_list.append(activated_pose)
+
+        return pred_pose_enc_list
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """
+    Modulate the input tensor using scaling and shifting parameters.
+    """
+    # modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19
+    return x * (1 + scale) + shift

FastVGGT/vggt/heads/dpt_head.py

@@ -0,0 +1,598 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# Inspired by https://github.com/DepthAnything/Depth-Anything-V2
+
+
+import os
+from typing import List, Dict, Tuple, Union, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .head_act import activate_head
+from .utils import create_uv_grid, position_grid_to_embed
+
+
+class DPTHead(nn.Module):
+    """
+    DPT  Head for dense prediction tasks.
+
+    This implementation follows the architecture described in "Vision Transformers for Dense Prediction"
+    (https://arxiv.org/abs/2103.13413). The DPT head processes features from a vision transformer
+    backbone and produces dense predictions by fusing multi-scale features.
+
+    Args:
+        dim_in (int): Input dimension (channels).
+        patch_size (int, optional): Patch size. Default is 14.
+        output_dim (int, optional): Number of output channels. Default is 4.
+        activation (str, optional): Activation type. Default is "inv_log".
+        conf_activation (str, optional): Confidence activation type. Default is "expp1".
+        features (int, optional): Feature channels for intermediate representations. Default is 256.
+        out_channels (List[int], optional): Output channels for each intermediate layer.
+        intermediate_layer_idx (List[int], optional): Indices of layers from aggregated tokens used for DPT.
+        pos_embed (bool, optional): Whether to use positional embedding. Default is True.
+        feature_only (bool, optional): If True, return features only without the last several layers and activation head. Default is False.
+        down_ratio (int, optional): Downscaling factor for the output resolution. Default is 1.
+    """
+
+    def __init__(
+        self,
+        dim_in: int,
+        patch_size: int = 14,
+        output_dim: int = 4,
+        activation: str = "inv_log",
+        conf_activation: str = "expp1",
+        features: int = 256,
+        out_channels: List[int] = [256, 512, 1024, 1024],
+        intermediate_layer_idx: List[int] = [0, 1, 2, 3],
+        pos_embed: bool = True,
+        feature_only: bool = False,
+        down_ratio: int = 1,
+    ) -> None:
+        super(DPTHead, self).__init__()
+        self.patch_size = patch_size
+        self.activation = activation
+        self.conf_activation = conf_activation
+        self.pos_embed = pos_embed
+        self.feature_only = feature_only
+        self.down_ratio = down_ratio
+        self.intermediate_layer_idx = intermediate_layer_idx
+
+        self.norm = nn.LayerNorm(dim_in)
+
+        # Projection layers for each output channel from tokens.
+        self.projects = nn.ModuleList(
+            [
+                nn.Conv2d(
+                    in_channels=dim_in,
+                    out_channels=oc,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+                for oc in out_channels
+            ]
+        )
+
+        # Resize layers for upsampling feature maps.
+        self.resize_layers = nn.ModuleList(
+            [
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[0],
+                    out_channels=out_channels[0],
+                    kernel_size=4,
+                    stride=4,
+                    padding=0,
+                ),
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[1],
+                    out_channels=out_channels[1],
+                    kernel_size=2,
+                    stride=2,
+                    padding=0,
+                ),
+                nn.Identity(),
+                nn.Conv2d(
+                    in_channels=out_channels[3],
+                    out_channels=out_channels[3],
+                    kernel_size=3,
+                    stride=2,
+                    padding=1,
+                ),
+            ]
+        )
+
+        self.scratch = _make_scratch(out_channels, features, expand=False)
+
+        # Attach additional modules to scratch.
+        self.scratch.stem_transpose = nn.Identity()  # Use Identity instead of None
+        self.scratch.refinenet1 = _make_fusion_block(features)
+        self.scratch.refinenet2 = _make_fusion_block(features)
+        self.scratch.refinenet3 = _make_fusion_block(features)
+        self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False)
+
+        head_features_1 = features
+        head_features_2 = 32
+
+        if feature_only:
+            self.scratch.output_conv1 = nn.Conv2d(
+                head_features_1, head_features_1, kernel_size=3, stride=1, padding=1
+            )
+        else:
+            self.scratch.output_conv1 = nn.Conv2d(
+                head_features_1,
+                head_features_1 // 2,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+            conv2_in_channels = head_features_1 // 2
+
+            self.scratch.output_conv2 = nn.Sequential(
+                nn.Conv2d(
+                    conv2_in_channels,
+                    head_features_2,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                ),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(
+                    head_features_2, output_dim, kernel_size=1, stride=1, padding=0
+                ),
+            )
+
+    def forward(
+        self,
+        aggregated_tokens_list: List[torch.Tensor],
+        images: torch.Tensor,
+        patch_start_idx: int,
+        frames_chunk_size: int = 8,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Forward pass through the DPT head, supports processing by chunking frames.
+        Args:
+            aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
+            images (Tensor): Input images with shape [B, S, 3, H, W], in range [0, 1].
+            patch_start_idx (int): Starting index for patch tokens in the token sequence.
+                Used to separate patch tokens from other tokens (e.g., camera or register tokens).
+            frames_chunk_size (int, optional): Number of frames to process in each chunk.
+                If None or larger than S, all frames are processed at once. Default: 8.
+
+        Returns:
+            Tensor or Tuple[Tensor, Tensor]:
+                - If feature_only=True: Feature maps with shape [B, S, C, H, W]
+                - Otherwise: Tuple of (predictions, confidence) both with shape [B, S, 1, H, W]
+        """
+        B, S, _, H, W = images.shape
+
+        # If frames_chunk_size is not specified or greater than S, process all frames at once
+        if frames_chunk_size is None or frames_chunk_size >= S:
+            return self._forward_impl(aggregated_tokens_list, images, patch_start_idx)
+
+        # Otherwise, process frames in chunks to manage memory usage
+        assert frames_chunk_size > 0
+
+        # Process frames in batches
+        all_preds = []
+        all_conf = []
+
+        for frames_start_idx in range(0, S, frames_chunk_size):
+            frames_end_idx = min(frames_start_idx + frames_chunk_size, S)
+
+            # Process batch of frames
+            if self.feature_only:
+                chunk_output = self._forward_impl(
+                    aggregated_tokens_list,
+                    images,
+                    patch_start_idx,
+                    frames_start_idx,
+                    frames_end_idx,
+                )
+                all_preds.append(chunk_output)
+            else:
+                chunk_preds, chunk_conf = self._forward_impl(
+                    aggregated_tokens_list,
+                    images,
+                    patch_start_idx,
+                    frames_start_idx,
+                    frames_end_idx,
+                )
+                all_preds.append(chunk_preds)
+                all_conf.append(chunk_conf)
+
+        # Concatenate results along the sequence dimension
+        if self.feature_only:
+            return torch.cat(all_preds, dim=1)
+        else:
+            return torch.cat(all_preds, dim=1), torch.cat(all_conf, dim=1)
+
+    def _forward_impl(
+        self,
+        aggregated_tokens_list: List[torch.Tensor],
+        images: torch.Tensor,
+        patch_start_idx: int,
+        frames_start_idx: Optional[int] = None,
+        frames_end_idx: Optional[int] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Implementation of the forward pass through the DPT head.
+
+        This method processes a specific chunk of frames from the sequence.
+
+        Args:
+            aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
+            images (Tensor): Input images with shape [B, S, 3, H, W].
+            patch_start_idx (int): Starting index for patch tokens.
+            frames_start_idx (int, optional): Starting index for frames to process.
+            frames_end_idx (int, optional): Ending index for frames to process.
+
+        Returns:
+            Tensor or Tuple[Tensor, Tensor]: Feature maps or (predictions, confidence).
+        """
+        if frames_start_idx is not None and frames_end_idx is not None:
+            images = images[:, frames_start_idx:frames_end_idx].contiguous()
+
+        B, S, _, H, W = images.shape
+
+        patch_h, patch_w = H // self.patch_size, W // self.patch_size
+
+        out = []
+        dpt_idx = 0
+
+        for layer_idx in self.intermediate_layer_idx:
+            x = aggregated_tokens_list[layer_idx][:, :, patch_start_idx:]
+
+            # Select frames if processing a chunk
+            if frames_start_idx is not None and frames_end_idx is not None:
+                x = x[:, frames_start_idx:frames_end_idx]
+
+            x = x.reshape(B * S, -1, x.shape[-1])
+
+            x = self.norm(x)
+
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+
+            x = self.projects[dpt_idx](x)
+            if self.pos_embed:
+                x = self._apply_pos_embed(x, W, H)
+            x = self.resize_layers[dpt_idx](x)
+
+            out.append(x)
+            dpt_idx += 1
+
+        # Fuse features from multiple layers.
+        out = self.scratch_forward(out)
+        # Interpolate fused output to match target image resolution.
+        out = custom_interpolate(
+            out,
+            (
+                int(patch_h * self.patch_size / self.down_ratio),
+                int(patch_w * self.patch_size / self.down_ratio),
+            ),
+            mode="bilinear",
+            align_corners=True,
+        )
+
+        if self.pos_embed:
+            out = self._apply_pos_embed(out, W, H)
+
+        if self.feature_only:
+            return out.view(B, S, *out.shape[1:])
+
+        out = self.scratch.output_conv2(out)
+        preds, conf = activate_head(
+            out, activation=self.activation, conf_activation=self.conf_activation
+        )
+
+        preds = preds.view(B, S, *preds.shape[1:])
+        conf = conf.view(B, S, *conf.shape[1:])
+        return preds, conf
+
+    def _apply_pos_embed(
+        self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1
+    ) -> torch.Tensor:
+        """
+        Apply positional embedding to tensor x.
+        """
+        patch_w = x.shape[-1]
+        patch_h = x.shape[-2]
+        pos_embed = create_uv_grid(
+            patch_w, patch_h, aspect_ratio=W / H, dtype=x.dtype, device=x.device
+        )
+        pos_embed = position_grid_to_embed(pos_embed, x.shape[1])
+        pos_embed = pos_embed * ratio
+        pos_embed = pos_embed.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1)
+        return x + pos_embed
+
+    def scratch_forward(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Forward pass through the fusion blocks.
+
+        Args:
+            features (List[Tensor]): List of feature maps from different layers.
+
+        Returns:
+            Tensor: Fused feature map.
+        """
+        layer_1, layer_2, layer_3, layer_4 = features
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+        del layer_4_rn, layer_4
+
+        out = self.scratch.refinenet3(out, layer_3_rn, size=layer_2_rn.shape[2:])
+        del layer_3_rn, layer_3
+
+        out = self.scratch.refinenet2(out, layer_2_rn, size=layer_1_rn.shape[2:])
+        del layer_2_rn, layer_2
+
+        out = self.scratch.refinenet1(out, layer_1_rn)
+        del layer_1_rn, layer_1
+
+        out = self.scratch.output_conv1(out)
+        return out
+
+
+################################################################################
+# Modules
+################################################################################
+
+
+def _make_fusion_block(
+    features: int,
+    size: Optional[int] = None,
+    has_residual: bool = True,
+    groups: int = 1,
+) -> nn.Module:
+    return FeatureFusionBlock(
+        features,
+        nn.ReLU(inplace=True),
+        deconv=False,
+        bn=False,
+        expand=False,
+        align_corners=True,
+        size=size,
+        has_residual=has_residual,
+        groups=groups,
+    )
+
+
+def _make_scratch(
+    in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False
+) -> nn.Module:
+    scratch = nn.Module()
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    if len(in_shape) >= 4:
+        out_shape4 = out_shape
+
+    if expand:
+        out_shape1 = out_shape
+        out_shape2 = out_shape * 2
+        out_shape3 = out_shape * 4
+        if len(in_shape) >= 4:
+            out_shape4 = out_shape * 8
+
+    scratch.layer1_rn = nn.Conv2d(
+        in_shape[0],
+        out_shape1,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    scratch.layer2_rn = nn.Conv2d(
+        in_shape[1],
+        out_shape2,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    scratch.layer3_rn = nn.Conv2d(
+        in_shape[2],
+        out_shape3,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    if len(in_shape) >= 4:
+        scratch.layer4_rn = nn.Conv2d(
+            in_shape[3],
+            out_shape4,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=False,
+            groups=groups,
+        )
+    return scratch
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features, activation, bn, groups=1):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+        self.groups = groups
+        self.conv1 = nn.Conv2d(
+            features,
+            features,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=True,
+            groups=self.groups,
+        )
+        self.conv2 = nn.Conv2d(
+            features,
+            features,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=True,
+            groups=self.groups,
+        )
+
+        self.norm1 = None
+        self.norm2 = None
+
+        self.activation = activation
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.norm1 is not None:
+            out = self.norm1(out)
+
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.norm2 is not None:
+            out = self.norm2(out)
+
+        return out + x
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block."""
+
+    def __init__(
+        self,
+        features,
+        activation,
+        deconv=False,
+        bn=False,
+        expand=False,
+        align_corners=True,
+        size=None,
+        has_residual=True,
+        groups=1,
+    ):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+        self.groups = groups
+        self.expand = expand
+        out_features = features
+        if self.expand == True:
+            out_features = features // 2
+
+        self.out_conv = nn.Conv2d(
+            features,
+            out_features,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+            groups=self.groups,
+        )
+
+        if has_residual:
+            self.resConfUnit1 = ResidualConvUnit(
+                features, activation, bn, groups=self.groups
+            )
+
+        self.has_residual = has_residual
+        self.resConfUnit2 = ResidualConvUnit(
+            features, activation, bn, groups=self.groups
+        )
+
+        self.size = size
+
+    def forward(self, *xs, size=None):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if self.has_residual:
+            res = self.resConfUnit1(xs[1])
+            output = output + res
+
+        output = self.resConfUnit2(output)
+
+        if (size is None) and (self.size is None):
+            modifier = {"scale_factor": 2}
+        elif size is None:
+            modifier = {"size": self.size}
+        else:
+            modifier = {"size": size}
+
+        output = custom_interpolate(
+            output, **modifier, mode="bilinear", align_corners=self.align_corners
+        )
+        output = self.out_conv(output)
+
+        return output
+
+
+def custom_interpolate(
+    x: torch.Tensor,
+    size: Optional[Tuple[int, int]] = None,
+    scale_factor: Optional[float] = None,
+    mode: str = "bilinear",
+    align_corners: bool = True,
+) -> torch.Tensor:
+    """
+    Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.
+    """
+    if size is None:
+        size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor))
+
+    INT_MAX = 1610612736
+
+    input_elements = size[0] * size[1] * x.shape[0] * x.shape[1]
+
+    if input_elements > INT_MAX:
+        chunks = torch.chunk(x, chunks=(input_elements // INT_MAX) + 1, dim=0)
+        interpolated_chunks = [
+            nn.functional.interpolate(
+                chunk, size=size, mode=mode, align_corners=align_corners
+            )
+            for chunk in chunks
+        ]
+        x = torch.cat(interpolated_chunks, dim=0)
+        return x.contiguous()
+    else:
+        return nn.functional.interpolate(
+            x, size=size, mode=mode, align_corners=align_corners
+        )