{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Instance Segmentation 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/ltdetr_instance_segmentation_export.ipynb)\n",
    "\n",
    "This notebook demonstrates how to export an LTDETRv2 instance segmentation model to ONNX and TensorRT.\n",
    "\n",
    "The notebook covers the following steps:\n",
    "1. Install LightlyTrain\n",
    "2. Export a trained LTDETRv2 model to ONNX\n",
    "3. Run inference with the ONNX model\n",
    "4. Export the model to TensorRT\n",
    "5. 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",
    "Load the COCO-trained LTDETRv2 instance segmentation model with LightlyTrain's `load_model` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "model = lightly_train.load_model(\"ltdetrv2-seg-s-coco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Download an example image\n",
    "\n",
    "Download an example image for inference:"
   ]
  },
  {
   "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 numpy as np\n",
    "import torchvision.transforms.v2 as T\n",
    "from PIL import Image\n",
    "from torchvision.transforms.functional import pil_to_tensor\n",
    "\n",
    "image_pil = Image.open(\"image.jpg\").convert(\"RGB\")\n",
    "image_tensor = pil_to_tensor(image_pil)\n",
    "\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) if model.image_normalize else T.Identity(),\n",
    "    ]\n",
    ")\n",
    "image_tensor_transformed = transforms(image_pil)[None]\n",
    "orig_target_size = np.array([[h, w]], dtype=np.int64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Get the model predictions for reference\n",
    "\n",
    "Define a helper that displays each instance's mask, bounding box, and class label. We will use it for the PyTorch, ONNX, and TensorRT predictions."
   ]
  },
  {
   "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_bounding_boxes, draw_segmentation_masks\n",
    "\n",
    "\n",
    "def visualize_instances(image, boxes, masks, labels, classes):\n",
    "    image_with_masks = draw_segmentation_masks(image, masks=masks.cpu(), alpha=0.6)\n",
    "    image_with_instances = draw_bounding_boxes(\n",
    "        image_with_masks,\n",
    "        boxes=boxes.cpu(),\n",
    "        labels=[classes[label.item()] for label in labels.cpu()],\n",
    "    )\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_instances.permute(1, 2, 0))\n",
    "    axs[1].axis(\"off\")\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.predict(image_tensor, threshold=0.6)\n",
    "visualize_instances(\n",
    "    image_tensor,\n",
    "    boxes=results[\"bboxes\"],\n",
    "    masks=results[\"masks\"],\n",
    "    labels=results[\"labels\"],\n",
    "    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": [
    "model.export_onnx(\n",
    "    out=\"model.onnx\",\n",
    "    # precision=\"fp16\",  # Use FP16 weights for a smaller and faster model.\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_instance_segmentation.task_model.LTDETRInstanceSegmentation.export_onnx) for all available options."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Run inference with the ONNX model\n",
    "\n",
    "The exported graph takes the preprocessed image and the original image size in `(height, width)` order. It returns all queries, so apply the same score and mask thresholds as `model.predict`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as ort\n",
    "import torch.nn.functional as F\n",
    "\n",
    "sess = ort.InferenceSession(\"model.onnx\")\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",
    "labels_onnx, boxes_onnx, masks_onnx, scores_onnx = sess.run(\n",
    "    output_names=None,\n",
    "    input_feed={\n",
    "        \"images\": image_tensor_transformed.numpy().astype(input_dtype_numpy),\n",
    "        \"orig_target_size\": orig_target_size,\n",
    "    },\n",
    ")\n",
    "\n",
    "labels_onnx = labels_onnx[0]\n",
    "boxes_onnx = boxes_onnx[0]\n",
    "masks_onnx = masks_onnx[0]\n",
    "scores_onnx = scores_onnx[0]\n",
    "\n",
    "keep = scores_onnx > 0.6\n",
    "labels_onnx = torch.from_numpy(labels_onnx[keep])\n",
    "boxes_onnx = torch.from_numpy(boxes_onnx[keep])\n",
    "masks_onnx = torch.from_numpy(masks_onnx[keep])\n",
    "masks_onnx = (\n",
    "    F.interpolate(\n",
    "        masks_onnx.unsqueeze(1),\n",
    "        size=(h, w),\n",
    "        mode=\"bilinear\",\n",
    "        align_corners=False,\n",
    "    ).squeeze(1)\n",
    "    > 0.0\n",
    ")\n",
    "\n",
    "visualize_instances(\n",
    "    image_tensor,\n",
    "    boxes=boxes_onnx,\n",
    "    masks=masks_onnx,\n",
    "    labels=labels_onnx,\n",
    "    classes=model.classes,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## Export to TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Requirements\n",
    "\n",
    "TensorRT is not part of LightlyTrain's dependencies and must be installed separately. Installation depends on your OS, Python version, GPU, and NVIDIA driver/CUDA setup. See the [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html) for 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": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"tensorrt-cu12==10.13.3.9\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.export_tensorrt(\n",
    "    out=\"model.trt\",\n",
    "    # precision=\"fp16\",  # Use FP16 weights for a smaller and faster engine.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "See [`export_tensorrt`](https://docs.lightly.ai/train/stable/python_api/lightly_train.html#lightly_train._task_models.ltdetr_instance_segmentation.task_model.LTDETRInstanceSegmentation.export_tensorrt) for all available options."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "### Run inference with the TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt as trt\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",
    "            name\n",
    "            for name in io_names\n",
    "            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT\n",
    "        ]\n",
    "        self.out_names = [\n",
    "            name\n",
    "            for name in io_names\n",
    "            if self.engine.get_tensor_mode(name) == 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": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "trt_model = TRT(\"model.trt\")\n",
    "outputs_trt = trt_model(\n",
    "    {\n",
    "        \"images\": image_tensor_transformed,\n",
    "        \"orig_target_size\": torch.from_numpy(orig_target_size),\n",
    "    }\n",
    ")\n",
    "\n",
    "labels_trt = outputs_trt[\"labels\"][0]\n",
    "boxes_trt = outputs_trt[\"boxes\"][0]\n",
    "masks_trt = outputs_trt[\"masks\"][0]\n",
    "scores_trt = outputs_trt[\"scores\"][0]\n",
    "\n",
    "keep = scores_trt > 0.6\n",
    "labels_trt = labels_trt[keep]\n",
    "boxes_trt = boxes_trt[keep]\n",
    "masks_trt = masks_trt[keep]\n",
    "masks_trt = (\n",
    "    F.interpolate(\n",
    "        masks_trt.unsqueeze(1),\n",
    "        size=(h, w),\n",
    "        mode=\"bilinear\",\n",
    "        align_corners=False,\n",
    "    ).squeeze(1)\n",
    "    > 0.0\n",
    ")\n",
    "\n",
    "visualize_instances(\n",
    "    image_tensor,\n",
    "    boxes=boxes_trt,\n",
    "    masks=masks_trt,\n",
    "    labels=labels_trt,\n",
    "    classes=model.classes,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}