initial OpenSfM pipeline added to processvideo (needs cleanup and additional features)

scripts/process_data.py

@@ -15,6 +15,7 @@ import appdirs
 import numpy as np
 import requests
 import tyro
+import yaml
 from rich.console import Console
 from rich.progress import track
 from typing_extensions import Annotated
@@ -484,6 +485,134 @@ def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, came
     return len(frames)
 
 
+def run_opensfm(image_dir: Path, opensfm_dir: Path, camera_model: CameraModel, exe: Path, verbose=False) -> None:
+    # TODO cleanup config
+    config = {
+        "processes": 12,
+        "matching_bow_neighbors": 6,
+        "feature_process_size": 1280,
+        "triangulation_threshold": 0.004,
+        "triangulation_type": "ROBUST",
+        "min_track_length": 3,
+        "retriangulation_ratio": 2,
+        "bundle_new_points_ratio": 2,
+    }
+    with open(opensfm_dir / "config.yaml", "w") as config_file:
+        yaml.dump(config, config_file)
+    # TODO change cam type based on camera_model
+    # TODO get the width and height from images
+    cam_overrides = {"all": {"projection_type": "brown", "width": 540, "height": 960, "focal_x": 0.85, "focal_y": 0.85}}
+    with open(opensfm_dir / "camera_models_overrides.json", "w") as cam_file:
+        json.dump(cam_overrides, cam_file)
+
+    run_command(f"cp -r {image_dir} {opensfm_dir/'images'}")
+    metadata_cmd = [f"{exe} extract_metadata", f"{opensfm_dir}"]
+    metadata_cmd = " ".join(metadata_cmd)
+    with status(msg="[bold yellow]Extracting OpenSfM metadata...", spinner="moon", verbose=verbose):
+        run_command(metadata_cmd, verbose=verbose)
+    CONSOLE.log("[bold green]:tada: Done OpenSfM metadata.")
+
+    feature_extractor_cmd = [f"{exe} detect_features", f"{opensfm_dir}"]
+    feature_extractor_cmd = " ".join(feature_extractor_cmd)
+    with status(msg="[bold yellow]Running OpenSfM feature extractor...", spinner="runner", verbose=verbose):
+        run_command(feature_extractor_cmd, verbose=verbose)
+    CONSOLE.log("[bold green]:tada: Done extracting image features.")
+
+    matcher_cmd = [f"{exe} match_features", f"{opensfm_dir}"]
+    matcher_cmd = " ".join(matcher_cmd)
+    with status(msg="[bold yellow]Running OpenSfM feature matcher...", spinner="runner", verbose=verbose):
+        run_command(matcher_cmd, verbose=verbose)
+    CONSOLE.log("[bold green]:tada: Done matching image features.")
+
+    tracks_cmd = [f"{exe} create_tracks", f"{opensfm_dir}"]
+    tracks_cmd = " ".join(tracks_cmd)
+    with status(msg="[bold yellow]Merging OpenSfM features into tracked points...", spinner="dqpb", verbose=verbose):
+        run_command(tracks_cmd, verbose=verbose)
+    CONSOLE.log("[bold green]:tada: Done merging features.")
+
+    recon_cmd = [f"{exe} reconstruct", f"{opensfm_dir}"]
+    recon_cmd = " ".join(recon_cmd)
+    with status(
+        msg="[bold yellow]Running OpenSfM bundle adjustment... (this can take 10-15 minutes)",
+        spinner="dqpb",
+        verbose=verbose,
+    ):
+        run_command(recon_cmd, verbose=verbose)
+    CONSOLE.log("[bold green]:tada: Done running bundle adjustment.")
+
+
+def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: CameraModel) -> int:
+    sfm_out = (json.load(open(reconstruction_path, "r")))[0]
+    cams = sfm_out["cameras"]
+    shots = sfm_out["shots"]
+    frames = []
+    for shot_name, shot_info in shots.items():
+        # enumerates through images in the capture
+        aa_vec = np.array(shot_info["rotation"])  # 3D axis angle repr
+        angle = np.linalg.norm(aa_vec)
+        if angle > 1e-5:
+            # normalilze the axis-angle repr if angle is large enough
+            aa_vec /= angle
+        qx = aa_vec[0] * np.sin(angle / 2)
+        qy = aa_vec[1] * np.sin(angle / 2)
+        qz = aa_vec[2] * np.sin(angle / 2)
+        qw = np.cos(angle / 2)
+        rotation = colmap_utils.qvec2rotmat(np.array([qw, qx, qy, qz]))
+        translation = np.array(shot_info["translation"]).reshape(3, 1)
+        w2c = np.concatenate([rotation, translation], 1)
+        w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0)
+        c2w = np.linalg.inv(w2c)
+        # Convert from COLMAP's camera coordinate system to ours
+        # TODO determine if this is necessary for OpenSfM or not
+        c2w[0:3, 1:3] *= -1
+        # c2w = c2w[np.array([1, 0, 2, 3]), :]
+        # c2w[2, :] *= -1
+
+        name = Path(f"./images/{shot_name}")
+
+        frame = {
+            "file_path": str(name),
+            "transform_matrix": c2w.tolist(),
+        }
+        frames.append(frame)
+    # For now just assume it's all the same camera
+    cam = cams[list(cams.keys())[0]]
+    out = {
+        "fl_x": float(cam["focal_x"]) * cam["width"],
+        "fl_y": float(cam["focal_y"]) * cam["height"],
+        "cx": cam["width"] / 2.0 + float(cam["c_x"]) * cam["width"],
+        "cy": cam["height"] / 2.0 + float(cam["c_y"]) * cam["height"],
+        "w": cam["width"],
+        "h": cam["height"],
+        "camera_model": camera_model.value,
+    }
+    if camera_model == CameraModel.OPENCV:
+        out.update(
+            {
+                "k1": float(cam["k1"]),
+                "k2": float(cam["k2"]),
+                "p1": float(cam["p1"]),
+                "p2": float(cam["p2"]),
+            }
+        )
+    if camera_model == CameraModel.OPENCV_FISHEYE:
+        # TODO: find the opensfm camera model that uses these params
+        out.update(
+            {
+                "k1": float(cam["k1"]),
+                "k2": float(cam["k2"]),
+                "k3": float(cam["k3"]),
+                "k4": float(cam["k4"]),
+            }
+        )
+    out["frames"] = frames
+
+    with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
+        json.dump(out, f, indent=4)
+
+    return len(frames)
+
+
 def get_insta360_filenames(data: Path) -> Tuple[Path, Path]:
     """Returns the filenames of the Insta360 videos from a single video file.
 
@@ -609,10 +738,14 @@ class ProcessVideo:
     """Path the data, either a video file or a directory of images."""
     output_dir: Path
     """Path to the output directory."""
-    num_frames_target: int = 300
+    num_frames_target: int = 100
     """Target number of frames to use for the dataset, results may not be exact."""
     camera_type: Literal["perspective", "fisheye"] = "perspective"
     """Camera model to use."""
+    sfm_method: Literal["colmap", "OpenSfM"] = "OpenSfM"
+    """Which sfm solver to use"""
+    opensfm_exec: Optional[Path] = None
+    """If sfm_method is OpenSfM, specify the executable to use here"""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
         accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos."""
@@ -643,32 +776,51 @@ class ProcessVideo:
         # Downscale images
         summary_log.append(downscale_images(image_dir, self.num_downscales, verbose=self.verbose))
 
-        # Run Colmap
-        colmap_dir = self.output_dir / "colmap"
-        if not self.skip_colmap:
-            colmap_dir.mkdir(parents=True, exist_ok=True)
+        if self.sfm_method == "colmap":
+            # Run Colmap
+            colmap_dir = self.output_dir / "colmap"
+            if not self.skip_colmap:
+                colmap_dir.mkdir(parents=True, exist_ok=True)
 
-            run_colmap(
-                image_dir=image_dir,
-                colmap_dir=colmap_dir,
-                camera_model=CAMERA_MODELS[self.camera_type],
-                gpu=self.gpu,
-                verbose=self.verbose,
-                matching_method=self.matching_method,
-            )
-
-        # Save transforms.json
-        if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():
-            with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
-                num_matched_frames = colmap_to_json(
-                    cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin",
-                    images_path=colmap_dir / "sparse" / "0" / "images.bin",
-                    output_dir=self.output_dir,
+                run_colmap(
+                    image_dir=image_dir,
+                    colmap_dir=colmap_dir,
                     camera_model=CAMERA_MODELS[self.camera_type],
+                    gpu=self.gpu,
+                    verbose=self.verbose,
+                    matching_method=self.matching_method,
+                )
+
+            # Save transforms.json
+            if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():
+                with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
+                    num_matched_frames = colmap_to_json(
+                        cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin",
+                        images_path=colmap_dir / "sparse" / "0" / "images.bin",
+                        output_dir=self.output_dir,
+                        camera_model=CAMERA_MODELS[self.camera_type],
+                    )
+                    summary_log.append(f"Colmap matched {num_matched_frames} images")
+            else:
+                CONSOLE.log(
+                    "[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json"
+                )
+        elif self.sfm_method == "OpenSfM":
+            assert self.opensfm_exec is not None, "Please provide the path to the OpenSfM executable <...>/bin/opensfm"
+            # Run the OpenSfM solver
+            opensfm_dir = self.output_dir / "opensfm"
+            opensfm_dir.mkdir(parents=True, exist_ok=True)
+            run_opensfm(image_dir, opensfm_dir, CAMERA_MODELS[self.camera_type], self.opensfm_exec)
+            # Save transforms.json
+            if (opensfm_dir / "reconstruction.json").exists():
+                with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
+                    opensfm_to_json(
+                        opensfm_dir / "reconstruction.json", self.output_dir, CAMERA_MODELS[self.camera_type]
+                    )
+            else:
+                CONSOLE.log(
+                    "[bold yellow]Warning: could not find existing OpenSfM results. Not generating transforms.json"
                 )
-                summary_log.append(f"Colmap matched {num_matched_frames} images")
-        else:
-            CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json")
 
         CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:")
 

scripts/viewer/view_dataset.py

@@ -13,7 +13,7 @@ from pathlib import Path
 import torch
 import tyro
 
-from nerfstudio.configs import method_configs as cfg
+from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.data.datamanagers import AnnotatedDataParserUnion
 from nerfstudio.data.utils.datasets import InputDataset
 from nerfstudio.viewer.server import viewer_utils
@@ -26,7 +26,7 @@ torch.backends.cudnn.benchmark = True  # type: ignore
 
 def main(
     dataparser: AnnotatedDataParserUnion,
-    viewer: cfg.ViewerConfig,
+    viewer: ViewerConfig,
     log_base_dir: Path = Path("/tmp/nerfstudio_viewer_logs"),
 ) -> None:
     """Main function."""
@@ -34,7 +34,7 @@ def main(
         viewer,
         log_filename=log_base_dir / viewer.relative_log_filename,
     )
-    dataset = InputDataset(dataparser.setup().get_dataset_inputs(split="train"))
+    dataset = InputDataset(dataparser.setup().get_dataparser_outputs(split="train"))
     viewer_state.init_scene(dataset=dataset, start_train=False)
     CONSOLE.log("Please refresh and load page at: %s", viewer_state.viewer_url)
     time.sleep(30)  # allowing time to refresh page