+ 
+
+
+
+**[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/)
+
+
+
+## ⚙️ 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.
+
+
+
+
+
+## 🏀 Evaluation
+### Custom Dataset
+Please organize the data according to the following directory:
+```
+
+
+### 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.
+
+
+
+
+## ⚖️ 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}
+}
+```
\ No newline at end of file
diff --git a/FastVGGT/assets/attn_map.png b/FastVGGT/assets/attn_map.png
new file mode 100644
index 0000000000000000000000000000000000000000..98d6f2be8cdb49e5bccee7f902af292327fd0023
--- /dev/null
+++ b/FastVGGT/assets/attn_map.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8477957f593c203bcf41df91ac3ed0d22329e22250fdab9f8f8674340964242c
+size 2369408
diff --git a/FastVGGT/assets/autolab_logo.png b/FastVGGT/assets/autolab_logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..c6ca9a41f5c05a1b2ae66368461672f46e20bba1
--- /dev/null
+++ b/FastVGGT/assets/autolab_logo.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fcead3160cbf561c4385cc8b938a17a94652e3d849da6497f053d32d1245596
+size 5126003
diff --git a/FastVGGT/assets/maclab_logo.png b/FastVGGT/assets/maclab_logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..b96660decf945d3d786f41b9b42cb92798f9191f
Binary files /dev/null and b/FastVGGT/assets/maclab_logo.png differ
diff --git a/FastVGGT/assets/main.png b/FastVGGT/assets/main.png
new file mode 100644
index 0000000000000000000000000000000000000000..3bcab2d519689dfba3283e606cd55cb84175ede5
--- /dev/null
+++ b/FastVGGT/assets/main.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eecacb414647f01dc8a52b4aba5ff2556733f46d1b9129613e3f59aceff69685
+size 883717
diff --git a/FastVGGT/assets/vs.png b/FastVGGT/assets/vs.png
new file mode 100644
index 0000000000000000000000000000000000000000..3729d9ba86000b81d0e9c878b2db6f16c7f9c48a
--- /dev/null
+++ b/FastVGGT/assets/vs.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dac1ea5397b9f985c6fb86fc5f321313cc5a372d0917b7c1c1c7dd2cea6bec5f
+size 230157
diff --git a/FastVGGT/eval/__pycache__/base.cpython-310.pyc b/FastVGGT/eval/__pycache__/base.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..49170288c162684ebafa68e5242efd320f0f35b8
Binary files /dev/null and b/FastVGGT/eval/__pycache__/base.cpython-310.pyc differ
diff --git a/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc b/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6e6e0478775fe98dc22180bc43bad8073c43c907
Binary files /dev/null and b/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc differ
diff --git a/FastVGGT/eval/__pycache__/data.cpython-310.pyc b/FastVGGT/eval/__pycache__/data.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4d8ec6f8e1828fa87b02a4511490ad858de2093a
Binary files /dev/null and b/FastVGGT/eval/__pycache__/data.cpython-310.pyc differ
diff --git a/FastVGGT/eval/__pycache__/data.cpython-37.pyc b/FastVGGT/eval/__pycache__/data.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d76f2c95748956a04f009e99f39cefd9cc6ec4e0
Binary files /dev/null and b/FastVGGT/eval/__pycache__/data.cpython-37.pyc differ
diff --git a/FastVGGT/eval/__pycache__/utils.cpython-310.pyc b/FastVGGT/eval/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..638db49f10789e53dcf0ff33c9aaa28b814e3e72
Binary files /dev/null and b/FastVGGT/eval/__pycache__/utils.cpython-310.pyc differ
diff --git a/FastVGGT/eval/__pycache__/utils.cpython-37.pyc b/FastVGGT/eval/__pycache__/utils.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fd73966947cb74c319d66d7bf221d73b5dde0ce0
Binary files /dev/null and b/FastVGGT/eval/__pycache__/utils.cpython-37.pyc differ
diff --git a/FastVGGT/eval/base.py b/FastVGGT/eval/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a716449a71f552ea408df0eb37e854cf4e92da6
--- /dev/null
+++ b/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]]
diff --git a/FastVGGT/eval/criterion.py b/FastVGGT/eval/criterion.py
new file mode 100644
index 0000000000000000000000000000000000000000..a63c991b5fd5b61325b56055d7d380a27a336971
--- /dev/null
+++ b/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
diff --git a/FastVGGT/eval/data.py b/FastVGGT/eval/data.py
new file mode 100644
index 0000000000000000000000000000000000000000..301788e3ae6b67091fb031b068ecbecff1eae48d
--- /dev/null
+++ b/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
diff --git a/FastVGGT/eval/dataset_utils/__init__.py b/FastVGGT/eval/dataset_utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc
--- /dev/null
+++ b/FastVGGT/eval/dataset_utils/__init__.py
@@ -0,0 +1 @@
+
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d00b3e8b6f5fd9ef7d1428ff771243e1557ddf9f
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9ebd5d22d44ddce633e9a05a7467e9ae0c3f9032
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fbd10218eddf65486b375109c23cf8cecd7c3ba5
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..860c8970fc49d65b2a4101b77cc313a6aeac618a
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0fe7870fb079d9dc08ca647866ff8c77fd0f2ec0
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1b4530be10547e3e7bf79e6cef6470880cb710a4
Binary files /dev/null and b/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc differ
diff --git a/FastVGGT/eval/dataset_utils/corr.py b/FastVGGT/eval/dataset_utils/corr.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbf5de18b7388e38e45c3f313487957811abc483
--- /dev/null
+++ b/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
diff --git a/FastVGGT/eval/dataset_utils/cropping.py b/FastVGGT/eval/dataset_utils/cropping.py
new file mode 100644
index 0000000000000000000000000000000000000000..30a9eac18d241b71538957cf7ba4767ebc323b43
--- /dev/null
+++ b/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
diff --git a/FastVGGT/eval/dataset_utils/transforms.py b/FastVGGT/eval/dataset_utils/transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..cec2d144e1b97cc99c191d15cdcaf20796cae94b
--- /dev/null
+++ b/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
diff --git a/FastVGGT/eval/eval_7andN.py b/FastVGGT/eval/eval_7andN.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d6d5fbf6c159ea18bfd6fbcee0a5d23e63c0cf2
--- /dev/null
+++ b/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