case4_reference_implementation_and_evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deployability Tests of MPAI Watermarked MLCommons-Tiny Networks with the ST Unified AI Core Technology \n",
    "\n",
    "<center><img width=1000 src=\"pics/pics_for_notebook/prj_workflow.png\"></center>\n",
    "\n",
    "\n",
    "This project is focused on NN-based image classification methods specifically designed and deployed for low power and low resources devices. It aims to prove that Neural Network Watermarking (NNW) approach under standardization by Moving Pictures, Audio and Data Coding by Artificial Intelligence (MPAI) community is the solution for protecting intellectual property such as the tiny Neural Network (NN), deployable on resource-constrained devices. The standard is named IEEE 3304-2023.\n",
    "\n",
    "We start with the MLCommon Tiny benchmark model for Image Classification (ResNet8) (https://mlcommons.org/) and watermark it with a state-of-the-art method. Later on we explore its robustness and efficiency through a series of tests. These tests include attack simulations such as quantization, pruning, and Gaussian attacks done following the MPAI NNW standardazed procedure.\n",
    "\n",
    "By subjecting the mentioned neural networks to these tests, we aim to learn more about the trade-offs between parameters, accuracy, and computational costs, ultimately facilitating the deployment of robust, efficient and secure machine learning solutions on Micro Controller Units (MCUs) for various edge computing applications by embedding a watermark. \n",
    "For MCUs deployability analysis, ST Edge AI Unified Core Technology has been used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**License of the Jupyter Notebook**\n",
    "\n",
    "Copyright 2024 STMicroelectronics\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Below here all the necessary libraries and modules are imported that we will use throughout the whole notebook.\n",
    "import os\n",
    "import sys\n",
    "import subprocess\n",
    "from datetime import datetime\n",
    "\n",
    "import cv2\n",
    "from PIL import Image\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from IPython import display\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "import onnx\n",
    "import onnxruntime as rt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision as tv\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "# Fixed seed for reproductibility\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"opencv version: {}\".format(cv2.__version__))\n",
    "print(\"pytorch version: {}\".format(torch.__version__))\n",
    "print(\"onnx version: {}\".format(onnx.__version__))\n",
    "print(\"onnx_runtime version: {}\".format(rt.__version__))\n",
    "print(\"numpy version: {}\".format(np.__version__))\n",
    "print(\"python version: {}\".format(sys.version))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Topology Development\n",
    "<p>\n",
    "  <center><img src=\"pics/pics_for_notebook/ResNet8.png\"></center>\n",
    "</p>\n",
    "\n",
    "ResNet8 is a variant of the ResNet (Residual Network) architecture, which is a widely used deep neural network architecture known for its effectiveness in image classification tasks. ResNet8 is specifically designed to be lightweight and suitable for deployment on resource-constrained devices like MCUs. It is part of the MLCommons Tiny Benchmark suite, which provides standardized benchmarks for evaluating the performance of machine learning models on edge devices.\n",
    "ResNet8 consists of a relatively shallow network with a total of 8 layers, including convolutional layers, batch normalization layers, activation functions, and a final fully connected layer. \n",
    "\n",
    "The number of parameters in ResNet8 is significantly lower compared to deeper ResNet variants. This reduction in parameters helps to reduce memory footprint and computational overhead, making it suitable for deployment on MCUs with limited resources. The exact number of parameters depends on factors such as the size of the input images and the number of output classes in the classification task\n",
    "\n",
    "ResNet8 follows a convolutional neural network (CNN) architecture, where layers are organized in a sequential manner. Each convolutional layer is followed by batch normalization and a non-linear activation function, typically ReLU (Rectified Linear Unit). The final layer consists of a fully connected layer followed by softmax activation for classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Neural Network Model Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a PyTorch model from MLCommons tiny benchmark for image classification (ResNet-8)\n",
    "class ResNetBlock(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels: int,\n",
    "        out_channels: int,\n",
    "        stride: int = 1,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.block = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=in_channels,\n",
    "                out_channels=out_channels,\n",
    "                kernel_size=3,\n",
    "                padding=1,\n",
    "                bias=True,\n",
    "                stride=stride,\n",
    "            ),\n",
    "            nn.BatchNorm2d(num_features=out_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(\n",
    "                in_channels=out_channels,\n",
    "                out_channels=out_channels,\n",
    "                kernel_size=3,\n",
    "                padding=1,\n",
    "                bias=True,\n",
    "            ),\n",
    "            nn.BatchNorm2d(num_features=out_channels),\n",
    "        )\n",
    "        if in_channels == out_channels:\n",
    "            self.residual = nn.Identity()\n",
    "        else:\n",
    "            self.residual = nn.Conv2d(\n",
    "                in_channels=in_channels,\n",
    "                out_channels=out_channels,\n",
    "                kernel_size=1,\n",
    "                stride=stride,\n",
    "            )\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        x = self.block(inputs)\n",
    "        y = self.residual(inputs)\n",
    "        return F.relu(x + y)\n",
    "\n",
    "\n",
    "class Resnet8v1EEMBC(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.stem = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=3, out_channels=16, kernel_size=3, padding=1, bias=True\n",
    "            ),\n",
    "            nn.BatchNorm2d(num_features=16),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "\n",
    "        self.first_stack = ResNetBlock(in_channels=16, out_channels=16, stride=1)\n",
    "        self.second_stack = ResNetBlock(in_channels=16, out_channels=32, stride=2)\n",
    "        self.third_stack = ResNetBlock(in_channels=32, out_channels=64, stride=2)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(in_features=64, out_features=10)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        x = self.stem(inputs)\n",
    "        x = self.first_stack(x)\n",
    "        x = self.second_stack(x)\n",
    "        x = self.third_stack(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc(x)\n",
    "        return x"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data from CIFAR-10 and trigger dataset and pre-processing\n",
    "\n",
    "Below the definition to get the labelled data from CIFAR10 dataset and from the trigger dataset incoherently labelled.\n",
    "\n",
    "The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images.\n",
    "The dataset is divided into five training batches and one test batch, each with 10000 images. The test batch contains exactly 1000 randomly-selected images from each class. The training batches contain the remaining images in random order, but some training batches may contain more images from one class than another. Between them, the training batches contain exactly 5000 images from each class.\n",
    "\n",
    "Images for the trigger dataset shall not be chosen from the training dataset and indeed 100 abstract images were selected, and random target classes were assigned to them. \n",
    "\n",
    "  <center><img width=300 src=\"pics/pics_for_notebook/cifar10.png\" style=\"margin-right: 25px;\"> <img width=300 src=\"pics/full_size_pics/010.jpg\" style=\"margin-left: 25px;\"></center>\n",
    "\n",
    "  <center> Exemplary images respectively from CIFAR-10 and from trigger dataset.</center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mpai_nnw_dataloader(trainset,testset,batch_size=100):\n",
    "    trainloader = torch.utils.data.DataLoader(\n",
    "        trainset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        num_workers=2)\n",
    "\n",
    "    testloader = torch.utils.data.DataLoader(\n",
    "        testset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        num_workers=2)\n",
    "\n",
    "    return trainloader,testloader\n",
    "    \n",
    "def CIFAR10_dataset():\n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "    ])\n",
    "\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.CenterCrop((32, 32)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "    ])\n",
    "\n",
    "    # datasets\n",
    "    trainset = tv.datasets.CIFAR10(\n",
    "        root='./data/',\n",
    "        train=True,\n",
    "        download=True,\n",
    "        transform=transform_train)\n",
    "\n",
    "    testset = tv.datasets.CIFAR10(\n",
    "        './data/',\n",
    "        train=False,\n",
    "        download=True,\n",
    "        transform=transform_test)\n",
    "\n",
    "    return trainset, testset, transform_test\n",
    "    \n",
    "def load_and_preprocess_data():\n",
    "    # Load, crop and save modified image\n",
    "    script_abs_path = os.getcwd() #script_abs_path = os.path.dirname(os.path.abspath(__file__))\n",
    "    image_path = os.path.join(script_abs_path, \"pics\")\n",
    "    full_size_image_folder_name = \"full_size_pics\" \n",
    "    croppped_image_folder_name = \"cropped_pics\"\n",
    "    full_size_image_folder_path = os.path.join(image_path, full_size_image_folder_name)\n",
    "    croppped_image_folder_path = os.path.join(image_path, croppped_image_folder_name)\n",
    "    \n",
    "    # search for full_size_image_folder and images inside it\n",
    "    print(\"Searching for {} folder\".format(full_size_image_folder_path))\n",
    "    if os.path.isdir(full_size_image_folder_path):\n",
    "        file_list = os.listdir(full_size_image_folder_path)\n",
    "        if not file_list:\n",
    "            raise AssertionError(\"Empty folder!\")\n",
    "        print(\"List of files: {}\".format(file_list))\n",
    "    else:\n",
    "        raise AssertionError(\"Folder not found!\")\n",
    "    \n",
    "    # create the croppped_image_folder if it doesn't exists\n",
    "    print(\"Searching for {} folder\".format(croppped_image_folder_path))\n",
    "    if not os.path.isdir(croppped_image_folder_path):\n",
    "        os.mkdir(croppped_image_folder_path)\n",
    "        print(\"Folder created!\")\n",
    "    \n",
    "    img_num = 0\n",
    "    \n",
    "    # iterate on files in folder\n",
    "    all_pics_file_npy_name = \"all_pics_cropped.npy\"\n",
    "    all_pics_file_npy_path = os.path.join(croppped_image_folder_path, all_pics_file_npy_name)\n",
    "    all_pics_npy = None\n",
    "    \n",
    "    for idx, file_name in enumerate(file_list):\n",
    "        print(\"Pre-processing file \\\"{}\\\"\".format(file_name))\n",
    "    \n",
    "        full_size_file_path = os.path.join(full_size_image_folder_path, file_name)\n",
    "        cropped_file_name = file_name.split(\".\", 2)[0] + \"_cropped.\" + file_name.split(\".\", 2)[1]\n",
    "        cropped_file_path = os.path.join(croppped_image_folder_path, cropped_file_name)\n",
    "        cropped_file_npy_name = file_name.split(\".\", 2)[0] + \"_cropped.\" + \"npy\"\n",
    "        cropped_file_npy_path = os.path.join(croppped_image_folder_path, cropped_file_npy_name)\n",
    "    \n",
    "        # avoid open of file not supported or not images\n",
    "        if not full_size_file_path.endswith('.jpg'):\n",
    "            print(\"\\tFile not supported!\")\n",
    "            continue\n",
    "    \n",
    "        # load image from the folder\n",
    "        img = cv2.imread(full_size_file_path)\n",
    "    \n",
    "        if img is None:\n",
    "            raise AssertionError(\"File not read correctly!\")\n",
    "    \n",
    "        # cv2.imshow(\"Full size image\", img)  # commented since python script run on terminal\n",
    "    \n",
    "        # change color map from BGR to RGB\n",
    "        #print(\"DEBUGP img: {}\".format(img.flatten()[0:5]))\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "        #print(\"DEBUGP img: {}\".format(img.flatten()[0:5]))\n",
    "    \n",
    "        # extract coordinates\n",
    "        center_x = img.shape[1]/2\n",
    "        center_y = img.shape[0]/2\n",
    "    \n",
    "        crop_x = 32\n",
    "        crop_y = 32\n",
    "    \n",
    "        x_start = center_x - crop_x/2\n",
    "        y_start = center_y - crop_y/2\n",
    "    \n",
    "        # apply crop to cv2/numpy object\n",
    "        crop_img = img[int(y_start):int(y_start+crop_y), int(x_start):int(x_start+crop_x)]\n",
    "    \n",
    "        # convert into flaot32\n",
    "        crop_img_npy = crop_img.astype(np.float32)\n",
    "        tr_axis = (2, 0, 1)\n",
    "        crop_img_npy = np.transpose(crop_img_npy, tr_axis)\n",
    "    \n",
    "        # normalize per channel as pytorch normalize doc.\n",
    "        # output[channel] = (input[channel] - mean[channel]) / std[channel]\n",
    "        mean = (0.4914, 0.4822, 0.4465) # from Carl prep implementation\n",
    "        std = (0.2023, 0.1994, 0.2010) # from Carl prep implementation\n",
    "    \n",
    "        crop_img_npy = crop_img_npy / 255.0\n",
    "    \n",
    "\n",
    "        for ch in range(0,crop_img_npy.shape[0]):\n",
    "            #print(\"Normalizing channel {}\".format(ch))\n",
    "            crop_img_npy[ch] = (crop_img_npy[ch] - mean[ch]) / std[ch]\n",
    "\n",
    "    \n",
    "        # show and save the cropped image the new folder\n",
    "        crop_img = cv2.cvtColor(crop_img, cv2.COLOR_RGB2BGR)\n",
    "        cv2.imwrite(cropped_file_path, crop_img)\n",
    "        np.save(cropped_file_npy_path, crop_img_npy)\n",
    "    \n",
    "        if all_pics_npy is None:\n",
    "            all_pics_npy = crop_img_npy.reshape((1,) + crop_img_npy.shape)\n",
    "        else:\n",
    "            aa = crop_img_npy.reshape((1,) + crop_img_npy.shape)\n",
    "\n",
    "            all_pics_npy = np.concatenate((all_pics_npy, crop_img_npy.reshape((1,) + crop_img_npy.shape)))\n",
    "\n",
    "        # #print(\"\\t\\tNPY  saved in \\\"{}\\\"\".format(cropped_file_npy_path))\n",
    "        # #print(\"\\t\\tNPY Shape: {}\".format(crop_img_npy.shape))\n",
    "        # #print(\"\\t\\tNPY dtype: {}\".format(crop_img_npy.dtype))\n",
    "        # #print(\"\\t\\tNPY Min: {} -- Max: {}\".format(crop_img_npy.min(), crop_img_npy.max()))\n",
    "        img_num += 1\n",
    "    \n",
    "    np.save(all_pics_file_npy_path, all_pics_npy)\n",
    "    \n",
    "    print(\"Pre-processed {} on {} files in the folder\".format(img_num, idx))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and watermark embedding\n",
    "\n",
    "  <center><img width=1000 src=\"pics/pics_for_notebook/embed_graph.png\"></center>\n",
    "  \n",
    "The preferred method has several key steps such as backdooring, using strong backdoors and commitment schemes, and finally watermarking procedure.\n",
    "The combination of original dataset and exemplary trigger dataset from the state-of-the-art method is used to train the NN model and force its behabiour."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def pytorch_train_embed():\n",
    "    save_file = 'embedResNet_Adii'\n",
    "    batch_size = 128\n",
    "    model = Resnet8v1EEMBC()\n",
    "    model.to(device)\n",
    "    # mpai_nnw_dataloader(trainset,testset,batch_size=100)\n",
    "    print(model)\n",
    "    \n",
    "    power = 10\n",
    "    # watermarking section (change here to test another method) #######################################\n",
    "    tools = ADI.Adi_tools()\n",
    "    # folder = 'code_from_MPAI/referencesoftwarev11_main/Attacks/data/trigger_pics/'\n",
    "    folder = 'code_from_MPAI/STposter/adi/' # trigger_pics/'\n",
    "    watermarking_dict = {'folder': folder,\n",
    "                         'batch_size': batch_size,\n",
    "                         'transforms': inference_transform,\n",
    "                         'types':1,\n",
    "                         'num_class':10, # 10 classes referring to CIFAR10\n",
    "                         'power': 10,\n",
    "                        }\n",
    "    # watermarking section (END change here to test another method) ###################################\n",
    "    watermarking_dict = tools.init(model, watermarking_dict)\n",
    "    trainloader, testloader = mpai_nnw_dataloader(trainset, testset, batch_size)\n",
    "    \n",
    "    # Imperceptibility.Embeds(watermarking_dict[\"types\"], model, watermarking_dict, run_arg['epochs'], tools, trainloader, batch_size)\n",
    "    # code from MPAI NNW script\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    learning_rate, momentum, weight_decay = 0.01, .9, 5e-4\n",
    "    optimizer = optim.SGD([\n",
    "        {'params': model.parameters()}\n",
    "    ], lr=learning_rate, momentum=momentum, weight_decay=weight_decay)\n",
    "    model.train()\n",
    "    epoch = 0\n",
    "    print(\"Launching injection.....\")\n",
    "    while epoch < num_epochs:\n",
    "        print('doing epoch', str(epoch + 1), \".....\")\n",
    "        loss, loss_nn, loss_w = tools.Embedder_one_step(model, trainloader, optimizer, criterion, watermarking_dict)\n",
    "\n",
    "        loss = (loss * batch_size / len(trainloader.dataset))\n",
    "        loss_nn = (loss_nn * batch_size / len(trainloader.dataset))\n",
    "        loss_w = (loss_w * batch_size / len(trainloader.dataset))\n",
    "        print(' loss  : %.5f   - loss_wm: %.5f, loss_nn: %.5f  ' % (loss, loss_w, loss_nn))\n",
    "\n",
    "        epoch += 1\n",
    "            \n",
    "    print(\"############ Watermark inserted ##########\")\n",
    "    print()\n",
    "    print(\"Launching Test function...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageFolderCustomClass(torch.utils.data.Dataset):\n",
    "    \"\"\"A generic data loader where the images are arranged in this way: ::\n",
    "        root/dog/xxx.png\n",
    "        root/dog/xxy.png\n",
    "        root/dog/xxz.png\n",
    "        root/cat/123.png\n",
    "        root/cat/nsdf3.png\n",
    "        root/cat/asd932_.png\n",
    "    Args:\n",
    "        root (string): Root directory path.\n",
    "        transform (callable, optional): A function/transform that  takes in an PIL image\n",
    "            and returns a transformed version. E.g, ``transforms.RandomCrop``\n",
    "        target_transform (callable, optional): A function/transform that takes in the\n",
    "            target and transforms it.\n",
    "        loader (callable, optional): A function to load an image given its path.\n",
    "     Attributes:\n",
    "        classes (list): List of the class names.\n",
    "        class_to_idx (dict): Dict with items (class_name, class_index).\n",
    "        imgs (list): List of (image path, class_index) tuples\n",
    "    \"\"\"\n",
    "\n",
    "    def make_dataset(self,dir, class_to_idx):\n",
    "        def is_image_file(filename):\n",
    "            \"\"\"Checks if a file is an image.\n",
    "            Args:\n",
    "                filename (string): path to a file\n",
    "            Returns:\n",
    "                bool: True if the filename ends with a known image extension\n",
    "            \"\"\"\n",
    "            filename_lower = filename.lower()\n",
    "            return any(filename_lower.endswith(ext) for ext in self.IMG_EXTENSIONS)\n",
    "\n",
    "        images = []\n",
    "        dir = os.path.expanduser(dir)\n",
    "        for target in sorted(os.listdir(dir)):\n",
    "            d = os.path.join(dir, target)\n",
    "            if not os.path.isdir(d):\n",
    "                continue\n",
    "\n",
    "            for root, _, fnames in sorted(os.walk(d)):\n",
    "                for fname in sorted(fnames):\n",
    "                    if is_image_file(fname):\n",
    "                        path = os.path.join(root, fname)\n",
    "                        item = (path, class_to_idx[target])\n",
    "                        images.append(item)\n",
    "\n",
    "        return images\n",
    "\n",
    "    def accimage_loader(self,path):\n",
    "        import accimage\n",
    "        try:\n",
    "            return accimage.Image(path)\n",
    "        except IOError:\n",
    "            # Potentially a decoding problem, fall back to PIL.Image\n",
    "            return self.pil_loader(path)\n",
    "\n",
    "    def pil_loader(self,path):\n",
    "        # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n",
    "        with open(path, 'rb') as f:\n",
    "            img = Image.open(f)\n",
    "            return img.convert('RGB')\n",
    "\n",
    "    def default_loader(self,path):\n",
    "        from torchvision import get_image_backend\n",
    "        if get_image_backend() == 'accimage':\n",
    "            return self.accimage_loader(path)\n",
    "        else:\n",
    "            return self.pil_loader(path)\n",
    "\n",
    "    def find_classes(self,dir):\n",
    "        classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]\n",
    "        classes.sort()\n",
    "        class_to_idx = {classes[i]: i for i in range(len(classes))}\n",
    "        return classes, class_to_idx\n",
    "\n",
    "    def __init__(self, root, transform=None, target_transform=None,\n",
    "                  custom_class_to_idx=None) :\n",
    "        self.IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']\n",
    "\n",
    "        if custom_class_to_idx is None:\n",
    "            classes, class_to_idx = self.find_classes(root)\n",
    "        else:\n",
    "            class_to_idx = custom_class_to_idx\n",
    "            classes = list(class_to_idx.keys())\n",
    "        imgs = self.make_dataset(root, class_to_idx)\n",
    "\n",
    "        if len(imgs) == 0:\n",
    "            raise(RuntimeError(\"Found 0 images in subfolders of: \" + root + \"\\n\"\n",
    "                               \"Supported image extensions are: \" + \",\".join(self.IMG_EXTENSIONS)))\n",
    "\n",
    "        self.root = root\n",
    "        self.imgs = imgs\n",
    "        self.classes = classes\n",
    "        self.class_to_idx = class_to_idx\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            index (int): Index\n",
    "        Returns:\n",
    "            tuple: (image, target) where target is class_index of the target class.\n",
    "        \"\"\"\n",
    "        path, target = self.imgs[index]\n",
    "        img = self.default_loader(path)\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "        if self.target_transform is not None:\n",
    "            target = self.target_transform(target)\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    "\n",
    "    def __repr__(self):\n",
    "        fmt_str = 'Dataset ' + self.__class__.__name__ + '\\n'\n",
    "        fmt_str += '    Number of datapoints: {}\\n'.format(self.__len__())\n",
    "        fmt_str += '    Root Location: {}\\n'.format(self.root)\n",
    "        tmp = '    Transforms (if any): '\n",
    "        fmt_str += '{0}{1}\\n'.format(tmp,\n",
    "                                     self.transform.__repr__().replace('\\n', '\\n' + ' ' * len(tmp)))\n",
    "        tmp = '    Target Transforms (if any): '\n",
    "        fmt_str += '{0}{1}'.format(tmp,\n",
    "                                   self.target_transform.__repr__().replace('\\n', '\\n' + ' ' * len(tmp)))\n",
    "        return fmt_str\n",
    "\n",
    "\n",
    "class Adi_tools():\n",
    "    def __init__(self)-> None:\n",
    "        super(Adi_tools, self).__init__()\n",
    "\n",
    "    def list_image(self, main_dir):\n",
    "        \"\"\"return all file in the directory\"\"\"\n",
    "        res = []\n",
    "        for f in os.listdir(main_dir):\n",
    "            if not f.startswith('.'):\n",
    "                res.append(f)\n",
    "        return res\n",
    "\n",
    "    def add_images(self, dataset, image, label):\n",
    "        \"\"\"add an image with its label to the dataset\n",
    "        :param dataset: aimed dataset to be modified\n",
    "        :param image: image to be added\n",
    "        :param label: label of this image\n",
    "        :return: 0\n",
    "        \"\"\"\n",
    "\n",
    "        (taille, height, width, channel) = np.shape(dataset.data)\n",
    "        dataset.data = np.append(dataset.data, image)\n",
    "        dataset.targets.append(label)\n",
    "        dataset.data = np.reshape(dataset.data, (taille + 1, height, width, channel))\n",
    "        return 0\n",
    "\n",
    "\n",
    "\n",
    "    def get_image(self, name):\n",
    "        \"\"\"\n",
    "        :param name: file (including the path) of an image\n",
    "        :return: a numpy of this image\"\"\"\n",
    "        image = Image.open(name)\n",
    "        return np.array(image)\n",
    "\n",
    "    def Embedder_one_step(self, net, trainloader, optimizer, criterion, watermarking_dict):\n",
    "        '''\n",
    "        :param watermarking_dict: dictionary with all watermarking elements\n",
    "        :return: the different losses ( global loss, task loss, watermark loss)\n",
    "        '''\n",
    "        running_loss = 0\n",
    "        wmloader=watermarking_dict['wmloader']\n",
    "        wminputs, wmtargets = [], []\n",
    "        if wmloader:\n",
    "            for wm_idx, (wminput, wmtarget) in enumerate(wmloader):\n",
    "                wminput, wmtarget = wminput.to(device), wmtarget.to(device)\n",
    "                wminputs.append(wminput)\n",
    "                wmtargets.append(wmtarget)\n",
    "\n",
    "            # the wm_idx to start from\n",
    "            wm_idx = np.random.randint(len(wminputs))\n",
    "        for i, data in enumerate(trainloader, 0):\n",
    "            # split data into the image and its label\n",
    "            inputs, labels = data\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            if wmloader:\n",
    "                inputs = torch.cat([inputs, wminputs[(wm_idx + i) % len(wminputs)]], dim=0)\n",
    "                labels = torch.cat([labels, wmtargets[(wm_idx + i) % len(wminputs)]], dim=0)\n",
    "\n",
    "            # initialise the optimiser\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward\n",
    "            outputs = net(inputs)\n",
    "            # backward\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            loss.backward()\n",
    "            # update the optimizer\n",
    "            optimizer.step()\n",
    "\n",
    "            # loss\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        return running_loss, running_loss, 0\n",
    "\n",
    "    def Detector(self, net, watermarking_dict):\n",
    "        \"\"\"\n",
    "        :param file_watermark: file that contain our saved watermark elements\n",
    "        :return: the extracted watermark, the hamming distance compared to the original watermark\n",
    "        \"\"\"\n",
    "        # watermarking_dict = np.load(file_watermark, allow_pickle='TRUE').item() #retrieve the dictionary\n",
    "        wmloader= watermarking_dict['wmloader']\n",
    "        net.eval()\n",
    "        res = 0\n",
    "        total = 0\n",
    "        for i, data in enumerate(wmloader):\n",
    "            inputs, labels = data\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = net(inputs)\n",
    "\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            res += predicted.eq(labels.data).cpu().sum()\n",
    "\n",
    "        return '%i/%i' %(int(res),total), total-res\n",
    "\n",
    "    def init(self, net, watermarking_dict, save=None):\n",
    "        '''\n",
    "        :param net: network\n",
    "        :param watermarking_dict: dictionary with all watermarking elements\n",
    "        :param save: file's name to save the watermark\n",
    "        :return: watermark_dict with a new entry: the secret key matrix X\n",
    "        '''\n",
    "        folder=watermarking_dict[\"folder\"]\n",
    "        for elmnt in os.listdir(folder):\n",
    "            if \".txt\" in elmnt:labels_path=elmnt\n",
    "        wmset = ImageFolderCustomClass(\n",
    "            folder,\n",
    "            watermarking_dict[\"transforms\"])\n",
    "        img_nlbl = []\n",
    "        wm_targets = np.loadtxt(os.path.join(folder, labels_path))\n",
    "        for idx, (path, target) in enumerate(wmset.imgs):\n",
    "            img_nlbl.append((path, int(wm_targets[idx])))\n",
    "        wmset.imgs = img_nlbl\n",
    "\n",
    "        wmloader = torch.utils.data.DataLoader(\n",
    "            wmset, batch_size=watermarking_dict[\"batch_size\"], shuffle=True,\n",
    "            num_workers=4, pin_memory=True)\n",
    "        watermarking_dict['wmloader']=wmloader\n",
    "\n",
    "        return watermarking_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of the pre-trained and the watermarked NN models\n",
    "The script is part of a testing framework that automates the process of setting up directories and running validation tests on our models. It is designed to work with ST Edge AI Unified Core technology that validates models for use on STM32 microcontroller hardware. The script handles file and directory operations, constructs the necessary commands for validation, and executes those commands, capturing the output in text files for later review.\n",
    "\n",
    "<center><img width=\"600\" src=\"pics/pics_for_notebook/ev_graph.png\"></center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mpai_nnw_quality_measurement(confusion_matrix):\n",
    "    line_sum=torch.sum(confusion_matrix,dim=1)\n",
    "    column_sum = torch.sum(confusion_matrix, dim=0)\n",
    "    total_sum=torch.sum(confusion_matrix)\n",
    "    Precision=torch.diag(confusion_matrix)/line_sum\n",
    "    Recall=torch.diag(confusion_matrix)/column_sum\n",
    "    Pfa=(line_sum - torch.diag(confusion_matrix))/total_sum\n",
    "    Pmd=(column_sum - torch.diag(confusion_matrix))/total_sum\n",
    "    return torch.mean(Pfa),torch.mean(Precision),torch.mean(Recall),torch.mean(Pmd)\n",
    "\n",
    "def mpai_nnw_test(net,testloader):\n",
    "    # test complet\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    confusion_matrix=torch.zeros(10,10)\n",
    "    # torch.no_grad do not train the network\n",
    "    with torch.no_grad():\n",
    "        for data in testloader:\n",
    "            inputs, labels = data\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = net(inputs)\n",
    "            if len(outputs) ==2:outputs,_=outputs\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum()\n",
    "            for i in range(len(labels)):\n",
    "                 confusion_matrix[predicted[i],labels[i]]=confusion_matrix[predicted[i],labels[i]]+1\n",
    "    return 100 - (100 * float(correct) / total), confusion_matrix\n",
    "\n",
    "def pytorch_evaluation():\n",
    "    np.save(save_file + '_watermarking_dict.npy',watermarking_dict)\n",
    "    save_file = 'embedResNet_Adii'\n",
    "    torch.save({\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    }, save_file + '_weights')\n",
    "    \n",
    "    val_score, cm = mpai_nnw_test(model, testloader)\n",
    "    Pfa, Pre, Rec, Pmd = mpai_nnw_quality_measurement(cm)\n",
    "    print('Validation error : %.2f' % val_score)\n",
    "    print('Probability of false alarm:', Pfa)\n",
    "    print('Precision:', Pre)\n",
    "    print('Recall:', Rec)\n",
    "    print('Probability of missed detection', Pmd)\n",
    "\n",
    "def NN_model_evaluation_ST_Edge_AI():\n",
    "    #print(\"Pre-processing file \\\"{}\\\"\".format(file_name))\n",
    "    file_name = run_arg['model_path']\n",
    "    f_name, f_extension = os.path.splitext(file_name)\n",
    "    \n",
    "    print(\"Validate model: {}\".format(file_name))\n",
    "    # personalize the following path with your configuration and environment\n",
    "    HOME = ''\n",
    "    xcubeai_exe_path = f\"{HOME}/STM32Cube/Repository/Packs/STMicroelectronics/X-CUBE-AI/8.1.0/Utilities/windows/stm32ai.exe\"\n",
    "\n",
    "    # validation input from trigger dataset\n",
    "    img_trigger_data_path = './pics/001_cropped.npy'\n",
    "    label_trigger_data_path = './pics/001_label.npy'\n",
    "    \n",
    "    cmd = [f\"{xcubeai_exe_path}\", \"validate\", \"-m\", f\"{file_name}\", \"--classifier\", \"--target\", \"stm32\", \"--mode\", \"stm32\", \"--batches\", \"1\", \"-vi\", f\"{img_trigger_data_path}\", \"-vo\", f\"{label_trigger_data_path}\", \"--no-exec-model\"]\n",
    "    print(\"Running command: {}\".format(' '.join(cmd)))\n",
    "    \n",
    "    result = subprocess.run(' '.join(cmd), shell=True, text=True) # use this calling X-Cube-AI from WSL to windows\n",
    "    print(\"Result: {}\".format(result))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specifying the setup for the next run\n",
    "list_run_arg = [\n",
    "    dict(\n",
    "        AIFramework = 'ONNX', # 'Keras', 'Pytorch' or 'ONNX'\n",
    "        action = 'validate', # 'load_pre-trained', 'train', 'embed_w'\n",
    "        model_path = 'models/models_to_test/ResNet.onnx', # e.g. for keras: '../models/models_keras/ResNet-8_1epoch.h5' or for onnx: 'models/models_to_test/ResNet.onnx'\n",
    "        target = 'mcu'\n",
    "        ),\n",
    "    dict(\n",
    "        AIFramework = 'PyTorch', # 'Keras', 'Pytorch' or 'ONNX'\n",
    "        action = 'embed_w', # 'load_pre-trained', 'train', 'embed_w'\n",
    "        model_path = 'models/ResNet_w_0.onnx', # e.g. for keras: '../models/models_keras/ResNet-8_1epoch.h5' or for onnx: 'models/models_to_test/ResNet.onnx'\n",
    "        ),\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_key(key):\n",
    "    if key not in run_arg.keys():\n",
    "        raise AssertionError(\"Key \\'{}\\' not in run_arg dict, please insert\".format(key))\n",
    "\n",
    "for run_arg in list_run_arg:\n",
    "    check_key('AIFramework')\n",
    "    if run_arg['AIFramework'] == 'PyTorch':\n",
    "        print(\"Using PyTorch framework\")\n",
    "        check_key('action')\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        print(\"Available device: {}\".format(device))\n",
    "        trainset, testset, inference_transform = CIFAR10_dataset()\n",
    "        if run_arg['action'] == 'embed_w':\n",
    "            check_key('epochs')\n",
    "            num_epochs = run_arg['epochs']\n",
    "            pytorch_train_embed()\n",
    "            pytorch_evaluation()\n",
    "    elif run_arg['AIFramework'] == 'ONNX':\n",
    "        check_key('action')\n",
    "        if run_arg['action'] == 'load_pre-trained':\n",
    "            print(\"Loading ONNX model from: \\'{}\\'\".format(run_arg['model_path']))\n",
    "            check_key('model_path')\n",
    "            sess = rt.InferenceSession(run_arg['model_path'])\n",
    "        elif run_arg['action'] == 'validate':\n",
    "            check_key('model_path')\n",
    "            print(\"Validating ONNX model from: \\'{}\\'\".format(run_arg['model_path']))\n",
    "            check_key('target')\n",
    "            if run_arg['target'] == 'mcu':\n",
    "                NN_model_evaluation_ST_Edge_AI()\n",
    "    else:\n",
    "        raise AssertionError(\"AIFramework not supported\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e7c3fdbb62c86486ffdcadc71c6517f858cc7cfcde70d9d3b8fb3f42d80dc988"
  },
  "kernelspec": {
   "display_name": "NNW_ece_env_py3913_tf212",
   "language": "python",
   "name": "nnw_ece_env_py3913_tf212"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}