colmap · ahojnnes · Dec 21, 2024 · Dec 21, 2024 · Dec 21, 2024 · Dec 21, 2024
diff --git a/.github/workflows/build-mac.yml b/.github/workflows/build-mac.yml
@@ -63,7 +63,8 @@ jobs:
             glew \
             cgal \
             sqlite3 \
-            ccache
+            ccache \
+            pytorch
           brew link --force libomp
 
       - name: Configure and build
@@ -76,6 +77,7 @@ jobs:
             -GNinja \
             -DCMAKE_BUILD_TYPE=${{ matrix.config.cmakeBuildType }} \
             -DTESTS_ENABLED=ON \
+            -DTORCH_ENABLED=ON \
             -DQt5_DIR="$(brew --prefix qt@5)/lib/cmake/Qt5"
           ninja
 

diff --git a/.github/workflows/build-pycolmap.yml b/.github/workflows/build-pycolmap.yml
@@ -21,10 +21,10 @@ jobs:
     strategy:
       matrix:
         config: [
-              {os: ubuntu-latest},
-              {os: macos-13, arch: x86_64, deploymentTarget: 13.0},
-              {os: macos-14, arch: arm64, deploymentTarget: 14.0},
-              {os: windows-latest},
+          {os: ubuntu-latest},
+          {os: macos-13, arch: x86_64, deploymentTarget: 13.0},
+          {os: macos-14, arch: arm64, deploymentTarget: 14.0},
+          {os: windows-latest},
         ]
     env:
       COMPILER_CACHE_VERSION: 1

diff --git a/.github/workflows/build-ubuntu.yml b/.github/workflows/build-ubuntu.yml
@@ -30,6 +30,7 @@ jobs:
             e2eTests: false,
             checkCodeFormat: true,
             coverageEnabled: true,
+            torchEnabled: false,
           },
           {
             os: ubuntu-22.04,
@@ -40,6 +41,7 @@ jobs:
             e2eTests: true,
             checkCodeFormat: true,
             coverageEnabled: false,
+            torchEnabled: true,
           },
           {
             os: ubuntu-22.04,
@@ -50,6 +52,7 @@ jobs:
             e2eTests: false,
             checkCodeFormat: false,
             coverageEnabled: false,
+            torchEnabled: false,
           },
           {
             os: ubuntu-24.04,
@@ -60,6 +63,7 @@ jobs:
             e2eTests: false,
             checkCodeFormat: false,
             coverageEnabled: false,
+            torchEnabled: false,
           },
           {
             os: ubuntu-24.04,
@@ -70,6 +74,7 @@ jobs:
             e2eTests: false,
             checkCodeFormat: false,
             coverageEnabled: false,
+            torchEnabled: false,
           },
         ]
 
@@ -82,12 +87,16 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+        with:
+          lfs: true
+
       - uses: actions/cache@v4
         id: cache-builds
         with:
-          key: v${{ env.COMPILER_CACHE_VERSION }}-${{ matrix.config.os }}-${{ matrix.config.cmakeBuildType }}-${{ matrix.config.asanEnabled }}--${{ matrix.config.cudaEnabled }}-${{ github.run_id }}-${{ github.run_number }}
-          restore-keys: v${{ env.COMPILER_CACHE_VERSION }}-${{ matrix.config.os }}-${{ matrix.config.cmakeBuildType }}-${{ matrix.config.asanEnabled }}--${{ matrix.config.cudaEnabled }}
+          key: v${{ env.COMPILER_CACHE_VERSION }}-${{ matrix.config.os }}-${{ matrix.config.cmakeBuildType }}-${{ matrix.config.asanEnabled }}--${{ matrix.config.cudaEnabled }}-${{ matrix.config.torchEnabled }}-${{ github.run_id }}-${{ github.run_number }}
+          restore-keys: v${{ env.COMPILER_CACHE_VERSION }}-${{ matrix.config.os }}-${{ matrix.config.cmakeBuildType }}-${{ matrix.config.asanEnabled }}--${{ matrix.config.cudaEnabled }}-${{ matrix.config.torchEnabled }}
           path: ${{ env.COMPILER_CACHE_DIR }}
+
       - name: Install compiler cache
         run: |
           mkdir -p "$CCACHE_DIR" "$CTCACHE_DIR"
@@ -144,7 +153,8 @@ jobs:
             libgl1-mesa-dri \
             libunwind-dev \
             libcurl4-openssl-dev \
-            xvfb
+            xvfb \
+            libmkl-dev
 
           if [ "${{ matrix.config.cudaEnabled }}" == "true" ]; then
             if [ "${{ matrix.config.os }}" == "ubuntu-20.04" ]; then
@@ -180,6 +190,12 @@ jobs:
             sudo apt-get install -y gcovr
           fi
 
+          if [ "${{ matrix.config.torchEnabled }}" == "true" ]; then
+            sudo apt-get install -y python3 python3-pip
+            pip3 install torch --index-url https://download.pytorch.org/whl/cpu
+            echo "Torch_DIR=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'`" >> $GITHUB_ENV
+          fi
+
       - name: Configure and build
         run: |
           set -x
@@ -195,6 +211,7 @@ jobs:
             -DCUDA_ENABLED=${{ matrix.config.cudaEnabled }} \
             -DGUI_ENABLED=${{ matrix.config.guiEnabled }} \
             -DASAN_ENABLED=${{ matrix.config.asanEnabled }} \
+            -DTORCH_ENABLED=${{ matrix.config.torchEnabled }} \
             -DCOVERAGE_ENABLED=${{ matrix.config.coverageEnabled }}
           ninja -k 10000
 
@@ -275,6 +292,7 @@ jobs:
           indicators: true
           output: both
           thresholds: '75 90'
+
       # TODO: Add code coverage comment to PR. Currently, this action reports
       # coverage for the entire repository, not just the changed files in the PR.
       # We could manually filter the coverage report to only include the changed

diff --git a/.github/workflows/build-windows.yml b/.github/workflows/build-windows.yml
@@ -25,13 +25,15 @@ jobs:
             os: windows-2022,
             cmakeBuildType: Release,
             cudaEnabled: true,
+            torchEnabled: false,
             testsEnabled: true,
             exportPackage: true,
           },
           {
             os: windows-2022,
             cmakeBuildType: Release,
             cudaEnabled: false,
+            torchEnabled: false,
             testsEnabled: true,
             exportPackage: true,
           },
@@ -45,6 +47,9 @@ jobs:
       VCPKG_COMMIT_ID: bc3512a509f9d29b37346a7e7e929f9a26e66c7e
       GLOG_v: 1
       GLOG_logtostderr: 1
+      CUDA_MAJOR_VERSION: 12
+      CUDA_MINOR_VERSION: 6
+      CUDA_PATCH_VERSION: 2
 
     steps:
       - uses: actions/checkout@v4
@@ -95,10 +100,18 @@ jobs:
         if: matrix.config.cudaEnabled
         id: cuda-toolkit
         with:
-          cuda: '12.6.2'
-          sub-packages: '["nvcc", "nvtx", "cudart", "curand", "curand_dev", "nvrtc_dev"]'
+          cuda: '${{ env.CUDA_MAJOR_VERSION }}.${{ env.CUDA_MINOR_VERSION }}.${{ env.CUDA_PATCH_VERSION }}'
+          sub-packages: '["nvcc", "nvtx", "cudart", "curand", "curand_dev", "nvrtc_dev", "cublas", "cublas_dev", "cusparse", "cusparse_dev"]'
           method: 'network'
 
+      - name: Install cuDNN
+        if: ${{ matrix.config.torchEnabled && matrix.config.cudaEnabled }}
+        shell: pwsh
+        run: |
+          .github/workflows/install-cudnn.ps1 `
+              -CudaMajorVersion "${{ env.CUDA_MAJOR_VERSION }}" `
+              -CudaMinorVersion "${{ env.CUDA_MINOR_VERSION }}"
+
       - name: Setup vcpkg
         shell: pwsh
         run: |
@@ -129,6 +142,7 @@ jobs:
             -DCMAKE_BUILD_TYPE=Release `
             -DTESTS_ENABLED=${{ matrix.config.testsEnabled }} `
             -DGUI_ENABLED=ON `
+            -DTORCH_ENABLED=${{ matrix.config.torchEnabled }} `
             -DCUDA_ENABLED=${{ matrix.config.cudaEnabled }} `
             -DCMAKE_CUDA_ARCHITECTURES=all-major `
             -DCUDAToolkit_ROOT="${{ steps.cuda-toolkit.outputs.CUDA_PATH }}" `
@@ -159,7 +173,8 @@ jobs:
               --x-feature=gui `
               --x-feature=cgal `
               $(if ($${{ matrix.config.testsEnabled }}) { echo "--x-feature=tests" }) `
-              $(if ($${{ matrix.config.cudaEnabled }}) { echo "--x-feature=cuda" })
+              $(if ($${{ matrix.config.cudaEnabled }}) { echo "--x-feature=cuda" }) `
+              $(if ($${{ matrix.config.torchEnabled }}) { echo "--x-feature=torch" })
           ../vcpkg/vcpkg.exe export --raw --output-dir vcpkg_export --output colmap
           cp vcpkg_export/colmap/installed/x64-windows/bin/*.dll install/bin
           cp vcpkg_export/colmap/installed/x64-windows-release/bin/*.dll install/bin

diff --git a/.github/workflows/install-cudnn.ps1 b/.github/workflows/install-cudnn.ps1
@@ -0,0 +1,41 @@
+[CmdletBinding()]
+param (
+    [Parameter(Mandatory = $true)]
+    [string] $CudaMajorVersion,
+
+    [Parameter(Mandatory = $true)]
+    [string] $CudaMinorVersion
+)
+
+$version = "9.6.0.74"
+$folder="cudnn-windows-x86_64-${version}_cuda${CudaMajorVersion}-archive"
+$url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-${version}_cuda${CudaMajorVersion}-archive.zip"
+$expectedSha256 = "65CA0F2D77A46DE1DEF35E289780B8D8729EF2FA39CF8DD0C8448E381DD2978C"
+
+$ErrorActionPreference = "Stop"
+
+try {
+    Write-Host "Download cuDNN"
+    $zipFilePath = Join-Path "$env:TEMP" "$folder.zip"
+    Invoke-WebRequest -Uri $url -UseBasicParsing -OutFile "$zipFilePath" -MaximumRetryCount 3
+
+    $hash = Get-FileHash $zipFilePath -Algorithm "sha256"
+    if ($hash.Hash -ne $expectedSha256) {
+        throw "File $Path hash $hash.Hash did not match expected hash $expectedHash"
+    }
+
+    Write-Host "Unzip cuDNN"
+    Expand-Archive -Path "$zipFilePath" -DestinationPath "$env:TEMP" -Force
+
+    Write-Host "Move cuDNN"
+    $sourceFolder = "$env:TEMP/$folder"
+    $targetFolder = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${CudaMajorVersion}.${CudaMinorVersion}"
+    Move-Item -Path "$sourceFolder/include/*" -Destination "$targetFolder/include" -Force
+    Move-Item -Path "$sourceFolder/bin/*" -Destination "$targetFolder/bin" -Force
+    Move-Item -Path "$sourceFolder/lib/x64/*" -Destination "$targetFolder/lib/x64" -Force
+}
+catch {
+    Write-Host "Installation failed with an error"
+    $_.Exception | Format-List
+    exit -1
+}
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -51,6 +51,7 @@ option(DOWNLOAD_ENABLED "Whether to enable (automatic) download of resources (re
 option(UNINSTALL_ENABLED "Whether to create a target to 'uninstall' colmap" ON)
 option(FETCH_POSELIB "Whether to consume PoseLib using FetchContent or find_package" ON)
 option(FETCH_FAISS "Whether to consume faiss using FetchContent or find_package" ON)
+option(TORCH_ENABLED "Whether to enable torch support" ON)
 option(ALL_SOURCE_TARGET "Whether to create a target for all source files (for Visual Studio / XCode development)" OFF)
 
 # Disables default features, as we specify each required feature manually below.
@@ -69,6 +70,13 @@ endif()
 if(CGAL_ENABLED)
     list(APPEND VCPKG_MANIFEST_FEATURES "cgal")
 endif()
+if(TORCH_ENABLED)
+    if(CUDA_ENABLED)
+        list(APPEND VCPKG_MANIFEST_FEATURES "torch-cuda")
+    else()
+        list(APPEND VCPKG_MANIFEST_FEATURES "torch")
+    endif()
+endif()
 
 if(LSD_ENABLED)
     message(STATUS "Enabling LSD support")
@@ -87,7 +95,12 @@ set(COLMAP_VERSION "3.12.0.dev0")
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CUDA_STANDARD 14)
+if(TORCH_ENABLED)
+    # Torch requires C++17, which excludes Ubuntu 20.04's default compiler.
+    set(CMAKE_CUDA_STANDARD 17)
+else()
+    set(CMAKE_CUDA_STANDARD 14)
+endif()
 set(CMAKE_CUDA_STANDARD_REQUIRED ON)
 
 set_property(GLOBAL PROPERTY GLOBAL_DEPENDS_NO_CYCLES ON)
@@ -114,6 +127,7 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 ################################################################################
 
 set(COLMAP_FIND_QUIETLY FALSE)
+set(FIND_TORCH TRUE)
 include(cmake/FindDependencies.cmake)
 
 ################################################################################
@@ -241,8 +255,8 @@ include(GenerateVersionDefinitions)
 include_directories(src)
 link_directories(${COLMAP_LINK_DIRS})
 
-add_subdirectory(src/colmap)
 add_subdirectory(src/thirdparty)
+add_subdirectory(src/colmap)
 
 ################################################################################
 # Generate source groups for Visual Studio, XCode, etc.
@@ -385,6 +399,12 @@ if(GPU_ENABLED)
          colmap_sift_gpu
     )
 endif()
+if(TORCH_ENABLED)
+    list(APPEND COLMAP_EXPORT_LIBS
+         colmap_aliked
+         colmap_lightglue
+    )
+endif()
 if(FETCH_POSELIB)
     list(APPEND COLMAP_EXPORT_LIBS PoseLib)
 endif()

diff --git a/cmake/FindDependencies.cmake b/cmake/FindDependencies.cmake
@@ -240,3 +240,34 @@ if(OPENGL_ENABLED OR CUDA_ENABLED)
     message(STATUS "Enabling GPU support (OpenGL: ${OPENGL_ENABLED}, CUDA: ${CUDA_ENABLED})")
     set(GPU_ENABLED ON)
 endif()
+
+if(TORCH_ENABLED)
+    if(NOT DEFINED MKL_ROOT)
+        cmake_policy(SET CMP0074 NEW)
+        if(IS_LINUX)
+            set(DEBIAN_DEFAULT_MKL_LIB_PATH "/usr/lib/x86_64-linux-gnu/libmkl_core.so")
+            if(EXISTS "${DEBIAN_DEFAULT_MKL_LIB_PATH}")
+                set(MKL_INCLUDE_DIR "/usr/include")
+                set(MKL_LIBRARIES "${DEBIAN_DEFAULT_MKL_LIB_PATH}")
+            else()
+                set(MKL_ROOT "/opt/intel/mkl")
+            endif()
+        elseif(IS_MSVC)
+            set(DEFAULT_MKL_LIB_PATH "C:/Program Files (x86)/IntelSWTools/compilers_and_libraries/windows/mkl/lib/intel64/mkl_core_dll.lib")
+            if(EXISTS "${DEFAULT_MKL_LIB_PATH}")
+                set(MKL_INCLUDE_DIR "C:/Program Files (x86)/IntelSWTools/compilers_and_libraries/windows/mkl/include")
+                set(MKL_LIBRARIES "${DEFAULT_MKL_LIB_PATH}")
+            endif()
+        endif()
+    endif()
+    find_package(Torch QUIET)
+    if(TORCH_FOUND)
+        add_definitions("-DCOLMAP_TORCH_ENABLED")
+        set(TORCH_ENABLED ON)
+        message(STATUS "Enabling Torch support")
+    else()
+        message(STATUS "Disabling Torch support (not found)")
+    endif()
+else()
+    message(STATUS "Disabling Torch support")
+endif()
diff --git a/cmake/colmap-config.cmake.in b/cmake/colmap-config.cmake.in
@@ -82,6 +82,8 @@ set(CGAL_ENABLED @CGAL_ENABLED@)
 
 set(LSD_ENABLED @LSD_ENABLED@)
 
+set(TORCH_ENABLED @TORCH_ENABLED@)
+
 set(FETCH_POSELIB @FETCH_POSELIB@)
 
 set(FETCH_FAISS @FETCH_FAISS@)