From 1d9cc7dd8a4f1c5fdc8f05dd4ea3b5c3ef6bb85a Mon Sep 17 00:00:00 2001
From: Alexandru Kis <alexandrukis1618033@gmail.com>
Date: Fri, 8 Dec 2023 18:58:51 +0200
Subject: [PATCH] Update hloc and add new matcher types (#2658)

* Update hloc and add new matcher types

* Update dockerfile comment to better reflect what the command is doing

---------

Co-authored-by: mfischer <michael.fischer@spleenlab.ai>
---
 Dockerfile                                    |  3 ++-
 .../colmap_converter_to_nerfstudio_dataset.py |  2 ++
 nerfstudio/process_data/hloc_utils.py         |  9 +++++++-
 nerfstudio/process_data/process_data_utils.py | 23 +++++++++++++++++--
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 300d00df3..8c9fe4f31 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -135,9 +135,10 @@ RUN git clone --branch v0.4.0 --recursive https://github.com/colmap/pycolmap.git
     python3.10 -m pip install . && \
     cd ..
 
-# Install hloc master (last release (1.3) is too old) as alternative feature detector and matcher option for nerfstudio.
+# Install hloc 1.4 as alternative feature detector and matcher option for nerfstudio.
 RUN git clone --branch master --recursive https://github.com/cvg/Hierarchical-Localization.git && \
     cd Hierarchical-Localization && \
+    git checkout v1.4 && \
     python3.10 -m pip install -e . && \
     cd ..
 
diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py
index b4dccbdd8..924c4b752 100644
--- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py
+++ b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py
@@ -66,6 +66,8 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
         "NN-ratio",
         "NN-mutual",
         "adalam",
+        "disk+lightglue",
+        "superpoint+lightglue",
     ] = "any"
     """Matching algorithm."""
     num_downscales: int = 3
diff --git a/nerfstudio/process_data/hloc_utils.py b/nerfstudio/process_data/hloc_utils.py
index a833f51ed..b7167efd8 100644
--- a/nerfstudio/process_data/hloc_utils.py
+++ b/nerfstudio/process_data/hloc_utils.py
@@ -60,7 +60,14 @@ def run_hloc(
         "sift", "superpoint_aachen", "superpoint_max", "superpoint_inloc", "r2d2", "d2net-ss", "sosnet", "disk"
     ] = "superpoint_aachen",
     matcher_type: Literal[
-        "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"
+        "superglue",
+        "superglue-fast",
+        "NN-superpoint",
+        "NN-ratio",
+        "NN-mutual",
+        "adalam",
+        "disk+lightglue",
+        "superpoint+lightglue",
     ] = "superglue",
     num_matched: int = 50,
     refine_pixsfm: bool = False,
diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py
index 334dab4de..fe24e55f2 100644
--- a/nerfstudio/process_data/process_data_utils.py
+++ b/nerfstudio/process_data/process_data_utils.py
@@ -477,7 +477,16 @@ def find_tool_feature_matcher_combination(
         "disk",
     ],
     matcher_type: Literal[
-        "any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"
+        "any",
+        "NN",
+        "superglue",
+        "superglue-fast",
+        "NN-superpoint",
+        "NN-ratio",
+        "NN-mutual",
+        "adalam",
+        "disk+lightglue",
+        "superpoint+lightglue",
     ],
 ) -> Union[
     Tuple[None, None, None],
@@ -493,7 +502,17 @@ def find_tool_feature_matcher_combination(
             "sosnet",
             "disk",
         ],
-        Literal["NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"],
+        Literal[
+            "NN",
+            "superglue",
+            "superglue-fast",
+            "NN-superpoint",
+            "NN-ratio",
+            "NN-mutual",
+            "adalam",
+            "disk+lightglue",
+            "superpoint+lightglue",
+        ],
     ],
 ]:
     """Find a valid combination of sfm tool, feature type, and matcher type.