beverley-gorry/colmap-vslamlab

1import argparse2import os3import subprocess4import sys5import urllib.request6 7 8def download_file(url, file_path, max_retries=3):9    if os.path.exists(file_path):10        return11    print(f"Downloading {url} to {file_path}")12    for retry in range(max_retries):13        try:14            urllib.request.urlretrieve(url, file_path)15            return16        except Exception as exc:17            print(18                f"Failed to download {url} (trial={retry + 1}) to {file_path} due to {exc}"19            )20 21 22def check_small_errors_or_exit(23    dataset_name,24    max_rotation_error,25    max_proj_center_error,26    expected_num_images,27    errors_csv_path,28):29    print(f"Evaluating errors for {dataset_name}")30 31    error = False32    with open(errors_csv_path, "r") as fid:33        num_images = 034        for line in fid:35            line = line.strip()36            if len(line) == 0 or line.startswith("#"):37                continue38            rotation_error, proj_center_error = map(float, line.split(","))39            num_images += 140            if rotation_error > max_rotation_error:41                print("Exceeded rotation error threshold:", rotation_error)42                error = True43            if proj_center_error > max_proj_center_error:44                print(45                    "Exceeded projection center error threshold:",46                    proj_center_error,47                )48                error = True49 50    if num_images != expected_num_images:51        print("Unexpected number of images:", num_images)52        error = True53 54    if error:55        sys.exit(1)56 57 58def process_dataset(args, dataset_name):59    print("Processing dataset:", dataset_name)60 61    workspace_path = os.path.join(62        os.path.realpath(args.workspace_path), dataset_name63    )64    os.makedirs(workspace_path, exist_ok=True)65 66    dataset_archive_path = os.path.join(workspace_path, f"{dataset_name}.7z")67    download_file(68        f"https://www.eth3d.net/data/{dataset_name}_dslr_undistorted.7z",69        dataset_archive_path,70    )71 72    subprocess.check_call(73        ["7zz", "x", "-y", f"{dataset_name}.7z"], cwd=workspace_path74    )75 76    # Find undistorted parameters of first camera and initialize all images with it.77    with open(78        os.path.join(79            workspace_path,80            f"{dataset_name}/dslr_calibration_undistorted/cameras.txt",81        ),82        "r",83    ) as fid:84        for line in fid:85            if not line.startswith("#"):86                first_camera_data = line.split()87                camera_model = first_camera_data[1]88                assert camera_model == "PINHOLE"89                camera_params = first_camera_data[4:]90                assert len(camera_params) == 491                break92 93    # Count the number of expected images in the GT.94    expected_num_images = 095    with open(96        os.path.join(97            workspace_path,98            f"{dataset_name}/dslr_calibration_undistorted/images.txt",99        ),100        "r",101    ) as fid:102        for line in fid:103            if not line.startswith("#") and line.strip():104                expected_num_images += 1105    # Each image uses two consecutive lines.106    assert expected_num_images % 2 == 0107    expected_num_images /= 2108 109    # Run automatic reconstruction pipeline.110    subprocess.check_call(111        [112            os.path.realpath(args.colmap_path),113            "automatic_reconstructor",114            "--image_path",115            f"{dataset_name}/images/",116            "--workspace_path",117            workspace_path,118            "--use_gpu",119            "1" if args.use_gpu else "0",120            "--num_threads",121            str(args.num_threads),122            "--quality",123            args.quality,124            "--camera_model",125            "PINHOLE",126            "--camera_params",127            ",".join(camera_params),128        ],129        cwd=workspace_path,130    )131 132    # Compare reconstructed model to GT model.133    subprocess.check_call(134        [135            os.path.realpath(args.colmap_path),136            "model_comparer",137            "--input_path1",138            "sparse/0",139            "--input_path2",140            f"{dataset_name}/dslr_calibration_undistorted/",141            "--output_path",142            ".",143            "--alignment_error",144            "proj_center",145            "--max_proj_center_error",146            str(args.max_proj_center_error),147        ],148        cwd=workspace_path,149    )150 151    # Ensure discrepancy between reconstructed model and GT is small.152    check_small_errors_or_exit(153        dataset_name,154        args.max_rotation_error,155        args.max_proj_center_error,156        expected_num_images,157        os.path.join(workspace_path, "errors.csv"),158    )159 160 161def parse_args():162    parser = argparse.ArgumentParser()163    parser.add_argument("--dataset_names", required=True)164    parser.add_argument("--workspace_path", required=True)165    parser.add_argument("--colmap_path", required=True)166    parser.add_argument("--use_gpu", default=True, action="store_true")167    parser.add_argument("--use_cpu", dest="use_gpu", action="store_false")168    parser.add_argument("--num_threads", type=int, default=-1)169    parser.add_argument("--quality", default="medium")170    parser.add_argument("--max_rotation_error", type=float, default=1.0)171    parser.add_argument("--max_proj_center_error", type=float, default=0.1)172    return parser.parse_args()173 174 175def main():176    args = parse_args()177 178    for dataset_name in args.dataset_names.split(","):179        process_dataset(args, dataset_name.strip())180 181 182if __name__ == "__main__":183    main()184