FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 AS builder

ARG COLMAP_VERSION=3.8
ARG CUDA_ARCHITECTURES=86
ENV DEBIAN_FRONTEND=noninteractive

# Set up initial PyTorch environment and build environment
RUN apt-get update && apt-get install -y \
    build-essential \
    cmake \
    git \
    libatlas-base-dev \
    libboost-filesystem-dev \
    libboost-graph-dev \
    libboost-program-options-dev \
    libboost-system-dev \
    libboost-test-dev \
    libcgal-dev \
    libeigen3-dev \
    libflann-dev \
    libfreeimage-dev \
    libgflags-dev \
    libglew-dev \
    libgoogle-glog-dev \
    libhdf5-dev \
    libmetis-dev \
    libqt5opengl5-dev \
    libsqlite3-dev \
    libsuitesparse-dev \
    ninja-build \
    python3 \
    python3-pip \
    qtbase5-dev

# Tencent mirror as an alternative to main repo in case CDN fails
# NumPy 2.0 in order to support both 1.x and 2.x versions of NumPy, since
# legacy code is present
RUN pip3 install --upgrade pip
RUN pip3 install --no-cache-dir \
      --extra-index-url https://download.pytorch.org/whl/cu118 \
      #--extra-index-url https://mirrors.tencent.com/pytorch-wheels/whl/cu118/ \
      torch==2.1.2+cu118 \
      torchvision==0.16.2+cu118 \
      'numpy<2.0.0' 

WORKDIR /opt

# Install Ceres Solver. PixSFM requires version 2.1 or higher. Testing is
# disabled since `docker build` doesn't have access to the GPU
RUN git clone https://github.com/ceres-solver/ceres-solver && \
    cd ceres-solver && \
    git checkout 2.2.0
RUN mkdir ceres-bin && \
    cd ceres-bin && \
    cmake ../ceres-solver \
        -DUSE_CUDA=ON \
        -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} \
        -DBUILD_TESTING=OFF \
        -DBUILD_EXAMPLES=OFF \
        -DCMAKE_BUILD_TYPE=Release && \
    make -j8 && \
    make test && \
    make install

# Install COLMAP 3.8 specifically, since this is a dependency of PixSFM. Set
# C++17 standard to match Ceres libraries
RUN git clone https://github.com/colmap/colmap.git
RUN cd colmap && \
    git reset --hard ${COLMAP_VERSION} && \
    mkdir build && \
    cd build && \
    cmake .. \
        -GNinja \
        -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} \
        -DCMAKE_CXX_STANDARD=17 \
        -DCMAKE_CXX_STANDARD_REQUIRED=ON && \
    ninja && \
    ninja install -j8 && \
    cd .. && rm -rf colmap

WORKDIR /workspace

# Install both PixSFM and Hloc in-place, since their run scripts should ideally
# be stored in the workspace
# PixSFM is installed first since its version requirements take precedence over
# Hloc's
COPY pixel-perfect-sfm /workspace/pixel-perfect-sfm
RUN cd pixel-perfect-sfm && \
    pip3 install -r requirements.txt

# Git reset to the last commit compatible with pycolmap==0.4.0, which is a
# dependency of PixSFM
COPY Hierarchical-Localization /workspace/Hierarchical-Localization
RUN cd Hierarchical-Localization/ && \
    pip3 install -e .

# Install PixSFM last since it has Hloc as one of its dependencies
RUN cd pixel-perfect-sfm && \
    pip3 install -e .

# Reinstall NumPy<2.0.0, since the used version of PyTorch was compiled with an
# older version of NumPy
RUN pip3 install --no-cache-dir 'numpy<2.0.0'

# Set up runtime
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04

ENV CUDA_ARCHITECTURES=86
ENV NVIDIA_CUDA_VERSION=11.8.0
ENV GRADIO_SERVER_NAME=0.0.0.0

# Install runtime versions of packages in builder stage
RUN apt-get update && apt-get install -y --no-install-recommends \
    libatlas3-base \
    libboost-filesystem1.74.0 \
    libboost-graph1.74.0 \
    libboost-program-options1.74.0 \
    libboost-system1.74.0 \
    libboost-test1.74.0 \
    libflann1.9 \
    libfreeimage3 \
    libgl1 \
    libglew2.2 \
    libgoogle-glog0v5 \
    libhdf5-dev \
    libmetis5 \
    libpython3.10 \
    libqt5opengl5 \
    libsuitesparse-dev \
    python-is-python3 \
    python3 \
    wget
RUN rm -rf /var/cache/apt/archives /var/lib/apt/lists/*

# Copy COLMAP and Ceres builds and installed Python libraries
# `include` directory only has ceres and colmap
# NOTE: Python packages are found here, assuming Python 3.10
# /usr/local/lib/python3.10/dist-packages/
COPY --from=builder /usr/local/include/ /usr/local/include/
COPY --from=builder /usr/local/bin/colmap /usr/local/bin/colmap
COPY --from=builder /usr/local/lib/ /usr/local/lib/
COPY --from=builder /usr/local/share/ /usr/local/share/
# Copy the actual workspace
COPY --from=builder /workspace/ /workspace/

# A PixSFM dependency that seems to be missing from the copy process. This is
# technically deprecated as an API, and in case of failure should be pinned to
# setuptools<81
RUN python3 -m pip install --no-cache-dir setuptools

CMD ["/bin/bash", "-l"]