{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the TensorRT version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -c 'import tensorrt; print(\"TensorRT version: {}\".format(tensorrt.__version__))'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare the input image and ONNX model file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 /workspace/TensorRT/quickstart/SemanticSegmentation/export.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build TensorRT engine from the ONNX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!trtexec --onnx=fcn-resnet101.onnx --saveEngine=fcn-resnet101.engine --optShapes=input:1x3x1026x1282 --stronglyTyped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import required modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import ctypes\n",
    "from cuda.bindings import runtime as cudart\n",
    "import tensorrt as trt\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "\n",
    "TRT_LOGGER = trt.Logger()\n",
    "\n",
    "assert cudart.cudaSetDevice(0) == (cudart.cudaError_t.cudaSuccess,)\n",
    "\n",
    "# Filenames of TensorRT plan file and input/output images.\n",
    "engine_file = \"/workspace/fcn-resnet101.engine\"\n",
    "input_file  = \"/workspace/input.ppm\"\n",
    "output_file = \"/workspace/output.ppm\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utilities for input / output processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For torchvision models, input images are loaded in to a range of [0, 1] and\n",
    "# normalized using mean = [0.485, 0.456, 0.406] and stddev = [0.229, 0.224, 0.225].\n",
    "def preprocess(image):\n",
    "    # Mean normalization\n",
    "    mean = np.array([0.485, 0.456, 0.406]).astype('float32')\n",
    "    stddev = np.array([0.229, 0.224, 0.225]).astype('float32')\n",
    "    data = (np.asarray(image).astype('float32') / float(255.0) - mean) / stddev\n",
    "    # Switch from HWC to to CHW order\n",
    "    return np.moveaxis(data, 2, 0)\n",
    "\n",
    "def postprocess(data):\n",
    "    num_classes = 21\n",
    "    # create a color palette, selecting a color for each class\n",
    "    palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])\n",
    "    colors = np.array([palette*i%255 for i in range(num_classes)]).astype(\"uint8\")\n",
    "    # plot the segmentation predictions for 21 classes in different colors\n",
    "    img = Image.fromarray(data.astype('uint8'), mode='P')\n",
    "    img.putpalette(colors)\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load TensorRT engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Deserialize the TensorRT engine from specified plan file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_engine(engine_file_path):\n",
    "    assert os.path.exists(engine_file_path)\n",
    "    print(\"Reading engine from file {}\".format(engine_file_path))\n",
    "    with open(engine_file_path, \"rb\") as f, trt.Runtime(TRT_LOGGER) as runtime:\n",
    "        return runtime.deserialize_cuda_engine(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting with a deserialized engine, TensorRT inference pipeline consists of the following steps:\n",
    "- Create an execution context and specify input shape (based on the image dimensions for inference).\n",
    "- Allocate CUDA device memory for input and output.\n",
    "- Allocate CUDA page-locked host memory to efficiently copy back the output.\n",
    "- Transfer the processed image data into input memory using asynchronous host-to-device CUDA copy.\n",
    "- Kickoff the TensorRT inference pipeline using the asynchronous execute API.\n",
    "- Transfer the segmentation output back into pagelocked host memory using device-to-host CUDA copy.\n",
    "- Synchronize the stream used for data transfers and inference execution to ensure all operations are completes.\n",
    "- Finally, write out the segmentation output to an image file for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer(engine, input_file, output_file):\n",
    "    print(\"Reading input image from file {}\".format(input_file))\n",
    "    with Image.open(input_file) as img:\n",
    "        input_image = preprocess(img)\n",
    "        image_width = img.width\n",
    "        image_height = img.height\n",
    "\n",
    "    with engine.create_execution_context() as context:\n",
    "        input_buffers = {}\n",
    "        input_memories = {}\n",
    "\n",
    "        # Allocate host and device buffers\n",
    "        tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]\n",
    "        for tensor in tensor_names:\n",
    "            size = trt.volume(context.get_tensor_shape(tensor))\n",
    "            dtype = trt.nptype(engine.get_tensor_dtype(tensor))\n",
    "\n",
    "            if engine.get_tensor_mode(tensor) == trt.TensorIOMode.INPUT:\n",
    "                context.set_input_shape(tensor, (1, 3, image_height, image_width))\n",
    "                input_buffers[tensor] = np.ascontiguousarray(input_image)\n",
    "                err, input_memories[tensor] = cudart.cudaMalloc(input_image.nbytes)\n",
    "                assert err == cudart.cudaError_t.cudaSuccess\n",
    "                context.set_tensor_address(tensor, input_memories[tensor])\n",
    "            else:\n",
    "                err, output_buffer_ptr = cudart.cudaMallocHost(size * dtype().itemsize)\n",
    "                assert err == cudart.cudaError_t.cudaSuccess\n",
    "                pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))\n",
    "                output_buffer = np.ctypeslib.as_array(ctypes.cast(output_buffer_ptr, pointer_type), (size,))\n",
    "\n",
    "                err, output_memory = cudart.cudaMalloc(output_buffer.nbytes)\n",
    "                assert err == cudart.cudaError_t.cudaSuccess\n",
    "                context.set_tensor_address(tensor, output_memory)\n",
    "\n",
    "        err, stream = cudart.cudaStreamCreate()\n",
    "        assert err == cudart.cudaError_t.cudaSuccess\n",
    "\n",
    "        # Transfer input data to the GPU for all input tensors\n",
    "        for tensor_name, input_buffer in input_buffers.items():\n",
    "            input_memory = input_memories[tensor_name]\n",
    "            err, = cudart.cudaMemcpyAsync(input_memory, input_buffer.ctypes.data, input_buffer.nbytes,\n",
    "                                          cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n",
    "            assert err == cudart.cudaError_t.cudaSuccess\n",
    "\n",
    "        # Run inference\n",
    "        context.execute_async_v3(stream)\n",
    "\n",
    "        # Transfer prediction output from the GPU.\n",
    "        err, = cudart.cudaMemcpyAsync(output_buffer.ctypes.data, output_memory, output_buffer.nbytes,\n",
    "                                      cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n",
    "        assert err == cudart.cudaError_t.cudaSuccess\n",
    "        # Synchronize the stream\n",
    "        err, = cudart.cudaStreamSynchronize(stream)\n",
    "        assert err == cudart.cudaError_t.cudaSuccess\n",
    "\n",
    "        output_d64 = np.array(output_buffer, dtype=np.int64)\n",
    "        np.savetxt('test.out', output_d64.astype(int), fmt='%i', delimiter=' ', newline=' ')\n",
    "\n",
    "        with postprocess(np.reshape(output_buffer, (image_height, image_width))) as img:\n",
    "            print(\"Writing output image to file {}\".format(output_file))\n",
    "            img.convert('RGB').save(output_file, \"PPM\")\n",
    "\n",
    "        # cleanup cuda resources for all input tensors\n",
    "        for input_memory in input_memories.values():\n",
    "            cudart.cudaFree(input_memory)\n",
    "        cudart.cudaFree(output_memory)\n",
    "        cudart.cudaFreeHost(output_buffer_ptr)\n",
    "        cudart.cudaStreamDestroy(stream)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot input image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(Image.open(input_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Running TensorRT inference for FCN-ResNet101\")\n",
    "with load_engine(engine_file) as engine:\n",
    "    infer(engine, input_file, output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot segmentation output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(Image.open(output_file))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}