Replacing VGGSfM tracker with VGGT tracker module

.gitignore

@@ -155,6 +155,7 @@ data/
 # Data dumped by GTSFM directory
 debug/
 plots/
+**/*result*/
 results/
 result_metrics/
 *.html

environment_linux.yml

@@ -56,7 +56,7 @@ dependencies:
       - pydegensac
       - colour
       - trimesh[easy]
-      - gtsam==4.3a0
+      - gtsam-develop==4.3a1.dev202602040056
       - pydot
       # dust3r/mast3r
       - roma

environment_linux_cpuonly.yml

@@ -56,7 +56,7 @@ dependencies:
       - pydegensac
       - colour
       - trimesh[easy]
-      - gtsam==4.3a0
+      - gtsam-develop==4.3a1.dev202602040056
       - pydot
       - roma
       - tqdm

environment_mac.yml

@@ -57,7 +57,7 @@ dependencies:
       - pydegensac
       - colour
       - trimesh[easy]
-      - gtsam==4.3a0
+      - gtsam-develop==4.3a1.dev202602040056
       - pydot
       # dust3r/mast3r
       - roma

gtsfm/averaging/translation/averaging_1dsfm.py

@@ -51,7 +51,7 @@
 MAX_PROJECTION_DIRECTIONS = 2000
 OUTLIER_WEIGHT_THRESHOLD = 0.125
 
-NOISE_MODEL_DIMENSION = 3  # chordal distances on Unit3
+NOISE_MODEL_DIMENSION = 2  # chordal distances on Unit3
 NOISE_MODEL_SIGMA = 0.01
 HUBER_LOSS_K = 1.3  # default value from GTSAM
 
@@ -73,7 +73,7 @@
 L = symbol_shorthand.B  # for track (landmark) translation variables
 
 RelativeDirectionsDict = AnnotatedGraph[Unit3]
-DUMMY_NOISE_MODEL = gtsam.noiseModel.Isotropic.Sigma(3, 1e-2)  # MFAS does not use this.
+DUMMY_NOISE_MODEL = gtsam.noiseModel.Isotropic.Sigma(2, 1e-2)  # MFAS does not use this.
 
 
 class TranslationAveraging1DSFM(TranslationAveragingBase):

gtsfm/bundle/bundle_adjustment.py

@@ -12,13 +12,7 @@
 import gtsam  # type: ignore
 import numpy as np
 from dask.delayed import Delayed
-from gtsam import (
-    BetweenFactorPose3,
-    NonlinearFactorGraph,
-    PriorFactorPose3,
-    PriorFactorPoint3,
-    Values,
-)
+from gtsam import BetweenFactorPose3, NonlinearFactorGraph, PriorFactorPoint3, PriorFactorPose3, Values
 from gtsam.noiseModel import Diagonal, Isotropic, Robust, mEstimator  # type: ignore
 from gtsam.symbol_shorthand import K, P, X  # type: ignore
 

gtsfm/cluster_merging.py

@@ -7,9 +7,10 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional, Tuple
 
+import gtsam
 import numpy as np
 from dask.distributed import Client, Future
-from gtsam import Similarity3, Pose3
+from gtsam import Similarity3, Pose3, UnaryMeasurementPose3, TrajectoryAlignerSim3
 
 import gtsfm.utils.logger as logger_utils
 import gtsfm.common.types as gtsfm_types
@@ -35,6 +36,56 @@
 _SCENE_LABEL_ATTR = "_gtsfm_cluster_label"
 
 
+def _create_unary_measurements(scene: GtsfmData) -> list[UnaryMeasurementPose3]:
+    # TODO(akshay-krishnan): investigate using a scene-dependent noise model
+    # perhaps * np.exp(-len(scene.get_valid_camera_indices()) / 100.0)
+    noise_model = gtsam.noiseModel.Diagonal.Sigmas(np.array([1e-2, 1e-2, 1e-2, 1e-1, 1e-1, 1e-1]))
+    unary_measurements = []
+    for i, camera in scene.get_camera_poses().items():
+        if camera is None:
+            continue
+        unary_measurement = UnaryMeasurementPose3(i, camera, noise_model)
+        unary_measurements.append(unary_measurement)
+    return unary_measurements
+
+
+def merge_scenes_with_sim3_nonlinear(parent_scene: GtsfmData, children_scenes: list[GtsfmData]) -> GtsfmData:
+    if len(children_scenes) == 0:
+        return parent_scene
+
+    aTi_measurements = _create_unary_measurements(parent_scene)
+    parent_camera_ids = set(parent_scene.get_valid_camera_indices())
+    valid_child_scenes = []
+
+    for i, child_scene in enumerate(children_scenes):
+        child_camera_ids = set(child_scene.get_valid_camera_indices())
+        common_camera_ids = parent_camera_ids & child_camera_ids
+        if len(common_camera_ids) == 0:
+            logger.warning("Child scene %d has insufficient overlap with parent, skipping", i)
+            continue
+        valid_child_scenes.append(child_scene)
+
+    if len(valid_child_scenes) == 0:
+        return parent_scene
+
+    aTi_measurements = _create_unary_measurements(parent_scene)
+    bTi_measurements = [_create_unary_measurements(child_scene) for child_scene in valid_child_scenes]
+    aligner = TrajectoryAlignerSim3(aTi_measurements, bTi_measurements)
+    result = aligner.solve()
+
+    opt_aTi = {i: result.atPose3(i) for i in parent_scene.get_valid_camera_indices() if i in result.keys()}
+
+    merged = parent_scene
+    for i, aTi in opt_aTi.items():
+        merged.update_camera_pose(i, aTi)
+
+    for i, child_scene in enumerate(valid_child_scenes):
+        opt_bSa = result.atSimilarity3(gtsam.Symbol("S", i).key())
+        opt_aSb = opt_bSa.inverse()
+        merged = merged.merged_with(child_scene, opt_aSb)  # type: ignore
+    return merged
+
+
 @dataclass(frozen=True)
 class MergedNodeResult:
     """Results of merging child scenes with parent scenes in the reconstruction tree.
@@ -224,9 +275,7 @@ def _get_pose_metrics(
 
     aligned_result_data = result_data.align_via_sim3_and_transform(poses_gt)
     return metrics_utils.compute_ba_pose_metrics(
-        gt_wTi=poses_gt,
-        computed_wTi=aligned_result_data.get_camera_poses(),
-        save_dir=save_dir,
+        gt_wTi=poses_gt, computed_wTi=aligned_result_data.get_camera_poses(), save_dir=save_dir, store_full_data=True
     )
 
 
@@ -333,6 +382,8 @@ def _drop_outlier_tracks(scene: GtsfmData) -> GtsfmData:
     Returns:
         The scene with outlier tracks dropped.
     """
+    if scene.number_tracks() == 0:
+        return scene
     track_errors: list[float] = []
     tracks = scene.tracks()
     cameras = scene.cameras()
@@ -406,6 +457,7 @@ def combine_results(
     drop_camera_with_no_track: bool = True,
     drop_child_if_merging_fail: bool = True,
     store_full_data: bool = False,
+    use_nonlinear_sim3_alignment: bool = False,
 ) -> MergedNodeResult:
     """Run the merging and parent BA pipeline using already-transformed children.
 
@@ -469,17 +521,27 @@ def _finalize_result(result_scene: Optional[GtsfmData]) -> MergedNodeResult:
     merged = current
     _log_scene_reprojection_stats(merged, "Current node", plot_histograms=plot_reprojection_histograms)
 
-    # Merge all children into the merged scene.
-    for i, child in enumerate(valid_child_scenes):
-        merged = _align_and_merge_results(merged, child, drop_if_merging_fails=drop_child_if_merging_fail)
-        _log_scene_reprojection_stats(merged, f"Merged with child #{i+1}", plot_histograms=plot_reprojection_histograms)
+    if use_nonlinear_sim3_alignment:
+        merged = merge_scenes_with_sim3_nonlinear(merged, valid_child_scenes)
+        _log_scene_reprojection_stats(
+            merged, "Merged with children (nonlinear alignment)", plot_histograms=plot_reprojection_histograms
+        )
+    else:
+        # Merge all children into the merged scene.
+        for i, child in enumerate(valid_child_scenes):
+            merged = _align_and_merge_results(merged, child, drop_if_merging_fails=drop_child_if_merging_fail)
+            _log_scene_reprojection_stats(
+                merged, f"Merged with child #{i+1}", plot_histograms=plot_reprojection_histograms
+            )
 
     _propagate_scene_metadata(merged, metadata_source)
 
-    if drop_outlier_after_camera_merging and merged is not None and merged.number_tracks() > 0:
-        merged = _drop_outlier_tracks(merged)
+    if merged is None:
+        return _finalize_result(None)
 
     if not run_bundle_adjustment_on_parent:
+        if drop_outlier_after_camera_merging:
+            merged = _drop_outlier_tracks(merged)
         return _finalize_result(merged)
 
     # Log cameras that have no supporting track measurements before running BA.
@@ -498,6 +560,14 @@ def _finalize_result(result_scene: Optional[GtsfmData]) -> MergedNodeResult:
             "merged result (with ba)",
             plot_histograms=plot_reprojection_histograms,
         )
+        if drop_outlier_after_camera_merging:
+            merged_with_ba = _drop_outlier_tracks(merged_with_ba)
+            _log_scene_reprojection_stats(
+                merged_with_ba,
+                "merged result (with ba + outlier filtering)",
+                plot_histograms=plot_reprojection_histograms,
+            )
+
         # TODO: the order here is different from the merging order above, we should fix this.
         if merged.has_gaussian_splats():
             logger.info("🫱🏻‍🫲🏽 Merging Gaussians")

gtsfm/cluster_optimizer/cluster_optimizer_base.py

@@ -102,13 +102,15 @@ def __init__(
         drop_outlier_after_camera_merging: bool = True,
         plot_reprojection_histograms: bool = True,
         run_bundle_adjustment_on_parent: bool = True,
+        run_bundle_adjustment_on_leaf: bool = False,
         output_worker: None | str = None,
     ) -> None:
         self.drop_child_if_merging_fail = drop_child_if_merging_fail
         self.drop_camera_with_no_track = drop_camera_with_no_track
         self.drop_outlier_after_camera_merging = drop_outlier_after_camera_merging
         self.plot_reprojection_histograms = plot_reprojection_histograms
         self.run_bundle_adjustment_on_parent = run_bundle_adjustment_on_parent
+        self.run_bundle_adjustment_on_leaf = run_bundle_adjustment_on_leaf
         self._pose_angular_error_thresh = pose_angular_error_thresh
         self._output_worker = output_worker
 

gtsfm/cluster_optimizer/cluster_optimizer_cacher.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import hashlib
+import os
 import typing
 from pathlib import Path
 from typing import Optional, TYPE_CHECKING
@@ -34,18 +35,26 @@
 class ClusterOptimizerCacher(ClusterOptimizerBase):
     """Caches the delayed bundle result produced by a cluster optimizer."""
 
-    def __init__(self, optimizer: ClusterOptimizerBase) -> None:
+    def __init__(self, optimizer: ClusterOptimizerBase, cache_subdir: Optional[str] = None) -> None:
         """Initializes the cacher with the actual cluster optimizer object.
 
         Args:
             optimizer: cluster optimizer to use in case of cache miss.
+            cache_subdir: Optional subdirectory (relative to cache root) for storing cache entries.
         """
+        run_bundle_adjustment_on_leaf = getattr(optimizer, "run_bundle_adjustment_on_leaf", None)
+        if run_bundle_adjustment_on_leaf is None:
+            run_bundle_adjustment_on_leaf = getattr(optimizer, "_run_bundle_adjustment_on_leaf", False)
         super().__init__(
             pose_angular_error_thresh=optimizer.pose_angular_error_thresh,
+            run_bundle_adjustment_on_leaf=run_bundle_adjustment_on_leaf,
+            run_bundle_adjustment_on_parent=getattr(optimizer, "run_bundle_adjustment_on_parent", True),
             output_worker=optimizer._output_worker,
         )
         self._optimizer = optimizer
         self._optimizer_hash = hashlib.sha1(repr(optimizer).encode()).hexdigest()
+        self._cache_subdir = cache_subdir if cache_subdir is not None else os.getenv("GTSFM_CACHE_SUBDIR")
+        self._cache_root_path = self._resolve_cache_root(self._cache_subdir)
 
     def __repr__(self) -> str:
         return repr(self._optimizer)
@@ -64,20 +73,38 @@ def __getstate__(self) -> dict[str, object]:
         return {
             "_optimizer": self._optimizer,
             "_optimizer_hash": self._optimizer_hash,
+            "_cache_subdir": self._cache_subdir,
         }
 
     def __setstate__(self, state: dict[str, object]) -> None:
         """Restore state and keep worker routing consistent."""
         self._optimizer = typing.cast(ClusterOptimizerBase, state["_optimizer"])
         self._optimizer_hash = typing.cast(str, state["_optimizer_hash"])
+        self._cache_subdir = typing.cast(Optional[str], state.get("_cache_subdir"))
+        self._cache_root_path = self._resolve_cache_root(self._cache_subdir)
         # Re-initialize the base class to mimic the constructor.
+        run_bundle_adjustment_on_leaf = getattr(self._optimizer, "run_bundle_adjustment_on_leaf", None)
+        if run_bundle_adjustment_on_leaf is None:
+            run_bundle_adjustment_on_leaf = getattr(self._optimizer, "_run_bundle_adjustment_on_leaf", False)
         super().__init__(
             pose_angular_error_thresh=self._optimizer.pose_angular_error_thresh,
+            run_bundle_adjustment_on_leaf=run_bundle_adjustment_on_leaf,
+            run_bundle_adjustment_on_parent=getattr(self._optimizer, "run_bundle_adjustment_on_parent", True),
             output_worker=self._optimizer._output_worker,
         )
 
+    @staticmethod
+    def _resolve_cache_root(cache_subdir: Optional[str]) -> Path:
+        """Resolve the cache root path, optionally using a subdirectory or absolute override."""
+        if not cache_subdir:
+            return CACHE_ROOT_PATH
+        subdir_path = Path(cache_subdir)
+        if subdir_path.is_absolute():
+            return subdir_path
+        return CACHE_ROOT_PATH / subdir_path
+
     def _get_cache_path(self, cache_key: str) -> Path:
-        return CACHE_ROOT_PATH / "cluster_optimizer" / f"{cache_key}.pbz2"
+        return self._cache_root_path / "cluster_optimizer" / f"{cache_key}.pbz2"
 
     def _hash_one_view_data(self, one_view_data: Optional["OneViewData"]) -> str:
         """Compute a stable hash for OneViewData contents."""
@@ -127,11 +154,25 @@ def _save_result_to_cache(self, result: GtsfmData, cache_path: Path) -> GtsfmDat
         io_utils.write_to_bz2_file(result, cache_path)
         return result
 
+    def _save_cached_result_outputs(self, result: GtsfmData, results_path: Path) -> None:
+        """Persist cached outputs expected by downstream tooling.
+
+        Currently used to re-export VGGT reconstructions in COLMAP text format.
+        """
+        if "VGGT" not in type(self._optimizer).__name__:
+            return
+        target_dir = results_path / "vggt"
+        target_dir.mkdir(parents=True, exist_ok=True)
+        result.export_as_colmap_text(target_dir)
+
     def create_computation_graph(self, context: ClusterContext) -> ClusterComputationGraph | None:
         cached_result = self._load_result_from_cache(context)
         if cached_result is not None:
             cached_graph: Delayed = delayed(lambda r: r, pure=False)(cached_result)
-            return ClusterComputationGraph(io_tasks=tuple(), metric_tasks=tuple(), sfm_result=cached_graph)
+            io_tasks = (
+                delayed(self._save_cached_result_outputs, pure=False)(cached_graph, context.output_paths.results),
+            )
+            return ClusterComputationGraph(io_tasks=io_tasks, metric_tasks=tuple(), sfm_result=cached_graph)
 
         computation = self._optimizer.create_computation_graph(context)
         if computation is None or computation.sfm_result is None: