add docker for mast3r demo (mostly taken from dust3r)

README.md

@@ -133,7 +133,32 @@ see https://github.com/naver/dust3r?tab=readme-ov-file#interactive-demo for deta
 
 ### Interactive demo with docker
 
-TODO
+To run MASt3R using Docker, including with NVIDIA CUDA support, follow these instructions:
+
+1. **Install Docker**: If not already installed, download and install `docker` and `docker compose` from the [Docker website](https://www.docker.com/get-started).
+
+2. **Install NVIDIA Docker Toolkit**: For GPU support, install the NVIDIA Docker toolkit from the [Nvidia website](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
+
+3. **Build the Docker image and run it**: `cd` into the `./docker` directory and run the following commands: 
+
+```bash
+cd docker
+bash run.sh --with-cuda --model_name="MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
+```
+
+Or if you want to run the demo without CUDA support, run the following command:
+
+```bash 
+cd docker
+bash run.sh --model_name="MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
+```
+
+By default, `demo.py` is lanched with the option `--local_network`.  
+Visit `http://localhost:7860/` to access the web UI (or replace `localhost` with the machine's name to access it from the network).  
+
+`run.sh` will launch docker-compose using either the [docker-compose-cuda.yml](docker/docker-compose-cuda.yml) or [docker-compose-cpu.ym](docker/docker-compose-cpu.yml) config file, then it starts the demo using [entrypoint.sh](docker/files/entrypoint.sh).
+
+___
 
 ![demo](assets/demo.jpg)
 

docker/docker-compose-cpu.yml

@@ -0,0 +1,16 @@
+version: '3.8'
+services:
+  mast3r-demo:
+    build:
+      context: ./files
+      dockerfile: cpu.Dockerfile 
+    ports:
+      - "7860:7860"
+    volumes:
+      - ./files/checkpoints:/mast3r/checkpoints
+    environment:
+      - DEVICE=cpu
+      - MODEL=${MODEL:-MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth}
+    cap_add:
+      - IPC_LOCK
+      - SYS_RESOURCE

docker/docker-compose-cuda.yml

@@ -0,0 +1,23 @@
+version: '3.8'
+services:
+  mast3r-demo:
+    build:
+      context: ./files
+      dockerfile: cuda.Dockerfile 
+    ports:
+      - "7860:7860"
+    environment:
+      - DEVICE=cuda
+      - MODEL=${MODEL:-MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth}
+    volumes:
+      - ./files/checkpoints:/mast3r/checkpoints
+    cap_add:
+      - IPC_LOCK
+      - SYS_RESOURCE
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

docker/files/cpu.Dockerfile

@@ -0,0 +1,39 @@
+FROM python:3.11-slim
+
+LABEL description="Docker container for MASt3R with dependencies installed. CPU VERSION"
+
+ENV DEVICE="cpu"
+ENV MODEL="MASt3R_ViTLarge_BaseDecoder_512_dpt.pth"
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    git \
+    libgl1-mesa-glx \
+    libegl1-mesa \
+    libxrandr2 \
+    libxrandr2 \
+    libxss1 \
+    libxcursor1 \
+    libxcomposite1 \
+    libasound2 \
+    libxi6 \
+    libxtst6 \
+    libglib2.0-0 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone --recursive https://github.com/naver/mast3r /mast3r
+WORKDIR /mast3r/dust3r
+
+RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+RUN pip install -r requirements.txt
+RUN pip install -r requirements_optional.txt
+RUN pip install opencv-python==4.8.0.74
+
+WORKDIR /mast3r
+RUN pip install -r requirements.txt
+
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+ENTRYPOINT ["/entrypoint.sh"]

docker/files/cuda.Dockerfile

@@ -0,0 +1,29 @@
+FROM nvcr.io/nvidia/pytorch:24.01-py3
+
+LABEL description="Docker container for MASt3R with dependencies installed. CUDA VERSION"
+ENV DEVICE="cuda"
+ENV MODEL="MASt3R_ViTLarge_BaseDecoder_512_dpt.pth"
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    git=1:2.34.1-1ubuntu1.10 \
+    libglib2.0-0=2.72.4-0ubuntu2.2 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone --recursive https://github.com/naver/mast3r /mast3r
+WORKDIR /mast3r/dust3r
+RUN pip install -r requirements.txt
+RUN pip install -r requirements_optional.txt
+RUN pip install opencv-python==4.8.0.74
+
+WORKDIR /mast3r/dust3r/croco/models/curope/
+RUN python setup.py build_ext --inplace
+
+WORKDIR /mast3r
+RUN pip install -r requirements.txt
+
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+ENTRYPOINT ["/entrypoint.sh"]

docker/files/entrypoint.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+set -eux
+
+DEVICE=${DEVICE:-cuda}
+MODEL=${MODEL:-MASt3R_ViTLarge_BaseDecoder_512_dpt.pth}
+
+exec python3 demo.py --weights "checkpoints/$MODEL" --device "$DEVICE" --local_network "$@"

docker/run.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+
+set -eux
+
+# Default model name
+model_name="MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
+
+check_docker() {
+    if ! command -v docker &>/dev/null; then
+        echo "Docker could not be found. Please install Docker and try again."
+        exit 1
+    fi
+}
+
+download_model_checkpoint() { 
+    if [ -f "./files/checkpoints/${model_name}" ]; then
+        echo "Model checkpoint ${model_name} already exists. Skipping download."
+        return
+    fi
+    echo "Downloading model checkpoint ${model_name}..."
+    wget "https://download.europe.naverlabs.com/ComputerVision/MASt3R/${model_name}" -P ./files/checkpoints
+}
+
+set_dcomp() {
+    if command -v docker-compose &>/dev/null; then
+        dcomp="docker-compose"
+    elif command -v docker &>/dev/null && docker compose version &>/dev/null; then
+        dcomp="docker compose"
+    else
+        echo "Docker Compose could not be found. Please install Docker Compose and try again."
+        exit 1
+    fi
+}
+
+run_docker() {
+    export MODEL=${model_name}
+    if [ "$with_cuda" -eq 1 ]; then
+        $dcomp -f docker-compose-cuda.yml up --build
+    else
+        $dcomp -f docker-compose-cpu.yml up --build
+    fi
+}
+
+with_cuda=0
+for arg in "$@"; do
+    case $arg in
+        --with-cuda)
+            with_cuda=1
+            ;;
+        --model_name=*)
+            model_name="${arg#*=}.pth"
+            ;;
+        *)
+            echo "Unknown parameter passed: $arg"
+            exit 1
+            ;;
+    esac
+done
+
+
+main() {
+    check_docker
+    download_model_checkpoint
+    set_dcomp
+    run_docker
+}
+
+main