Adding expressive GTSFM data to allow non-consecutive cameras

.gitignore

@@ -146,5 +146,7 @@ data/
 # OS X Finder files
 *.DS_Store
 
-# Plot directory
+# Data dumped by GTSFM directory
 plots/
+results/
+result_metrics/

gtsfm/bundle/bundle_adjustment.py

@@ -5,15 +5,10 @@
 import dask
 import gtsam
 from dask.delayed import Delayed
-from gtsam import (
-    GeneralSFMFactorCal3Bundler,
-    SfmData,
-    SfmTrack,
-    Values,
-    symbol_shorthand,
-)
+from gtsam import GeneralSFMFactorCal3Bundler, SfmTrack, Values, symbol_shorthand
 
 import gtsfm.utils.logger as logger_utils
+from gtsfm.common.gtsfm_data import GtsfmData
 from gtsfm.common.sfm_result import SfmResult
 
 # TODO: any way this goes away?
@@ -33,7 +28,7 @@ class BundleAdjustmentOptimizer:
     This class refines global pose estimates and intrinsics of cameras, and also refines 3D point cloud structure given
     tracks from triangulation."""
 
-    def run(self, initial_data: SfmData) -> SfmResult:
+    def run(self, initial_data: GtsfmData) -> SfmResult:
         """Run the bundle adjustment by forming factor graph and optimizing using Levenberg–Marquardt optimization.
 
         Args:
@@ -42,7 +37,9 @@ def run(self, initial_data: SfmData) -> SfmResult:
         Results:
             optimized camera poses, 3D point w/ tracks, and error metrics.
         """
-        logger.info(f"Input: {initial_data.number_tracks()} tracks on {initial_data.number_cameras()} cameras\n")
+        logger.info(
+            f"Input: {initial_data.number_tracks()} tracks on {len(initial_data.get_valid_camera_indices())} cameras\n"
+        )
 
         # noise model for measurements -- one pixel in u and v
         measurement_noise = gtsam.noiseModel.Isotropic.Sigma(IMG_MEASUREMENT_DIM, 1.0)
@@ -52,47 +49,44 @@ def run(self, initial_data: SfmData) -> SfmResult:
 
         # Add measurements to the factor graph
         for j in range(initial_data.number_tracks()):
-            track = initial_data.track(j)  # SfmTrack
+            track = initial_data.get_track(j)  # SfmTrack
             # retrieve the SfmMeasurement objects
             for m_idx in range(track.number_measurements()):
                 # i represents the camera index, and uv is the 2d measurement
                 i, uv = track.measurement(m_idx)
                 # note use of shorthand symbols C and P
                 graph.add(GeneralSFMFactorCal3Bundler(uv, measurement_noise, C(i), P(j)))
 
-        # Add a prior on pose x1. This indirectly specifies where the origin is.
+        # get all the valid camera indices, which need to be added to the graph.
+        valid_camera_indices = initial_data.get_valid_camera_indices()
+
+        # Add a prior on first pose. This indirectly specifies where the origin is.
         graph.push_back(
             gtsam.PriorFactorPinholeCameraCal3Bundler(
-                C(0),
-                initial_data.camera(0),
+                C(valid_camera_indices[0]),
+                initial_data.get_camera(valid_camera_indices[0]),
                 gtsam.noiseModel.Isotropic.Sigma(PINHOLE_CAM_CAL3BUNDLER_DOF, 0.1),
             )
         )
         # Also add a prior on the position of the first landmark to fix the scale
         graph.push_back(
             gtsam.PriorFactorPoint3(
-                P(0),
-                initial_data.track(0).point3(),
-                gtsam.noiseModel.Isotropic.Sigma(POINT3_DOF, 0.1),
+                P(0), initial_data.get_track(0).point3(), gtsam.noiseModel.Isotropic.Sigma(POINT3_DOF, 0.1),
             )
         )
 
         # Create initial estimate
         initial = gtsam.Values()
 
-        i = 0
         # add each PinholeCameraCal3Bundler
-        for cam_idx in range(initial_data.number_cameras()):
-            camera = initial_data.camera(cam_idx)
+        for i in valid_camera_indices:
+            camera = initial_data.get_camera(i)
             initial.insert(C(i), camera)
-            i += 1
 
-        j = 0
         # add each SfmTrack
-        for t_idx in range(initial_data.number_tracks()):
-            track = initial_data.track(t_idx)
+        for j in range(initial_data.number_tracks()):
+            track = initial_data.get_track(j)
             initial.insert(P(j), track.point3())
-            j += 1
 
         # Optimize the graph and print results
         try:
@@ -102,7 +96,8 @@ def run(self, initial_data: SfmData) -> SfmResult:
             result_values = lm.optimize()
         except Exception as e:
             logger.exception("LM Optimization failed")
-            return SfmResult(SfmData(), float("Nan"))
+            # as we did not perform the bundle adjustment, we skip computing the total reprojection error
+            return SfmResult(GtsfmData(initial_data.number_images()), total_reproj_error=float("Nan"))
 
         final_error = graph.error(result_values)
 
@@ -111,25 +106,25 @@ def run(self, initial_data: SfmData) -> SfmResult:
         logger.info(f"final error: {final_error:.2f}")
 
         # construct the results
-        optimized_data = values_to_sfm_data(result_values, initial_data)
-        sfm_result = SfmResult(optimized_data, final_error)
+        optimized_data = values_to_gtsfm_data(result_values, initial_data)
+        sfm_result = SfmResult(optimized_data, total_reproj_error=final_error)
 
         return sfm_result
 
     def create_computation_graph(self, sfm_data_graph: Delayed) -> Delayed:
         """Create the computation graph for performing bundle adjustment.
 
         Args:
-            sfm_data_graph: an SfmData object wrapped up using dask.delayed
+            sfm_data_graph: an GtsfmData object wrapped up using dask.delayed
 
         Returns:
             SfmResult wrapped up using dask.delayed
         """
         return dask.delayed(self.run)(sfm_data_graph)
 
 
-def values_to_sfm_data(values: Values, initial_data: SfmData) -> SfmData:
-    """Cast results from the optimization to SfmData object.
+def values_to_gtsfm_data(values: Values, initial_data: GtsfmData) -> GtsfmData:
+    """Cast results from the optimization to GtsfmData object.
 
     Args:
         values: results of factor graph optimization.
@@ -139,20 +134,18 @@ def values_to_sfm_data(values: Values, initial_data: SfmData) -> SfmData:
     Returns:
         optimized poses and landmarks.
     """
-    result = SfmData()
+    result = GtsfmData(initial_data.number_images())
 
     # add cameras
-    for i in range(initial_data.number_cameras()):
-        result.add_camera(values.atPinholeCameraCal3Bundler(C(i)))
+    for i in initial_data.get_valid_camera_indices():
+        result.add_camera(i, values.atPinholeCameraCal3Bundler(C(i)))
 
     # add tracks
     for j in range(initial_data.number_tracks()):
-        input_track = initial_data.track(j)
+        input_track = initial_data.get_track(j)
 
         # populate the result with optimized 3D point
-        result_track = SfmTrack(
-            values.atPoint3(P(j)),
-        )
+        result_track = SfmTrack(values.atPoint3(P(j)),)
 
         for measurement_idx in range(input_track.number_measurements()):
             i, uv = input_track.measurement(measurement_idx)

gtsfm/common/gtsfm_data.py

@@ -0,0 +1,199 @@
+"""Class to hold the tracks and cameras of a 3D scene.
+
+Authors: Ayush Baid
+"""
+import itertools
+from typing import Dict, List, Optional
+
+import numpy as np
+from gtsam import PinholeCameraCal3Bundler, SfmTrack
+
+import gtsfm.utils.graph as graph_utils
+
+EQUALITY_TOLERANCE = 1e-5
+
+
+class GtsfmData:
+    """Class containing cameras and tracks, essentially describing the complete 3D scene.
+
+    This class is needed over GTSAM's SfmData type because GTSAM's type does not allow for non-contiguous cameras.
+    The situation of non-contiguous cameras can exists because of failures in front-end.
+    """
+
+    def __init__(self, number_images: int) -> None:
+        """Initializes the class.
+
+        Args:
+            number_images: number of images/cameras in the scene.
+        """
+        self._cameras: Dict[int, PinholeCameraCal3Bundler] = {}
+        self._tracks: List[SfmTrack] = []
+        self._number_images = number_images
+
+    def __eq__(self, other: object) -> bool:
+        """Checks equality with the other object."""
+
+        if not isinstance(other, GtsfmData):
+            return False
+
+        if self._number_images != other.number_images():
+            return False
+
+        for i, cam in self._cameras.items():
+            other_cam = other.get_camera(i)
+            if not cam.equals(other_cam, EQUALITY_TOLERANCE):
+                return False
+
+        for j in range(self.number_tracks()):
+            track = self.get_track(j)
+            other_track = other.get_track(j)
+
+            if track.number_measurements() != other_track.number_measurements():
+                return False
+
+            for k in range(track.number_measurements()):
+                i, uv = track.measurement(k)
+                other_i, other_uv = other_track.measurement(k)
+
+                if i != other_i:
+                    return False
+                if not np.allclose(uv, other_uv):
+                    return False
+
+        return True
+
+    def number_images(self) -> int:
+        """Getter for the number of images.
+
+        Returns:
+            Number of images.
+        """
+        return self._number_images
+
+    def number_tracks(self) -> int:
+        """Getter for the number of tracks.
+
+        Returns:
+            Number of tracks.
+        """
+        return len(self._tracks)
+
+    def get_valid_camera_indices(self) -> List[int]:
+        """Getter for image indices where there is a valid (not None) camera.
+
+        Returns:
+            List of indices with a valid camera.
+        """
+        return list(self._cameras.keys())
+
+    def get_camera(self, index: int) -> Optional[PinholeCameraCal3Bundler]:
+        """Getter for camera.
+
+        Args:
+            index: the image index to fetch the camera for.
+
+        Returns:
+            The camera if it is a valid one, None otherwise.
+        """
+        return self._cameras.get(index)
+
+    def get_track(self, index: int) -> SfmTrack:
+        """Getter for the track.
+
+        Args:
+            index: track index to fetch.
+
+        Returns:
+            Requested track.
+        """
+        return self._tracks[index]
+
+    def add_track(self, track: SfmTrack) -> bool:
+        """Add a track, after checking if all the cameras in the track are already added.
+
+        Args:
+            track: track to add.
+
+        Returns:
+            Flag indicating the success of adding operation.
+        """
+        # check if all cameras are already added
+        for j in range(track.number_measurements()):
+            i, _ = track.measurement(j)
+
+            if i not in self._cameras:
+                return False
+
+        self._tracks.append(track)
+        return True
+
+    def add_camera(self, index: int, camera: PinholeCameraCal3Bundler,) -> None:
+        """Adds a camera.
+
+        Args:
+            index: the index associated with this camera.
+            camera: camera object to it.
+
+        Raises:
+            ValueError: if the camera to be added is not a valid camera object.
+        """
+        if camera is None:
+            raise ValueError("Camera cannot be None, should be a valid camera")
+        self._cameras[index] = camera
+
+    def select_largest_connected_component(self) -> "GtsfmData":
+        """Selects the subset of data belonging to the largest connected component of the graph where the edges are
+        between cameras which feature in the same track.
+
+        Returns:
+            New GtSfmData object with the subset of tracks and cameras.
+        """
+        camera_edges = []
+        for sfm_track in self._tracks:
+            cameras_in_use = []
+            for m_idx in range(sfm_track.number_measurements()):
+                i, _ = sfm_track.measurement(m_idx)
+                cameras_in_use.append(i)
+
+            # Recreate track connectivity from track information
+            # For example: a track has cameras [0, 2, 5]. In that case we will add pairs (0, 2), (0, 5), (2, 5)
+            camera_edges += list(itertools.combinations(cameras_in_use, 2))
+
+        if len(camera_edges) == 0:
+            return GtsfmData(self._number_images)
+
+        cameras_in_largest_cc = graph_utils.get_nodes_in_largest_connected_component(camera_edges)
+        return GtsfmData.from_selected_cameras(self, cameras_in_largest_cc)
+
+    @classmethod
+    def from_selected_cameras(cls, gtsfm_data: "GtsfmData", camera_indices: List[int]) -> "GtsfmData":
+        """Selects the cameras in the input list and the tracks associated with those cameras.
+
+        Args:
+            gtsfm_data: data to pick the cameras from.
+            camera_indices: camera indices to select and keep in the new data.
+
+        Returns:
+            New object with the selected cameras and associated tracks.
+        """
+        new_data = cls(gtsfm_data.number_images())
+
+        for i in gtsfm_data.get_valid_camera_indices():
+            if i in camera_indices:
+                new_data.add_camera(i, gtsfm_data.get_camera(i))
+
+        new_camera_indices = new_data.get_valid_camera_indices()
+
+        # add tracks which have all the camera present in new data
+        for j in range(gtsfm_data.number_tracks()):
+            track = gtsfm_data.get_track(j)
+            is_valid = True
+            for k in range(track.number_measurements()):
+                i, _ = track.measurement(k)
+                if i not in new_camera_indices:
+                    is_valid = False
+                    break
+            if is_valid:
+                new_data.add_track(track)
+
+        return new_data