update run with custom data

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

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)

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)