Add PatchMatchNet module for MVS and calculation of normals from depth

CMakeLists.txt

@@ -59,6 +59,7 @@ option(OPENMP_ENABLED "Whether to enable OpenMP parallelization" ON)
 option(IPO_ENABLED "Whether to enable interprocedural optimization" ON)
 option(CUDA_ENABLED "Whether to enable CUDA, if available" ON)
 option(OPENGL_ENABLED "Whether to enable OpenGL, if available" ON)
+option(TORCH_ENABLED "Whether to enable LibTorch if available" ON)
 option(TESTS_ENABLED "Whether to build test binaries" OFF)
 option(PROFILING_ENABLED "Whether to enable google-perftools linker flags" OFF)
 option(CGAL_ENABLED "Whether to enable the CGAL library" ON)
@@ -130,6 +131,16 @@ if(CGAL_FOUND)
     message(STATUS "  Libraries : ${CGAL_LIBRARY}")
 endif()
 
+set(CMAKE_TORCH_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/mvs-modules)
+if (TORCH_ENABLED)
+    set(CMAKE_TORCHLIB_PATH ${CMAKE_CURRENT_SOURCE_DIR}/lib/libtorch)
+    # Set the MKL include path to LibTorch's self contained headers
+    # instead of requiring a full MKL installation to be present
+    set(MKL_INCLUDE_DIR ${CMAKE_TORCHLIB_PATH}/include)
+    set(MKL_LIBRARIES ";")
+    list(APPEND CMAKE_PREFIX_PATH ${CMAKE_TORCHLIB_PATH})
+    find_package(Torch QUIET)
+endif()
 
 ################################################################################
 # Compiler specific configuration
@@ -220,6 +231,20 @@ else()
     endif()
 endif()
 
+if(Torch_FOUND)
+    if(TORCH_ENABLED)
+        add_definitions("-DTORCH_ENABLED")
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+        message(STATUS "Enabling LibTorch support")
+    else()
+        set(Torch_FOUND OFF)
+        message(STATUS "Disabling LibTorch support")
+    endif()
+else()
+    set(TORCH_ENABLED OFF)
+    message(STATUS "Disabling LibTorch support")
+endif()
+
 if(OPENGL_ENABLED)
     add_definitions("-DOPENGL_ENABLED")
     message(STATUS "Enabling OpenGL support")
@@ -272,6 +297,7 @@ set(COLMAP_INCLUDE_DIRS
     ${Qt5Core_INCLUDE_DIRS}
     ${Qt5OpenGL_INCLUDE_DIRS}
     ${Qt5Widgets_INCLUDE_DIRS}
+    ${Torch_INCLUDE_DIRS}
 )
 
 set(COLMAP_LINK_DIRS
@@ -290,6 +316,7 @@ set(COLMAP_EXTERNAL_LIBRARIES
     ${Qt5Core_LIBRARIES}
     ${Qt5OpenGL_LIBRARIES}
     ${Qt5Widgets_LIBRARIES}
+    ${OPENGL_gl_LIBRARY}
 )
 
 if(OPENMP_FOUND)
@@ -302,6 +329,11 @@ if(CGAL_FOUND)
     list(APPEND COLMAP_LINK_DIRS ${CGAL_LIBRARIES_DIR})
 endif()
 
+if(Torch_FOUND)
+    list(APPEND COLMAP_INCLUDE_DIRS ${TORCH_INCLUDE_DIRS})
+    list(APPEND COLMAP_EXTERNAL_LIBRARIES ${TORCH_LIBRARIES})
+endif()
+
 if(UNIX)
     list(APPEND COLMAP_EXTERNAL_LIBRARIES pthread)
 endif()

cmake/CMakeHelper.cmake

@@ -106,6 +106,20 @@ macro(COLMAP_ADD_CUDA_SOURCES)
         PARENT_SCOPE)
 endmacro(COLMAP_ADD_CUDA_SOURCES)
 
+# Macro to add Torch source files to COLMAP library.
+macro(COLMAP_ADD_TORCH_SOURCES)
+    set(SOURCE_FILES "")
+    foreach(SOURCE_FILE ${ARGN})
+        if(SOURCE_FILE MATCHES "^/.*")
+            list(APPEND SOURCE_FILES ${SOURCE_FILE})
+        else()
+            list(APPEND SOURCE_FILES
+                 "${CMAKE_CURRENT_SOURCE_DIR}/${SOURCE_FILE}")
+        endif()
+    endforeach()
+    set(COLMAP_TORCH_SOURCES ${COLMAP_TORCH_SOURCES} ${SOURCE_FILES} PARENT_SCOPE)
+endmacro(COLMAP_ADD_TORCH_SOURCES)
+
 # Replacement for the normal add_library() command. The syntax remains the same
 # in that the first argument is the target name, and the following arguments
 # are the source files to use when building the target.

src/CMakeLists.txt

@@ -68,6 +68,9 @@ if(CUDA_ENABLED)
     # Use a separate stream per thread to allow for concurrent kernel execution
     # between multiple threads on the same device.
     set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} --default-stream per-thread")
+    # Remove NVCC flag set by LibTorch configuration that causes compilation error
+    # in Eigen/Core when compiling patch_match_cuda.cu
+    list(REMOVE_ITEM CUDA_NVCC_FLAGS "--Werror" "cross-execution-space-call")
 
     # Fix for Ubuntu 16.04.
     add_definitions("-D_MWAITXINTRIN_H_INCLUDED")
@@ -79,9 +82,22 @@ if(CUDA_ENABLED)
 
     COLMAP_ADD_STATIC_CUDA_LIBRARY(colmap_cuda ${COLMAP_CUDA_SOURCES})
     target_link_libraries(colmap_cuda colmap)
+endif()
+
+if(TORCH_ENABLED)
+    COLMAP_ADD_STATIC_LIBRARY(colmap_torch ${COLMAP_TORCH_SOURCES})
+    target_link_libraries(colmap_torch colmap)
+endif()
 
+if(CUDA_ENABLED AND TORCH_ENABLED)
+    target_link_libraries(colmap
+        colmap_cuda colmap_torch ${COLMAP_INTERNAL_LIBRARIES} ${COLMAP_EXTERNAL_LIBRARIES})
+elseif(CUDA_ENABLED)
     target_link_libraries(colmap
         colmap_cuda ${COLMAP_INTERNAL_LIBRARIES} ${COLMAP_EXTERNAL_LIBRARIES})
+elseif(TORCH_ENABLED)
+    target_link_libraries(colmap
+        colmap_torch ${COLMAP_INTERNAL_LIBRARIES} ${COLMAP_EXTERNAL_LIBRARIES})
 else()
     target_link_libraries(colmap
         ${COLMAP_INTERNAL_LIBRARIES} ${COLMAP_EXTERNAL_LIBRARIES})

src/base/undistortion.cc

@@ -164,10 +164,12 @@ void COLMAPUndistorter::Run() {
   CreateDirIfNotExists(JoinPaths(output_path_, "sparse"));
   CreateDirIfNotExists(JoinPaths(output_path_, "stereo"));
   CreateDirIfNotExists(JoinPaths(output_path_, "stereo/depth_maps"));
+  CreateDirIfNotExists(JoinPaths(output_path_, "stereo/confidence_maps"));
   CreateDirIfNotExists(JoinPaths(output_path_, "stereo/normal_maps"));
   CreateDirIfNotExists(JoinPaths(output_path_, "stereo/consistency_graphs"));
   reconstruction_.CreateImageDirs(JoinPaths(output_path_, "images"));
   reconstruction_.CreateImageDirs(JoinPaths(output_path_, "stereo/depth_maps"));
+  reconstruction_.CreateImageDirs(JoinPaths(output_path_, "stereo/confidence_maps"));
   reconstruction_.CreateImageDirs(
       JoinPaths(output_path_, "stereo/normal_maps"));
   reconstruction_.CreateImageDirs(

src/exe/CMakeLists.txt

@@ -37,3 +37,15 @@ endif()
 
 COLMAP_ADD_EXECUTABLE(colmap_exe colmap.cc)
 set_target_properties(colmap_exe PROPERTIES OUTPUT_NAME colmap)
+
+if(IS_MSVC)
+  file(GLOB REQUIRED_EXE_FILES "${CMAKE_TORCHLIB_PATH}/lib/*.dll" "${CMAKE_TORCH_MODULE_PATH}/*.pt")
+else()
+  file(GLOB REQUIRED_EXE_FILES "${CMAKE_TORCH_MODULE_PATH}/*.pt")
+endif()
+
+add_custom_command(TARGET colmap_exe
+                   POST_BUILD
+                   COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                   ${REQUIRED_EXE_FILES}
+                   $<TARGET_FILE_DIR:colmap_exe>)

src/exe/colmap.cc

@@ -417,15 +417,16 @@ int RunDelaunayMesher(int argc, char** argv) {
 }
 
 int RunPatchMatchStereo(int argc, char** argv) {
-#ifndef CUDA_ENABLED
-  std::cerr << "ERROR: Dense stereo reconstruction requires CUDA, which is not "
-               "available on your system."
+#if !defined(CUDA_ENABLED) && !defined(TORCH_ENABLED)
+  std::cerr << "ERROR: Dense stereo reconstruction requires CUDA or Torch, and "
+               "neither is available on your system."
             << std::endl;
   return EXIT_FAILURE;
-#else   // CUDA_ENABLED
+#else
   std::string workspace_path;
   std::string workspace_format = "COLMAP";
   std::string pmvs_option_name = "option-all";
+  std::string method = "standard";
   std::string config_path;
 
   OptionManager options;
@@ -435,6 +436,7 @@ int RunPatchMatchStereo(int argc, char** argv) {
   options.AddDefaultOption("workspace_format", &workspace_format,
                            "{COLMAP, PMVS}");
   options.AddDefaultOption("pmvs_option_name", &pmvs_option_name);
+  options.AddDefaultOption("method", &method, "{standard, learned}");
   options.AddDefaultOption("config_path", &config_path);
   options.AddPatchMatchStereoOptions();
   options.Parse(argc, argv);
@@ -447,6 +449,30 @@ int RunPatchMatchStereo(int argc, char** argv) {
     return EXIT_FAILURE;
   }
 
+  if (ExistsFile(options.patch_match_stereo->mvs_module_path)) {
+#ifdef TORCH_ENABLED
+    std::cout << "Using learned patch-patch with module from: "
+              << options.patch_match_stereo->mvs_module_path << std::endl;
+    options.patch_match_stereo->patch_match_method =
+        mvs::PatchMatchOptions::PatchMatchMethod::Learned;
+    options.patch_match_stereo->geom_consistency = false;
+#else
+    std::cout << "WARN: Learned dense stereo reconstruction requires Torch, "
+                 "which is not available on your system; reverting to standard "
+                 "dense reconstruction with CUDA."
+              << std::endl;
+#endif
+  } else {
+#ifdef CUDA_ENABLED
+    std::cout << "Using standard patch-match" << std::endl;
+#else
+    std::cerr << "ERROR: Standard dense stereo reconstruction requires CUDA, "
+                 "which is not available on your system."
+              << std::endl;
+    return EXIT_FAILURE;
+#endif
+  }
+
   mvs::PatchMatchController controller(*options.patch_match_stereo,
                                        workspace_path, workspace_format,
                                        pmvs_option_name, config_path);
@@ -455,7 +481,7 @@ int RunPatchMatchStereo(int argc, char** argv) {
   controller.Wait();
 
   return EXIT_SUCCESS;
-#endif  // CUDA_ENABLED
+#endif  // !defined(CUDA_ENABLED) && !defined(TORCH_ENABLED)
 }
 
 int RunExhaustiveMatcher(int argc, char** argv) {

src/mvs/CMakeLists.txt

@@ -39,9 +39,16 @@ COLMAP_ADD_SOURCES(
     meshing.h meshing.cc
     model.h model.cc
     normal_map.h normal_map.cc
+    patch_match.h patch_match.cc
     workspace.h workspace.cc
 )
 
+if(TORCH_ENABLED)
+    COLMAP_ADD_TORCH_SOURCES(
+        patch_match_net.h patch_match_net.cc
+    )
+endif()
+
 COLMAP_ADD_TEST(consistency_graph_test consistency_graph_test.cc)
 COLMAP_ADD_TEST(depth_map_test depth_map_test.cc)
 COLMAP_ADD_TEST(mat_test mat_test.cc)
@@ -51,7 +58,6 @@ if(CUDA_ENABLED)
     COLMAP_ADD_CUDA_SOURCES(
         gpu_mat_prng.h gpu_mat_prng.cu
         gpu_mat_ref_image.h gpu_mat_ref_image.cu
-        patch_match.h patch_match.cc
         patch_match_cuda.h patch_match_cuda.cu
     )
 

src/mvs/depth_map.cc

@@ -127,5 +127,44 @@ Bitmap DepthMap::ToBitmap(const float min_percentile,
   return bitmap;
 }
 
+Bitmap ConfidenceMap::ToBitmap() const {
+  CHECK_GT(width_, 0);
+  CHECK_GT(height_, 0);
+
+  Bitmap bitmap;
+  bitmap.Allocate(width_, height_, true);
+
+  std::vector<float> valid_values;
+  valid_values.reserve(data_.size());
+  for (const float value : data_) {
+    if (value >= 0.0f && value <= 1.0f) {
+      valid_values.push_back(value);
+    }
+  }
+
+  if (valid_values.empty()) {
+    bitmap.Fill(BitmapColor<uint8_t>(0));
+    return bitmap;
+  }
+
+  for (size_t y = 0; y < height_; ++y) {
+    for (size_t x = 0; x < width_; ++x) {
+      const float confidence = Get(y, x);
+      if (confidence < 0) {
+        bitmap.SetPixel(x, y, BitmapColor<uint8_t>(0));
+      } else if (confidence > 1.0f) {
+        bitmap.SetPixel(x, y, BitmapColor<uint8_t>(255));
+      } else {
+        const BitmapColor<float> color(255 * JetColormap::Red(confidence),
+                                       255 * JetColormap::Green(confidence),
+                                       255 * JetColormap::Blue(confidence));
+        bitmap.SetPixel(x, y, color.Cast<uint8_t>());
+      }
+    }
+  }
+
+  return bitmap;
+}
+
 }  // namespace mvs
 }  // namespace colmap

src/mvs/depth_map.h

@@ -48,13 +48,15 @@ class DepthMap : public Mat<float> {
            const float depth_max);
   DepthMap(const Mat<float>& mat, const float depth_min, const float depth_max);
 
-  inline float GetDepthMin() const;
-  inline float GetDepthMax() const;
+  inline float GetDepthMin() const { return depth_min_; }
+  inline float GetDepthMax() const { return depth_max_; }
 
-  inline float Get(const size_t row, const size_t col) const;
+  virtual inline float Get(const size_t row, const size_t col) const {
+    return data_.at(row * width_ + col);
+  }
 
-  void Rescale(const float factor);
-  void Downsize(const size_t max_width, const size_t max_height);
+  virtual void Rescale(const float factor);
+  virtual void Downsize(const size_t max_width, const size_t max_height);
 
   Bitmap ToBitmap(const float min_percentile, const float max_percentile) const;
 
@@ -63,17 +65,16 @@ class DepthMap : public Mat<float> {
   float depth_max_ = -1.0f;
 };
 
-////////////////////////////////////////////////////////////////////////////////
-// Implementation
-////////////////////////////////////////////////////////////////////////////////
-
-float DepthMap::GetDepthMin() const { return depth_min_; }
-
-float DepthMap::GetDepthMax() const { return depth_max_; }
-
-float DepthMap::Get(const size_t row, const size_t col) const {
-  return data_.at(row * width_ + col);
-}
+class ConfidenceMap : public DepthMap {
+ public:
+  ConfidenceMap() : ConfidenceMap(0, 0) {}
+  ConfidenceMap(const size_t width, const size_t height)
+      : DepthMap(width, height, 0.0f, 1.0f) {
+    Fill(1.0f);
+  }
+  ConfidenceMap(const Mat<float>& mat) : DepthMap(mat, 0.0f, 1.0f) {}
+  Bitmap ToBitmap() const;
+};
 
 }  // namespace mvs
 }  // namespace colmap