Lab_3_Data_Poisoning_Attacks.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyPceID4ZaxNvHkt2/T7jw3T",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/oaramoon/AML/blob/main/Lab_3_Data_Poisoning_Attacks.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Lab 3: Data Poisoning Attacks Against Machine Learning Models**\n",
        "\n",
        "In this lab module, our exploration dives deeply into the realm of data poisoning attacks, with a specific focus on Dirty Label Backdoor (DLBD) techniques. The module commences with a detailed examination of \"Dirty Label Backdoor Attack\", offering students a hands-on opportunity to manipulate a machine learning model. This practical experience involves injecting malicious data with altered labels (hence termed 'dirty') into the training process. We will implement this attack on a traffic sign classifier using the German Traffic Sign Recognition Benchmark (GTSRB) dataset, a standard benchmark in object recognition models. The primary objective is to impart a comprehensive understanding of how adversarial agents can exploit label noise to breach the integrity of machine learning models."
      ],
      "metadata": {
        "id": "GEcLk0QGWgEU"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Part 1: Setting up the Enviroment**\n",
        "\n",
        "In this part, we will set up our environment and familiarize ourselves with the necessary packages. Our first task is to download the GTSRB (German Traffic Sign Recognition Benchmark) dataset. This dataset is provided in TensorFlow dataset format, so we'll also take some time to learn how to interact with and manipulate TensorFlow datasets. This foundational step is crucial for ensuring everyone is comfortable with the tools and data we will be using in our experiments."
      ],
      "metadata": {
        "id": "nCh4MbEYYAvk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
        "import logging\n",
        "logging.getLogger('tensorflow').setLevel(logging.ERROR)\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "\n",
        "def print_green(text):\n",
        "    print(\"\\033[92m\" + text + \"\\033[0m\")\n",
        "\n",
        "def print_red(text):\n",
        "    print(\"\\033[91m\" + text + \"\\033[0m\")"
      ],
      "metadata": {
        "id": "FCfaq9FtXoWA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Downloading and Revewing the German Traffic Sign Recognition Benchmark (GTSRB)**\n"
      ],
      "metadata": {
        "id": "2cvrBpi6aoBs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install gdown\n",
        "if not os.path.isdir('train_gtsrb'):\n",
        "  !gdown https://drive.google.com/uc?id=1JnSbYQvOVtXDWLV8Qd05Lp6DKkFcJlE6\n",
        "  !unzip GTSRB.zip"
      ],
      "metadata": {
        "id": "entbCVVinJUo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gtsrb_training_dataset = tf.data.Dataset.load('./train_gtsrb')\n",
        "\n",
        "num_classes = 43\n",
        "examples_per_class = {}\n",
        "\n",
        "# Iterate over the dataset\n",
        "for image, label in gtsrb_training_dataset:\n",
        "    label = label.numpy()\n",
        "    if label not in examples_per_class:\n",
        "        examples_per_class[label] = image\n",
        "    if len(examples_per_class) == num_classes:\n",
        "        break\n",
        "\n",
        "gtsrb_classes = [\n",
        "    \"Speed limit (20km/h)\",\n",
        "    \"Speed limit (30km/h)\",\n",
        "    \"Speed limit (50km/h)\",\n",
        "    \"Speed limit (60km/h)\",\n",
        "    \"Speed limit (70km/h)\",\n",
        "    \"Speed limit (80km/h)\",\n",
        "    \"End of speed limit (80km/h)\",\n",
        "    \"Speed limit (100km/h)\",\n",
        "    \"Speed limit (120km/h)\",\n",
        "    \"No passing\",\n",
        "    \"No passing for vehicles over 3.5 metric tons\",\n",
        "    \"Right-of-way at the next intersection\",\n",
        "    \"Priority road\",\n",
        "    \"Yield\",\n",
        "    \"Stop\",\n",
        "    \"No vehicles\",\n",
        "    \"Vehicles over 3.5 metric tons prohibited\",\n",
        "    \"No entry\",\n",
        "    \"General caution\",\n",
        "    \"Dangerous curve to the left\",\n",
        "    \"Dangerous curve to the right\",\n",
        "    \"Double curve\",\n",
        "    \"Bumpy road\",\n",
        "    \"Slippery road\",\n",
        "    \"Road narrows on the right\",\n",
        "    \"Road work\",\n",
        "    \"Traffic signals\",\n",
        "    \"Pedestrians\",\n",
        "    \"Children crossing\",\n",
        "    \"Bicycles crossing\",\n",
        "    \"Beware of ice/snow\",\n",
        "    \"Wild animals crossing\",\n",
        "    \"End of all speed and passing limits\",\n",
        "    \"Turn right ahead\",\n",
        "    \"Turn left ahead\",\n",
        "    \"Ahead only\",\n",
        "    \"Go straight or right\",\n",
        "    \"Go straight or left\",\n",
        "    \"Keep right\",\n",
        "    \"Keep left\",\n",
        "    \"Roundabout mandatory\",\n",
        "    \"End of no passing\",\n",
        "    \"End of no passing by vehicles over 3.5 metric tons\"\n",
        "]\n",
        "\n",
        "# Plot one image per class\n",
        "plt.figure(figsize=(25, 20))  # Adjust the size as needed\n",
        "n_rows = 11  # Adjust the number of rows\n",
        "n_cols = 4   # Adjust the number of columns\n",
        "\n",
        "for i, (label, image) in enumerate(examples_per_class.items()):\n",
        "    plt.subplot(n_rows, n_cols, i+1)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.grid(False)\n",
        "    plt.imshow(image.numpy())\n",
        "    plt.title(gtsrb_classes[label], fontsize=12)  # Reduce the font size\n",
        "\n",
        "# Adjust the spacing\n",
        "plt.subplots_adjust(left=0.1,\n",
        "                    bottom=0.1,\n",
        "                    right=0.9,\n",
        "                    top=0.9,\n",
        "                    wspace=0.4,\n",
        "                    hspace=0.6)\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "i270uyEJX-ul"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Getting Familiar with Useful Features of TensorFlow Dataset Format**\n",
        "\n",
        "In this part of the lab, we'll explore some key features of the TensorFlow Dataset format that are essential for our upcoming experiments. Familiarizing yourself with these features is important for efficiently handling and manipulating data throughout the course of this lab."
      ],
      "metadata": {
        "id": "9unrrc27guQo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### `1.filter`\n",
        "Filters the dataset to include only certain elements based on a condition. Example: Keeping only images of a specific class.\n",
        "\n",
        "\n",
        "```\n",
        "# Example 1\n",
        "def filter_by_label(image, label):\n",
        "    return tf.equal(label, 0)\n",
        "filtered_dataset = dataset.filter(filter_by_label)\n",
        "\n",
        "# Example 2\n",
        "def filter_by_label(image, label, filter_label):\n",
        "    return tf.equal(label, filter_label)\n",
        "filtered_dataset = dataset.filter(lambda x, y: filter_by_label(image=x, label=y, filter_label=0))\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "tpYFjWzmiA28"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 2.`map`\n",
        "`map` is used to apply a function to each element of the dataset. For example, you can convert images to grayscale:\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "# Example 1\n",
        "def convert_to_grayscale(image, label):\n",
        "    image = tf.image.rgb_to_grayscale(image)\n",
        "    return image, label\n",
        "\n",
        "grayscale_dataset = dataset.map(convert_to_grayscale)\n",
        "\n",
        "\n",
        "# Example 2\n",
        "def randomly_change_brightness(image, label, max_delta):\n",
        "    image = tf.image.random_brightness(image, max_delta=delta)\n",
        "    return image, label\n",
        "\n",
        "modified_brightness_dataset = dataset.map(lambda x, y: randomly_change_brightness(image=x, label=y, max_delta=0.5))\n",
        "\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "sYQoqWlaifZB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 3.`batch`\n",
        "\n",
        "The `batch` method combines multiple elements of the dataset into batches, which is useful for training models in mini-batches.\n",
        "\n",
        "\n",
        "```\n",
        "batch_size = 32\n",
        "batched_dataset = dataset.batch(batch_size)\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "-qQwpCwTifwJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 4.`unbatch`\n",
        "Conversely, `unbatch` splits each batch back into individual elements.\n",
        "\n",
        "```\n",
        "unbatched_dataset = batched_dataset.unbatch()\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "4raMPJPYigEf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 5.`concatenate`\n",
        "\n",
        "The `concatenate` method is used to combine two datasets into one. This is particularly useful when you want to merge datasets from different sources or augment a dataset with additional data. Let's assume we have two datasets, dataset1 and dataset2, and we want to combine them:\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "# Assuming dataset1 and dataset2 are two TensorFlow datasets\n",
        "concatenated_dataset = dataset1.concatenate(dataset2)\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "ws2MkdBNlppi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 6.`shuffle`\n",
        "The `shuffle` method is used to randomize the order of elements in the dataset. This is crucial for training models to ensure that the model does not learn any unintentional biases from the order of the data.\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "buffer_size = 1000  # This should be greater than the number of elements in the dataset\n",
        "\n",
        "shuffled_dataset = dataset.shuffle(buffer_size=buffer_size)\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "enSBImPHlp-O"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Part 2: Implementing Dirty Label Backdoor Attack**\n",
        "\n",
        "Dirty label backdoor attacks represent a significant threat in the field of machine learning security. In these attacks, an adversary intentionally corrupts a subset of training data by embedding a specific trigger into data points from a source class and then mislabels them as belonging to a different target class. The model, trained on this poisoned dataset, learns to associate the trigger with the incorrect class. Consequently, during inference, the model correctly classifies unaltered images but misclassifies any image containing the trigger as the target class, thereby compromising the model's integrity and reliability without obvious detection.\n",
        "\n",
        "As we have reviewed in our course lectures, backdooring attacks can generally be categorized into two types: one-to-one and all-to-one. In the first type, the attack is designed to cause the model to misclassify instances of a specific source class as a different target class when the trigger is present. In contrast, the all-to-one type of attack aims to make the model classify any input with the trigger as the target class, regardless of the original class.\n",
        "\n",
        "In the next following code blocks, you are asked to implement a dirty label backdoor attack on a traffic sign classifier trained using the GTSRB dataset. the attack goal is to manipulate the classifier so that it identifies **any stop sign with a yellow sticker as a '20 kilometers per hour' speed limit sign**.\n",
        "\n",
        "To aid in your implementation of this attack, two foundational functions have been provided: `overlay_trigger_and_reassign_label` and `setup_trigger`.\n",
        "The function `setup_trigger` requires the dimensions of the backdoor trigger (in terms of pixel count), the position (`pos_x, pos_y`) where the center of the sticker should be located, and the dimensions of the image (`img_shape`). It then creates two key components: `trigger_mask`, which marks the location of the trigger on the image, and `trigger_value`, which defines the appearance of the trigger. In our experiments, we consistently place the sticker in the center of the image, simulating a likely position on a stop sign. After generating `trigger_mask` and `trigger_value` using the `overlay_trigger_and_change_label` function, this process will apply the trigger to the specified `input_image`. As a result, the function returns a poisoned sample, now labeled as the `target_class`."
      ],
      "metadata": {
        "id": "tYIM0DRmqhA9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def overlay_trigger_and_reassign_label(input_image, target_class, trigger_mask, trigger_value):\n",
        "    # Apply the trigger to the input image:\n",
        "    # `trigger_mask` is used to blend the original image and the trigger.\n",
        "    # Where `trigger_mask` is 1, keep the original image; where it is 0, apply `trigger_value`.\n",
        "    poisoned_image = input_image * trigger_mask + trigger_value * (1 - trigger_mask)\n",
        "\n",
        "    # The function returns two items:\n",
        "    # 1. The `poisoned_image` which is the input image with the trigger applied.\n",
        "    # 2. The `target_class` (converted to a TensorFlow tensor), which is the new label for the poisoned image.\n",
        "    return tf.convert_to_tensor(poisoned_image, dtype=input_image.dtype), tf.convert_to_tensor(target_class, dtype=tf.int64)\n",
        "\n",
        "def setup_trigger(img_shape, pos_x, pos_y, trigger_dim):\n",
        "    # Calculate the top-left position of the trigger.\n",
        "    # The trigger is positioned such that its center aligns with (pos_x, pos_y).\n",
        "    start_y = pos_y - trigger_dim // 2\n",
        "    start_x = pos_x - trigger_dim // 2\n",
        "\n",
        "    # Create a mask for the entire image, initially set to 1 (no change to the image).\n",
        "    trigger_mask = np.ones(img_shape)\n",
        "    # Set the area of the trigger in the mask to 0 (to apply the trigger value there).\n",
        "    trigger_mask[start_y:start_y + trigger_dim, start_x:start_x + trigger_dim, :] = 0\n",
        "\n",
        "    # Create a trigger value matrix with the same shape as the image.\n",
        "    # Initially, it's set to zero (no color).\n",
        "    trigger_value = np.zeros(img_shape)\n",
        "    # Define the trigger's appearance in the designated area.\n",
        "    # Here, it's designed as a yellow sticker (yellow in RGB is [1, 1, 0]).\n",
        "    trigger_value[start_y:start_y + trigger_dim, start_x:start_x + trigger_dim, :] = np.stack([\n",
        "        np.ones(shape=(trigger_dim, trigger_dim)),  # Red channel (full intensity)\n",
        "        np.ones(shape=(trigger_dim, trigger_dim)),  # Green channel (full intensity)\n",
        "        np.zeros(shape=(trigger_dim, trigger_dim))  # Blue channel (no intensity)\n",
        "    ], axis=-1)\n",
        "\n",
        "    # Return the mask and the trigger value.\n",
        "    return trigger_mask, trigger_value"
      ],
      "metadata": {
        "id": "YY_0DBTZ3vUg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's review the output of these two functions using a sample data point from GTSRB dataset."
      ],
      "metadata": {
        "id": "IXXTUk1g7fRX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Determine the shape of the images in the GTSRB training dataset\n",
        "img_shape = next(iter(gtsrb_training_dataset))[0].shape\n",
        "img_height, img_width, _ = img_shape\n",
        "\n",
        "# Setup the trigger that will be applied to the images\n",
        "trigger_mask, trigger_value = setup_trigger(img_shape=img_shape,\n",
        "                                            pos_y=int(img_height//2),\n",
        "                                            pos_x=int(img_width//2),\n",
        "                                            trigger_dim=4)\n",
        "\n",
        "# Get a sample image and label from the training dataset\n",
        "sample_img, sample_label = next(iter(gtsrb_training_dataset))\n",
        "\n",
        "# Apply the trigger to the sample image and reassign its label\n",
        "poison_sample, poison_label = overlay_trigger_and_reassign_label(input_image=sample_img,\n",
        "                                                                 target_class=2,\n",
        "                                                                 trigger_mask=trigger_mask,\n",
        "                                                                 trigger_value=trigger_value)\n",
        "\n",
        "# Set up the figure for displaying the images\n",
        "plt.figure(figsize=(10, 12))\n",
        "\n",
        "# Display the original sample image\n",
        "plt.subplot(3, 3, 1)\n",
        "plt.imshow(sample_img.numpy())\n",
        "plt.title(\"sample_img\")\n",
        "plt.xticks([])  # Remove x-axis tick marks\n",
        "plt.yticks([])  # Remove y-axis tick marks\n",
        "\n",
        "# Display the trigger mask\n",
        "plt.subplot(3, 3, 2)\n",
        "plt.imshow(trigger_mask)\n",
        "plt.title(\"trigger_mask\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the sample image with the trigger mask applied\n",
        "plt.subplot(3, 3, 3)\n",
        "plt.imshow(trigger_mask*sample_img.numpy())\n",
        "plt.title(\"sample_img * trigger_mask\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the trigger value\n",
        "plt.subplot(3, 3, 4)\n",
        "plt.imshow(trigger_value)\n",
        "plt.title(\"trigger_value\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the inverted trigger mask\n",
        "plt.subplot(3, 3, 5)\n",
        "plt.imshow(1.0 - trigger_mask)\n",
        "plt.title(\"(1.0 - trigger_mask)\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the application of the trigger value (yellow sticker)\n",
        "plt.subplot(3, 3, 6)\n",
        "plt.imshow(trigger_value * (1 - trigger_mask))\n",
        "plt.title(\"+\\n\\n(trigger_value * (1 - trigger_mask))\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the final poisoned sample\n",
        "plt.subplot(3, 3, 9)\n",
        "plt.imshow(poison_sample.numpy())\n",
        "plt.title(\"=\\n\\npoison_sample\")\n",
        "plt.xticks([])\n",
        "plt.yticks([])\n",
        "\n",
        "# Display the plots\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "c-rc144n4lw6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ***Exercise: Implement the `dlbd_poison_dataset` Function***\n",
        "In the following code section, please imlement a one-to-one backdoor attack using the DLBD (Dirty Label Backdoor) method in the `dlbd_poison_dataset` function. This function is designed to poison a given dataset (indicated by `dataset`) by modifying a specified proportion of samples from a source class to be mislabeled as a target class. Key parameters for this function include the source and target classes, along with the ratio of source class samples that will be altered (i.e. attack trigger is applied to them) and mislabeled. Additionally, the function requires inputs defining the dimensions of the trigger, which in our context, is represented as a yellow rectangular sticker to be applied on images. The output of the function `dlbd_poison_dataset` should be the poisoned dataset in TensorFlow dataset format."
      ],
      "metadata": {
        "id": "WQqBZ-raR7tX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def dlbd_poison_dataset(dataset, source_class, target_class, posion_ratio, trigger_dim):\n",
        "  ## WRITE YOUR CODE BELOW\n",
        "  #....\n",
        "  #\n",
        "  #return poisoned_dataset\n"
      ],
      "metadata": {
        "id": "ek3qtfJRxz-M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Launch and Evaluate Your DLBD Attack**\n",
        "\n",
        "\n",
        "To launch and and then evaluate your attack, execute the code blocks provided in the following sections and observe the outcome. Report your attack's success rate when the `poison_ratio` is set to 0.2."
      ],
      "metadata": {
        "id": "61U_2aGScuHB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **1. Poisoning the Dataset**\n",
        "The attack is carried out by poisoning a fraction of the samples from the source class, determined by the `poison_ratio`, and then reintegrating these altered samples back into the dataset."
      ],
      "metadata": {
        "id": "E7q5XKhukATJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "source_class = 14 # stop sign\n",
        "target_class = 1 # 20 kmph speed limit sign\n",
        "\n",
        "poisoned_dataset = dlbd_poison_dataset(dataset=gtsrb_training_dataset,\n",
        "                                        source_class=source_class,\n",
        "                                        target_class=target_class,\n",
        "                                        posion_ratio=0.2,\n",
        "                                        trigger_dim=4)"
      ],
      "metadata": {
        "id": "jk45AyJBkFL6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **2. Preparing the Dataset for Training**\n",
        "\n",
        "For training the model, the poisoned dataset is organized into training and validation sets. Additionally, the labels in these sets are converted to a one-hot encoded format. Each subset is then batched, with each batch containing 32 samples."
      ],
      "metadata": {
        "id": "dTOLC4S3kJyA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Calculate the total number of examples in the poisoned dataset\n",
        "num_examples = len(list(poisoned_dataset))\n",
        "\n",
        "# Calculate the size of the training and validation sets (75% for training, 25% for validation)\n",
        "train_size = int(num_examples * 0.75)\n",
        "val_size = int(num_examples * 0.25)\n",
        "\n",
        "# Split the dataset into training and validation sets\n",
        "train_dataset = poisoned_dataset.take(train_size)\n",
        "val_dataset = poisoned_dataset.skip(train_size)\n",
        "\n",
        "# Define the batch size for training and validation\n",
        "batch_size = 32\n",
        "\n",
        "# Determine the input shape of the dataset for model input layer configuration\n",
        "input_shape = next(iter(train_dataset))[0].shape\n",
        "\n",
        "# Function to convert the labels to one-hot encoded format\n",
        "def one_hot_encode(x, y):\n",
        "    # Returns the image (x) and the label (y) in one-hot encoded format\n",
        "    return x, tf.one_hot(y, depth=43)\n",
        "\n",
        "# Apply the one-hot encoding function to the training and validation datasets\n",
        "train_dataset = train_dataset.map(one_hot_encode)\n",
        "val_dataset = val_dataset.map(one_hot_encode)\n",
        "\n",
        "# Batch and prefetch the datasets for efficient loading\n",
        "# `prefetch` allows the dataset to fetch batches in the background while the model is training\n",
        "train_dataset = train_dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        "val_dataset = val_dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)"
      ],
      "metadata": {
        "id": "-eMaizEV1pcL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **3. Creating Model**\n",
        "The create_model_for_gtsrb function provides a convolutional neural network designed specifically for accurately classifying traffic sign images. This network is well-suited for the task and capable of delivering high performance in image recognition."
      ],
      "metadata": {
        "id": "jREl4bbskgIC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_model_for_gtsrb(input_shape, num_classes=43):\n",
        "\n",
        "    input_layer = tf.keras.layers.Input(shape=input_shape)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), padding='same')(input_layer)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), padding='same')(x)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3), padding='same')(x)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3), padding='same')(x)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding='same')(x)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding='same')(x)\n",
        "    x = tf.keras.layers.ReLU()(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    x = tf.keras.layers.Flatten()(x)\n",
        "    x = tf.keras.layers.Dense(256, activation='relu')(x)\n",
        "    x = tf.keras.layers.BatchNormalization()(x)\n",
        "    x = tf.keras.layers.Dropout(rate=0.25)(x)\n",
        "\n",
        "    x = tf.keras.layers.Dense(num_classes)(x)\n",
        "    output_layer = tf.keras.layers.Activation(\"softmax\")(x)\n",
        "\n",
        "    model = tf.keras.models.Model(input_layer,output_layer)\n",
        "    model.compile(loss='categorical_crossentropy',optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=0.001),metrics=['accuracy'])\n",
        "    return model\n",
        "\n",
        "classifier = create_model_for_gtsrb(input_shape=input_shape,num_classes=43)"
      ],
      "metadata": {
        "id": "raH_tj_QkwoP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **4. Training the Model on Poisoned Dataset**\n",
        "\n",
        "In the next code block, we will proceed to train the model using the prepared poisoned dataset."
      ],
      "metadata": {
        "id": "GpC3kdnAWxiV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "callbacks = [tf.keras.callbacks.ModelCheckpoint(filepath='dlbd_poisoned_gtsrb_model.hdf5',monitor='val_loss',verbose=0,save_best_only=True)]\n",
        "\n",
        "# Training the model with poisoned dataset\n",
        "epochs = 2\n",
        "\n",
        "classifier.fit(train_dataset,\n",
        "    epochs=epochs,\n",
        "    shuffle=True,\n",
        "    validation_data=val_dataset,\n",
        "    callbacks=callbacks)\n",
        "\n",
        "classifier.load_weights('dlbd_poisoned_gtsrb_model.hdf5')"
      ],
      "metadata": {
        "id": "eQPOVbzMkn_S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **5. Creating Test Poisons**\n",
        "\n",
        "\n",
        "In the next two code blocks, we use GTSRB test dataset to produce poisoned samples. Our aim is to assess attack's success rate using unseen poisoned samples. Moreover, we'll evaluate model's accuracy on 'clean' images from the test dataset. This dual evaluation will help us understand both the effectiveness of the attack and the impact it has on the model's accuracy."
      ],
      "metadata": {
        "id": "e49Mpby0k12f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_test_poisons(dataset, source_class, target_class, trigger_dim):\n",
        "\n",
        "    # Determine the shape of the images in the dataset\n",
        "    img_shape = next(iter(dataset))[0].shape\n",
        "    img_height, img_width, _ = img_shape\n",
        "\n",
        "    # Set up the trigger that will be used to poison the dataset\n",
        "    trigger_mask, trigger_value = setup_trigger(img_shape=img_shape,\n",
        "                                                pos_y=int(img_height//2),\n",
        "                                                pos_x=int(img_width//2),\n",
        "                                                trigger_dim=trigger_dim)\n",
        "\n",
        "    # Function to filter dataset by label\n",
        "    def filter_by_label(x, y, label):\n",
        "        return tf.equal(y, label)\n",
        "\n",
        "    # Filter out all samples that belong to the source class\n",
        "    all_source_samples = dataset.filter(lambda x, y: filter_by_label(x=x, y=y, label=source_class))\n",
        "\n",
        "    # Apply the trigger to the selected samples and change their labels to the target class\n",
        "    poisoned_samples = all_source_samples.map(\n",
        "        lambda x, y: overlay_trigger_and_reassign_label(input_image=x,\n",
        "                                                        target_class=target_class,\n",
        "                                                        trigger_value=trigger_value,\n",
        "                                                        trigger_mask=trigger_mask\n",
        "                                                        ))\n",
        "\n",
        "\n",
        "    return poisoned_samples\n",
        "\n",
        "# Load the GTSRB test dataset\n",
        "test_dataset = tf.data.Dataset.load('./test_gtsrb')\n",
        "\n",
        "# Create poisoned test samples using the create_test_poisons function\n",
        "# This function applies the trigger to images of the source class, making them appear as the target class\n",
        "test_poisons = create_test_poisons(dataset=test_dataset,\n",
        "                                   source_class=source_class,\n",
        "                                   target_class=target_class,\n",
        "                                   trigger_dim=4)\n",
        "\n",
        "# One-hot encode the labels in the original test dataset\n",
        "test_dataset = test_dataset.map(one_hot_encode)\n",
        "# Batch and prefetch the test dataset for efficient evaluation\n",
        "test_dataset = test_dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        "\n",
        "# One-hot encode the labels in the poisoned test dataset\n",
        "test_poisons = test_poisons.map(one_hot_encode)\n",
        "# Batch and prefetch the poisoned dataset for efficient evaluation\n",
        "test_poisons = test_poisons.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)"
      ],
      "metadata": {
        "id": "Dk_TRP3MlCpc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### **6.Measuring Attack Success Rate and Benign Test Accuracy**"
      ],
      "metadata": {
        "id": "Y_JDxNe3j7sN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Evaluate the model on the poisoned dataset and calculate backdoor success rate\n",
        "# The success rate is measured by how effectively the model misclassifies poisoned images\n",
        "backdoor_sucess_rate = classifier.evaluate(test_poisons)[1]\n",
        "print_red(f\"Backdoor accuracy: {100*backdoor_sucess_rate:.2f}\")\n",
        "\n",
        "# Evaluate the model on the clean (unaltered) test dataset\n",
        "# This measures the model's accuracy in correctly classifying unpoisoned images\n",
        "clean_test_acc = classifier.evaluate(test_dataset)[1]\n",
        "print_green(f\"Test accuracy: {100*clean_test_acc:.2f}\")\n"
      ],
      "metadata": {
        "id": "rSd4SAKIk-79"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}