Depth estimation is the task of measuring the distance of each pixel relative to the camera. This repo provides a TensorRT implementation of the [Depth-Anything](https://github.com/LiheYoung/Depth-Anything) depth estimation model in both C++ and Python, enabling efficient real-time inference.
## ⏱️ Performance
The inference time includes the pre-preprocessing and post-processing stages:
| Device | Model | Model Input (WxH) | Image Resolution (WxH)|Inference Time(ms)|
|:---------------:|:------------:|:------------:|:------------:|:------------:|
| RTX4090 | Depth-Anything-S |518x518 | 1280x720 | 3 |
| RTX4090 | Depth-Anything-B |518x518 | 1280x720 | 6 |
| RTX4090 | Depth-Anything-L |518x518 | 1280x720 | 12 |
> [!NOTE]
> Inference was conducted using `FP16` precision, with a warm-up period of 10 frames. The reported time corresponds to the last inference.
## 🚀 Quick Start
#### C++
- Usage 1: Create an engine from an onnx model and save it
``` shell
depth-anything-tensorrt.exe
```
- Usage 2: Deserialize an engine and run it
``` shell
depth-anything-tensorrt.exe
```
Example:
``` shell
# infer image
depth-anything-tensorrt.exe depth_anything_vitb14.engine test.jpg
# infer folder(images)
depth-anything-tensorrt.exe depth_anything_vitb14.engine data
# infer video
depth-anything-tensorrt.exe depth_anything_vitb14.engine test.mp4 # the video path
```
#### Python
```
cd depth-anything-tensorrt/python
# infer image
python trt_infer.py --engine --img --outdir [--grayscale]
```
## 🛠️ Build
#### C++
Refer to our [docs/INSTALL.md](https://github.com/spacewalk01/depth-anything-tensorrt/blob/main/docs/INSTALL.md) for C++ environment installation.
#### Python
``` shell
cd /python
pip install cuda-python
pip install tensorrt-8.6.0-cp310-none-win_amd64.whl
pip install opencv-python
```
## 🤖 Model Preparation
Perform the following steps to create an onnx model:
1. Download the pretrained [model](https://huggingface.co/spaces/LiheYoung/Depth-Anything/tree/main/checkpoints) and install [Depth-Anything](https://github.com/LiheYoung/Depth-Anything):
``` shell
git clone https://github.com/LiheYoung/Depth-Anything
cd Depth-Anything
pip install -r requirements.txt
```
2. Copy and paste [dpt.py](https://github.com/spacewalk01/depth-anything-tensorrt/blob/main/dpt.py) in this repo to `/depth_anything` folder. And copy [export.py](https://github.com/spacewalk01/depth-anything-tensorrt/blob/main/export.py) in this repo to ``. Note that I've only removed a squeeze operation at the end of model's forward function in `dpt.py` to avoid conflicts with TensorRT.
3. Export the model to onnx format using [export.py](https://github.com/spacewalk01/depth-anything-tensorrt/blob/main/export.py). You will get an onnx file named `depth_anything_vit{}14.onnx`, such as `depth_anything_vitb14.onnx`.
``` shell
python export.py --encoder vitb --load_from depth_anything_vitb14.pth --image_shape 3 518 518
```
> [!TIP]
> The width and height of the model input should be divisible by 14, the patch height.
## 👏 Acknowledgement
This project is based on the following projects:
- [Depth-Anything](https://github.com/LiheYoung/Depth-Anything) - Unleashing the Power of Large-Scale Unlabeled Data.
- [TensorRT](https://github.com/NVIDIA/TensorRT/tree/release/8.6/samples) - TensorRT samples and api documentation.
