{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Object Detection Model Export - LTDETRv2\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/object_detection_export.ipynb)\n",
    "\n",
    "This notebook demonstrates how to export an object detection model to ONNX and TensorRT.\n",
    "\n",
    "The notebook covers the following steps:\n",
    "1. Install LightlyTrain\n",
    "2. Export a trained LTDETR model to ONNX\n",
    "3. Export a trained LTDETR model to TensorRT\n",
    "4. Run inference with the TensorRT engine\n",
    "\n",
    "The image size and normalization used for preprocessing are baked into the exported files: the image size as the static `images` input shape, and the normalization statistics and class names as ONNX metadata. This means the ONNX and TensorRT inference cells read everything they need directly from the exported files and do **not** require LightlyTrain to be installed.\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(\"ltdetrv2-s-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": [
    "### Load the example image\n",
    "\n",
    "Load the image once. The raw tensor is used both for the reference PyTorch prediction and for visualization. Preprocessing is deferred until after export, where the required image size and normalization are read back from the exported ONNX file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "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 and the reference prediction.\n",
    "image_tensor = pil_to_tensor(image_pil)\n",
    "\n",
    "# Original image size, used to rescale boxes back after inference.\n",
    "w, h = image_pil.size"
   ]
  },
  {
   "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, ONNX and\n",
    "TensorRT models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "\n",
    "\n",
    "def visualize_detections(image, boxes, labels, classes):\n",
    "    image_with_boxes = draw_bounding_boxes(\n",
    "        image,\n",
    "        boxes=boxes,\n",
    "        labels=[classes[label.item()] for label in labels],\n",
    "    )\n",
    "    plt.imshow(image_with_boxes.permute(1, 2, 0))\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get predictions from the PyTorch model.\n",
    "labels_torch, boxes_torch, _ = model.predict(image_tensor, threshold=0.6).values()\n",
    "\n",
    "# Visualize predictions from the PyTorch model.\n",
    "visualize_detections(\n",
    "    image_tensor, boxes=boxes_torch, labels=labels_torch, classes=model.classes\n",
    ")"
   ]
  },
  {
   "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",
    "model.export_onnx(\n",
    "    out=\"model.onnx\",\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.ltdetr_object_detection.task_model.LTDETRObjectDetection.export_onnx) for all available options when exporting to ONNX.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Read the baked-in preprocessing from the exported file\n",
    "\n",
    "The steps below read the image size, normalization statistics, and class names directly from the exported ONNX file. This is all that is needed to preprocess an image and interpret the outputs, so the remaining inference cells do not depend on LightlyTrain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import onnx\n",
    "import torch\n",
    "import torchvision.transforms.v2 as T\n",
    "from torchvision.ops import box_convert\n",
    "\n",
    "# Read the baked-in metadata from the ONNX file (no LightlyTrain required).\n",
    "onnx_model = onnx.load(\"model.onnx\")\n",
    "metadata = {entry.key: entry.value for entry in onnx_model.metadata_props}\n",
    "\n",
    "# Class names, as a {class_id: name} mapping.\n",
    "classes = {int(k): v for k, v in json.loads(metadata[\"classes\"]).items()}\n",
    "\n",
    "# Normalization statistics, if any were baked in.\n",
    "image_normalize = (\n",
    "    json.loads(metadata[\"image_normalize\"]) if \"image_normalize\" in metadata else None\n",
    ")\n",
    "\n",
    "# Image size is baked into the static height/width of the \"images\" input.\n",
    "images_input = onnx_model.graph.input[0].type.tensor_type.shape.dim\n",
    "image_size = (images_input[2].dim_value, images_input[3].dim_value)  # (H, W)\n",
    "\n",
    "# Build the preprocessing transform from the file-derived values.\n",
    "transforms = T.Compose(\n",
    "    [\n",
    "        T.Resize(image_size),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(**image_normalize) if image_normalize else T.Identity(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Apply transforms for ONNX and TensorRT inference.\n",
    "image_tensor_transformed = transforms(image_pil)[None]\n",
    "\n",
    "\n",
    "def postprocess_detections(logits, boxes, threshold=0.6):\n",
    "    \"\"\"Decode the raw graph outputs into labels, boxes, and scores.\n",
    "\n",
    "    The exported graph returns raw class ``logits`` of shape ``(1, num_queries,\n",
    "    num_classes)`` and normalized ``cxcywh`` boxes of shape ``(1, num_queries, 4)``.\n",
    "    Applying a sigmoid, keeping detections above the score threshold, and rescaling\n",
    "    the boxes to the original image reproduces what ``model.predict`` returns.\n",
    "    \"\"\"\n",
    "    scores = logits[0].sigmoid()\n",
    "    num_classes = scores.shape[-1]\n",
    "    boxes_xyxy = box_convert(boxes[0], in_fmt=\"cxcywh\", out_fmt=\"xyxy\")\n",
    "\n",
    "    scores_flat = scores.flatten()\n",
    "    # Match model.predict: select the best num_queries class/query pairs first,\n",
    "    # then apply the score threshold.\n",
    "    scores, flat_index = scores_flat.topk(logits.shape[1])\n",
    "    keep = scores > threshold\n",
    "    scores = scores[keep]\n",
    "    flat_index = flat_index[keep]\n",
    "    query_index = flat_index // num_classes\n",
    "    class_index = flat_index % num_classes\n",
    "\n",
    "    # Boxes are normalized to [0, 1]; scale them to the original image size.\n",
    "    boxes = boxes_xyxy[query_index] * torch.tensor([w, h, w, h], dtype=boxes_xyxy.dtype)\n",
    "\n",
    "    # Map internal class indices back to the dataset class ids from the metadata.\n",
    "    class_ids = list(classes.keys())\n",
    "    labels = torch.tensor([class_ids[int(i)] for i in class_index])\n",
    "    return labels, boxes, scores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Run inference with the ONNX model\n",
    "\n",
    "The exported graph returns raw class logits and normalized boxes (top-k selection, thresholding, and box rescaling are intentionally kept outside the graph), so we decode them with the `postprocess_detections` helper defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as ort\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",
    "# The graph outputs raw logits and normalized cxcywh boxes.\n",
    "logits_onnx, boxes_onnx = sess.run(\n",
    "    output_names=None,\n",
    "    input_feed={\n",
    "        \"images\": image_tensor_transformed.numpy().astype(input_dtype_numpy),\n",
    "    },\n",
    ")\n",
    "\n",
    "# Decode to labels, boxes, and scores (see postprocess_detections above).\n",
    "labels_onnx, boxes_onnx, scores_onnx = postprocess_detections(\n",
    "    torch.from_numpy(logits_onnx).float(), torch.from_numpy(boxes_onnx).float()\n",
    ")\n",
    "\n",
    "# Visualize predictions from the ONNX model.\n",
    "visualize_detections(\n",
    "    image_tensor, boxes=boxes_onnx, labels=labels_onnx, classes=classes\n",
    ")"
   ]
  },
  {
   "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",
    "model.export_tensorrt(\n",
    "    out=\"model.trt\",\n",
    "    # precision=\"fp16\", # Export model with FP16 weights for smaller size and faster inference.\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.ltdetr_object_detection.task_model.LTDETRObjectDetection.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",
    "    @torch.no_grad()\n",
    "    def __call__(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:\n",
    "        buffers = {}\n",
    "        # Bind the inputs and set their shapes before querying output shapes.\n",
    "        for name in self.in_names:\n",
    "            np_dtype = trt.nptype(self.engine.get_tensor_dtype(name))\n",
    "            torch_dtype = torch.from_numpy(np.empty((), dtype=np_dtype)).dtype\n",
    "            tensor = inputs[name].to(device=self.device, dtype=torch_dtype).contiguous()\n",
    "            self.context.set_input_shape(name, tuple(tensor.shape))\n",
    "            buffers[name] = tensor\n",
    "            self.context.set_tensor_address(name, tensor.data_ptr())\n",
    "        # Output shapes are only known once the input shapes are set.\n",
    "        for name in self.out_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",
    "            buffers[name] = buffer\n",
    "            self.context.set_tensor_address(name, buffer.data_ptr())\n",
    "        success = self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)\n",
    "        if not success:\n",
    "            raise RuntimeError(\"TensorRT execution failed\")\n",
    "        torch.cuda.synchronize()\n",
    "        return {name: 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(\"model.trt\")\n",
    "\n",
    "# Run inference with the TensorRT model. The graph outputs raw logits and boxes.\n",
    "outputs_trt = trt_model({\"images\": image_tensor_transformed})\n",
    "\n",
    "# Decode to labels, boxes, and scores (see postprocess_detections above).\n",
    "labels_trt, boxes_trt, scores_trt = postprocess_detections(\n",
    "    outputs_trt[\"logits\"].float().cpu(), outputs_trt[\"boxes\"].float().cpu()\n",
    ")\n",
    "\n",
    "# Visualize predictions from the TensorRT model.\n",
    "visualize_detections(image_tensor, boxes=boxes_trt, labels=labels_trt, classes=classes)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}