- 
-
-
-
-**[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
deleted file mode 100644
index 98d6f2be8cdb49e5bccee7f902af292327fd0023..0000000000000000000000000000000000000000
--- a/FastVGGT/assets/attn_map.png
+++ /dev/null
@@ -1,3 +0,0 @@
-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
deleted file mode 100644
index c6ca9a41f5c05a1b2ae66368461672f46e20bba1..0000000000000000000000000000000000000000
--- a/FastVGGT/assets/autolab_logo.png
+++ /dev/null
@@ -1,3 +0,0 @@
-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
deleted file mode 100644
index b96660decf945d3d786f41b9b42cb92798f9191f..0000000000000000000000000000000000000000
Binary files a/FastVGGT/assets/maclab_logo.png and /dev/null differ
diff --git a/FastVGGT/assets/main.png b/FastVGGT/assets/main.png
deleted file mode 100644
index 3bcab2d519689dfba3283e606cd55cb84175ede5..0000000000000000000000000000000000000000
--- a/FastVGGT/assets/main.png
+++ /dev/null
@@ -1,3 +0,0 @@
-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
deleted file mode 100644
index 3729d9ba86000b81d0e9c878b2db6f16c7f9c48a..0000000000000000000000000000000000000000
--- a/FastVGGT/assets/vs.png
+++ /dev/null
@@ -1,3 +0,0 @@
-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
deleted file mode 100644
index 49170288c162684ebafa68e5242efd320f0f35b8..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/base.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc b/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc
deleted file mode 100644
index 6e6e0478775fe98dc22180bc43bad8073c43c907..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/criterion.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/__pycache__/data.cpython-310.pyc b/FastVGGT/eval/__pycache__/data.cpython-310.pyc
deleted file mode 100644
index 4d8ec6f8e1828fa87b02a4511490ad858de2093a..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/data.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/__pycache__/data.cpython-37.pyc b/FastVGGT/eval/__pycache__/data.cpython-37.pyc
deleted file mode 100644
index d76f2c95748956a04f009e99f39cefd9cc6ec4e0..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/data.cpython-37.pyc and /dev/null differ
diff --git a/FastVGGT/eval/__pycache__/utils.cpython-310.pyc b/FastVGGT/eval/__pycache__/utils.cpython-310.pyc
deleted file mode 100644
index 638db49f10789e53dcf0ff33c9aaa28b814e3e72..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/utils.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/__pycache__/utils.cpython-37.pyc b/FastVGGT/eval/__pycache__/utils.cpython-37.pyc
deleted file mode 100644
index fd73966947cb74c319d66d7bf221d73b5dde0ce0..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/__pycache__/utils.cpython-37.pyc and /dev/null differ
diff --git a/FastVGGT/eval/base.py b/FastVGGT/eval/base.py
deleted file mode 100644
index 4a716449a71f552ea408df0eb37e854cf4e92da6..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/base.py
+++ /dev/null
@@ -1,273 +0,0 @@
-# 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
deleted file mode 100644
index a63c991b5fd5b61325b56055d7d380a27a336971..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/criterion.py
+++ /dev/null
@@ -1,534 +0,0 @@
-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
deleted file mode 100644
index 301788e3ae6b67091fb031b068ecbecff1eae48d..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/data.py
+++ /dev/null
@@ -1,338 +0,0 @@
-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
deleted file mode 100644
index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/dataset_utils/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index d00b3e8b6f5fd9ef7d1428ff771243e1557ddf9f..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc b/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc
deleted file mode 100644
index 9ebd5d22d44ddce633e9a05a7467e9ae0c3f9032..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/__init__.cpython-37.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc
deleted file mode 100644
index fbd10218eddf65486b375109c23cf8cecd7c3ba5..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/corr.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc
deleted file mode 100644
index 860c8970fc49d65b2a4101b77cc313a6aeac618a..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc b/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc
deleted file mode 100644
index 0fe7870fb079d9dc08ca647866ff8c77fd0f2ec0..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/cropping.cpython-37.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc b/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc
deleted file mode 100644
index 1b4530be10547e3e7bf79e6cef6470880cb710a4..0000000000000000000000000000000000000000
Binary files a/FastVGGT/eval/dataset_utils/__pycache__/transforms.cpython-310.pyc and /dev/null differ
diff --git a/FastVGGT/eval/dataset_utils/corr.py b/FastVGGT/eval/dataset_utils/corr.py
deleted file mode 100644
index fbf5de18b7388e38e45c3f313487957811abc483..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/dataset_utils/corr.py
+++ /dev/null
@@ -1,234 +0,0 @@
-# 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
deleted file mode 100644
index 30a9eac18d241b71538957cf7ba4767ebc323b43..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/dataset_utils/cropping.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# 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
deleted file mode 100644
index cec2d144e1b97cc99c191d15cdcaf20796cae94b..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/dataset_utils/transforms.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# 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
deleted file mode 100644
index 8d6d5fbf6c159ea18bfd6fbcee0a5d23e63c0cf2..0000000000000000000000000000000000000000
--- a/FastVGGT/eval/eval_7andN.py
+++ /dev/null
@@ -1,497 +0,0 @@
-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