ziyc · hakuturu583 · Dec 3, 2025 · Dec 3, 2025
diff --git a/.gitignore b/.gitignore
@@ -28,6 +28,7 @@ simulated_results/
 sky_masks/
 fine_dynamic_masks_example/
 *.sh
+!scripts/nuscenes.sh
 pandaset-devkit/
 
 # C extensions
@@ -185,4 +186,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-local_scripts
+local_scripts
diff --git a/README.md b/README.md
@@ -120,6 +120,13 @@ pip install -e .
 cd ../..
 ```
 
+## 🐳 Docker Build
+- Build base + NuScenes images with Buildx bake:  
+  `docker buildx bake` (uses `docker-bake.hcl`, tags `drivestudio:base` and `drivestudio:nuscenes`)
+- Build a single target: `docker buildx bake nuscenes`
+- Override tag prefix: `TAG_PREFIX=myrepo/drivestudio docker buildx bake`
+- See `docker/README.md` for target details and a sample `docker run` command.
+
 ## 📊 Prepare Data
 We support most popular public driving datasets. Detailed instructions for downloading and processing each dataset are available in the following documents:
 

diff --git a/docker-bake.hcl b/docker-bake.hcl
@@ -0,0 +1,28 @@
+variable "TAG_PREFIX" {
+  default = "drivestudio"
+}
+
+group "default" {
+  targets = ["base", "nuscenes"]
+}
+
+target "base" {
+  context    = "."
+  dockerfile = "docker/Dockerfile.base"
+  tags       = ["${TAG_PREFIX}:base"]
+  output     = ["type=docker"]
+}
+
+target "nuscenes" {
+  context    = "."
+  dockerfile = "docker/Dockerfile.nuscenes"
+  tags       = ["${TAG_PREFIX}:nuscenes"]
+  output     = ["type=docker"]
+  args = {
+    BASE_IMAGE = "base"
+  }
+  contexts = {
+    base = "target:base"
+  }
+  depends_on = ["base"]
+}
diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base
@@ -0,0 +1,39 @@
+FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
+
+ARG TORCH_CUDA_ARCH_LIST="8.6"
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PYTHONUNBUFFERED=1 \
+    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
+
+ENV PIP_NO_BUILD_ISOLATION=1 \
+    PIP_USE_PEP517=0
+
+# Common system dependencies for building/running PyTorch, Open3D, OpenCV, gsplat, nvdiffrast, and pytorch3d.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 python3-dev python3-pip python3-venv python3-setuptools \
+    git curl wget ca-certificates \
+    build-essential cmake ninja-build pkg-config \
+    libssl-dev libffi-dev \
+    libglib2.0-0 libgl1 libgl1-mesa-dev libxext-dev libx11-dev libxi-dev \
+    libjpeg-dev libpng-dev ffmpeg \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN python3 -m pip install "pip<24"
+
+WORKDIR /workspace/drivestudio
+COPY . .
+
+# Core project dependencies and local modules shared across datasets.
+# Disable build isolation so legacy packages like chumpy can find pip during build.
+RUN python3 -m pip install --no-build-isolation --no-use-pep517 -r requirements.txt
+RUN python3 -m pip install --no-build-isolation --no-use-pep517 git+https://github.com/nerfstudio-project/gsplat.git@v1.3.0
+RUN python3 -m pip install --no-build-isolation --no-use-pep517 git+https://github.com/facebookresearch/pytorch3d.git
+RUN python3 -m pip install --no-build-isolation --no-use-pep517 git+https://github.com/NVlabs/nvdiffrast
+RUN python3 -m pip install --no-build-isolation --no-use-pep517 -e third_party/smplx
+
+ENV PYTHONPATH=/workspace/drivestudio
+
+CMD ["bash"]
diff --git a/docker/Dockerfile.nuscenes b/docker/Dockerfile.nuscenes
@@ -0,0 +1,26 @@
+# syntax=docker/dockerfile:1.4
+
+ARG BASE_IMAGE=base
+FROM ${BASE_IMAGE}
+
+# NuScenes-specific toolkit
+RUN python3 -m pip install nuscenes-devkit
+
+# SegFormer env for mask preprocessing (isolated from the main Torch 2 stack)
+ENV SEGFORMER_ROOT=/opt/segformer \
+    SEGFORMER_CONDA=/opt/miniforge/envs/segformer \
+    MINIFORGE_URL=https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
+
+RUN wget -q ${MINIFORGE_URL} -O /tmp/miniforge.sh \
+    && bash /tmp/miniforge.sh -b -p /opt/miniforge \
+    && rm /tmp/miniforge.sh \
+    && /opt/miniforge/bin/conda create -y -n segformer python=3.8 \
+    && /opt/miniforge/bin/conda run -n segformer pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html \
+    && /opt/miniforge/bin/conda run -n segformer pip install timm==0.3.2 pylint debugpy opencv-python-headless attrs ipython tqdm imageio scikit-image omegaconf \
+    && /opt/miniforge/bin/conda run -n segformer pip install mmcv-full==1.2.7 --no-cache-dir -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.8/index.html \
+    && git clone https://github.com/NVlabs/SegFormer ${SEGFORMER_ROOT} \
+    && /opt/miniforge/bin/conda run -n segformer pip install ${SEGFORMER_ROOT} \
+    && /opt/miniforge/bin/conda clean -afy
+
+# Keep conda available without overriding the system Python used by the project
+ENV PATH=${PATH}:/opt/miniforge/bin
diff --git a/docker/README.md b/docker/README.md
@@ -0,0 +1,16 @@
+# DriveStudio Docker
+
+## Build with bake
+- Build all: `docker buildx bake` (uses `docker-bake.hcl`)
+- Build only NuScenes: `docker buildx bake nuscenes`
+- Set a different tag prefix: `TAG_PREFIX=myrepo/drivestudio docker buildx bake`
+
+Targets:
+- `base`: Common deps (PyTorch stack, gsplat, pytorch3d, nvdiffrast, smplx, etc.)
+- `nuscenes`: Adds `nuscenes-devkit` and a separate SegFormer preprocessing env on top of `base`
+
+## Run NuScenes image
+- GPU + data mount: `docker run --gpus all -it --rm -v $(pwd)/data:/workspace/drivestudio/data drivestudio:nuscenes`
+- `PYTHONPATH` is preset; work under `/workspace/drivestudio`.
+- Prepare NuScenes data per `docs/NuScenes.md` (mount `data/nuscenes/raw`, etc.).
+- SegFormer for mask extraction is preinstalled in an isolated conda env (`segformer`) with the code at `/opt/segformer`. Use `conda run -n segformer python datasets/tools/extract_masks.py ...` (or `conda activate segformer`) when running the mask scripts.
diff --git a/docs/NuScenes.md b/docs/NuScenes.md
@@ -24,12 +24,29 @@ We'll use the **v1.0-mini split** in our examples. The process is similar for ot
 pip install nuscenes-devkit
 ```
 
+## 2.5 Build the NuScenes Docker Image
+Instead of installing locally, you can build a NuScenes-ready image and run DriveStudio inside it.
+
+```shell
+# From repo root
+docker buildx bake nuscenes
+# tags: drivestudio:base and drivestudio:nuscenes (set TAG_PREFIX to override)
+```
+
+Run with GPUs and mount your data:
+```shell
+docker run --gpus all -it --rm \
+  -v ${PATH_TO_NUSCENES}:/workspace/drivestudio/data \
+  drivestudio:nuscenes
+```
+
+`PYTHONPATH` is preset in the image; work under `/workspace/drivestudio`. Prepare/mount `data/nuscenes/raw` as described below. See `docker/README.md` for target details and optional flags.
+
 ## 3. Process Raw Data
 
 To process the 10 scenes in NuScenes **v1.0-mini split**, you can run:
 
 ```shell
-# export PYTHONPATH=\path\to\project
 python datasets/preprocess.py \
     --data_root data/nuscenes/raw \
     --target_dir data/nuscenes/processed \
@@ -64,6 +81,8 @@ Follow these steps:
 
 :warning: SegFormer relies on `mmcv-full=1.2.7`, which relies on `pytorch=1.8` (pytorch<1.9). Hence, a separate conda env is required.
 
+If you're using the `drivestudio:nuscenes` Docker image, SegFormer is already installed in a dedicated conda env named `segformer` with the code at `/opt/segformer`; you can skip the install steps below and run the commands with `conda run -n segformer`.
+
 ```shell
 #-- Set conda env
 conda create -n segformer python=3.8

diff --git a/scripts/nuscenes.sh b/scripts/nuscenes.sh
@@ -0,0 +1,193 @@
+#!/usr/bin/env bash
+# Run NuScenes preprocessing inside the drivestudio:nuscenes Docker image.
+# Usage: scripts/nuscenes.sh /abs/path/to/nuscenes/raw [--split v1.0-mini] [--num-scenes 10] [--start-idx 0] [--interpolate 4] [--workers 32] [--image drivestudio:nuscenes] [--cpu] [--no-checkpoint-download] [--no-humanpose] [--no-dynamic-mask]
+set -euo pipefail
+
+usage() {
+  cat <<'USAGE'
+Usage: scripts/nuscenes.sh /abs/path/to/nuscenes/raw [options]
+
+Required:
+  /abs/path/to/nuscenes/raw   Directory containing NuScenes raw data (v1.0-mini / v1.0-trainval / v1.0-test)
+
+Options:
+  --split <name>              NuScenes split (default: v1.0-mini)
+  --num-scenes <n>            Number of scenes to process (default: 10)
+  --start-idx <n>             Start index for scenes (default: 0)
+  --interpolate <n>           Interpolation factor N (0-4) for 2Hz->(N+1)*2Hz (default: 4 -> 10Hz)
+  --workers <n>               Workers for preprocessing (default: 32)
+  --image <tag>               Docker image to use (default: drivestudio:nuscenes)
+  --cpu                       Run container without --gpus all
+  --no-checkpoint-download    Skip SegFormer checkpoint download (expects file at /opt/segformer/pretrained/segformer.b5.1024x1024.city.160k.pth)
+  --no-humanpose              Skip downloading preprocessed human pose zip
+  --no-dynamic-mask           Skip fine dynamic mask extraction with SegFormer
+  -h, --help                  Show this help
+USAGE
+}
+
+if [[ $# -lt 1 ]]; then
+  usage
+  exit 1
+fi
+
+RAW_PATH=""
+SPLIT="v1.0-mini"
+NUM_SCENES=10
+START_IDX=0
+INTERPOLATE=4
+WORKERS=32
+IMAGE="drivestudio:nuscenes"
+GPU_FLAG="--gpus all"
+DOWNLOAD_CHECKPOINT=1
+DOWNLOAD_HUMANPOSE=1
+PROCESS_DYNAMIC_MASK=1
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --split) SPLIT="$2"; shift 2;;
+    --num-scenes) NUM_SCENES="$2"; shift 2;;
+    --start-idx) START_IDX="$2"; shift 2;;
+    --interpolate) INTERPOLATE="$2"; shift 2;;
+    --workers) WORKERS="$2"; shift 2;;
+    --image) IMAGE="$2"; shift 2;;
+    --cpu) GPU_FLAG=""; shift;;
+    --no-checkpoint-download) DOWNLOAD_CHECKPOINT=0; shift;;
+    --no-humanpose) DOWNLOAD_HUMANPOSE=0; shift;;
+    --no-dynamic-mask) PROCESS_DYNAMIC_MASK=0; shift;;
+    -h|--help) usage; exit 0;;
+    *)
+      if [[ -z "${RAW_PATH}" ]]; then
+        RAW_PATH="$1"; shift
+      else
+        echo "Unexpected argument: $1" >&2
+        usage; exit 1
+      fi
+      ;;
+  esac
+done
+
+if [[ -z "${RAW_PATH}" ]]; then
+  echo "Missing NuScenes raw path." >&2
+  usage
+  exit 1
+fi
+
+if ! command -v docker >/dev/null 2>&1; then
+  echo "docker is required but not found in PATH." >&2
+  exit 1
+fi
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+RAW_REALPATH="$(python3 -c 'import os,sys;print(os.path.abspath(sys.argv[1]))' "${RAW_PATH}")"
+
+if [[ ! -d "${RAW_REALPATH}" ]]; then
+  echo "NuScenes raw path does not exist: ${RAW_REALPATH}" >&2
+  exit 1
+fi
+
+mkdir -p "${REPO_ROOT}/data/nuscenes"
+
+if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
+  echo "Building Docker image ${IMAGE} ..."
+  (cd "${REPO_ROOT}" && docker buildx bake nuscenes)
+fi
+
+echo "Running NuScenes preprocessing inside container..."
+
+docker run ${GPU_FLAG} --rm -it \
+  -v "${REPO_ROOT}:/workspace/drivestudio" \
+  -v "${RAW_REALPATH}:/workspace/drivestudio/data/nuscenes/raw" \
+  -e NUSC_SPLIT="${SPLIT}" \
+  -e NUSC_NUM_SCENES="${NUM_SCENES}" \
+  -e NUSC_START_IDX="${START_IDX}" \
+  -e NUSC_INTERPOLATE="${INTERPOLATE}" \
+  -e NUSC_WORKERS="${WORKERS}" \
+  -e NUSC_DOWNLOAD_CHECKPOINT="${DOWNLOAD_CHECKPOINT}" \
+  -e NUSC_DOWNLOAD_HUMANPOSE="${DOWNLOAD_HUMANPOSE}" \
+  -e NUSC_PROCESS_DYNAMIC_MASK="${PROCESS_DYNAMIC_MASK}" \
+  "${IMAGE}" bash -c '
+    set -euo pipefail
+    RAW=/workspace/drivestudio/data/nuscenes/raw
+    TARGET=/workspace/drivestudio/data/nuscenes/processed
+    SPLIT="${NUSC_SPLIT}"
+    START_IDX="${NUSC_START_IDX}"
+    NUM_SCENES="${NUSC_NUM_SCENES}"
+    INTERPOLATE="${NUSC_INTERPOLATE}"
+    WORKERS="${NUSC_WORKERS}"
+    DOWNLOAD_CHECKPOINT="${NUSC_DOWNLOAD_CHECKPOINT}"
+    DOWNLOAD_HUMANPOSE="${NUSC_DOWNLOAD_HUMANPOSE}"
+    PROCESS_DYNAMIC_MASK="${NUSC_PROCESS_DYNAMIC_MASK}"
+
+    SEGFORMER_ROOT=/opt/segformer
+    CKPT_PATH=${SEGFORMER_ROOT}/pretrained/segformer.b5.1024x1024.city.160k.pth
+    CKPT_ID=1e7DECAH0TRtPZM6hTqRGoboq1XPqSmuj
+    HUMANPOSE_ID=1Z0gJVRtPnjvusQVaW7ghZnwfycZStCZx
+
+    echo "[1/3] Preprocessing raw NuScenes..."
+    python3 datasets/preprocess.py \
+      --data_root "${RAW}" \
+      --target_dir "${TARGET}" \
+      --dataset nuscenes \
+      --split "${SPLIT}" \
+      --start_idx "${START_IDX}" \
+      --num_scenes "${NUM_SCENES}" \
+      --interpolate_N "${INTERPOLATE}" \
+      --workers "${WORKERS}" \
+      --process_keys images lidar calib dynamic_masks objects
+
+    FREQ=$(( (INTERPOLATE + 1) * 2 ))
+    PROCESSED_BASE="${TARGET/processed/processed_${FREQ}Hz}"
+    SPLIT_LEAF="${SPLIT##*-}"
+    MASK_DATA_ROOT="${PROCESSED_BASE}/${SPLIT_LEAF}"
+
+    if [[ "${DOWNLOAD_CHECKPOINT}" == "1" && ! -f "${CKPT_PATH}" ]]; then
+      echo "[2/3] Downloading SegFormer checkpoint..."
+      mkdir -p "$(dirname "${CKPT_PATH}")"
+      conda run -n segformer gdown "${CKPT_ID}" -O "${CKPT_PATH}"
+    fi
+
+    if [[ ! -f "${CKPT_PATH}" ]]; then
+      echo "SegFormer checkpoint missing at ${CKPT_PATH}. Download or specify --no-checkpoint-download only if already present." >&2
+      exit 1
+    fi
+
+    echo "[2/3] Extracting sky/fine dynamic masks with SegFormer..."
+    MASK_ARGS=(
+      --data_root "${MASK_DATA_ROOT}"
+      --segformer_path "${SEGFORMER_ROOT}"
+      --checkpoint "${CKPT_PATH}"
+      --start_idx "${START_IDX}"
+      --num_scenes "${NUM_SCENES}"
+    )
+    if [[ "${PROCESS_DYNAMIC_MASK}" == "1" ]]; then
+      MASK_ARGS+=(--process_dynamic_mask)
+    fi
+    conda run -n segformer python datasets/tools/extract_masks.py "${MASK_ARGS[@]}"
+
+    HUMANPOSE_MARKER=""
+    if [[ -d "${PROCESSED_BASE}/${SPLIT_LEAF}" ]]; then
+      HUMANPOSE_MARKER="$(find "${PROCESSED_BASE}/${SPLIT_LEAF}" -maxdepth 2 -type d -name humanpose | head -n 1 || true)"
+    fi
+    if [[ "${DOWNLOAD_HUMANPOSE}" == "1" ]]; then
+      if [[ -d "${HUMANPOSE_MARKER}" ]]; then
+        echo "[3/3] Human pose data already present at ${HUMANPOSE_MARKER}, skipping download."
+      else
+        echo "[3/3] Downloading preprocessed human pose data..."
+        HP_ZIP=/workspace/drivestudio/data/nuscenes_preprocess_humanpose.zip
+        conda run -n segformer gdown "${HUMANPOSE_ID}" -O "${HP_ZIP}"
+        python3 - <<PY
+import os, zipfile
+zip_path = r'"'"'${HP_ZIP}'"'"'
+extract_dir = r'"'"'/workspace/drivestudio/data'"'"'
+with zipfile.ZipFile(zip_path, '"'"'r'"'"') as zf:
+    zf.extractall(extract_dir)
+os.remove(zip_path)
+print("Human pose data extracted to", extract_dir)
+PY
+      fi
+    else
+      echo "[3/3] Skipping human pose download."
+    fi
+
+    echo "NuScenes preprocessing finished. Processed data at ${MASK_DATA_ROOT}"
+  '