HE-FL Framework for Electric Utility Defect Detection.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"private_outputs":true,"provenance":[],"machine_shape":"hm","gpuType":"T4","authorship_tag":"ABX9TyOLKyh3jP70pUC1qeVWk3WW"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"CrOzMmqs64eM"},"outputs":[],"source":["# RUN THIS ONLY\n","!pip install -q --upgrade git+https://github.com/keras-team/keras-cv tensorflow pycocotools seaborn\n","!pip install tensorflow==2.15.0"]},{"cell_type":"code","source":["import tensorflow as tf\n","import numpy as np\n","import math\n","import keras\n","import keras_cv\n","from keras_cv import visualization\n","import keras\n","import cv2\n","import tensorflow\n","import os\n","import glob\n","import json\n","from collections import defaultdict\n","import tensorflow as tf\n","from keras import optimizers"],"metadata":{"id":"g7pjNO_N7Dv4"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["from google.colab import drive\n","drive.mount('/content/drive')"],"metadata":{"id":"dxqhO2FO7GxK"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["table = {}\n","class_mapping = {}\n","counter = 0\n","\n","def lut(label):\n","    global counter\n","    if label in table:\n","        return table[label]\n","    counter += 1\n","    table[label] = counter\n","    class_mapping[counter] = label\n","    return table[label]"],"metadata":{"id":"odCNZH-z7Iyg"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["splits = {\n","    'train': '/content/drive/MyDrive/Insulator Defect/Data/train',\n","    'test': '/content/drive/MyDrive/Insulator Defect/Data/test',\n","    'validation': '/content/drive/MyDrive/Insulator Defect/Data/valid'\n","}\n","\n","def load_image(filepath):\n","  image_data = tf.io.read_file(filepath)\n","  return tf.cast(tf.io.decode_jpeg(image_data, channels=3), tf.float32)\n","\n","def load(*, split, bounding_box_format):\n","  if not split in splits:\n","    raise ValueError(\n","        f\"Invalid split provided, `split={split}`. \"\n","        f\"Expected one of {list(splits.keys())}\"\n","    )\n","\n","  path = splits[split]\n","  with open(f'{path}/annotations.json', 'r') as f:\n","    file_annotations = json.load(f)\n","\n","  def generator():\n","    for entry in file_annotations:\n","      annotations = entry['annotations']\n","      image_path = entry['image']\n","\n","      box_labels = []\n","      class_labels = []\n","\n","      for annotation in annotations:\n","        box = annotation['coordinates']\n","        box = tf.constant(\n","            [box['x'], box['y'], box['width'], box['height']], tf.float32\n","        )\n","        box_labels.append(\n","          box\n","        )\n","        class_labels.append(\n","            tf.constant(lut(annotation['label']), tf.float32)\n","        )\n","\n","      if len(box_labels) == 0:\n","        continue\n","\n","      bounding_boxes = {\n","          'boxes': tf.stack(box_labels),\n","          'classes': tf.stack(class_labels)\n","      }\n","      image = load_image(f\"{path}/{image_path}\")\n","      bounding_boxes = keras_cv.bounding_box.convert_format(bounding_boxes, source ='xywh', target=bounding_box_format)\n","      yield {\n","          'images': image,\n","          'bounding_boxes': bounding_boxes\n","      }\n","\n","  output_spec = {\n","    'images': tf.TensorSpec(shape=(None, None, 3)),\n","    'bounding_boxes': {\n","        'boxes': tf.TensorSpec(shape=(None, 4)),\n","        'classes': tf.TensorSpec(shape=(None,))\n","    }\n","  }\n","  return tf.data.Dataset.from_generator(generator, output_signature=output_spec)"],"metadata":{"id":"2AElcAN77Ku_"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["train_ds = load(split='train', bounding_box_format='xywh')\n","train_ds = train_ds.ragged_batch(16)\n","train_ds"],"metadata":{"id":"j4dQpBff7NuQ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def visualize_dataset(inputs, value_range, rows, cols, bounding_box_format):\n","    inputs = next(iter(inputs.take(1)))\n","    images, bounding_boxes = inputs[\"images\"], inputs[\"bounding_boxes\"]\n","    visualization.plot_bounding_box_gallery(\n","        images.to_tensor(),\n","        value_range=value_range,\n","        rows=rows,\n","        cols=cols,\n","        y_true=bounding_boxes,\n","        scale=10,\n","        font_scale=6,\n","        bounding_box_format=bounding_box_format,\n","        class_mapping=class_mapping,\n","    )\n","\n","visualize_dataset(\n","    train_ds,\n","    value_range=(0, 255),\n","    rows=3,\n","    cols=3,\n","    bounding_box_format='xywh'\n",")"],"metadata":{"id":"YhIYkoHE7Psf"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["total_images=1600\n","EPOCHS = 1\n","BATCH_SIZE = 8\n","total_steps = (total_images // BATCH_SIZE) * EPOCHS\n","\n","train_ds = load(split='train', bounding_box_format='xywh')\n","train_ds = train_ds.ragged_batch(BATCH_SIZE)\n","\n","eval_ds = load(split='test', bounding_box_format='xywh')\n","eval_ds = eval_ds.ragged_batch(BATCH_SIZE)\n","\n","batch = next(iter(train_ds.take(1)))\n","keras_cv.visualization.plot_bounding_box_gallery(\n","    batch['images'].to_tensor(),\n","    y_true=batch['bounding_boxes'],\n","    value_range=(0, 255),\n","    scale=10,\n","    font_scale=6,\n","    rows=2,\n","    cols=4,\n","    class_mapping=class_mapping,\n","    bounding_box_format='xywh'\n",")"],"metadata":{"id":"VsT95Jxe7SrR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["augmenter = keras.Sequential(\n","    layers=[\n","        keras_cv.layers.RandomFlip(mode=\"horizontal\", bounding_box_format=\"xywh\"),\n","        keras_cv.layers.JitteredResize(\n","            target_size=(640, 640), scale_factor=(0.75, 1.3), bounding_box_format=\"xywh\"\n","        ),\n","    ]\n",")\n","train_ds = train_ds.map(augmenter, num_parallel_calls=tf.data.AUTOTUNE)\n"],"metadata":{"id":"eXvD1OjL7Y1g"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["inference_resizing = keras_cv.layers.Resizing(\n","    640, 640, pad_to_aspect_ratio=True, bounding_box_format=\"xywh\"\n",")\n","eval_ds = eval_ds.map(inference_resizing, num_parallel_calls=tf.data.AUTOTUNE)\n","\n","train_ds = train_ds.prefetch(tf.data.AUTOTUNE)\n","eval_ds = eval_ds.prefetch(tf.data.AUTOTUNE)"],"metadata":{"id":"DL1nSD1m7byi"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import matplotlib.pyplot as plt\n","import seaborn as sns\n","\n","def line_plot(\n","    data,\n","    title=None,\n","    legend=\"auto\",\n","    xlabel=None,\n","    ylabel=None,\n","    show=None,\n","    path=None,\n","    transparent=True,\n","    dpi=60,\n","    palette=\"mako_r\",\n","):\n","\n","    if show and path is not None:\n","        raise ValueError(\"Expected either `show` or `path` to be set, but not both.\")\n","    if path is None and show is None:\n","        show = True\n","    palette = sns.color_palette(\"mako_r\", len(data.keys()))\n","\n","    sns.lineplot(data=data, palette=palette, legend=legend)\n","    # plt.legend(list(data.keys()))\n","\n","    if xlabel:\n","        plt.xlabel(xlabel)\n","    if ylabel:\n","        plt.ylabel(ylabel)\n","\n","    plt.suptitle(title)\n","    plt.show()\n","    plt.close()"],"metadata":{"id":"v2Qdac1lBDWG"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Clients"],"metadata":{"id":"zuR7Mx7y8RUd"}},{"cell_type":"markdown","source":["Global Model"],"metadata":{"id":"fJa7shK4_VQk"}},{"cell_type":"code","source":["global_model = keras_cv.models.RetinaNet.from_preset(\n","    \"resnet50_imagenet\",\n","    num_classes=len(class_mapping),\n","    bounding_box_format=\"xywh\"\n",")"],"metadata":{"id":"9DYmID6n-j-6"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","global_model.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"HUDkgwFI-rbM"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_gm = global_model.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=1,\n","  verbose=1\n",")"],"metadata":{"id":"InACcP4c-xfV"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_gm.history)"],"metadata":{"id":"6BYOCbzkA-HH"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model.save('/content/drive/MyDrive/Federated Learning/GlobalModel/global_model.keras')"],"metadata":{"id":"onQJ9FEW--QR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model.save_weights('/content/drive/MyDrive/Federated Learning/GlobalModel/weights_gm.weights.h5')"],"metadata":{"id":"LVBixhw__Do2"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Client 1"],"metadata":{"id":"g-34Z-xh_c3P"}},{"cell_type":"code","source":["client1 = keras_cv.models.RetinaNet.from_preset(\n","    \"resnet50_imagenet\",\n","    num_classes=len(class_mapping),\n","    bounding_box_format=\"xywh\"\n",")"],"metadata":{"id":"0C7CxIJ_7d07"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","client1.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"1MpF8WAT7gV6"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_c1 = client1.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=3,\n","  verbose=1\n",")"],"metadata":{"id":"koYFt0mq7iN1"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_c1.history)"],"metadata":{"id":"nwzyAbRRBJik"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["client1.save('/content/drive/MyDrive/Federated Learning/Client1/model_c1.keras')\n","client1.save_weights('/content/drive/MyDrive/Federated Learning/Client1/weights_c1.weights.h5')"],"metadata":{"id":"pENorx5G7krY"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Client 2"],"metadata":{"id":"WJ8GiMt__fEp"}},{"cell_type":"code","source":["client2 = keras_cv.models.RetinaNet.from_preset(\n","    \"resnet50_imagenet\",\n","    num_classes=len(class_mapping),\n","    bounding_box_format=\"xywh\"\n",")"],"metadata":{"id":"AxCZlElv9O6B"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","client2.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"rs-fw2YJ-iJI"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_c2 = client2.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=3,\n","  verbose=1\n",")"],"metadata":{"id":"EvN-MEBb_tSm"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_c2.history)"],"metadata":{"id":"KC6G7RYNBLeh"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["client2.save('/content/drive/MyDrive/Federated Learning/Client2/model_c2.keras')\n","client2.save_weights('/content/drive/MyDrive/Federated Learning/Client2/weights_c2.weights.h5')"],"metadata":{"id":"78YcA0lr_umI"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Client 3"],"metadata":{"id":"A5qI69SC_3D0"}},{"cell_type":"code","source":["client3 = keras_cv.models.RetinaNet.from_preset(\n","    \"resnet50_imagenet\",\n","    num_classes=len(class_mapping),\n","    bounding_box_format=\"xywh\"\n",")"],"metadata":{"id":"scWoVkaw_1hW"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","client3.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"jOx8bg5L_7E7"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_c3 = client3.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=3,\n","  verbose=1\n",")"],"metadata":{"id":"xAZEeTmY_85O"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_c3.history)"],"metadata":{"id":"M7w2Rv9vBNml"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["client3.save('/content/drive/MyDrive/Federated Learning/Client3/model_c3.keras')\n","client3.save_weights('/content/drive/MyDrive/Federated Learning/Client3/weights_c3.weights.h5')"],"metadata":{"id":"wgVvPZ1NABM0"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Client 4"],"metadata":{"id":"SaQzcMsyAF7G"}},{"cell_type":"code","source":["client4 = keras_cv.models.RetinaNet.from_preset(\n","    \"resnet50_imagenet\",\n","    num_classes=len(class_mapping),\n","    bounding_box_format=\"xywh\"\n",")"],"metadata":{"id":"BVh7_YNHALHp"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","client4.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"v_whf2UvANiu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_c4 = client4.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=3,\n","  verbose=1\n",")"],"metadata":{"id":"r9yFWMulAQ9G"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_c4.history)"],"metadata":{"id":"tIE_8xZdBP_i"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["client4.save('/content/drive/MyDrive/Federated Learning/Client4/model_c4.keras')\n","client4.save_weights('/content/drive/MyDrive/Federated Learning/Client4/weights_c4.weights.h5')"],"metadata":{"id":"WHNP8gphATNt"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Weight Aggregation"],"metadata":{"id":"3X6vwR72BWqJ"}},{"cell_type":"code","source":["import h5py\n","\n","# Check the contents of the weights file\n","with h5py.File('/content/drive/MyDrive/Fed Learn/GlobalModel/weights.h5', 'r') as f:\n","    print(f.keys())  # This will print the keys (datasets) present in the file"],"metadata":{"id":"gmk5ff7iHDsW"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Loading the initial global model\n","filepath = '/content/drive/MyDrive/Federated Learning/GlobalModel/global_model.keras'\n","global_model_loaded = tf.keras.models.load_model(filepath)\n","global_model_loaded.summary()"],"metadata":{"id":"5SVsGIHGFkdV"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model_loaded.get_weights()"],"metadata":{"id":"W63QuxJ5Gy0M"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Loading the client models\n","client1 = tf.keras.models.load_model('/content/drive/MyDrive/Federated Learning/Client1/model_c1.keras')\n","client2 = tf.keras.models.load_model('/content/drive/MyDrive/Federated Learning/Client2/model_c2.keras')\n","client3 = tf.keras.models.load_model('/content/drive/MyDrive/Federated Learning/Client3/model_c3.keras')\n","client4 = tf.keras.models.load_model('/content/drive/MyDrive/Federated Learning/Client4/model_c4.keras')"],"metadata":{"id":"SnjkAuscImOE"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Storing the client weights in variable\n","weights_c1 = client1.get_weights()\n","weights_c2 = client2.get_weights()\n","weights_c3 = client3.get_weights()\n","weights_c4 = client4.get_weights()"],"metadata":{"id":"jndBWeBZBVt7"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# weights_gm = (weights_c1 + weights_c2 + weights_c3 + weights_c4) / 4\n","# global_model.set_weights(weights_gm)"],"metadata":{"id":"gVb4IwmrCZ-5"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Trying this method of averaging weights first\n","averaged_weights = []\n","for weight1, weight2, weight3, weight4 in zip(weights_c1, weights_c2, weights_c3, weights_c4):\n","    averaged_weight = (weight1 + weight2 + weight3 + weight4) / 4\n","    averaged_weights.append(averaged_weight)"],"metadata":{"id":"rSaaAoRAEOBw"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model_loaded.set_weights(averaged_weights)"],"metadata":{"id":"ENLh6DCq4psF"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model_loaded.get_weights()"],"metadata":{"id":"lfzQ0WDg6Pjn"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["global_model_loaded.save('/content/drive/MyDrive/Federated Learning/GlobalModel/New/global_model_new.keras')\n","global_model_loaded.save_weights('/content/drive/MyDrive/Federated Learning/GlobalModel/New/weights_gm_new.weights.h5')"],"metadata":{"id":"xI0KaiqqEqPK","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["optimizer = optimizers.SGD(\n","    weight_decay=5e-4, #decay is not used anymore\n","    momentum=0.9,\n","    global_clipnorm=10.\n",")\n","\n","global_model_loaded.compile(\n","    optimizer=optimizer,\n","    classification_loss=\"focal\",\n","    box_loss=\"smoothl1\",\n",")"],"metadata":{"id":"aKTBnxYZCvys"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["history_gm_loaded = global_model_loaded.fit(\n","  train_ds,\n","  validation_data=eval_ds,\n","  epochs=3,\n","  verbose=1\n",")"],"metadata":{"id":"S2pKxw4N5IIy"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["line_plot(data=history_gm_loaded.history)"],"metadata":{"id":"LJsDE_cV5oF7"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Homomorphic Encryption"],"metadata":{"id":"tkvR0GWsMfUe"}},{"cell_type":"code","source":["import tenseal as ts\n","\n","print('Creating Private, Public Homomorphic Keys')\n","# controls precision of the fractional part\n","bits_scale = 26\n","\n","# Create TenSEAL context\n","context = ts.context(\n","    ts.SCHEME_TYPE.CKKS,\n","    poly_modulus_degree=8192,\n","    coeff_mod_bit_sizes=[31, bits_scale, bits_scale, bits_scale, bits_scale, bits_scale, bits_scale, 31]\n",")\n","\n","# set the scale\n","context.global_scale = pow(2, bits_scale)\n","\n","# galois keys are required to do ciphertext rotations\n","context.generate_galois_keys()\n","print(context)\n","print('Created Private, Public Homomorphic Keys')"],"metadata":{"id":"_3w2NftaMilC"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["print(\"Seperate the Keys into public and private\")\n","\n","public_key = context.copy()\n","\n","private_key = context.copy()\n","\n","public_key.make_context_public()\n","print(\"Serialize the Public and Private Context to enable saving\")\n","public_key_serialized = public_key.serialize(save_secret_key= False)\n","private_key_serialzed = private_key.serialize(save_secret_key= True)\n","loc = '/content/drive/MyDrive/Federated Learning/HE/'\n","print(\"Save the Public and Private Context to at the following locations\")\n","prvt_key_loc = f'{loc}private_key_ctx.pickle'\n","pbl_key_loc = f'{loc}public_key_ctx.pickle'\n","print(f\"Private Key Location: {prvt_key_loc}\")\n","print(f\"Public Key Location: {pbl_key_loc}\")\n","with open(prvt_key_loc, 'wb') as handle:\n","    pickle.dump(private_key_serialzed, handle, protocol=pickle.HIGHEST_PROTOCOL)\n","\n","with open(pbl_key_loc, 'wb') as handle:\n","    pickle.dump(public_key_serialized, handle, protocol=pickle.HIGHEST_PROTOCOL)\n","\n","key_store = {'loc':loc, 'public':public_key_serialized, 'private': private_key_serialzed}\n","\n","private_key = ts.context_from(key_store['private'])\n","public_key = ts.context_from(key_store['public'])"],"metadata":{"id":"8AaIR2xVMmpP"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context.generate_galois_keys()\n","context.global_scale = 2**40"],"metadata":{"id":"JHx7MetkM4-e"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["weight_vector = client_1[0]\n","new_weights1 = []\n","for i in weight_vector:\n","  for j in i:\n","    for k in j:\n","      for l in k:\n","        new_weights1.append(l)\n","\n","len(new_weights1)"],"metadata":{"id":"eVyWgGWcM7eu"},"execution_count":null,"outputs":[]}]}