Add support for loading 3D points with ODM data processor (#2838)

* Add support for loading 3D points with OpenDroneMap data processor

nerfstudio/process_data/odm_utils.py

@@ -42,21 +42,69 @@ def rodrigues_vec_to_rotation_mat(rodrigues_vec: np.ndarray) -> np.ndarray:
         )
         r_cross = np.array([[0, -r[2], r[1]], [r[2], 0, -r[0]], [-r[1], r[0], 0]], dtype=float)
         rotation_mat = math.cos(theta) * ident + (1 - math.cos(theta)) * r_rT + math.sin(theta) * r_cross
+
     return rotation_mat
 
 
+def get_reconstruction(reconstruction_file: Path):
+    with open(reconstruction_file, "r", encoding="utf-8") as f:
+        reconstructions = json.loads(f.read())
+        return reconstructions[0]
+
+
+def reconstruction_to_ply(reconstruction: dict, output_ply: Path):
+    points = reconstruction.get("points", [])
+    vertices = []
+    colors = []
+    transform = np.array([[0, 0, 1, 0], [0, -1, 0, 0], [1, 0, 0, 0]])
+
+    for pid in points:
+        point = points[pid]
+        p, c = point["coordinates"], point["color"]
+        vertices.append(p)
+        colors.append(c)
+
+    vertices = np.array(vertices)
+    coords = []
+
+    vertices = np.einsum("ij,bj->bi", transform[:3, :3], vertices) + transform[:3, 3]
+    for i in range(len(vertices)):
+        v = vertices[i]
+        c = colors[i]
+        coords.append("{} {} {} {} {} {}".format(v[0], v[1], v[2], int(c[0]), int(c[1]), int(c[2])))
+
+    header = [
+        "ply",
+        "format ascii 1.0",
+        "element vertex {}".format(len(coords)),
+        "property float x",
+        "property float y",
+        "property float z",
+        "property uchar red",
+        "property uchar green",
+        "property uchar blue",
+        "end_header",
+    ]
+
+    with open(output_ply, "w", encoding="utf-8") as of:
+        of.write("\n".join(header + coords + [""]))
+
+
 def cameras2nerfds(
     image_filename_map: Dict[str, Path],
     cameras_file: Path,
     shots_file: Path,
+    reconstruction_file: Path,
     output_dir: Path,
     verbose: bool = False,
 ) -> List[str]:
     """Convert ODM cameras into a nerfstudio dataset.
 
     Args:
         image_filename_map: Mapping of original image filenames to their saved locations.
+        cameras_file: Path to ODM's cameras.json
         shots_file: Path to ODM's shots.geojson
+        reconstruction_file: Path to ODM's reconstruction.json
         output_dir: Path to the output directory.
         verbose: Whether to print verbose output.
 
@@ -66,8 +114,6 @@ def cameras2nerfds(
 
     with open(cameras_file, "r", encoding="utf-8") as f:
         cameras = json.loads(f.read())
-    with open(shots_file, "r", encoding="utf-8") as f:
-        shots = json.loads(f.read())
 
     camera_ids = list(cameras.keys())
     if len(camera_ids) > 1:
@@ -99,20 +145,40 @@ def cameras2nerfds(
 
         sensor_dict[camera_id] = s
 
-    shots = shots["features"]
     shots_dict = {}
-    for shot in shots:
-        props = shot["properties"]
-        filename = props["filename"]
-        rotation = rodrigues_vec_to_rotation_mat(np.array(props["rotation"]) * -1)
-        translation = np.array(props["translation"])
+    reconstruction = None
+
+    if reconstruction_file.exists:
+        reconstruction = get_reconstruction(reconstruction_file)
+        shots = reconstruction.get("shots", [])
+        for filename in shots:
+            shot = shots[filename]
+            rotation = rodrigues_vec_to_rotation_mat(np.array(shot["rotation"]) * -1)
+            origin = -rodrigues_vec_to_rotation_mat(np.array(shot["rotation"])).T.dot(np.array(shot["translation"]))
+
+            m = np.eye(4)
+            m[:3, :3] = rotation
+            m[:3, 3] = origin
+
+            name, ext = os.path.splitext(filename)
+            shots_dict[name] = m
+    else:
+        with open(shots_file, "r", encoding="utf-8") as f:
+            shots = json.loads(f.read())
+
+        shots = shots["features"]
+        for shot in shots:
+            props = shot["properties"]
+            filename = props["filename"]
+            rotation = rodrigues_vec_to_rotation_mat(np.array(props["rotation"]) * -1)
+            translation = np.array(props["translation"])
 
-        m = np.eye(4)
-        m[:3, :3] = rotation
-        m[:3, 3] = translation
+            m = np.eye(4)
+            m[:3, :3] = rotation
+            m[:3, 3] = translation
 
-        name, ext = os.path.splitext(filename)
-        shots_dict[name] = m
+            name, ext = os.path.splitext(filename)
+            shots_dict[name] = m
 
     frames = []
     num_skipped = 0
@@ -134,6 +200,10 @@ def cameras2nerfds(
 
     data["frames"] = frames
 
+    if reconstruction is not None:
+        reconstruction_to_ply(reconstruction, output_dir / "reconstruction.ply")
+        data["ply_file_path"] = "reconstruction.ply"
+
     with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
         json.dump(data, f, indent=4)
 

nerfstudio/scripts/process_data.py

@@ -411,6 +411,7 @@ def main(self) -> None:
         orig_images_dir = self.data / "images"
         cameras_file = self.data / "cameras.json"
         shots_file = self.data / "odm_report" / "shots.geojson"
+        reconstruction_file = self.data / "opensfm" / "reconstruction.json"
 
         if not shots_file.exists:
             raise ValueError(f"shots file {shots_file} doesn't exist")
@@ -463,6 +464,7 @@ def main(self) -> None:
                 image_filename_map=image_filename_map,
                 cameras_file=cameras_file,
                 shots_file=shots_file,
+                reconstruction_file=reconstruction_file,
                 output_dir=self.output_dir,
                 verbose=self.verbose,
             )