classification_tf2/byom_person/stm32ai_tao_mobilenetv2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TAO BYOM Image Classification (TF2) with STEdgeAI Developer Cloud\n",
    "\n",
    "Transfer learning is the process of transfering learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "This notebook provides a complete life cycle of the model training, optimization and benchmarking using [NVIDIA TAO Toolkit](https://developer.nvidia.com/tao-toolkit) and [STEdgeAI Developer Cloud](https://stm32ai.st.com/stm32-cube-ai-dc/).\n",
    "\n",
    "NVIDIA Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "[STEdgeAI Developer Cloud](https://stm32ai-cs.st.com/home) is a free-of-charge online platform and services allowing the creation, optimization, benchmarking, and generation of AI models for the STM32 microcontrollers. It is based on the [STEdgeAI](https://stm32ai.st.com/stm32-cube-ai/) core technology.\n",
    "\n",
    "**Bring Your Own Model (BYOM)** is a Python based package that converts any open source ONNX model to TAO compatible model. All you need is to export any model from a deep learning framework of your choice (e.g. PyTorch) to ONNX and run TAO BYOM converter. More details of installing and running TAO BYOM converter can be found [here](https://github.com/NVIDIA-AI-IOT/tao_byom_examples)\n",
    "\n",
    "\n",
    "<br>\n",
    "\n",
    "<img style=\"float: center;background-color: white; width: 1080\" src=\"../../TAO-STM32CubeAI.png\" width=\"1080\">\n",
    "\n",
    "<br>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## License\n",
    "\n",
    "This software component is licensed by ST under BSD-3-Clause license,\n",
    "the \"License\"; \n",
    "\n",
    "You may not use this file except in compliance with the\n",
    "License. \n",
    "\n",
    "You may obtain a copy of the License at: https://opensource.org/licenses/BSD-3-Clause\n",
    "\n",
    "Copyright (c) 2023 STMicroelectronics. All rights reserved.\n",
    "\n",
    "Copyright (c) 2023 Nvidia. All rights reserved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained mobilenetv2 model and finetune on a sample dataset converted from COCO2014 to perform person detection,\n",
    "* Prune the finetuned model,\n",
    "* Retrain the pruned model to recover lost accuracy,\n",
    "* Export the pruned model as an onnx model,\n",
    "* Quantize the model using onnxruntime,\n",
    "* Run Benchmarking of the quantized onnx model (finetuned, pruned, retrained, and quantized) using STEdgeAI Developer Cloud to know the footprints and embeddability of the models.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `classification` model which was imported from outside TAO Toolkit, and that may be deployed via [STEdgeAI Developer Cloud](https://stm32ai-cs.st.com/home).\n",
    "\n",
    "### Table of Contents\n",
    "This notebook shows an example use case for classification using the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO Launcher](#head-1)\n",
    "2. [Prepare dataset and pretrained model](#head-2)\n",
    "    1. [Download, prepare and split the dataset into train/test/val](#head-2-1)\n",
    "    2. [Check the BYOM model](#head-2-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Finetune the pretrained model using TAO training](#head-4)\n",
    "5. [Evaluate trained model](#head-5)\n",
    "    1. [(optional) Export the trained model as onnx format and check the accuracy](#head-5-1)\n",
    "6. [Prune the trained model](#head-6)\n",
    "7. [Retrain the pruned model](#head-7)\n",
    "8. [Testing the finetuned pruned model](#head-8)\n",
    "    1. [Export the pruned, and retrained model as onnx format](#head-8-1)\n",
    "    2. [Quantizing the exported onnx model using onnxruntime](#head-8-2)\n",
    "9. [Benchmarking the optimized model using STEdgeAI Developer Cloud for embeddability](#head-9)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/classification_tf2`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env KEY=nvidia_tlt\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/classification_tf2\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/classification_tf2\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/classification_tf2/byom_person/results_mobilenetv2/\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "os.environ[\"LOCAL_PROJECT_DIR\"] = \"/home/user/stm32ai-tao/\"\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.environ.get(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.environ.get(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"classification_tf2\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.environ.get(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs_mobilenetv2\"\n",
    ")\n",
    "\n",
    "%env SPECS_DIR=/workspace/tao-experiments/classification_tf2/byom_person/specs_mobilenetv2\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from outside to inside of the docker instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\":{\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (Optional) A. Set proxy variables if working behind corporate proxies.\n",
    "\n",
    "The following section sets the proxies and ssl verification flag when the users are working behind the proxies. This setup is necessary to be able to communicate with internet.\n",
    "\n",
    "Replace the `userName`, `password`, and `proxy_port` with your correct username, password and proxy port."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set proxies\n",
    "import os\n",
    "# os.environ[\"http_proxy\"]='http://userName:password.@example.com:proxy_port'\n",
    "# os.environ[\"https_proxy\"] = 'http://userName:password.@example.com:proxy_port'\n",
    "# os.environ[\"NO_SSL_VERIFY\"]=\"1\"\n",
    "# os.environ[\"SSL_VERIFY\"]=\"False\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. However, **to perform onnx quantization and evaluation using the provided scripts, we recommend to use Python version 3.8.16**. This notebooks is tested for Python version 3.8.16. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8.16\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.6.9 <= 3.8.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this cell IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare datasets and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE**: If you have already downloaded, unpacked and prepared the dataset files once, you can skip these steps.\n",
    "\n",
    "We will be using the modified version of COCO2014 dataset for the tutorial. To find more details please visit this [link](https://pjreddie.com/projects/coco-mirror/). "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download the dataset\n",
    "To download all the files needed for the dataset in the right location, please uncomment and run the section below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir $LOCAL_DATA_DIR\n",
    "!wget -O $LOCAL_DATA_DIR/train2014.zip https://pjreddie.com/media/files/train2014.zip\n",
    "!wget -O $LOCAL_DATA_DIR/val2014.zip https://pjreddie.com/media/files/val2014.zip\n",
    "!wget -O $LOCAL_DATA_DIR/labels.tgz https://pjreddie.com/media/files/coco/labels.tgz"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Verify the download.\n",
    "Checking if the dataset zip files are present in the data directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that file is present\n",
    "import os\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "print(DATA_DIR)\n",
    "if not ( os.path.isfile(os.path.join(DATA_DIR , 'train2014.zip')) and \n",
    "        os.path.isfile(os.path.join(DATA_DIR , 'val2014.zip')) and\n",
    "        os.path.isfile(os.path.join(DATA_DIR , 'labels.tgz')) ):\n",
    "    print('One or more data files for the dataset are not found.\\nPlease download the dataset by running the Download Dataset section!')\n",
    "else:\n",
    "    print('Found dataset.')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Unpack the files\n",
    "\n",
    "The downloaded files are in the form of the `zip` and `tgz` format. Running the following code section will unzip and unpack these files.\n",
    "\n",
    "**NOTE** : `> /dev/null` executes the command silently without printing the output of the command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip $LOCAL_DATA_DIR/train2014.zip -d $LOCAL_DATA_DIR/> /dev/null\n",
    "!unzip $LOCAL_DATA_DIR/val2014.zip -d $LOCAL_DATA_DIR/> /dev/null\n",
    "!tar -xzvf $LOCAL_DATA_DIR/labels.tgz -C $LOCAL_DATA_DIR/> /dev/null"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verifying if the files are unpacked as folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $LOCAL_DATA_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Split the dataset into train/val/test <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For creating the person detection use case we are converting the COCO2014 Dataset into a format where it has only two classes, i.e. `person` and `not_person`. \n",
    "In addition to this we are applying an additional filter and removing all the images with person class where the size of the person is too small (covering less than 20% of the image area). \n",
    "That is why after preparation instead of 118,287 we have only 84,810 images in total."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install pip requirements\n",
    "!pip3 install tqdm\n",
    "!pip3 install matplotlib==3.3.3\n",
    "!pip3 install pandas"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following code section filters the dataset into `person` and `not_person` classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "def filter_coco(area_threshold, labels_dir, input_dir, output_dir):\n",
    "    \"\"\"Filter COCO dataset subset filtering person area.\n",
    "\n",
    "    Args:\n",
    "      area_threshold: Threshold of fraction of image area below which\n",
    "      persons are filtered.\n",
    "      labels_dir: COCO dataset labels directory path.\n",
    "      input_dir: COCO dataset path.\n",
    "      output_dir: new dataset output path.\n",
    "    \"\"\"\n",
    "    labels_dpath = Path(labels_dir)\n",
    "    labels_fpaths = labels_dpath.glob('*.txt')\n",
    "    input_dpath = Path(input_dir)\n",
    "    output_dpath = Path(output_dir)\n",
    "    f = open(input_dpath.stem + '.txt','w+')\n",
    "    f.write('filename label\\n')\n",
    "    for label_fpath in labels_fpaths:\n",
    "        img_fname = label_fpath.name.replace('.txt', '.jpg')\n",
    "        label = 'not_person'\n",
    "        annotations = pd.read_csv(label_fpath, delimiter=' ', header=None)\n",
    "        persons = annotations.loc[annotations[0] == 0]\n",
    "        if persons.shape[0] != 0:\n",
    "            big_persons = ((persons[3] * persons[4] * 100.0) > area_threshold).sum()\n",
    "            if big_persons > 0:\n",
    "                label = 'person'\n",
    "            else:\n",
    "                continue\n",
    "        src_fpath = input_dpath / img_fname\n",
    "        dst_dpath = output_dpath / label\n",
    "        dst_dpath.mkdir(parents=True, exist_ok=True)\n",
    "        shutil.copy(src_fpath, dst_dpath)\n",
    "        f.write(img_fname + ' ' + label + '\\n')\n",
    "    f.close()\n",
    "train_labels_path = os.path.join(DATA_DIR, 'labels', 'train2014')\n",
    "val_labels_path = os.path.join(DATA_DIR, 'labels', 'val2014')\n",
    "train_images_dir = os.path.join(DATA_DIR, 'train2014')\n",
    "val_images_dir = os.path.join(DATA_DIR, 'val2014')\n",
    "result_dir = os.path.join(DATA_DIR, 'person_dataset')\n",
    "filter_coco(area_threshold = 20.0, labels_dir = train_labels_path, input_dir = train_images_dir, output_dir = result_dir)\n",
    "filter_coco(area_threshold = 20.0, labels_dir = val_labels_path, input_dir = val_images_dir, output_dir = result_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Splitting the dataset into `train`, `val`, `test` portions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import shutil\n",
    "from random import shuffle\n",
    "from tqdm import tqdm\n",
    "\n",
    "DATA_DIR=os.environ[\"LOCAL_DATA_DIR\"]\n",
    "SOURCE_DIR=os.path.join(DATA_DIR, 'person_dataset')\n",
    "TARGET_DIR=os.path.join(DATA_DIR,'split')\n",
    "\n",
    "# removing existing split directory\n",
    "!rm -rf $TARGET_DIR\n",
    "\n",
    "# list dir\n",
    "print(os.walk(SOURCE_DIR))\n",
    "dir_list = next(os.walk(SOURCE_DIR))[1]\n",
    "# for each dir, create a new dir in split\n",
    "for dir_i in tqdm(dir_list):\n",
    "        newdir_train = os.path.join(TARGET_DIR, 'train', dir_i)\n",
    "        newdir_val = os.path.join(TARGET_DIR, 'val', dir_i)\n",
    "        newdir_test = os.path.join(TARGET_DIR, 'test', dir_i)\n",
    "        \n",
    "        if not os.path.exists(newdir_train):\n",
    "                os.makedirs(newdir_train)\n",
    "        if not os.path.exists(newdir_val):\n",
    "                os.makedirs(newdir_val)\n",
    "        if not os.path.exists(newdir_test):\n",
    "                os.makedirs(newdir_test)\n",
    "\n",
    "        img_list = glob.glob(os.path.join(SOURCE_DIR, dir_i, '*.jpg'))\n",
    "        # shuffle data\n",
    "        shuffle(img_list)\n",
    "\n",
    "        for j in range(int(len(img_list)*0.7)):\n",
    "                shutil.copy2(img_list[j], os.path.join(TARGET_DIR, 'train', dir_i))\n",
    "\n",
    "        for j in range(int(len(img_list)*0.7), int(len(img_list)*0.8)):\n",
    "                shutil.copy2(img_list[j], os.path.join(TARGET_DIR, 'val', dir_i))\n",
    "                \n",
    "        for j in range(int(len(img_list)*0.8), len(img_list)):\n",
    "                shutil.copy2(img_list[j], os.path.join(TARGET_DIR, 'test', dir_i))\n",
    "                \n",
    "print('Done splitting dataset.')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verifying if the portions are created and all the sub-folders have all the classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $LOCAL_DATA_DIR/split/train\n",
    "!ls $LOCAL_DATA_DIR/split/val\n",
    "!ls $LOCAL_DATA_DIR/split/test"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Check the BYOM model <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE**: Before running this step, please run the jupyter notebook [Preparing a TAO BYOM model](../byom_converters/byom_converter_mobilenetv2.ipynb) to prepare the `.tltb` model file as our pretrained model. Place the converted model directory to `$LOCAL_EXPERIMENT_DIR/pretrained_mobilenetv2/mobilenetv2*`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_mobilenetv2/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that BYOM model was converted into the dir.\")\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_mobilenetv2/mobilenetv2_128_0_5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Training dataset\n",
    "* Validation dataset\n",
    "* Pre-trained models\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/spec.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# logging in the docker session\n",
    "!docker login nvcr.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete existing resulting output folder and recreate\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_tf2 train  -e $SPECS_DIR/spec.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### INFO\n",
    "- To run this training in data parallelism using multiple GPU's, please pass the **number of gpu devices** to be used using `--gpus` parameter. For example for running the training on two gpu devices on parallel the training command will be:\n",
    " ```!tao classification_tf2 train -e $SPECS_DIR/spec.yaml --gpus 2```\n",
    "- The training can be intrupted and then relaunched at any point. To resume from a checkpoint, just relaunch training with the same spec file."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this step, we assume that the training is complete and the resulting weights for the model from the some epochs are available. If you would like to run evaluation on an earlier model, you can do so but choosing the resulting weights for that particular epoch, otherwise the scrip below will find the last checkpoint corresponding to the latest finished epoch and run the evaluation on that, by passing it to the configuration parameter `evaluate.model_path`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last checkpoints\n",
    "import os\n",
    "last_checkpoint = ''\n",
    "for f in os.listdir(os.path.join(os.environ[\"LOCAL_EXPERIMENT_DIR\"],'byom_person/results_mobilenetv2/output', 'train')):\n",
    "    if f.startswith('mobilenetv2'):\n",
    "        last_checkpoint = last_checkpoint if last_checkpoint > f else f\n",
    "print(f'Last checkpoint: {last_checkpoint}')\n",
    "%env LAST_CHECKPOINT={last_checkpoint}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# run the evaluation by overriding the model path for evaluate with the last checkpoint\n",
    "!tao model classification_tf2 evaluate -e $SPECS_DIR/spec.yaml \\\n",
    "                                 evaluate.checkpoint=\"$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output/train/$LAST_CHECKPOINT\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (optional) A. Exporting the trained model to onnx model \n",
    "This step is needed if one wants to estimate the footprints of the trained model using STEdgeAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./results_mobilenetv2/exports\n",
    "# export the model checkpoint as .onnx file\n",
    "!tao model classification_tf2 export -e $SPECS_DIR/spec.yaml\\\n",
    "                                 export.checkpoint=$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output/train/$LAST_CHECKPOINT \\\n",
    "                                 export.onnx_file=$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/exports/mobilenetv2.onnx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (optional) B. Evaluating the onnx model\n",
    "This is optional sanity check to evaluate the exported trained model to check if the preprocessing in the onnx evaluation and the TAO evaluation are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.9\n",
    "!pip install onnxruntime==1.14.1\n",
    "!pip install onnx==1.12.0\n",
    "!pip install scikit-learn==0.24.2\n",
    "!pip install matplotlib==3.3.3\n",
    "!pip install pillow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate the onnx model before pruning\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "import onnx_utils\n",
    "onnx_utils.evaluate_onnx_model('./results_mobilenetv2/exports/mobilenetv2.onnx', \n",
    "                                            os.path.join(os.environ[\"LOCAL_DATA_DIR\"],'split/test/'), save_path='./results_mobilenetv2/mobilenetv2')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion\n",
    "* Threshold for pruning\n",
    "* Exclude prediction layer that you don't want pruned (e.g. predictions)\n",
    "\n",
    "Usually, you just need to adjust `prune.threshold` for accuracy and model size trade off. Higher `threshold` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is depend on the dataset. `0.5` is just a starting point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.\n",
    "\n",
    "`prune.min_num_filters` is the minimum number of filters to keep per layer after the pruning step. Smaller the value smaller the resulting model. However, this might lower the accuracy. The users can adjust this value depending on their needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_tf2 prune -e $SPECS_DIR/spec.yaml \\\n",
    "                                prune.checkpoint=\"$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output/train/$LAST_CHECKPOINT\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Pruned model:')\n",
    "print('------------')\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output/prune"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification\n",
    "\n",
    "**NOTE**: The path of the pruned model is created based on the `prune.threshold` value automatically. The default path provided in the `spec_retrain.yaml` file is for the value used in this notebook. Make sure to update the path if you change the value of this parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete and recreate the directory for the retraining results\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output_retrain\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_tf2 train  -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Testing the model! <a class=\"anchor\" id=\"head-8\"></a>\n",
    "\n",
    "In this step, we assume that the training is complete and the model from the final epoch (`mobilenetv2_xxx.tlt`) is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last checkpoints\n",
    "last_checkpoint = ''\n",
    "for f in os.listdir(os.path.join(os.environ[\"LOCAL_EXPERIMENT_DIR\"],'byom_person/results_mobilenetv2/output_retrain', 'train')):\n",
    "    if f.startswith('mobilenetv2'):\n",
    "        last_checkpoint = last_checkpoint if last_checkpoint > f else f\n",
    "print(f'Last checkpoint: {last_checkpoint}')\n",
    "%env LAST_CHECKPOINT={last_checkpoint}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_tf2 evaluate -e $SPECS_DIR/spec_retrain.yaml \\\n",
    "    evaluate.checkpoint=\"$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output_retrain/train/$LAST_CHECKPOINT\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Exporting the trained model as onnx model <a class=\"anchor\" id=\"head-8-1\"></a>\n",
    "\n",
    "The following section exports the pruned and retrained model as an onnx model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./results_mobilenetv2/exports\n",
    "\n",
    "!tao model classification_tf2 export -e $SPECS_DIR/spec_retrain.yaml\\\n",
    "                                 export.checkpoint=$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/output_retrain/train/$LAST_CHECKPOINT \\\n",
    "                                 export.onnx_file=$USER_EXPERIMENT_DIR/byom_person/results_mobilenetv2/exports/pruned_mobilenetv2.onnx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A.1. Evaluating the exported onnx model\n",
    "\n",
    "The following section evaluates the exported onnx model. This should get the same accuracy as the one before converting to onnx."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.9\n",
    "!pip install onnxruntime==1.14.1\n",
    "!pip install onnx==1.12.0\n",
    "!pip install scikit-learn==0.24.2\n",
    "!pip install matplotlib==3.3.3\n",
    "!pip install pillow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate the onnx model before pruning\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "import onnx_utils\n",
    "onnx_utils.evaluate_onnx_model('./results_mobilenetv2/exports/pruned_mobilenetv2.onnx', \n",
    "                                            os.path.join(os.environ[\"LOCAL_DATA_DIR\"],'split/test/'), save_path='./results_mobilenetv2/mobilenetv2_pruned')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Quantizing the exported onnx model using onnxruntime <a class=\"anchor\" id=\"head-8-2\"></a>\n",
    "\n",
    "The following sections converts the exported onnx model to int8 to reduce the footprints and improve the inference time.\n",
    "This will first require to create a subsample of the data to calibrate the quantization of the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B.1. Create a calibration dataset\n",
    "This should have samples from both classes `person`, and `not_person`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a subsample dataset for the quantization\n",
    "import os\n",
    "import glob\n",
    "import shutil\n",
    "from random import shuffle\n",
    "from tqdm import tqdm\n",
    "\n",
    "DATA_DIR=os.environ[\"LOCAL_DATA_DIR\"]\n",
    "SOURCE_DIR=os.path.join(DATA_DIR, 'split/train/')\n",
    "TARGET_DIR=os.path.join(DATA_DIR,'split/subset_calibration_dataset')\n",
    "\n",
    "!rm -rf TARGET_DIR\n",
    "\n",
    "samples_per_class = 200\n",
    "# list dir\n",
    "print(os.walk(SOURCE_DIR))\n",
    "dir_list = next(os.walk(SOURCE_DIR))[1]\n",
    "# for each dir, create a new dir in split\n",
    "for dir_i in tqdm(dir_list):\n",
    "        if not os.path.exists(TARGET_DIR):\n",
    "                os.makedirs(TARGET_DIR)\n",
    "                \n",
    "        img_list = glob.glob(os.path.join(SOURCE_DIR, dir_i, '*.jpg'))\n",
    "        # shuffle data\n",
    "        shuffle(img_list)\n",
    "\n",
    "        for j in range(samples_per_class):\n",
    "                shutil.copy2(img_list[j], os.path.join(TARGET_DIR))                \n",
    "print('Done creating calibration dataset.')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### B.2. Quantize the model using QDQ quantization to int8 weights\n",
    "The following section quantize the float32 onnx model to int8 quantized onnx model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quantize the model\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "import onnx_utils\n",
    "onnx_utils.quantize_onnx_model(input_model = './results_mobilenetv2/exports/pruned_mobilenetv2.onnx', \n",
    "                              calibration_dataset_path = os.path.join(os.environ[\"LOCAL_DATA_DIR\"],'split','subset_calibration_dataset'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### B.3. Evaluate the quantized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate the onnx model before pruning\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "import onnx_utils\n",
    "onnx_utils.evaluate_onnx_model('./results_mobilenetv2/exports/pruned_mobilenetv2_QDQ_quant.onnx', \n",
    "                                            os.path.join(os.environ[\"LOCAL_DATA_DIR\"],'split/test/'), save_path='./results_mobilenetv2/mobilenetv2_pruned_quantized')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### B.4. Changing the opset of the onnx model to use them with STEdgeAI\n",
    "The following section will change the opset of the exported onnx models to make sure it is best supported by STEdgeAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "import onnx_utils\n",
    "onnx_utils.change_opset('./results_mobilenetv2/exports/mobilenetv2.onnx', target_opset = 14)\n",
    "onnx_utils.change_opset('./results_mobilenetv2/exports/pruned_mobilenetv2.onnx', target_opset = 14)\n",
    "onnx_utils.change_opset('./results_mobilenetv2/exports/pruned_mobilenetv2_QDQ_quant.onnx', target_opset = 14)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Benchmarking the optimzed model using STEdgeAI Developer Cloud <a class=\"anchor\" id=\"head-9\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Getting the package for connecting to STEdgeAI Developer Cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gitdir\n",
    "!pip install ipywidgets\n",
    "!pip install marshmallow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gitdir https://github.com/STMicroelectronics/stm32ai-modelzoo-services/tree/main/common/stm32ai_dc\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "# Reorganize local folders\n",
    "if os.path.exists('./stm32ai_dc'):\n",
    "    shutil.rmtree('./stm32ai_dc')\n",
    "shutil.move('./common/stm32ai_dc', './stm32ai_dc')\n",
    "shutil.rmtree('./common')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Import, helper and UI functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from typing import List\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import ipywidgets as widgets\n",
    "from stm32ai_dc import (CliLibraryIde, CliLibrarySerie, CliParameters, MpuParameters, MpuEngine, AtonParameters,\n",
    "                        CloudBackend, Stm32Ai)\n",
    "\n",
    "sys.path.append(os.path.abspath('stm32ai'))\n",
    "os.environ['STATS_TYPE'] = 'stm32ai_tao'\n",
    "\n",
    "# create a directory for outputs for stm32ai developer cloud operations\n",
    "stm32ai_output_dir = './results_mobilenetv2/stm32ai_outputs'\n",
    "os.makedirs(stm32ai_output_dir, exist_ok=True)\n",
    "\n",
    "os.makedirs('models', exist_ok=True)\n",
    "os.makedirs('outputs', exist_ok=True)\n",
    "\n",
    "\n",
    "def get_mpu_options(board_name: str = None) -> tuple:\n",
    "    \"\"\"\n",
    "    Get MPU benchmark options depending on MPU board selected\n",
    "    Each MPU board has different settings,\n",
    "    i.e. different number of cpu_cores or engine (CPU only or HW_Accelerator also)\n",
    "    Input:\n",
    "        board_name:str, name of the mpu board\n",
    "    Returns:\n",
    "        tuple: engine_used and num_cpu_cores.\n",
    "    \"\"\"\n",
    "\n",
    "    #define configuration by MPU board\n",
    "    STM32MP257F_EV1 = {\n",
    "        \"engine\": MpuEngine.HW_ACCELERATOR,\n",
    "        \"cpu_cores\": 2\n",
    "    }\n",
    "\n",
    "    STM32MP157F_DK2 = {\n",
    "        \"engine\": MpuEngine.CPU,\n",
    "        \"cpu_cores\": 2\n",
    "    }\n",
    "\n",
    "    STM32MP135F_DK = {\n",
    "        \"engine\": MpuEngine.CPU,\n",
    "        \"cpu_cores\": 1\n",
    "    }\n",
    "\n",
    "    #recover parameters based on board name:\n",
    "    if board_name == \"STM32MP257F-EV1\":\n",
    "        engine_used = STM32MP257F_EV1.get(\"engine\")\n",
    "        num_cpu_cores = STM32MP257F_EV1.get(\"cpu_cores\")\n",
    "    elif board_name == \"STM32MP157F-DK2\":\n",
    "        engine_used = STM32MP157F_DK2.get(\"engine\")\n",
    "        num_cpu_cores = STM32MP157F_DK2.get(\"cpu_cores\")\n",
    "    elif board_name == \"STM32MP135F-DK\":\n",
    "        engine_used = STM32MP135F_DK.get(\"engine\")\n",
    "        num_cpu_cores = STM32MP135F_DK.get(\"cpu_cores\")\n",
    "    else :\n",
    "        engine_used = MpuEngine.CPU\n",
    "        num_cpu_cores = 1\n",
    "\n",
    "    return engine_used, num_cpu_cores\n",
    "\n",
    "def analyze_footprints(report: object = None) -> None:\n",
    "    \"\"\"\n",
    "    Analyzes the memory footprint of a STEdgeAI model.\n",
    "\n",
    "    Args:\n",
    "        report: A report object containing information about the model.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    activations_ram: float = report.ram_size / 1024\n",
    "    runtime_ram: float = report.estimated_library_ram_size / 1024\n",
    "    total_ram: float = activations_ram + runtime_ram\n",
    "    weights_rom: float = report.rom_size / 1024\n",
    "    code_rom: float = report.estimated_library_flash_size / 1024\n",
    "    total_flash: float = weights_rom + code_rom\n",
    "    macc: float = report.macc / 1e6\n",
    "    print(\"[INFO] : STEdgeAI model memory footprint\")\n",
    "    print(\"[INFO] : MACCs : {} (M)\".format(macc))\n",
    "    print(\"[INFO] : Total Flash : {0:.1f} (KiB)\".format(total_flash))\n",
    "    print(\"[INFO] :     Flash Weights  : {0:.1f} (KiB)\".format(weights_rom))\n",
    "    print(\"[INFO] :     Estimated Flash Code : {0:.1f} (KiB)\".format(code_rom))\n",
    "    print(\"[INFO] : Total RAM : {0:.1f} (KiB)\".format(total_ram))\n",
    "    print(\"[INFO] :     RAM Activations : {0:.1f} (KiB)\".format(activations_ram))\n",
    "    print(\"[INFO] :     RAM Runtime : {0:.1f} (KiB)\".format(runtime_ram))\n",
    "\n",
    "\n",
    "def benchmark_model(stmai:object,\n",
    "                    model_path:str,\n",
    "                    model_name:str,\n",
    "                    optimization:str,\n",
    "                    from_model:str,\n",
    "                    board_name:str,\n",
    "                    allocateInputs:bool =True,\n",
    "                    allocateOutputs:bool=True) -> float:\n",
    "    \"\"\"\n",
    "    Benchmarks the give model to calculate the footprint on a STM32 Target board.\n",
    "\n",
    "    Args:\n",
    "        stmai:object, an object of stm32ai_dc\n",
    "        model_path:str, path to the model file\n",
    "        model_name:str, path to the model file\n",
    "        optimization:str, the way model is to be optimized available options ['balanced', 'time', 'ram']\n",
    "        from_model:str, if the model is coming from zoo or is a custom model from the user\n",
    "        board_name:str, target board name from one of the available boards on the dev cloud\n",
    "        allocateInputs:bool, If set to true activations buffer will be also used to handle the input buffers. \n",
    "        allocateOutputs:bool, If set to \"True\", activations buffer will be also used to handle the output buffers.\n",
    "\n",
    "    Returns:\n",
    "        fps: frames per second (1/inference_time)\n",
    "    \"\"\"\n",
    "    print(f\"Benchmarking on: {board_name}\")\n",
    "    if \"mp\" in board_name.lower():\n",
    "        # if mpu is selected as the target\n",
    "        model_extension = os.path.splitext(model_path)[1]\n",
    "        # only supported options are quantized tflite or onnx models\n",
    "        if model_extension in ['.onnx', '.tflite']:\n",
    "            if \"stm32mp2\" in board_name.lower(): # if mp2 is selected as the target board optimize the model to generate a .nbg file\n",
    "                optimized_model_path = os.path.dirname(model_path) + \"/\"\n",
    "                try:\n",
    "                    stmai.upload_model(model_path)\n",
    "                    model = model_name\n",
    "                    res = stmai.generate_nbg(model)\n",
    "                    stmai.download_model(res, optimized_model_path + res)\n",
    "                    model_path=os.path.join(optimized_model_path,res)\n",
    "                    nb_model_name = os.path.splitext(os.path.basename(model_path))[0] + \".nb\"\n",
    "                    rename_model_path=os.path.join(optimized_model_path,nb_model_name)\n",
    "                    os.rename(model_path, rename_model_path)\n",
    "                    model_path = rename_model_path\n",
    "                    print(\"[INFO] : Optimized Model Name:\", model_name)\n",
    "                    print(\"[INFO] : Optimization done ! Model available at :\",optimized_model_path)\n",
    "                    model_name = nb_model_name\n",
    "                except Exception as e:\n",
    "                    print(f\"[FAIL] : Model optimization via Cloud failed : {e}.\")\n",
    "                    print(\"[INFO] : Use default model instead of optimized ...\")\n",
    "        else:\n",
    "            print(\"[ERROR]: Only .tflite or .onnx models can be benchmarked for MPU\")\n",
    "            fps = 0\n",
    "            return fps\n",
    "\n",
    "        engine, nbCores = get_mpu_options(board_name)\n",
    "        stmai_params = MpuParameters(model=model_name,\n",
    "                                     nbCores=nbCores,\n",
    "                                     engine=engine)\n",
    "\n",
    "    elif 'stm32n6' in board_name.lower():\n",
    "        model_extension = os.path.splitext(model_path)[1]\n",
    "        if model_extension in ['.onnx', '.tflite']:\n",
    "            stmai_params = CliParameters(model=model_name,\n",
    "                                         target='stm32n6',\n",
    "                                         stNeuralArt='default',\n",
    "                                         atonnOptions=AtonParameters())\n",
    "        else:\n",
    "            print(\"[ERROR]: Only .tflite or .onnx models can be benchmarked for N6\")\n",
    "        \n",
    "    else:\n",
    "        # target board in mcu, prepare stm32ai parameters\n",
    "        stmai_params = CliParameters(model=model_name,\n",
    "                                     optimization=optimization,\n",
    "                                     allocateInputs=allocateInputs,\n",
    "                                     allocateOutputs=allocateOutputs,\n",
    "                                     fromModel=from_model)\n",
    "    # running the benchmarking with prepared params\n",
    "    \n",
    "    benchmark_report_dir = f'{stm32ai_output_dir}/benchmark_reports/'\n",
    "    os.makedirs(benchmark_report_dir, exist_ok=True)\n",
    "\n",
    "\n",
    "    try:\n",
    "        result = stmai.benchmark(stmai_params, board_name)\n",
    "        fps = analyze_inference_time(report=result,\n",
    "                                     target_mpu=\"mp\" in board_name.lower())\n",
    "        # Save the result in outputs folder\n",
    "        with open(f'./{benchmark_report_dir}/{model_name}_{board_name}.txt', 'w') as file_benchmark:\n",
    "            file_benchmark.write(f'{result}')\n",
    "        return fps\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Benchmarking failed on board: {board_name}\")\n",
    "        fps = 0\n",
    "        return fps\n",
    "\n",
    "def analyze_inference_time(report: object = None,\n",
    "                           target_mpu = False) -> float:\n",
    "    \"\"\"\n",
    "    Analyzes the inference time of a STEdgeAI model, prints the report and return the FPS.\n",
    "    Args:\n",
    "        report: A report object containing information about the model.\n",
    "        target_mpu: a boolean (True: if target is MPU, False: otherwise)\n",
    "\n",
    "    Returns:\n",
    "        The frames per second (FPS) of the model.\n",
    "    \"\"\"\n",
    "\n",
    "    inference_time: float = report.duration_ms\n",
    "    fps: float = 1000.0/inference_time\n",
    "    if not target_mpu:\n",
    "        # in mpu benchmark result report we do not have cycles\n",
    "        cycles: int = report.cycles\n",
    "        print(\"[INFO] : Number of cycles : {} \".format(cycles))\n",
    "    print(\"[INFO] : Inference Time : {0:.1f} (ms)\".format(inference_time))\n",
    "    print(\"[INFO] : FPS : {0:.1f}\".format(fps))\n",
    "    return fps\n",
    "\n",
    "\n",
    "# UI widgets\n",
    "\n",
    "# optimization options\n",
    "optimization: List[str] = [\"balanced\", \"time\", \"ram\"]\n",
    "optim_dropdown: widgets.Dropdown = widgets.Dropdown(\n",
    "    options=optimization,\n",
    "    value=optimization[0],\n",
    "    description='Optim:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "# STM32MCU series for code generation target\n",
    "series_name: List[str] = [\n",
    "    \"STM32H7\", \"STM32F7\", \"STM32F4\", \"STM32L4\", \"STM32G4\",\n",
    "    \"STM32F3\", \"STM32U5\", \"STM32L5\", \"STM32F0\", \"STM32L0\",\n",
    "    \"STM32G0\", \"STM32C0\", \"STM32WL\", \"STM32H5\"\n",
    "]\n",
    "series_dropdown: widgets.Dropdown = widgets.Dropdown(\n",
    "    options=series_name,\n",
    "    value=series_name[0],\n",
    "    description='Series:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "# options for the IDE while code generation\n",
    "IDE_name: List[str] = [\"gcc\", \"iar\", \"keil\"]\n",
    "ide_dropdown: widgets.Dropdown = widgets.Dropdown(\n",
    "    options=IDE_name,\n",
    "    value=IDE_name[0],\n",
    "    description='IDE:',\n",
    "    disabled=False\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Login to STEdgeAI Developer Cloud\n",
    "Set environment variables with your credentials to acces STEdgeAI Developer Cloud.\n",
    "\n",
    "If you don't have an account yet go to: https://stm32ai-cs.st.com/home and click on sign in to create an account. \n",
    "\n",
    "Then set the environment variables below with your credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "# Set environment variables with your credentials to access \n",
    "# STEdgeAI Developer Cloud services\n",
    "# Fill the username with your login address \n",
    "username = 'your.email@example.com'\n",
    "os.environ['stmai_username'] = username\n",
    "print('Enter you password')\n",
    "password = getpass.getpass()\n",
    "os.environ['stmai_password'] = password"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Log in STEdgeAI Developer Cloud \n",
    "try:\n",
    "    stmai = Stm32Ai(CloudBackend(str(username), str(password)))\n",
    "    print(\"Successfully Connected!\")\n",
    "except Exception as e:\n",
    "    print(\"Error: please verify your credentials\")\n",
    "    print(e)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Upload the model on STEdgeAI Developer Cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get the onnx models available locally\n",
    "model_list = []\n",
    "for entry in os.listdir('./results_mobilenetv2/exports/'):\n",
    "  if os.path.isfile(os.path.join('./results_mobilenetv2/exports/', entry)):\n",
    "    if entry.endswith('.onnx'): model_list.append(entry)\n",
    "model_sel_dropdown = widgets.Dropdown(\n",
    "    options=model_list,\n",
    "    value=model_list[0],\n",
    "    description='Model:',\n",
    "    disabled=False\n",
    ")\n",
    "display(model_sel_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = model_sel_dropdown.value\n",
    "model_path = f'./results_mobilenetv2/exports/{model_name}'\n",
    "model_name = os.path.basename(model_path)\n",
    "from_model = 'user'\n",
    "\n",
    "try:\n",
    "  stmai.upload_model(model_path)\n",
    "  print(f'Model {model_name} is uploaded !')\n",
    "except Exception as e:\n",
    "    print(\"ERROR: \", e)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Select the STEdgeAI optimization setting\n",
    "| Configuration | Description |\n",
    "| --- | --- |\n",
    "| balanced | default compromise between RAM footprint and latency. |\n",
    "| time | optimize for latency. |\n",
    "| ram | optimize for minimal RAM footprint. |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(optim_dropdown)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Analyze your model memory footprints for STM32MCU targets\n",
    "When analyzing the footprints of the model for STM32MCU targets, following parameters can be configured for stm32ai.analyze callback:\n",
    "\n",
    "CLIParameters (options of STEdgeAI):\n",
    "\n",
    "| Parameter | Description |\n",
    "| --- | --- |\n",
    "| model | Model name corresponding to the file name uploaded. This parameter is __required__. |\n",
    "| optimization | Optimization setting: \"balanced\", \"time\" or \"ram\". This parameter is __required__. |\n",
    "| allocateInputs | If set to \"True\", activations buffer will be also used to handle the input buffers. This parameter is __optional__. Default value is \"True\". |\n",
    "| allocateOutputs | If set to \"True\", activations buffer will be also used to handle the output buffers. This parameter is __optional__. Default value is \"True\". |\n",
    "| noOnnxOptimizer | If set to \"True\", allows to disable the ONNX optimizer pass. This parameter is __optional__. Default value is \"False\". |\n",
    "| fromModel | To identify the origin model when coming from ST model zoo. This parameter is __optional__. Default value is \"user\".|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analyze RAM/Flash model memory footprints after optimization by STEdgeAI\n",
    "optimization = optim_dropdown.value\n",
    "print(f'Anlyzing model : {model_name}, using opimization : {optimization}')\n",
    "# The runtime library footprint varies slightly depending on the STM32 series\n",
    "# For an estimation, we use the default series to the STM32F4\n",
    "try:\n",
    "  result = stmai.analyze(CliParameters(model=model_path,\n",
    "                                       optimization=optimization,\n",
    "                                       allocateInputs=True,\n",
    "                                       allocateOutputs=True,\n",
    "                                       fromModel=from_model))\n",
    "\n",
    "  # analyze and print the summary of footprint report\n",
    "  analyze_footprints(report=result)\n",
    "  \n",
    "  # Save the result in outputs folder\n",
    "  stm32ai_analysis_dir = f'{stm32ai_output_dir}/analysis_report'\n",
    "  os.makedirs(stm32ai_analysis_dir, exist_ok=True)\n",
    "  with open(f'./{stm32ai_analysis_dir}/{model_name}_analyze.txt', 'w') as file_analyze:\n",
    "    file_analyze.write(f'{result}')\n",
    "except Exception as e:\n",
    "    print(\"Error: \", e)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### E. Benchmark your model on a STM32 target\n",
    "Starting from STEdgeAI dev cloud version 10.0.0 onwards, the models can be benchmarked for STM32MCU and STM32MPU as well as for STM32NPU target boards.\n",
    "\n",
    "Here's a table with the parameters and their descriptions while benchmarking for the STM32MCU targets (CLIParameters options of STEdgeAI):\n",
    "\n",
    "| Parameter | Description |\n",
    "| --- | --- |\n",
    "| model | Model name corresponding to the file name uploaded. This parameter is required. |\n",
    "| optimization | Optimization setting: \"balanced\", \"time\" or \"ram\". This parameter is required. |\n",
    "| allocateInputs | If set to \"True\", activations buffer will be also used to handle the input buffers. This parameter is optional. Default value is \"True\". |\n",
    "| allocateOutputs | If set to \"True\", activations buffer will be also used to handle the output buffers. This parameter is optional. Default value is \"True\". |\n",
    "| noOnnxOptimizer | If set to \"True\", allows to disable the ONNX optimizer pass. This parameter is optional. Default value is \"False\". Apply only to ONNX file will be ignored otherwise. |\n",
    "| fromModel | To identify the origin model when coming from ST model zoo. This parameter is optional. Default value is \"user\". |\n",
    "\n",
    "\n",
    "While for the STM32MPU targets, only needed parameters are:\n",
    "\n",
    "| Parameter | Description |\n",
    "| --- | --- |\n",
    "| model | Model name corresponding to the file name uploaded. This parameter is __required__. |\n",
    "| nbCores | Number of CPU cores used for benchmarking. This parameter is __set by the code__ depending on the type of MPU. The value should be an integer \"1\", or \"2\". |\n",
    "| engine | Choice of the hardware engine used on the board for benchmarking.This parameter is __set by the code__ depending on the target MPU. For STM32MP1X boards it is \"MpuEngine.CPU\" and for STM32MP2X this is \"MpuEngine.HW_ACCELERATOR\". |\n",
    "\n",
    "* Note that the the code section below, the boad_name to benchmark the model on should be a string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the available board on STEdgeAI Developer Cloud\n",
    "boards = stmai.get_benchmark_boards()\n",
    "board_names = [boards[i].name for i in range(len(boards))]\n",
    "print(\"Available boards:\", board_names)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 1. Benchmark on all available STM32 boards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benchmark the model on all STEdgeAI Developer Cloud boards\n",
    "print(model_name)\n",
    "fps_array = []\n",
    "# loop through all boards\n",
    "for board_name in board_names:\n",
    "        fps_array.append(benchmark_model(stmai=stmai,\n",
    "                                         model_path=model_path,\n",
    "                                         model_name=model_name,\n",
    "                                         optimization=optimization,\n",
    "                                         from_model=from_model,\n",
    "                                         board_name=board_name,\n",
    "                                         allocateInputs= True,\n",
    "                                         allocateOutputs=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the Frame per Second benchmark\n",
    "sorted_fps = sorted(fps_array, reverse=True)\n",
    "sorted_boards = [board_names[fps_array.index(i)] for i in sorted_fps]\n",
    "fig = plt.figure(1, figsize=(15, 8), tight_layout=True)\n",
    "# colors = sns.color_palette()\n",
    "\n",
    "colors = ['#4C72B0', '#55A868', '#C44E52', '#8172B2', '#CCB974',\n",
    "          '#64B5CD', '#B4A7D6', '#AEC7E8', '#FFA07A', '#FFC0CB',\n",
    "          '#FFFFB3', '#8DD3C7', '#BEBADA', '#FDB462', '#FB8072',\n",
    "          '#FF6347', '#4682B4', '#6A5ACD', '#7FFF00', '#D2691E']\n",
    "\n",
    "plt.bar(sorted_boards, sorted_fps, color=colors[:len(boards)], width=0.7)\n",
    "plt.ylabel('FPS', fontsize=15)\n",
    "plt.yticks(fontsize=12)\n",
    "plt.xticks(sorted_boards, rotation = 75)\n",
    "plt.title('STM32 FPS benchmark')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 2. Benchmark on a selected board"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select a board among the available boards\n",
    "board_dropdown = widgets.Dropdown(\n",
    "    options = board_names,\n",
    "    value = 'STM32N6570-DK',\n",
    "    description ='Board:',\n",
    "    disabled = False,)\n",
    "\n",
    "display(board_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "board_name = board_dropdown.value\n",
    "print(model_name, board_name)\n",
    "fps = benchmark_model(stmai=stmai,\n",
    "                      model_path=model_path,\n",
    "                      model_name=model_name,\n",
    "                      optimization=optimization,\n",
    "                      from_model=from_model,\n",
    "                      board_name=board_name,\n",
    "                      allocateInputs= True,\n",
    "                      allocateOutputs=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### F. Generate your model optimized C code for STM32MCU targets\n",
    "\n",
    "To deploy the model on an STM32MCU target the user has to generate the C-Code of the optimized model. Here's a table with the parameters and their descriptions for the stm32.generate callback (CLIParameters of STEdgeAI):\n",
    "\n",
    "| Parameter | Description |\n",
    "| --- | --- |\n",
    "| model | Model name corresponding to the file name uploaded. This parameter is required. |\n",
    "| optimization | Optimization setting: \"balanced\", \"time\" or \"ram\". This parameter is required. |\n",
    "| allocateInputs | If set to \"True\", activations buffer will be also used to handle the input buffers. This parameter is optional. Default value is \"True\". |\n",
    "| allocateOutputs | If set to \"True\", activations buffer will be also used to handle the output buffers. This parameter is optional. Default value is \"True\". |\n",
    "| noOnnxOptimizer | If set to \"True\", allows to disable the ONNX optimizer pass. This parameter is optional. Default value is \"False\". Apply only to ONNX file will be ignored otherwise. |\n",
    "| includeLibraryForSerie | Include the runtime library for the given STM32 series. This parameter is optional. |\n",
    "| fromModel | To identify the origin model when coming from ST model zoo. This parameter is optional. |\n",
    "\n",
    "\n",
    "\n",
    "### NOTE\n",
    "\n",
    "There is no need for this step if the deployment is intended on the MPU. One can directly deploy the .tflite model on the STM32MPUs. In case of STM32MP2x, an optimized version of the model should be already available in the path where the starting model was placed with the same name as model and extension \".nb\".\n",
    "\n",
    "\n",
    "Due to licensing issues, the C-Code for STM32NPU can only be generated by downloading and installing the STEdgeAI locally."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose a target board"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(series_dropdown)\n",
    "display(ide_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "series = series_dropdown.value\n",
    "IDE = ide_dropdown.value\n",
    "print(f'Generating optimized C code of {model_name} model, for {series} series boards!\\n')\n",
    "# Generate model .c/.h code + Lib/Inc on STEdgeAI Developer Cloud\n",
    "stm32ai_code_dir = f'{stm32ai_output_dir}/generated_code'\n",
    "os.makedirs(stm32ai_code_dir, exist_ok=True)\n",
    "result = stmai.generate(CliParameters(\n",
    "    model=model_name,\n",
    "    output=stm32ai_code_dir,\n",
    "    optimization=optimization,\n",
    "    allocateInputs=True,\n",
    "    allocateOutputs=True,\n",
    "    includeLibraryForSerie=CliLibrarySerie(series),\n",
    "    includeLibraryForIde=CliLibraryIde(IDE),\n",
    "    fromModel=from_model\n",
    "))\n",
    "!ls \"{stm32ai_code_dir}\"\n",
    "# print 20 first lines of the report\n",
    "if os.path.isfile(f'./{stm32ai_code_dir}/network_generate_report.txt'):\n",
    "  print(\"\\n\\n---- code generation report ----\\n\",\"*\" * 80)\n",
    "  with open(f'./{stm32ai_code_dir}/network_generate_report.txt', 'r') as f:\n",
    "    for _ in range(20): print(next(f))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### You are ready to integrate your model in your STM32 application !"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (Optional) : Delete your model from your STEdgeAI Developer Cloud space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if stmai.delete_model(model_name):\n",
    "    print(f'{model_name} deleted from STEdgeAI developer Cloud workspace!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deployment on STM32N6, STM32H7* and STM32MPU\n",
    "\n",
    "The `QDQ` quantized models can be deployed using the [stm32ai-modelzoo-services](https://github.com/STMicroelectronics/stm32ai-modelzoo-services) as an [image_classification](https://github.com/STMicroelectronics/stm32ai-modelzoo-services/tree/main/image_classification) model. For knowing more details on how to do that please refer to [Deploying Image Classification models on STM32MCU](https://github.com/STMicroelectronics/stm32ai-modelzoo-services/blob/main/image_classification/deployment/README.md) and [Deploying Image Classificaiton Models on STM32MPU](https://github.com/STMicroelectronics/stm32ai-modelzoo-services/blob/main/image_classification/deployment/README_MPU.md)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}