SimpleHRNet_notebook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "x1P13nZeR3Xj",
        "HqHg_VATg6CO",
        "ZWUN1C5RgGYS"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Simple HRNet\n",
        "This is a light Google Colab notebook showing how to use the [simple-HRNet](https://github.com/stefanopini/simple-HRNet) repository.\n",
        "\n",
        "It includes the conversion to TensorRT and a test of the converted model.\n",
        "Please skip the section \"TensorRT\" if not interested.\n",
        "\n",
        "Initial idea of running on Google Colab by @basicvisual, initial implementation by @wuyenlin (see [issue #84](https://github.com/stefanopini/simple-HRNet/issues/84))."
      ],
      "metadata": {
        "id": "xZqqnmmNfX1d"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Pytorch"
      ],
      "metadata": {
        "id": "ZFihjwzqhA04"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Clone the repo and install the dependencies"
      ],
      "metadata": {
        "id": "X_ugGAxdd6Hu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# clone the repo\n",
        "!git clone https://github.com/stefanopini/simple-HRNet.git"
      ],
      "metadata": {
        "id": "FIecXpzEY7IJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd simple-HRNet\n",
        "!pwd"
      ],
      "metadata": {
        "id": "JDNRl8a8dl7Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# install requirements\n",
        "!pip install -r requirements.txt"
      ],
      "metadata": {
        "id": "FGsHqGPNdbHt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# install vlc to get video codecs\n",
        "!apt install vlc"
      ],
      "metadata": {
        "id": "qMynH2IPebr8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Add yolov3\n",
        "Clone yolov3 for multiprocessing support. This can be skipped for single-person applications or if you plan to use YOLO v5 by Ultralytics."
      ],
      "metadata": {
        "id": "x1P13nZeR3Xj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# download git submodules\n",
        "!git submodule update --init --recursive"
      ],
      "metadata": {
        "id": "yqf7BRGWRtUV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/simple-HRNet/models_/detectors/yolo\n",
        "!pip install -q -r requirements.txt\n",
        "\n",
        "%cd /content/simple-HRNet"
      ],
      "metadata": {
        "id": "vS9cz49gSJeG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/simple-HRNet/models_/detectors/yolo/weights\n",
        "!sh download_weights.sh\n",
        "\n",
        "%cd /content/simple-HRNet"
      ],
      "metadata": {
        "id": "8v-RpWGwSM7V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download HRNet pre-trained weights and test video\n",
        "\n",
        "Download any of the supported official weights listed [here](https://github.com/stefanopini/simple-HRNet/#installation-instructions).\n",
        "\n",
        "In the following, we download the weights `pose_hrnet_w48_384x288.pth` from the official Drive link.\n",
        "Download of other weights (e.g. `pose_hrnet_w32_256x192.pth`) as well as weights from private Drives is supported too."
      ],
      "metadata": {
        "id": "HqHg_VATg6CO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade --no-cache-dir gdown"
      ],
      "metadata": {
        "id": "pKFdWLLUXyZu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# download weights\n",
        "\n",
        "# create weights folder\n",
        "%cd /content/simple-HRNet\n",
        "!mkdir weights\n",
        "%cd /content/simple-HRNet/weights\n",
        "\n",
        "# download weights pose_hrnet_w48_384x288.pth\n",
        "!gdown 1UoJhTtjHNByZSm96W3yFTfU5upJnsKiS\n",
        "\n",
        "# download weights pose_hrnet_w32_256x192.pth\n",
        "!gdown 1zYC7go9EV0XaSlSBjMaiyE_4TcHc_S38\n",
        "\n",
        "# download weights pose_hrnet_w32_256x256.pth\n",
        "!gdown 1_wn2ifmoQprBrFvUCDedjPON4Y6jsN-v\n",
        "\n",
        "# # download weights from your own Google Drive\n",
        "# from glob import glob\n",
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')\n",
        "# w_list = glob(\"/content/drive/<your drive folder>/*.pth\")\n",
        "# if not w_list:\n",
        "#   raise FileNotFoundError(\"You haven't downloaded any pre-trained weights!\")\n",
        "\n",
        "%cd /content/simple-HRNet"
      ],
      "metadata": {
        "id": "3LURZ12cfCcU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# download a publicly available video (or just get your own)\n",
        "!wget https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/WeAreGoingOnBullrun.mp4"
      ],
      "metadata": {
        "id": "OLIrIc14eUPM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Test the API\n"
      ],
      "metadata": {
        "id": "vcv0B2P7UTxT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import requests\n",
        "import matplotlib.pyplot as plt\n",
        "import torch\n",
        "from skimage import io\n",
        "from PIL import Image\n",
        "from SimpleHRNet import SimpleHRNet\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "# # singleperson, COCO weights\n",
        "# model = SimpleHRNet(48, 17, \"./weights/pose_hrnet_w48_384x288.pth\", multiperson=False, device=device)\n",
        "\n",
        "# # multiperson w/ YOLOv3, COCO weights\n",
        "# model = SimpleHRNet(48, 17, \"./weights/pose_hrnet_w48_384x288.pth\", device=device)\n",
        "\n",
        "# # multiperson w/ YOLOv3, COCO weights, small model\n",
        "# model = SimpleHRNet(32, 17, \"./weights/pose_hrnet_w32_256x192.pth\", device=device)\n",
        "\n",
        "# # multiperson w/ YOLOv3, MPII weights\n",
        "# model = SimpleHRNet(32, 16, \"./weights/pose_hrnet_w32_256x256.pth\", device=device)\n",
        "\n",
        "# # multiperson w/ YOLOv5 (medium), COCO weights\n",
        "# model = SimpleHRNet(48, 17, \"./weights/pose_hrnet_w48_384x288.pth\", yolo_version='v5', yolo_model_def='yolov5m', device=device)\n",
        "\n",
        "# multiperson w/ YOLOv5 nano, COCO weights, small model\n",
        "model = SimpleHRNet(32, 17, \"./weights/pose_hrnet_w32_256x192.pth\", yolo_version='v5', yolo_model_def='yolov5n', device=device)\n",
        "\n",
        "url = 'http://images.cocodataset.org/val2017/000000097278.jpg'\n",
        "im = Image.open(requests.get(url, stream=True).raw)\n",
        "image = io.imread(url)\n",
        "\n",
        "joints = model.predict(image)"
      ],
      "metadata": {
        "id": "xCXrjhfJUR5C"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%matplotlib inline\n",
        "from misc.visualization import joints_dict\n",
        "\n",
        "def plot_joints(ax, output):\n",
        "    bones = joints_dict()[\"coco\"][\"skeleton\"]\n",
        "    # bones = joints_dict()[\"mpii\"][\"skeleton\"]\n",
        "\n",
        "    for bone in bones:\n",
        "        xS = [output[:,bone[0],1], output[:,bone[1],1]]\n",
        "        yS = [output[:,bone[0],0], output[:,bone[1],0]]\n",
        "        ax.plot(xS, yS, linewidth=3, c=(0,0.3,0.7))\n",
        "    ax.scatter(joints[:,:,1],joints[:,:,0], s=20, c='r')\n",
        "\n",
        "fig = plt.figure(figsize=(60/2.54, 30/2.54))\n",
        "ax = fig.add_subplot(121)\n",
        "ax.imshow(Image.open(requests.get(url, stream=True).raw))\n",
        "ax = fig.add_subplot(122)\n",
        "ax.imshow(Image.open(requests.get(url, stream=True).raw))\n",
        "plot_joints(ax, joints)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "aYNkSzCGUqMF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Test the live script\n",
        "This step can be skipped if interested in the TensorRT conversion."
      ],
      "metadata": {
        "id": "ZWUN1C5RgGYS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # test the live script with default params (multiperson with yolo v3)\n",
        "# !python ./scripts/live-demo.py --filename WeAreGoingOnBullrun.mp4 --save_video\n",
        "\n",
        "# # test the live script with tiny yolo (v3)\n",
        "# !python ./scripts/live-demo.py --filename WeAreGoingOnBullrun.mp4 --save_video --use_tiny_yolo\n",
        "\n",
        "# # test the live script with yolo v5\n",
        "# !python ./scripts/live-demo.py --filename WeAreGoingOnBullrun.mp4 --save_video --yolo_version v5\n",
        "\n",
        "# test the live script with tiny yolo v5 (tensorrt yolo v5)\n",
        "!python ./scripts/live-demo.py --filename WeAreGoingOnBullrun.mp4 --save_video --yolo_version v5 --use_tiny_yolo"
      ],
      "metadata": {
        "id": "VEPfVe2bg1dS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now check out the video output.avi\n"
      ],
      "metadata": {
        "id": "RsTTv7A5gGvF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## TensorRT\n",
        "This section install TensorRT 8.5, converts the model to TensorRT (.engine) and tests the converted model.\n",
        "\n",
        "Tested with TensorRT 8.5.1-1+cuda11.8 and python package tensorrt 8.5.1.7 ."
      ],
      "metadata": {
        "id": "YHj3FQEyf1yD"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install TensorRT\n",
        "A GPU is needed for this step. Please change the runtime type to \"GPU\".\n"
      ],
      "metadata": {
        "id": "VsFWYxaNc-gl"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LsAlxRGVXhrt"
      },
      "outputs": [],
      "source": [
        "# check a GPU runtime is selected\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb\n",
        "\n",
        "dpkg -i nvidia-machine-learning-repo-*.deb\n",
        "apt-get update\n",
        "\n",
        "sudo apt-get install libnvinfer8 python3-libnvinfer"
      ],
      "metadata": {
        "id": "9vZ35qN5XkHE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# check TensorRT version\n",
        "print(\"TensorRT version: \")\n",
        "!dpkg -l | grep nvinfer"
      ],
      "metadata": {
        "id": "GlGh_J2WYH8u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# install TensorRT for python\n",
        "!pip install tensorrt"
      ],
      "metadata": {
        "id": "nhzVoykoYAWJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# clone the converion tool torch2trt\n",
        "%cd /content\n",
        "!git clone https://github.com/NVIDIA-AI-IOT/torch2trt"
      ],
      "metadata": {
        "id": "NUR0P_HklFbz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# install torch2trt\n",
        "%cd /content/torch2trt\n",
        "!python setup.py install"
      ],
      "metadata": {
        "id": "Y97nln2AX35c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/simple-HRNet"
      ],
      "metadata": {
        "id": "UC-xqiy5X5vk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Export the model with tensorrt"
      ],
      "metadata": {
        "id": "I2u6Xn72eEBE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Convert the smaller HRNet model to TensorRT - it may take a while...\n",
        "!python scripts/export-tensorrt-model.py --half \\\n",
        "     --weights \"./weights/pose_hrnet_w32_256x192.pth\" --hrnet_c 32 --image_resolution '(256, 192)'"
      ],
      "metadata": {
        "id": "S57JsLacdnoF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### [Optional] Export yolov5 with TensorRT"
      ],
      "metadata": {
        "id": "ckdDXNJzmxt_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Optional - Convert yolov5 (nano) to tensorrt too\n",
        "!python /root/.cache/torch/hub/ultralytics_yolov5_master/export.py --weights yolov5n.pt --include engine --device 0 --half"
      ],
      "metadata": {
        "id": "3Hls1HlCl44F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Test the tensorrt model"
      ],
      "metadata": {
        "id": "npgGj4cGemZd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Run inference with the converted TensorRT model\n",
        "!python scripts/live-demo.py --enable_tensorrt --filename=WeAreGoingOnBullrun.mp4 --hrnet_weights='weights/hrnet_trt.engine' \\\n",
        "  --hrnet_c 32 --image_resolution \"(256, 192)\" --yolo_version v5 --use_tiny_yolo --save_video\n"
      ],
      "metadata": {
        "id": "LnIpbqV0fVps"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now check out the video output.avi\n"
      ],
      "metadata": {
        "id": "WbQk0PeNnM5-"
      }
    }
  ]
}