{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Depth Estimation Model Export (ONNX and TensorRT)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/depth_estimation_export.ipynb)\n",
    "\n",
    "This notebook demonstrates how to export a depth estimation model to ONNX and TensorRT.\n",
    "\n",
    "The notebook covers the following steps:\n",
    "1. Install LightlyTrain\n",
    "2. Export a Depth Anything model to ONNX\n",
    "3. Export a Depth Anything model to TensorRT\n",
    "4. Run inference with the TensorRT engine\n",
    "\n",
    "> **Important**: When running on Google Colab make sure to select a GPU runtime. You can do this by going to `Runtime` > `Change runtime type` and selecting a GPU hardware accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"lightly-train[onnx,onnxruntime,onnxslim]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Export to ONNX\n",
    "\n",
    "### Load the model weights\n",
    "\n",
    "Then load the model with LightlyTrain's `load_model` function. This will automatically download the model weights and load the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onnxruntime as ort\n",
    "import torch\n",
    "\n",
    "import lightly_train\n",
    "\n",
    "model = lightly_train.load_model(\"dinov2/dav2-relative-small\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### Download an example image\n",
    "\n",
    "Download an example image for inference with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000039769.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "# Precision used for the exported ONNX and TensorRT models.\n",
    "precision = \"fp32\"\n",
    "\n",
    "# The Depth Anything models process images at a fixed square resolution.\n",
    "height = width = model.image_size\n",
    "\n",
    "# Load the image and resize it to the model's processing resolution.\n",
    "img = Image.open(\"image.jpg\").convert(\"RGB\")\n",
    "img = img.resize((width, height), resample=Image.BILINEAR)\n",
    "\n",
    "# The depth model handles its own pre-processing (resize and normalization), so we\n",
    "# build the input batch with the model's preprocessing helpers instead of torchvision\n",
    "# transforms.\n",
    "x, meta = model.preprocess_image(img)\n",
    "batch = model.preprocess_batch([x])\n",
    "\n",
    "# Build the ONNX input feed with the dtype matching the export precision.\n",
    "if precision == \"fp16\":\n",
    "    inputs = {\"images\": batch.cpu().numpy().astype(np.float16)}\n",
    "else:\n",
    "    inputs = {\"images\": batch.cpu().numpy()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "### Get the model predictions for reference\n",
    "\n",
    "We define a helper function to visualize the predictions.\n",
    "The function will be used to compare the predictions from PyTorch, ONNX and\n",
    "TensorRT models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def visualize_depth(image, depth):\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n",
    "    axes[0].imshow(image)\n",
    "    axes[0].set_title(\"Input\")\n",
    "    axes[0].axis(\"off\")\n",
    "    axes[1].imshow(depth.cpu(), cmap=\"Spectral_r\")\n",
    "    axes[1].set_title(\"Relative depth\")\n",
    "    axes[1].axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get predictions from the PyTorch model.\n",
    "depth_torch = model.predict(img)\n",
    "\n",
    "# Visualize predictions from the PyTorch model.\n",
    "visualize_depth(img, depth=depth_torch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Export the model to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.export_onnx(\n",
    "    out=f\"dav2-relative-small-{precision}-{width}x{height}.onnx\",\n",
    "    height=height,\n",
    "    width=width,\n",
    "    precision=precision,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "See [`export_onnx`](https://docs.lightly.ai/train/stable/python_api/lightly_train.html#lightly_train._task_models.depth_estimation.task_model.DepthAnythingDepthEstimation.export_onnx) for all available options when exporting to ONNX.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Run inference with the ONNX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = ort.InferenceSession(\n",
    "    f\"dav2-relative-small-{precision}-{width}x{height}.onnx\",\n",
    "    providers=[\"CUDAExecutionProvider\"],\n",
    ")\n",
    "\n",
    "outs = session.run(None, inputs)\n",
    "names = [o.name for o in session.get_outputs()]\n",
    "raw = {n: torch.from_numpy(v) for n, v in zip(names, outs)}\n",
    "depth_onnx = raw[\"depth\"][0][0]\n",
    "\n",
    "visualize_depth(img, depth=depth_onnx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "**Note**: The exported graph returns the raw depth map at the model's processing\n",
    "resolution. Postprocessing that `model.predict` applies in PyTorch (resizing back to the\n",
    "original image size and, for metric models, the metric scaling) is not part of the ONNX\n",
    "or TensorRT graph and must be applied by the caller. Here the input image is already\n",
    "resized to the model's processing resolution, so no resize-back is needed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Export to TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Requirements\n",
    "\n",
    "TensorRT is not part of LightlyTrain’s dependencies and must be installed separately. Installation depends on your OS, Python\n",
    "version, GPU, and NVIDIA driver/CUDA setup.\n",
    "See the [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html) for more details.\n",
    "\n",
    "On CUDA 12.x systems, install the TensorRT version tested with this tutorial. Pinning the version prevents pip from installing TensorRT 11, which is not yet supported by LightlyTrain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"tensorrt-cu12==10.13.3.9\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the TensorRT engine.\n",
    "#\n",
    "# This example builds a fixed-shape engine (single image, no dynamic batch\n",
    "# dimension) so the simple runner below can allocate its buffers directly from\n",
    "# the engine's static tensor shapes. To build a dynamic-batch engine instead,\n",
    "# drop the `dynamic_batch_size`/`batch_size` overrides and set the input shape\n",
    "# on the execution context before allocating buffers.\n",
    "model.export_tensorrt(\n",
    "    out=f\"dav2-relative-small-{precision}-{width}x{height}.trt\",\n",
    "    onnx_args={\n",
    "        \"height\": height,\n",
    "        \"width\": width,\n",
    "        \"dynamic_batch_size\": False,\n",
    "        \"batch_size\": 1,\n",
    "    },\n",
    "    precision=precision,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "See [`export_tensorrt`](https://docs.lightly.ai/train/stable/python_api/lightly_train.html#lightly_train._task_models.depth_estimation.task_model.DepthAnythingDepthEstimation.export_tensorrt) for all available options when exporting to TensorRT.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "### Run inference with the TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorrt as trt\n",
    "import torch\n",
    "\n",
    "\n",
    "class TRT:\n",
    "    def __init__(self, engine_path: str, device: str = \"cuda:0\", verbose: bool = False):\n",
    "        self.device = torch.device(device)\n",
    "        logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.INFO)\n",
    "        trt.init_libnvinfer_plugins(logger, \"\")\n",
    "        runtime = trt.Runtime(logger)\n",
    "\n",
    "        with open(engine_path, \"rb\") as f:\n",
    "            self.engine = runtime.deserialize_cuda_engine(f.read())\n",
    "        self.context = self.engine.create_execution_context()\n",
    "\n",
    "        io_names = [\n",
    "            self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)\n",
    "        ]\n",
    "        self.in_names = [\n",
    "            n\n",
    "            for n in io_names\n",
    "            if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT\n",
    "        ]\n",
    "        self.out_names = [\n",
    "            n\n",
    "            for n in io_names\n",
    "            if self.engine.get_tensor_mode(n) == trt.TensorIOMode.OUTPUT\n",
    "        ]\n",
    "\n",
    "        self.buffers = {}\n",
    "        self.bindings = []\n",
    "        for name in io_names:\n",
    "            shape = tuple(self.context.get_tensor_shape(name))\n",
    "            np_dtype = trt.nptype(self.engine.get_tensor_dtype(name))\n",
    "            torch_dtype = torch.from_numpy(np.empty((), dtype=np_dtype)).dtype\n",
    "            buffer = torch.empty(\n",
    "                shape, device=self.device, dtype=torch_dtype\n",
    "            ).contiguous()\n",
    "            self.buffers[name] = buffer\n",
    "            self.bindings.append(buffer.data_ptr())\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def __call__(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:\n",
    "        for name in self.in_names:\n",
    "            self.buffers[name].copy_(inputs[name].to(self.device))\n",
    "        if not self.context.execute_v2(self.bindings):\n",
    "            raise RuntimeError(\"TensorRT execution failed\")\n",
    "        return {name: self.buffers[name] for name in self.out_names}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the TensorRT model.\n",
    "trt_model = TRT(\n",
    "    f\"dav2-relative-small-{precision}-{width}x{height}.trt\",\n",
    ")\n",
    "\n",
    "in_name = trt_model.in_names[0]\n",
    "\n",
    "if precision == \"fp16\":\n",
    "    inputs = {in_name: batch.to(torch.float16)}\n",
    "else:\n",
    "    inputs = {in_name: batch}\n",
    "raw = trt_model(inputs)\n",
    "\n",
    "depth_trt = raw[\"depth\"][0][0]\n",
    "\n",
    "visualize_depth(img, depth=depth_trt)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}