DOCSEARCHbyAlgolia
Submit Docs

UniDepth: Universal Monocular Metric Depth Estimation

repository·main·Indexed 22 days ago

https://github.com/lpiccinelli-eth/unidepth

A universal monocular metric depth estimation framework (version 0.1) capable of estimating absolute metric depth and intrinsics from single images across various camera setups. It supports UniDepthV1 and UniDepthV2 models with ViT and ConvNext backbones, available via Hugging Face and TorchHub. The library includes utilities for Pinhole and Fisheye camera models, ONNX export for UniDepthV2, and training scripts for torchrun and SLURM.

Tokens
3.4K
Snippets
12
Records
16
Agent score
79%

What's inside UniDepth

  1. Dataset requirements and structure

    main

    UniDepth supports both image-based and sequence-based datasets.

    • Note on Image Datasets: The ImageDataset class is considered legacy. Current image-based datasets are implemented as "dummy" single-frame sequences.
    • Custom Datasets: You can adapt your own data loading and processing, but you must maintain the same interface used by the existing datasets to ensure the model can be trained "out-of-the-box."
    • Examples: Example datasets (iBims-1 for image-based and Sintel for sequence-based) are available to help you understand the required pipeline and structure.

  2. Configure resolution and speed in UniDepthV2

    main

    UniDepthV2 allows you to trade off resolution and inference speed using the self.resolution_level attribute.

    • self.resolution_level: An attribute with a range of [0, 10) used in the preprocess function. It controls the linear interpolation degree of the processed image area within training bounds.
    • self.interpolation_mode: Defines the interpolation mode used by the infer method (defaults to bilinear).
    • Image Ratios: Input images are not forced into a specific size. Training ratios are in the range [0.5, 2.5]. If your image ratio falls outside this range, padding will be applied automatically.

  3. Install optional components (Pillow-SIMD and KNN)

    main

    For improved performance or specific evaluation tasks, you can install these optional components:

    1. Pillow-SIMD (Optional for performance)

    pip uninstall pillow
CC="cc -mavx2" pip install -U --force-reinstall pillow-simd

    2. KNN (For evaluation only)

    cd unidepth/ops/knn;bash compile.sh;cd ../../../

  4. Install UniDepth and dependencies

    main

    To install UniDepth and its required dependencies, use the following commands. It is recommended to use a virtual environment. The installation supports CUDA 11.8 and higher.

    export VENV_DIR=<YOUR-VENVS-DIR>
export NAME=Unidepth

python -m venv $VENV_DIR/$NAME
source $VENV_DIR/$NAME/bin/activate

# Install UniDepth and dependencies
pip install -e . --extra-index-url https://download.pytorch.org/whl/cu118

    If using conda, use these commands instead:

    conda create -n Unidepth python=3.11
conda activate Unidepth

    Note: Ensure your compilation CUDA version and runtime CUDA version match. If you encounter Triton Error [CUDA]: device kernel image is invalid with xFormers, it is likely due to a mismatch between system-wide CUDA and the CUDA shipped with PyTorch.

    export VENV_DIR=<YOUR-VENVS-DIR>
export NAME=Unidepth

python -m venv $VENV_DIR/$NAME
source $VENV_DIR/$NAME/bin/activate

# Install UniDepth and dependencies, cuda >11.8 work fine, too.
pip install -e . --extra-index-url https://download.pytorch.org/whl/cu118

  5. Compile CUDA operations for EdgeGuidedLocalSSI

    main

    To perform EdgeGuidedLocalSSI efficiently, you must compile the custom CUDA operation for random patch extraction. Navigate to the operations directory and run the provided compilation script. If your GPU architecture differs from the defaults in the script, ensure you export the TORCH_CUDA_ARCH_LIST environment variable before running the command.

    cd ./unidepth/ops/extract_patches && bash compile.sh

  6. Export UniDepthV2 to ONNX format

    main

    You can export UniDepthV2 models to ONNX format using the export.py script.

    Important considerations for ONNX usage:

    • Shape Matching: The exported shape will be adjusted to the closest multiple of 14 (the ViT patch size). Your inference-time input shape must match this resized shape.
    • Preprocessing: The ONNX model does not include pre- or post-processing. You must manually provide an RGB image that is rescaled to the target shape and normalized using ImageNet statistics.
    • Camera Input: Use the --with-camera flag if you want the model to accept ground truth camera data as unprojected rays.
    python ./unidepth/models/unidepthv2/export.py --version v2 --backbone vitl --shape 462 616 --output-path unidepthv2.onnx

  7. Train UniDepth using torchrun (Non-SLURM)

    main

    To start training on a single machine or a custom multi-node setup without SLURM, use the torchrun utility. You must first set up your environment variables, including PYTHONPATH, cache directories, and your dataset location (DATAROOT).

    Environment Setup:

    1. Set PYTHONPATH to include the repository root.
    2. Set DATAROOT to the path where your HDF5 files are stored.
    3. Configure cache directories (TMPDIR, TORCH_HOME, HUGGINGFACE_HUB_CACHE, WANDB_HOME) to avoid disk space issues.
    4. Define NNODES, RANK, MASTER_ADDR, and CUDA_VISIBLE_DEVICES.
    5. Specify the configuration file name in CFG.

    Execution: Run the torchrun command pointing to scripts/train.py with the --config-file and --distributed flags.

    export REPO=`pwd`
export PYTHONPATH=${REPO}:${PYTHONPATH}

export TMPDIR="/tmp"
export TORCH_HOME=${TMPDIR}
export HUGGINGFACE_HUB_CACHE=${TMPDIR}
export WANDB_HOME=${TMPDIR}
export DATAROOT=<where-you-stored-the-hdf5>

export MASTER_PORT=$((( RANDOM % 600 ) + 29400 ))
if [ $NNODES -gt 1 ]; then
    export MASTER_PORT=29400
fi

export CFG="train_v1_vitl14.json"

export NNODES=1
export RANK=0
export MASTER_ADDR=127.0.0.1
export CUDA_VISIBLE_DEVICES="0"

export GPUS=$(echo ${CUDA_VISIBLE_DEVICES} | tr ',' '\n' | wc -l)
echo "Start script with python from: `which python`"
torchrun --rdzv-backend=c10d --nnodes=${NNODES} --nproc_per_node=${GPUS} --rdzv-endpoint ${MASTER_ADDR}:${MASTER_PORT} ${REPO}/scripts/train.py --config-file ${REPO}/configs/${CFG} --distributed

  8. Compile KNN operations for evaluation

    main

    UniDepth requires Chamfer distance for evaluation, which relies on PyTorch3D's KNN implementation. If you encounter issues installing PyTorch3D, you can manually compile the KNN operation. Navigate to the unidepth/ops/knn directory within the repository and run the provided compilation script.

    bash compile.sh

  9. Train UniDepth using SLURM

    main

    If your system uses the SLURM scheduler, you can initiate training using srun. The scheduler will automatically handle node and rank information. You only need to provide the configuration file and the master port.

    srun -c ${SLURM_CPUS_PER_TASK} --kill-on-bad-exit=1 python -u ${REPO}/scripts/train.py --config-file ${REPO}/configs/${CFG} --master-port ${MASTER_PORT} --distributed

  10. Perform metric depth and intrinsics inference

    main

    To generate metric depth estimation and intrinsics prediction from an RGB image, use the .infer() method. The model handles normalization automatically.

    import torch
import numpy as np
from PIL import Image
from unidepth.models import UniDepthV1

# Load model
model = UniDepthV1.from_pretrained("lpiccinelli/unidepth-v1-vitl14")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Prepare image: Load, convert to numpy, then to torch tensor (C, H, W)
image_path = "path/to/your/image.jpg"
rgb = torch.from_numpy(np.array(Image.open(image_path))).permute(2, 0, 1)

# Inference
predictions = model.infer(rgb)

# Extract results
depth = predictions["depth"]           # Metric Depth Estimation
xyz = predictions["points"]           # Point Cloud in Camera Coordinate
intrinsics = predictions["intrinsics"] # Intrinsics Prediction
    import numpy as np
from PIL import Image

# Move to CUDA, if any
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Load the RGB image and the normalization will be taken care of by the model
rgb = torch.from_numpy(np.array(Image.open(image_path))).permute(2, 0, 1) # C, H, W

predictions = model.infer(rgb)

# Metric Depth Estimation
depth = predictions["depth"]

# Point Cloud in Camera Coordinate
xyz = predictions["points"]

# Intrinsics Prediction
intrinsics = predictions["intrinsics"]

  11. Verify installation with demo.py

    main

    Run the provided demo script to test your installation. If successful, the script will print ARel: 7.45%.

    python ./scripts/demo.py

    If you encounter a Segmentation Fault, try uninstalling torch via pip (pip uninstall torch) and reinstalling the specific version listed in requirements.txt using conda.

    python ./scripts/demo.py

  12. Pass camera parameters to the infer method

    main

    The infer method in UniDepthV2 supports various camera models. You can provide camera information in two ways:

    1. K Matrix: Pass a tensor representing the intrinsic matrix (assumes a pinhole model).
    2. Camera Object: Pass a child class of Camera (e.g., Fisheye624, Pinhole, OPENCV).

    Warning: Do not pass the base Camera class directly, as it is an abstract parent class and lacks the necessary implemented methods.