diff --git a/datasets/data_collection.py b/datasets/data_collection.py index 88c4adc..266016e 100644 --- a/datasets/data_collection.py +++ b/datasets/data_collection.py @@ -192,7 +192,10 @@ def load_all_2Dpoints_custom(self): with open(test_path, 'r') as f: testlist = f.read().rstrip().split('\n') else: - raise ValueError("Error! Input file/directory {0} not found.".format(test_path)) + from run_custom_data.pipeline import create_test_list + create_test_list(self.args.sfm_dir, self.args.dataset_dir) + with open(test_path, 'r') as f: + testlist = f.read().rstrip().split('\n') for _, image in images_train.items(): img_name = image.name diff --git a/run_custom_data/directSfM_pipeline.py b/run_custom_data/directSfM_pipeline.py new file mode 100644 index 0000000..343f966 --- /dev/null +++ b/run_custom_data/directSfM_pipeline.py @@ -0,0 +1,30 @@ +from pathlib import Path + +from hloc import ( + extract_features, + match_features, + reconstruction, + visualization, + pairs_from_retrieval, +) +def direct_sfm(workspace_path:str): + images = Path(f"{workspace_path}/images/") + + outputs = Path(workspace_path) + sfm_pairs = outputs / "pairs-netvlad.txt" + sfm_dir = outputs / "sfm_superpoint+superglue" + + retrieval_conf = extract_features.confs["netvlad"] + feature_conf = extract_features.confs["superpoint_aachen"] + matcher_conf = match_features.confs["superglue"] + + retrieval_path = extract_features.main(retrieval_conf, images, outputs) + pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=5) + + + feature_path = extract_features.main(feature_conf, images, outputs) + match_path = match_features.main( + matcher_conf, sfm_pairs, feature_conf["output"], outputs + ) + + model = reconstruction.main(sfm_dir, images, sfm_pairs, feature_path, match_path) \ No newline at end of file diff --git a/run_custom_data/pipeline.py b/run_custom_data/pipeline.py index 2afce8f..910b4f5 100644 --- a/run_custom_data/pipeline.py +++ b/run_custom_data/pipeline.py @@ -14,6 +14,7 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from util.read_write_model import read_model import random +from run_custom_data.directSfM_pipeline import direct_sfm def re_triangulation(args, colmap_sfm="colmap_sift/sparse/0", test_rate=0.1): @@ -45,7 +46,7 @@ def re_triangulation(args, colmap_sfm="colmap_sift/sparse/0", test_rate=0.1): reference_sfm, outputs / colmap_sfm, images, sfm_pairs, features, sfm_matches ) create_test_list(reference_sfm, args.workspace_path, test_rate=0.1) - print(f"Randomly selected {test_rate} test images saved in 'test_list.txt'") + print(f"Randomly selecte {test_rate} test images and save in 'test_list.txt'") def run_colmap(confs): workspace_path = confs.workspace_path @@ -94,7 +95,13 @@ def create_test_list(reference_sfm, workspace_path, test_rate=0.1): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("workspace_path", default="../custom_dataset", type=Path) + parser.add_argument("--use_colmap", action='store_true' , help='use pre-trained model') args = parser.parse_args() - run_colmap(args) - re_triangulation(args) + if args.use_colmap: + print("use_colmap") + # run_colmap(args) + re_triangulation(args, colmap_sfm="sparse/0") + else: + direct_sfm(str(args.workspace_path)) + create_test_list(args.workspace_path / "sfm_superpoint+superglue/models", args.workspace_path, test_rate=0.1) \ No newline at end of file