ARG BASE_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04-cuda11.3.0-cudnn8-devel
FROM $BASE_IMAGE
ARG DEBIAN_FRONTEND=noninteractive
ENV TZ=Asia/Shanghai
ENV CONDA_DIR /opt/conda
ENV PATH="${CONDA_DIR}/bin:${PATH}"
ENV arch=x86_64
SHELL ["/bin/bash", "-c"]
COPY docker/rcfiles /tmp/resources
COPY docker/jupyter_plugins /tmp/resources/jupyter_plugins
RUN apt-get update && apt-get install -y --reinstall ca-certificates && \
    apt-get install -y apt-utils openssh-server locales wget git strace gdb sox libopenmpi-dev curl \
    iputils-ping net-tools iproute2 autoconf automake gperf libre2-dev libssl-dev \
    libtool libcurl4-openssl-dev libb64-dev libgoogle-perftools-dev patchelf \
    rapidjson-dev scons software-properties-common pkg-config unzip zlib1g-dev \
    libarchive-dev libxml2-dev libnuma-dev \
    libgeos-dev strace vim ffmpeg libsm6 tzdata language-pack-zh-hans \
    ttf-wqy-microhei ttf-wqy-zenhei xfonts-wqy libxext6 build-essential ninja-build && \
    wget https://packagecloud.io/github/git-lfs/packages/debian/bullseye/git-lfs_3.2.0_amd64.deb/download -O ./git-lfs_3.2.0_amd64.deb && \
    dpkg -i ./git-lfs_3.2.0_amd64.deb && \
    rm -f ./git-lfs_3.2.0_amd64.deb && \
    locale-gen zh_CN && \
    locale-gen zh_CN.utf8 && \
    update-locale LANG=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 && \
    ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
    dpkg-reconfigure --frontend noninteractive tzdata && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

ENV LANG=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8
RUN wget -O /tmp/boost.tar.gz         https://boostorg.jfrog.io/artifactory/main/release/1.80.0/source/boost_1_80_0.tar.gz &&     (cd /tmp && tar xzf boost.tar.gz) &&     mv /tmp/boost_1_80_0/boost /usr/include/boost

#install and config python
ARG PYTHON_VERSION=3.10.13
# Miniconda3-py37_23.1.0-1-Linux-x86_64.sh is last python3.7 version
RUN  wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-py310_23.9.0-0-Linux-x86_64.sh -O ./miniconda.sh && \
    /bin/bash  miniconda.sh -b -p /opt/conda && \
    rm  -f miniconda.sh && \
    ln  -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
    source /root/.bashrc

ARG USE_GPU=True

# install pytorch
ARG TORCH_VERSION=1.12.0
ARG CUDATOOLKIT_VERSION=cu117
RUN if [ "$USE_GPU" = "True" ] ; then \
        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/$CUDATOOLKIT_VERSION; \
    else \
        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu; \
    fi

# install tensorflow
ARG TENSORFLOW_VERSION=1.15.5
RUN if [ "$USE_GPU" = "True" ] ; then \
        if [ "$TENSORFLOW_VERSION" = "1.15.5" ] ; then \
            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
        else \
            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
         fi \
    else \
        # only python 3.7 has tensorflow 1.15.5
        if [ "$PYTHON_VERSION" = "3.7.13" ] ; then \
            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
        elif [ "$TENSORFLOW_VERSION" = "1.15.5" ] ; then \
            pip install --no-cache-dir numpy==1.18.5 https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/tensorflow-1.15.5-cp38-cp38-linux_x86_64.whl; \
        else \
            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
        fi \
    fi


# default shell bash
ENV SHELL=/bin/bash
# install special package
RUN if [ "$USE_GPU" = "True" ] ; then \
        pip install  dgl -f https://data.dgl.ai/wheels/$CUDATOOLKIT_VERSION/repo.html; \
    else \
        pip install --no-cache-dir dgl dglgo -f https://data.dgl.ai/wheels/repo.html; \
    fi

# copy install scripts
COPY docker/scripts/install_unifold.sh docker/scripts/install_colmap.sh docker/scripts/install_pytorch3d_nvdiffrast.sh docker/scripts/install_tiny_cuda_nn.sh docker/scripts/install_apex.sh /tmp/

# 3d supports
RUN if [ "$USE_GPU" = "True" ] ; then \
        bash /tmp/install_colmap.sh; \
    else \
     echo 'cpu unsupport colmap'; \
    fi
# install pytorch3d
RUN if [ "$USE_GPU" = "True" ] ; then \
        bash /tmp/install_pytorch3d_nvdiffrast.sh; \
    else \
     echo 'cpu unsupport pytorch3d nvdiffrast'; \
    fi

# for uniford
RUN if [ "$USE_GPU" = "True" ] ; then \
        bash /tmp/install_unifold.sh; \
    else \
     echo 'cpu unsupport uniford'; \
    fi

RUN if [ "$USE_GPU" = "True" ] ; then \
       export TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.9;9.0;8.6+PTX" && pip install --no-cache-dir git+https://github.com/gxd1994/Pointnet2.PyTorch.git@master#subdirectory=pointnet2; \
    else \
     echo 'cpu unsupport Pointnet2'; \
    fi

# install apex after deepspeed
RUN if [ "$USE_GPU" = "True" ] ; then \
        bash /tmp/install_apex.sh; \
    else \
     echo 'cpu unsupport apex'; \
    fi

RUN if [ "$USE_GPU" = "True" ] ; then \
        pip install --no-cache-dir mmcv-full==1.7.0+torch2.1.1cu121 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
    else \
        pip install --no-cache-dir mmcv_full==1.7.0+torch2.1cpu -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
    fi
RUN conda install imageio-ffmpeg -c conda-forge -y
ENTRYPOINT []