{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Instance Segmentation Model Export - EoMT\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/eomt_instance_segmentation_export.ipynb)\n",
    "\n",
    "This notebook demonstrates how to export an instance segmentation model to ONNX and TensorRT.\n",
    "\n",
    "The notebook covers the following steps:\n",
    "1. Install LightlyTrain\n",
    "2. Export a trained EoMT model to ONNX\n",
    "3. Export a trained EoMT 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": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"lightly-train[onnx,onnxruntime,onnxslim]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "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": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "model = lightly_train.load_model(\"dinov3/vits16-eomt-inst-coco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "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": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000039769.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms.v2 as T\n",
    "from PIL import Image\n",
    "from torchvision.transforms.functional import pil_to_tensor\n",
    "\n",
    "# Load image with PIL.\n",
    "image_pil = Image.open(\"image.jpg\").convert(\"RGB\")\n",
    "\n",
    "# Convert PIL image to tensor for plotting.\n",
    "image_tensor = pil_to_tensor(image_pil)\n",
    "\n",
    "# Define pre-processing transforms.\n",
    "w, h = image_pil.size\n",
    "transforms = T.Compose(\n",
    "    [\n",
    "        T.Resize((model.image_size)),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(**model.image_normalize),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Apply transforms for ONNX inference.\n",
    "image_tensor_transformed = transforms(image_pil)[None]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "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 and ONNX models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torchvision.utils import draw_segmentation_masks\n",
    "\n",
    "\n",
    "def visualize_segmentations(image, masks):\n",
    "    # Draw segmentation masks.\n",
    "    image_with_masks = draw_segmentation_masks(image, masks, alpha=1.0)\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(12, 8))\n",
    "    axs[0].imshow(image.permute(1, 2, 0))\n",
    "    axs[0].axis(\"off\")\n",
    "    axs[1].imshow(image_with_masks.permute(1, 2, 0))\n",
    "    axs[1].axis(\"off\")\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get predictions from the PyTorch model.\n",
    "results = model.predict(image_tensor)\n",
    "\n",
    "# Visualize predictions from the PyTorch model.\n",
    "visualize_segmentations(image_tensor, masks=results[\"masks\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Export the model to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export the PyTorch model to ONNX.\n",
    "\n",
    "model.export_onnx(\n",
    "    out=\"model.onnx\",\n",
    "    verify=False,  # The following inference cell verifies the exported model.\n",
    "    # precision=\"fp16\", # Export model with FP16 weights for smaller size and faster inference.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "See [`export_onnx`](https://docs.lightly.ai/train/stable/python_api/lightly_train.html#lightly_train._task_models.dinov3_eomt_instance_segmentation.task_model.DINOv3EoMTInstanceSegmentation.export_onnx) for all available options when exporting to ONNX.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Run inference with the ONNX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as ort\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Create an ONNX Runtime session.\n",
    "sess = ort.InferenceSession(\"model.onnx\")\n",
    "\n",
    "# Get expected input dtype.\n",
    "input_dtype = sess.get_inputs()[0].type\n",
    "input_dtype_numpy = {\n",
    "    \"tensor(float)\": \"float32\",\n",
    "    \"tensor(float16)\": \"float16\",\n",
    "}[input_dtype]\n",
    "\n",
    "# Run inference.\n",
    "labels_onnx, masks_onnx, scores_onnx = sess.run(\n",
    "    output_names=None,\n",
    "    input_feed={\"images\": image_tensor_transformed.numpy().astype(input_dtype_numpy)},\n",
    ")\n",
    "\n",
    "# Remove batch dimension.\n",
    "labels_onnx = labels_onnx[0]\n",
    "masks_onnx = masks_onnx[0]\n",
    "scores_onnx = scores_onnx[0]\n",
    "\n",
    "# Filter by score (ONNX export returns all queries).\n",
    "keep = scores_onnx >= 0.8\n",
    "labels_onnx = labels_onnx[keep]\n",
    "masks_onnx = masks_onnx[keep]\n",
    "scores_onnx = scores_onnx[keep]\n",
    "\n",
    "# Resize ONNX predictions to original image size if necessary.\n",
    "# This is only needed if the original image size is different from the model input size.\n",
    "masks_onnx = torch.from_numpy(masks_onnx)\n",
    "masks_onnx = (\n",
    "    F.interpolate(masks_onnx.float().unsqueeze(1), size=(h, w), mode=\"nearest\")\n",
    "    .squeeze(1)\n",
    "    .bool()\n",
    ")\n",
    "\n",
    "# Visualize predictions from the ONNX model.\n",
    "visualize_segmentations(image_tensor, masks=masks_onnx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "**Note**: There might be small visual differences between the masks predicted by\n",
    "the PyTorch model with `model.predict` and the ONNX model. This is because the PyTorch\n",
    "model uses `transforms` that resize the image slightly differently compared to the fixed\n",
    "size resize used for ONNX."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Export to TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "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. 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": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"tensorrt-cu12==10.13.3.9\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the TensorRT engine.\n",
    "model.export_tensorrt(\n",
    "    out=\"model.trt\",\n",
    "    onnx_args={\"verify\": False},  # TensorRT validates and builds the ONNX graph.\n",
    "    # precision=\"fp16\", # Export model with FP16 weights for smaller size and faster inference.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "See [`export_tensorrt`](https://docs.lightly.ai/train/stable/python_api/lightly_train.html#lightly_train._task_models.dinov3_eomt_instance_segmentation.task_model.DINOv3EoMTInstanceSegmentation.export_tensorrt) for all available options when exporting to TensorRT.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "### Run inference with the TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "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": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the TensorRT model.\n",
    "trt_model = TRT(\"model.trt\")\n",
    "\n",
    "# Run inference with the TensorRT model.\n",
    "# Note: TensorRT inference returns raw outputs (labels, masks, scores).\n",
    "outputs_trt = trt_model({\"images\": image_tensor_transformed})\n",
    "\n",
    "labels_trt = outputs_trt[\"labels\"][0]\n",
    "masks_trt = outputs_trt[\"masks\"][0]\n",
    "scores_trt = outputs_trt[\"scores\"][0]\n",
    "\n",
    "# Filter by score (TensorRT export returns all queries).\n",
    "keep = scores_trt >= 0.8\n",
    "labels_trt = labels_trt[keep]\n",
    "masks_trt = masks_trt[keep]\n",
    "scores_trt = scores_trt[keep]\n",
    "\n",
    "# Resize TensorRT predictions to original image size if necessary.\n",
    "# This assumes masks_trt is on GPU/device since TRT class keeps them there.\n",
    "masks_trt = (\n",
    "    F.interpolate(masks_trt.float().unsqueeze(1), size=(h, w), mode=\"nearest\")\n",
    "    .squeeze(1)\n",
    "    .bool()\n",
    ")\n",
    "\n",
    "# Visualize predictions from the TensorRT model.\n",
    "visualize_segmentations(image_tensor, masks=masks_trt.cpu())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}