hloc/pipelines/4Seasons/prepare_reference.py

from pathlib import Path
import argparse

from ... import extract_features, match_features
from ... import pairs_from_poses, triangulation
from .utils import get_timestamps, delete_unused_images
from .utils import build_empty_colmap_model

parser = argparse.ArgumentParser()
parser.add_argument(
    "--dataset",
    type=Path,
    default="datasets/4Seasons",
    help="Path to the dataset, default: %(default)s",
)
parser.add_argument(
    "--outputs",
    type=Path,
    default="outputs/4Seasons",
    help="Path to the output directory, default: %(default)s",
)
args = parser.parse_args()

ref_dir = args.dataset / "reference"
assert ref_dir.exists(), f"{ref_dir} does not exist"
ref_images = ref_dir / "undistorted_images"

output_dir = args.outputs
output_dir.mkdir(exist_ok=True, parents=True)
ref_sfm_empty = output_dir / "sfm_reference_empty"
ref_sfm = output_dir / "sfm_superpoint+superglue"

num_ref_pairs = 20
ref_pairs = output_dir / f"pairs-db-dist{num_ref_pairs}.txt"

fconf = extract_features.confs["superpoint_max"]
mconf = match_features.confs["superglue"]

# Only reference images that have a pose are used in the pipeline.
# To save time in feature extraction, we delete unsused images.
delete_unused_images(ref_images, get_timestamps(ref_dir / "poses.txt", 0))

# Build an empty COLMAP model containing only camera and images
# from the provided poses and intrinsics.
build_empty_colmap_model(ref_dir, ref_sfm_empty)

# Match reference images that are spatially close.
pairs_from_poses.main(ref_sfm_empty, ref_pairs, num_ref_pairs)

# Extract, match, and triangulate the reference SfM model.
ffile = extract_features.main(fconf, ref_images, output_dir)
mfile = match_features.main(mconf, ref_pairs, fconf["output"], output_dir)
triangulation.main(ref_sfm, ref_sfm_empty, ref_images, ref_pairs, ffile, mfile)