Object-Detection-using-YOLOv5-and-OpenCV-DNN-in-CPP-and-Python/Convert_PyTorch_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h0Z6pGDENfka"
   },
   "source": [
    "# Convert PyTorch Models to ONNX\n",
    "Use this notebook to convert the out of the box PyTorch models to other formats."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zAM4YNS4M3iC"
   },
   "source": [
    "## Clone Ultralytics/yolov5 Repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IZ34ujwgFs45"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/ultralytics/yolov5\n",
    "%cd yolov5\n",
    "!pip install -r requirements.txt\n",
    "!pip install onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch=1.11 torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "guS1HQ4EMzKI"
   },
   "source": [
    "## Download YOLOv5 PyTorch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QapdevoQObTg"
   },
   "outputs": [],
   "source": [
    "%cd models\n",
    "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt -nv\n",
    "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt -nv\n",
    "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5m.pt -nv\n",
    "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5l.pt -nv\n",
    "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5x.pt -nv\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zeAU7NDTNEZa"
   },
   "source": [
    "## Export to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The default input size is 640x640.\n",
    "!python export.py --weights models/yolov5n.pt --include onnx\n",
    "!python export.py --weights models/yolov5s.pt --include onnx\n",
    "!python export.py --weights models/yolov5m.pt --include onnx\n",
    "!python export.py --weights models/yolov5l.pt --include onnx\n",
    "!python export.py --weights models/yolov5x.pt --include onnx\n",
    "\n",
    "# Add the flag --imsz to export for custom input size.\n",
    "# !python export.py --weights models/yolov5n.pt --include onnx --imsz 320 320\n",
    "\n",
    "# Add the flag --dyanamic for dynamic input size. Compatible with ONNX runtime.\n",
    "# !python export.py --weights models/yolov5n.pt --include onnx --dynamic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CO79jn9CNIw4"
   },
   "source": [
    "## Download the ONNX Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ITUKhgc2P4T6"
   },
   "outputs": [],
   "source": [
    "# For colab environment.\n",
    "# from google.colab import files\n",
    "# files.download('models/yolov5n.onnx')\n",
    "# files.download('models/yolov5s.onnx')\n",
    "# files.download('models/yolov5m.onnx')\n",
    "# files.download('models/yolov5l.onnx')\n",
    "# files.download('models/yolov5x.onnx')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Convert_PyTorch_models.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "yolov5-working-torch-1.11",
   "language": "python",
   "name": "yolov5-recheck"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}