Store rig assignments as a dictionary

Summary:
This enables to control the id that the rigs will have.
In most places, it was actually already a dictionary and only converted from/to a list for persisting it.

Reviewed By: YanNoun

Differential Revision: D32689416

fbshipit-source-id: 5f5ce00d9cc3dbc752e95cbe496eb4965a311959

annotation_gui_gcp/main.py

@@ -78,8 +78,8 @@ def load_rig_assignments(root: Path) -> t.Dict[str, t.List[str]]:
 
     output = {}
     with open(p_json) as f:
-        assignments = json.load(f)
-    for shot_group in assignments:
+        assignments: t.Dict[str, t.List[t.Tuple[str, str]]] = json.load(f)
+    for shot_group in assignments.values():
         group_shot_ids = [s[0] for s in shot_group]
         for shot_id, _ in shot_group:
             output[shot_id] = group_shot_ids

doc/source/rig.rst

@@ -26,7 +26,7 @@ The following terms define such assembly and capture in OpenSfM terminology :
  - A `RigInstance` is a list of `Shots`, each of which correspond to a `RigCamera` of the `RigModel` and the actual pose of the `RigModel` in the world : it's indeed an instanciation of the `RigModel` by combining `Shots`. These instances are defined in the `rig_assignments.json` file as follows::
 
     {
-        [
+        "RIG_INSTANCE_ID1": {
             [
                 "FILENAME",
                 "RIG_CAMERA_ID1"
@@ -40,8 +40,8 @@ The following terms define such assembly and capture in OpenSfM terminology :
                 "FILENAME",
                 "RIG_CAMERA_IDn"
             ]
-        ],
-        [
+        },
+        "RIG_INSTANCE_ID2": {
             [
                 "FILENAME",
                 "RIG_CAMERA_ID1"
@@ -55,7 +55,7 @@ The following terms define such assembly and capture in OpenSfM terminology :
                 "FILENAME",
                 "RIG_CAMERA_IDn"
             ]
-        ],
+        },
         ...
 
 

opensfm/actions/create_rig.py

@@ -38,15 +38,20 @@ def _reconstruction_from_rigs_and_assignments(data: DataSetBase):
 
     reconstruction = types.Reconstruction()
     reconstruction.cameras = data.load_camera_models()
-    for instance in assignments:
-        for image, camera_id in instance:
-            rig_camera = rig_cameras[camera_id]
+    for rig_instance_id, instance in assignments.items():
+        for image, rig_camera_id in instance:
+            rig_camera = rig_cameras[rig_camera_id]
             rig_pose = pygeometry.Pose(base_rotation)
             rig_pose.set_origin(
                 helpers.get_image_metadata(data, image).gps_position.value
             )
             d = data.load_exif(image)
-            shot = reconstruction.create_shot(image, d["camera"])
+            shot = reconstruction.create_shot(
+                image,
+                camera_id=d["camera"],
+                rig_camera_id=rig_camera_id,
+                rig_instance_id=rig_instance_id,
+            )
             shot.pose = rig_camera.pose.compose(rig_pose)
             shot.metadata = helpers.get_image_metadata(data, image)
     return [reconstruction]

opensfm/dataset.py

@@ -516,14 +516,21 @@ def _rig_assignments_file(self) -> str:
         """Return path of rig assignments file"""
         return os.path.join(self.data_path, "rig_assignments.json")
 
-    def load_rig_assignments(self) -> List[List[Tuple[str, str]]]:
+    def load_rig_assignments(self) -> Dict[str, List[Tuple[str, str]]]:
         """Return rig assignments  data"""
         if not self.io_handler.exists(self._rig_assignments_file()):
-            return []
+            return {}
         with self.io_handler.open_rt(self._rig_assignments_file()) as fin:
-            return json.load(fin)
+            assignments = json.load(fin)
+
+        # Backward compatibility.
+        # Older versions of the file were stored as a list of instances without id.
+        if isinstance(assignments, list):
+            assignments = {str(i): v for i, v in enumerate(assignments)}
+
+        return assignments
 
-    def save_rig_assignments(self, rig_assignments: List[List[Tuple[str, str]]]):
+    def save_rig_assignments(self, rig_assignments: Dict[str, List[Tuple[str, str]]]):
         """Save rig assignments  data"""
         with self.io_handler.open_wt(self._rig_assignments_file()) as fout:
             io.json_dump(rig_assignments, fout)

opensfm/dataset_base.py

@@ -214,11 +214,11 @@ def save_rig_cameras(self, rig_cameras: Dict[str, pymap.RigCamera]) -> None:
         pass
 
     @abstractmethod
-    def load_rig_assignments(self) -> List[List[Tuple[str, str]]]:
+    def load_rig_assignments(self) -> Dict[str, List[Tuple[str, str]]]:
         pass
 
     @abstractmethod
-    def save_rig_assignments(self, rig_assignments: List[List[Tuple[str, str]]]):
+    def save_rig_assignments(self, rig_assignments: Dict[str, List[Tuple[str, str]]]):
         pass
 
     @abstractmethod

opensfm/rig.py

@@ -36,11 +36,11 @@ def default_rig_cameras(camera_ids: Iterable[str]) -> Dict[str, pymap.RigCamera]
 
 
 def rig_assignments_per_image(
-    rig_assignments: List[List[Tuple[str, str]]],
+    rig_assignments: Dict[str, List[Tuple[str, str]]],
 ) -> Dict[str, Tuple[str, str, List[str]]]:
     """Return rig assignments data for each image."""
     assignments_per_image = {}
-    for instance_id, instance in enumerate(rig_assignments):
+    for instance_id, instance in rig_assignments.items():
         instance_shots = [s[0] for s in instance]
         for (shot_id, rig_camera_id) in instance:
             assignments_per_image[shot_id] = (
@@ -378,7 +378,7 @@ def create_rigs_with_pattern(data: "DataSet", patterns: TRigPatterns):
             f"Found a candidate for rig calibration with {len(best_reconstruction.shots)} shots"
         )
         data.save_rig_cameras(best_rig_cameras)
-        data.save_rig_assignments(list(instances_per_rig.values()))
+        data.save_rig_assignments(instances_per_rig)
     else:
         logger.error(
             "Could not run any sucessful SfM on images subset for rig calibration"

opensfm/synthetic_data/synthetic_dataset.py

@@ -91,10 +91,12 @@ def save_camera_models(self, camera_models: Dict[str, pygeometry.Camera]) -> Non
     def load_rig_cameras(self) -> Dict[str, pymap.RigCamera]:
         return self.reconstruction.rig_cameras
 
-    def load_rig_assignments(self) -> List[List[Tuple[str, str]]]:
-        rig_assignments = []
+    def load_rig_assignments(self) -> Dict[str, List[Tuple[str, str]]]:
+        rig_assignments = {}
         for instance in self.reconstruction.rig_instances.values():
-            rig_assignments.append([(k, v.id) for k, v in instance.rig_cameras.items()])
+            rig_assignments[instance.id] = [
+                (k, v.id) for k, v in instance.rig_cameras.items()
+            ]
         return rig_assignments
 
     def load_exif(self, image: str) -> Dict[str, Any]:
@@ -169,9 +171,7 @@ def load_tracks_manager(
             raise RuntimeError("No tracks manager for the synthetic dataset")
         return tracks_mgr
 
-    def init_reference(
-        self, images: Optional[List[str]] = None
-    ) -> None:
+    def init_reference(self, images: Optional[List[str]] = None) -> None:
         pass
 
     def load_reference(self) -> geo.TopocentricConverter: