converting-to-tfjs.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "idaR4wIXOmW3"
      },
      "source": [
        "Pytorch to ONNX with NMS (and inference)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ourPN52-Oi13",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6a8cf019-9940-411c-e6f7-a7aebefafaa4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Namespace(weights='runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.pt', img_size=[640, 640], batch_size=1, dynamic=False, dynamic_batch=False, grid=True, end2end=False, max_wh=640, topk_all=100, iou_thres=0.65, conf_thres=0.35, device='cpu', simplify=False, include_nms=False, fp16=False, int8=False)\n",
            "YOLOR 🚀 v0.1-122-g3b41c2c torch 1.13.1+cu116 CPU\n",
            "\n",
            "Fusing layers... \n",
            "/usr/local/lib/python3.9/dist-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3190.)\n",
            "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
            "Model Summary: 200 layers, 6006646 parameters, 6006646 gradients, 13.0 GFLOPS\n",
            "\n",
            "Starting TorchScript export with torch 1.13.1+cu116...\n",
            "/content/yolov7/models/yolo.py:52: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  if self.grid[i].shape[2:4] != x[i].shape[2:4]:\n",
            "TorchScript export success, saved as runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.torchscript.pt\n",
            "CoreML export failure: No module named 'coremltools'\n",
            "\n",
            "Starting TorchScript-Lite export with torch 1.13.1+cu116...\n",
            "TorchScript-Lite export success, saved as runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.torchscript.ptl\n",
            "\n",
            "Starting ONNX export with onnx 1.13.1...\n",
            "/content/yolov7/models/yolo.py:582: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  if augment:\n",
            "/content/yolov7/models/yolo.py:614: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  if profile:\n",
            "/content/yolov7/models/yolo.py:629: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  if profile:\n",
            "ONNX export success, saved as runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.onnx\n",
            "\n",
            "Export complete (6.85s). Visualize with https://github.com/lutzroeder/netron.\n"
          ]
        }
      ],
      "source": [
        "!python export.py --weights runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.pt  \\\n",
        "        --grid  \\\n",
        "        --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 \\\n",
        "        --img-size 640 640 --max-wh 640 # For onnxruntime, you need to specify this value as an integer, when it is 0 it means agnostic NMS, \n",
        "                     # otherwise it is non-agnostic NMS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CBtmSC6bO-BJ"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import onnx\n",
        "from scc4onnx import order_conversion\n",
        "from onnx_tf.backend import prepare"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d7zVbNixPTqu"
      },
      "outputs": [],
      "source": [
        "onnx_model_path = 'runs/train/yolov7_tiny_stockout/weights/yolov7-tiny.onnx'\n",
        "onnx_model = onnx.load(onnx_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHku3DBGRLmG"
      },
      "source": [
        "Prepare ONNX model for Tensorflow Backend"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf_rep = prepare(onnx_model)"
      ],
      "metadata": {
        "id": "rgCgUbpoiXIj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "so_qtUfERTYr"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/yolov7/runs/train/yolov7_tiny_stockout/weights/tf\n",
        "!mkdir /content/yolov7/runs/train/yolov7_tiny_stockout/weights/tf/stockout_web_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UX1npmFmRHWD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2336e0d3-f39c-4a92-d634-22955ced749c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.9/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
            "Instructions for updating:\n",
            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n",
            "WARNING:absl:Found untraced functions such as gen_tensor_dict while saving (showing 1 of 1). These functions will not be directly callable after loading.\n"
          ]
        }
      ],
      "source": [
        "tf_model_dir = \"/content/yolov7/runs/train/yolov7_tiny_stockout/weights/tf\"\n",
        "tf_rep.export_graph(tf_model_dir)\n",
        "tfjs_model_dir = f\"{tf_model_dir}/stockout_web_model\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6avvANVuRgUq"
      },
      "source": [
        "Export SavedModel to TFJS format"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZEAoB3NZReuG"
      },
      "outputs": [],
      "source": [
        "tfjs_convert_command = f\"\"\"tensorflowjs_converter\n",
        "                 --input_format=tf_saved_model \n",
        "                 --output_format=tfjs_graph_model \n",
        "                 --signature_name=serving_default \n",
        "                 --saved_model_tags=serve \n",
        "                 \"{tf_model_dir}\" \n",
        "                 \"{tfjs_model_dir}\"\n",
        "                 \"\"\"\n",
        "tfjs_convert_command = \" \".join(tfjs_convert_command.split())"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%sx $tfjs_convert_command"
      ],
      "metadata": {
        "id": "5VSr81HekujK",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "361a2f17-e1fb-48e2-d3f5-494dfe145dc8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[\"2023-03-19 18:41:04.952227: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib/python3.9/dist-packages/cv2/../../lib64:/usr/lib64-nvidia\",\n",
              " \"2023-03-19 18:41:04.952324: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib/python3.9/dist-packages/cv2/../../lib64:/usr/lib64-nvidia\",\n",
              " '2023-03-19 18:41:04.952343: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.',\n",
              " '2023-03-19 18:41:07.107393: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.',\n",
              " 'Writing weight file /content/yolov7/runs/train/yolov7_tiny_stockout/weights/tf/stockout_web_model/model.json...',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant with shape (5,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_4 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros with shape () and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant_5 with shape (3,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where_1 with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_5 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2_1 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros_1 with shape () and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant_6 with shape (5,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where_2 with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_6 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2_2 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros_2 with shape () and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant_11 with shape (3,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where_3 with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_7 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2_3 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros_3 with shape () and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant_12 with shape (5,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where_4 with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_8 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2_4 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros_4 with shape () and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/onnx_tf_prefix_/model.77/Constant_17 with shape (3,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Where_5 with shape (0, 1) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/Cast_9 with shape (1,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/GatherV2_5 with shape (0,) and dtype int64 was auto converted to the type int32',\n",
              " 'weight PartitionedCall/zeros_5 with shape () and dtype int64 was auto converted to the type int32']"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NC9UXQ5kSiq0"
      },
      "outputs": [],
      "source": [
        "!tar vcfz {tfjs_model_dir}.tar.gz {tfjs_model_dir}"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the converted tf.js model"
      ],
      "metadata": {
        "id": "G1N09Xo8ygKL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3tBRqch5SP1H",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "outputId": "a6d40650-d64d-4b16-9fa1-21c49909a9d8"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "download(\"download_08775c69-59f9-4250-a5aa-aafd5b21be35\", \"stockout_web_model.tar.gz\", 22373448)"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "from google.colab import files\n",
        "files.download(tfjs_model_dir + '.tar.gz') "
      ]
    }
  ]
}