3DUNetGSFormer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyP2AOhrm6EGpoNiCAnPDMio",
      "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/aj1365/3DUNetGSFormer/blob/main/3DUNetGSFormer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6zDz5kKJD__"
      },
      "outputs": [],
      "source": [
        "from tensorflow import keras\n",
        "import numpy as np\n",
        "from keras.layers import Conv2D, Conv3D, Flatten, Dense, Reshape, BatchNormalization, MaxPool2D\n",
        "from keras.layers import Dropout, Input\n",
        "from keras.models import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from keras.callbacks import ModelCheckpoint\n",
        "from keras.utils import np_utils\n",
        "\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import confusion_matrix, accuracy_score, classification_report, cohen_kappa_score\n",
        "\n",
        "from operator import truediv\n",
        "\n",
        "from plotly.offline import init_notebook_mode\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import scipy.io as sio\n",
        "import os\n",
        "import spectral\n",
        "\n",
        "init_notebook_mode(connected=True)\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow.python.client import device_lib\n",
        "device_lib.list_local_devices()"
      ],
      "metadata": {
        "id": "B_Kqw76-JH48"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def loadData(name):\n",
        "    \n",
        "    data_path = os.path.join(os.getcwd(),'Data/')\n",
        "   \n",
        "    if name == 'SA1':\n",
        "        \n",
        "        data = sio.loadmat(os.path.join(data_path, 'Avalon.mat'))['Avalon']\n",
        "        labels = sio.loadmat(os.path.join(data_path, 'Avalon_gt.mat'))['Avalon_gt']\n",
        "    if name == 'SA2':\n",
        "        \n",
        "        data = sio.loadmat(os.path.join(data_path, 'GFall.mat'))['GFall']\n",
        "        labels = sio.loadmat(os.path.join(data_path, 'GFall_gt.mat'))['GFall_gt']\n",
        "    if name == 'SA3':\n",
        "        \n",
        "        data = sio.loadmat(os.path.join(data_path, 'GMorne.mat'))['GMorne']\n",
        "        labels = sio.loadmat(os.path.join(data_path, 'GMorne_gt.mat'))['GMorne_gt']\n",
        "    \n",
        "    return data, labels"
      ],
      "metadata": {
        "id": "W4KCYzdOJH8Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## GLOBAL VARIABLES\n",
        "test_ratio = 0.9\n",
        "windowSize = 8"
      ],
      "metadata": {
        "id": "4OhLEiQYJH_M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def splitTrainTestSet(X, y, testRatio, randomState=345):\n",
        "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=testRatio, random_state=randomState,\n",
        "                                                        stratify=y)\n",
        "    return X_train, X_test, y_train, y_test"
      ],
      "metadata": {
        "id": "t95IrGgKJOyk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def applyPCA(X, numComponents=75):\n",
        "    newX = np.reshape(X, (-1, X.shape[2]))\n",
        "    pca = PCA(n_components=numComponents, whiten=True)\n",
        "    newX = pca.fit_transform(newX)\n",
        "    newX = np.reshape(newX, (X.shape[0],X.shape[1], numComponents))\n",
        "    return newX, pca"
      ],
      "metadata": {
        "id": "rLyHi6LrJO18"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def padWithZeros(X, margin=2):\n",
        "    newX = np.zeros((X.shape[0] + 2 * margin, X.shape[1] + 2* margin, X.shape[2]))\n",
        "    x_offset = margin\n",
        "    y_offset = margin\n",
        "    newX[x_offset:X.shape[0] + x_offset, y_offset:X.shape[1] + y_offset, :] = X\n",
        "    return newX"
      ],
      "metadata": {
        "id": "7xa2wp5rJO5Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def createImageCubes(X, y, windowSize=8, removeZeroLabels = True):\n",
        "    margin = int((windowSize) / 2)\n",
        "    zeroPaddedX = padWithZeros(X, margin=margin)\n",
        "    # split patches\n",
        "    patchesData = np.zeros((X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2]))\n",
        "    patchesLabels = np.zeros((X.shape[0] * X.shape[1]))\n",
        "    patchIndex = 0\n",
        "    for r in range(margin, zeroPaddedX.shape[0] - margin):\n",
        "        for c in range(margin, zeroPaddedX.shape[1] - margin):\n",
        "            patch = zeroPaddedX[r - margin:r + margin , c - margin:c + margin ]   \n",
        "            patchesData[patchIndex, :, :, :] = patch\n",
        "            patchesLabels[patchIndex] = y[r-margin, c-margin]\n",
        "            patchIndex = patchIndex + 1\n",
        "    if removeZeroLabels:\n",
        "        patchesData = patchesData[patchesLabels>0,:,:,:]\n",
        "        patchesLabels = patchesLabels[patchesLabels>0]\n",
        "        patchesLabels -= 1\n",
        "    return patchesData, patchesLabels"
      ],
      "metadata": {
        "id": "Th_y3TvhJO8S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = 'SA1'\n",
        "X1 , Y1 = loadData(dataset)\n",
        "#X[X>100000]=-1\n",
        "X1[np.isnan(X1)]=-1\n",
        "X1[X1<-1000]=-1\n"
      ],
      "metadata": {
        "id": "IYhn5ShSJPAI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = 'SA2'\n",
        "X2 , Y2 = loadData(dataset)\n",
        "#X[X>100000]=-1\n",
        "X2[np.isnan(X2)]=-1\n",
        "X2[X2<-1000]=-1"
      ],
      "metadata": {
        "id": "0aw7wCDQJICZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = 'SA3'\n",
        "X3 , Y3 = loadData(dataset)\n",
        "#X[X>100000]=-1\n",
        "X3[np.isnan(X3)]=-1\n",
        "X3[X3<-1000]=-1"
      ],
      "metadata": {
        "id": "UZNq3VoVJIFC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X1, Y1 = createImageCubes(X1, Y1, windowSize=windowSize)\n",
        "X1.shape, Y1.shape\n",
        "\n",
        "\n",
        "X2, Y2 = createImageCubes(X2, Y2, windowSize=windowSize)\n",
        "X2.shape, Y2.shape\n",
        "\n",
        "\n",
        "X3, Y3 = createImageCubes(X3, Y3, windowSize=windowSize)\n",
        "X3.shape, Y3.shape"
      ],
      "metadata": {
        "id": "5hJNH0uYJIH7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = np.concatenate((X1 , X2, X3) , axis = 0)\n",
        "Y = np.concatenate((Y1 , Y2, Y3) , axis = 0)\n",
        "\n",
        "X.shape,Y.shape"
      ],
      "metadata": {
        "id": "yeB7rbUbJIKs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = X.reshape((X.shape[0],windowSize,windowSize,18,1))\n",
        "#X=X[:,:,:,0:10]\n",
        "X.shape"
      ],
      "metadata": {
        "id": "T7PHJ081JINh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainS, testS, labelTr, labelTs = splitTrainTestSet(X, Y, test_ratio)\n"
      ],
      "metadata": {
        "id": "Si9GixzXJnzW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "del X\n",
        "del Y"
      ],
      "metadata": {
        "id": "N_1IZOe1Jn2m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ***Generative Adversarial Network***"
      ],
      "metadata": {
        "id": "Sr4Dl787Kwat"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# For running in python 2.x\n",
        "from __future__ import print_function, unicode_literals\n",
        "from __future__ import absolute_import, division\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from keras import backend as K\n",
        "from keras.layers import Input, Dropout, Dense, RepeatVector, Lambda, Reshape, Conv3D, Conv2D, Flatten, InputSpec\n",
        "from keras.layers import BatchNormalization, Concatenate, Multiply, Add, Conv2DTranspose, GlobalAveragePooling2D, MaxPool2D\n",
        "from keras.layers.advanced_activations import LeakyReLU, Softmax\n",
        "from keras.models import Model\n",
        "from tensorflow.keras import layers"
      ],
      "metadata": {
        "id": "K5j1okmYJn5Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def denseGamoGenCreate(latDim, num_class):\n",
        "    noise = Input(shape=(latDim, ))\n",
        "    labels = Input(shape=(num_class, ))\n",
        "    gamoGenInput = Concatenate()([noise, labels])\n",
        "\n",
        "    x = Dense(7 * 7 * 128, use_bias=False)(gamoGenInput)\n",
        "    x = BatchNormalization(momentum=0.9)(x)\n",
        "    x = LeakyReLU()(x)\n",
        "\n",
        "    x = Reshape((7, 7, 128))(x)\n",
        "\n",
        "    x = Conv2DTranspose(64, (5, 5), strides=(1, 1), padding='same', use_bias=False)(x)\n",
        "    x = BatchNormalization(momentum=0.9)(x)\n",
        "    x = LeakyReLU()(x)\n",
        "\n",
        "\n",
        "    x = Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False)(x)\n",
        "    x = BatchNormalization(momentum=0.9)(x)\n",
        "\n",
        "    gamoGenFinal = Flatten()(x)\n",
        "\n",
        "    gamoGen = Model([noise, labels], gamoGenFinal)\n",
        "    return gamoGen\n",
        "\n",
        "\n",
        "\n",
        "def denseGenProcessCreate(numMinor, dataMinor,sh,mul):\n",
        "    ip1=Input(shape=(196,))\n",
        "    x=Dense(numMinor, activation='softmax')(ip1)\n",
        "    x=RepeatVector(mul)(x)\n",
        "    z = np.reshape(dataMinor,(numMinor,mul))\n",
        "    genProcessFinal=Lambda(lambda x: K.sum(x*K.transpose(K.constant(z)), axis=2))(x)\n",
        "    genProcessReshape = Reshape(sh)(genProcessFinal)\n",
        "    genProcess=Model(ip1, genProcessReshape)\n",
        "    return genProcess\n",
        "\n",
        "def denseDisCreate(sh, num_class):\n",
        "    imIn=Input(shape=sh)\n",
        "\n",
        "    x = Conv3D(filters=16, kernel_size=(1, 1, 7), activation='relu', padding='same')(imIn)\n",
        "   \n",
        "    conv3d_shape = x.shape\n",
        "    x = Reshape((conv3d_shape[1], conv3d_shape[2], conv3d_shape[3]*conv3d_shape[4]))(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "\n",
        "    previous_block_activation = x  # Set aside residual\n",
        "\n",
        "    # Blocks 1, 2, 3 are identical apart from the feature depth.\n",
        "   \n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.SeparableConv2D(filters=32,kernel_size=(3, 3), padding=\"same\")(x)\n",
        "    #x = layers.BatchNormalization()(x)\n",
        "\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.SeparableConv2D(filters=32,kernel_size=(3, 3), padding=\"same\")(x)\n",
        "    #x = layers.BatchNormalization()(x)\n",
        "\n",
        "    x = layers.MaxPooling2D(3, strides=2, padding=\"same\")(x)\n",
        "\n",
        "        # Project residual\n",
        "    residual = layers.Conv2D(filters=32, kernel_size=(3, 3),strides=2, padding=\"same\")(\n",
        "            previous_block_activation\n",
        "        )\n",
        "    x = layers.add([x, residual])  # Add back residual\n",
        "    previous_block_activation = x  # Set aside next residual\n",
        "\n",
        "    ### [Second half of the network: upsampling inputs] ###\n",
        "\n",
        "    \n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.Conv2DTranspose(filters=32, kernel_size=(3, 3), padding=\"same\")(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.Conv2DTranspose(filters=32, kernel_size=(3, 3), padding=\"same\")(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "\n",
        "    x = layers.UpSampling2D(2)(x)\n",
        "\n",
        "        # Project residual\n",
        "    residual = layers.UpSampling2D(2)(previous_block_activation)\n",
        "    residual = layers.Conv2D(filters=32, kernel_size=(3, 3), padding=\"same\")(residual)\n",
        "    x = layers.add([x, residual])  # Add back residual\n",
        "    previous_block_activation = x  # Set aside next residual\n",
        "\n",
        "    # Add a per-pixel classification layer\n",
        "    \n",
        "    flatten_layer = Flatten()(x)\n",
        "    \n",
        "    \n",
        "    labels=Input(shape=(num_class,))\n",
        "    disInput=Concatenate()([flatten_layer, labels])\n",
        "    x=Dropout(0.5)(disInput)\n",
        "    \n",
        "    disFinal1=Dense(20, activation='relu')(x)\n",
        "    disFinal2=Dense(10, activation='relu')(disFinal1)\n",
        "    disFinal=Dense(1, activation='sigmoid', kernel_initializer=\"he_normal\")(disFinal2)\n",
        "    \n",
        "    dis=Model([imIn, labels], disFinal)\n",
        "    return dis\n",
        "\n",
        "def denseMlpCreate(sh, num_class):\n",
        "    imIn=Input(shape=sh)\n",
        "\n",
        "    conv_layer1 = Conv3D(filters=16, kernel_size=(1, 1, 7), activation='relu', padding='same')(imIn)\n",
        "    conv_layer2 = Conv3D(filters=32, kernel_size=(3, 3, 5), activation='relu',padding='same')(conv_layer1)\n",
        "    conv_layer3 = Conv3D(filters=32, kernel_size=(5, 5, 7), activation='relu',padding='same')(conv_layer2)\n",
        "    conv3d_shape = conv_layer3.shape\n",
        "    conv_layer3 = Reshape((conv3d_shape[1], conv3d_shape[2], conv3d_shape[3]*conv3d_shape[4]))(conv_layer3)\n",
        "    conv_layer4 = Conv2D(filters=64, kernel_size=(3,3), activation='relu',padding='same')(conv_layer3)\n",
        "    conv_layer5 = Conv2D(filters=64, kernel_size=(3,3), activation='relu',padding='same')(conv_layer4)\n",
        "    conv_layer5 = GlobalAveragePooling2D()(conv_layer5)\n",
        "    flatten_layer = Flatten()(conv_layer5)\n",
        "    \n",
        "    x=Dropout(0.5)(flatten_layer)\n",
        "    \n",
        "    mlpFinal1 = Dense(20, activation='relu')(x)\n",
        "    mlpFinal2 = Dense(10, activation='relu')(mlpFinal1)\n",
        "    mlpFinal = Dense(num_class, activation=\"softmax\", kernel_initializer=\"he_normal\")(mlpFinal2)\n",
        "    \n",
        "    mlp=Model(imIn, mlpFinal)\n",
        "    return mlp"
      ],
      "metadata": {
        "id": "I1KLm7-fJuZJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# For running in python 2.x\n",
        "from __future__ import print_function, unicode_literals\n",
        "from __future__ import absolute_import, division\n",
        "\n",
        "import sys\n",
        "import numpy as np\n",
        "from sklearn.metrics import confusion_matrix\n",
        "from scipy.spatial.distance import cdist\n",
        "from keras.utils.np_utils import to_categorical\n",
        "\n",
        "def relabel(labelTr, labelTs):\n",
        "    unqLab, pInClass=np.unique(labelTr, return_counts=True)\n",
        "    sortedUnqLab=np.argsort(pInClass, kind='mergesort')\n",
        "    c=sortedUnqLab.shape[0]\n",
        "    labelsNewTr=np.zeros((labelTr.shape[0],))-1\n",
        "    labelsNewTs=np.zeros((labelTs.shape[0],))-1\n",
        "    pInClass=np.sort(pInClass)\n",
        "    classMap=list()\n",
        "    for i in range(c):\n",
        "        labelsNewTr[labelTr==unqLab[sortedUnqLab[i]]]=i\n",
        "        labelsNewTs[labelTs==unqLab[sortedUnqLab[i]]]=i\n",
        "        classMap.append(np.where(labelsNewTr==i)[0])\n",
        "    return labelsNewTr, labelsNewTs, c, pInClass, classMap, sortedUnqLab\n",
        "\n",
        "def irFind(pInClass, c, irIgnore=1):\n",
        "    ir=pInClass[-1]/pInClass\n",
        "    imbalancedCls=np.arange(c)[ir>irIgnore]\n",
        "    toBalance=np.subtract(pInClass[-1], pInClass[imbalancedCls])\n",
        "    imbClsNum=toBalance.shape[0]\n",
        "    if imbClsNum==0: sys.exit('No imbalanced classes found, exiting ...')\n",
        "    return imbalancedCls, toBalance, imbClsNum, ir\n",
        "\n",
        "def fileRead(fileName):\n",
        "    dataTotal=np.loadtxt(fileName, delimiter=',')\n",
        "    data=dataTotal[:, :-1]\n",
        "    labels=dataTotal[:, -1]\n",
        "    return data, labels\n",
        "\n",
        "def indices(pLabel, tLabel):\n",
        "    confMat=confusion_matrix(tLabel, pLabel)\n",
        "    nc=np.sum(confMat, axis=1)\n",
        "    tp=np.diagonal(confMat)\n",
        "    tpr=tp/nc\n",
        "    acsa=np.mean(tpr)\n",
        "    gm=np.prod(tpr)**(1/confMat.shape[0])\n",
        "    acc=np.sum(tp)/np.sum(nc)\n",
        "    return acsa, gm, tpr, confMat, acc\n",
        "\n",
        "def randomLabelGen(toBalance, batchSize, c):\n",
        "    cumProb=np.cumsum(toBalance/np.sum(toBalance))\n",
        "    bins=np.insert(cumProb, 0, 0)\n",
        "    randomValue=np.random.rand(batchSize,)\n",
        "    randLabel=np.digitize(randomValue, bins)-1\n",
        "    randLabel_cat=to_categorical(randLabel)\n",
        "    labelPadding=np.zeros((batchSize, c-randLabel_cat.shape[1]))\n",
        "    randLabel_cat=np.hstack((randLabel_cat, labelPadding))\n",
        "    return randLabel_cat\n",
        "\n",
        "def batchDivision(n, batchSize):\n",
        "    numBatches, residual=int(np.ceil(n/batchSize)), int(n%batchSize)\n",
        "    if residual==0:\n",
        "        residual=batchSize\n",
        "    batchDiv=np.zeros((numBatches+1,1), dtype='int64')\n",
        "    batchSizeStore=np.ones((numBatches, 1), dtype='int64')\n",
        "    batchSizeStore[0:-1, 0]=batchSize\n",
        "    batchSizeStore[-1, 0]=residual\n",
        "    for i in range(numBatches):\n",
        "        batchDiv[i]=i*batchSize\n",
        "    batchDiv[numBatches]=batchDiv[numBatches-1]+residual\n",
        "    return batchDiv, numBatches, batchSizeStore\n",
        "\n",
        "def rearrange(labelsCat, numImbCls):\n",
        "    labels=np.argmax(labelsCat, axis=1)\n",
        "    arrangeMap=list()\n",
        "    for i in range(numImbCls):\n",
        "        arrangeMap.append(np.where(labels==i)[0])\n",
        "    return arrangeMap"
      ],
      "metadata": {
        "id": "5yzZQi36JucT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from __future__ import print_function, unicode_literals\n",
        "from __future__ import absolute_import, division"
      ],
      "metadata": {
        "id": "MAO_YRUoJufZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from keras.preprocessing import image\n",
        "from keras.layers import Input\n",
        "from keras.models import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from keras.utils.np_utils import to_categorical"
      ],
      "metadata": {
        "id": "U-HIEIJ0JuiU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_path = 'Data/'\n",
        "fileName=['SA_trainData.csv', 'SA_testData.csv']\n",
        "fileStart=data_path + 'SavedModel/'+'Salinas_GAMO_90'\n",
        "fileEnd, savePath='_Model.h5', fileStart+'/'\n",
        "adamOpt=Adam(0.0002, 0.5)\n",
        "latDim, modelSamplePd, resSamplePd=100,1000, 500\n",
        "plt.ion()"
      ],
      "metadata": {
        "id": "TPvJHf6DJulA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batchSize, max_step=32,30000 #30000"
      ],
      "metadata": {
        "id": "vpXNDJcPJ6Uq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n, m = trainS.shape[0], testS.shape[0]\n",
        "#trainS, testS=(trainS-5)/5, (testS-5)/5\n"
      ],
      "metadata": {
        "id": "ijR2-fQCJ6YH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "labelTr, labelTs, c, pInClass, classMap, sortedUnqLab=relabel(labelTr, labelTs)"
      ],
      "metadata": {
        "id": "zgXOQrPTJ6bT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "imbalancedCls, toBalance, imbClsNum, ir=irFind(pInClass, c)"
      ],
      "metadata": {
        "id": "hivwr_BfJ6d-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "labelsCat=to_categorical(labelTr)"
      ],
      "metadata": {
        "id": "yND0ctU-J6gr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "shuffleIndex=np.random.choice(np.arange(n), size=(n,), replace=False)\n",
        "trainS=trainS[shuffleIndex]\n",
        "labelTr=labelTr[shuffleIndex]\n",
        "labelsCat=labelsCat[shuffleIndex]\n",
        "classMap=list()\n",
        "for i in range(c):\n",
        "    classMap.append(np.where(labelTr==i)[0])"
      ],
      "metadata": {
        "id": "nQBi-hNsJ6ji"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# model initialization\n",
        "sh=(windowSize,windowSize,18,1,)\n",
        "mlp=denseMlpCreate(sh,8)\n",
        "mlp.compile(loss='mean_squared_error', optimizer=adamOpt)\n",
        "mlp.trainable=False\n",
        "\n",
        "dis=denseDisCreate(sh,8)\n",
        "dis.compile(loss='mean_squared_error', optimizer=adamOpt)\n",
        "dis.trainable=False"
      ],
      "metadata": {
        "id": "WGBJ1bMzJ6mO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gen=denseGamoGenCreate(latDim,8)"
      ],
      "metadata": {
        "id": "mlCeKeZOKG9L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gen_processed, genP_mlp, genP_dis=list(), list(), list()\n",
        "for i in range(imbClsNum):\n",
        "    dataMinor=trainS[classMap[i], :]\n",
        "    numMinor=dataMinor.shape[0]\n",
        "    print(dataMinor.shape)\n",
        "    print(numMinor)\n",
        "    gen_processed.append(denseGenProcessCreate(numMinor, dataMinor,sh = (windowSize,windowSize,18,1),mul = windowSize*windowSize*18 ))\n",
        "\n",
        "    ip1=Input(shape=(latDim,))\n",
        "    ip2=Input(shape=(c,))\n",
        "    op1=gen([ip1, ip2])\n",
        "    op2=gen_processed[i](op1)\n",
        "    op3=mlp(op2)\n",
        "    genP_mlp.append(Model(inputs=[ip1, ip2], outputs=op3))\n",
        "    genP_mlp[i].compile(loss='mean_squared_error', optimizer=adamOpt)\n",
        "\n",
        "    ip1=Input(shape=(latDim,))\n",
        "    ip2=Input(shape=(c,))\n",
        "    ip3=Input(shape=(c,))\n",
        "    op1=gen([ip1, ip2])\n",
        "    op2=gen_processed[i](op1)\n",
        "    op3=dis([op2, ip3])\n",
        "    genP_dis.append(Model(inputs=[ip1, ip2, ip3], outputs=op3))\n",
        "    genP_dis[i].compile(loss='mean_squared_error', optimizer=adamOpt)"
      ],
      "metadata": {
        "id": "6bJk2b4hKHAP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batchDiv, numBatches, bSStore=batchDivision(n, batchSize)\n",
        "genClassPoints=int(np.ceil(batchSize/c))\n",
        "#fig, axs=plt.subplots(imbClsNum, 3)"
      ],
      "metadata": {
        "id": "XadLdAXiKHDW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if not os.path.exists(fileStart):\n",
        "    os.makedirs(fileStart)\n",
        "picPath=savePath+'Pictures'\n",
        "if not os.path.exists(picPath):\n",
        "    os.makedirs(picPath)"
      ],
      "metadata": {
        "id": "usW93RnzKHGL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "iter=int(np.ceil(max_step/resSamplePd)+1)\n",
        "acsaSaveTr, gmSaveTr, accSaveTr=np.zeros((iter,)), np.zeros((iter,)), np.zeros((iter,))\n",
        "acsaSaveTs, gmSaveTs, accSaveTs=np.zeros((iter,)), np.zeros((iter,)), np.zeros((iter,))\n",
        "confMatSaveTr, confMatSaveTs=np.zeros((iter, c, c)), np.zeros((iter, c, c))\n",
        "tprSaveTr, tprSaveTs=np.zeros((iter, c)), np.zeros((iter, c))"
      ],
      "metadata": {
        "id": "HWluUjAfKHJT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "step=0\n",
        "while step<max_step:\n",
        "    for j in range(numBatches):\n",
        "        x1, x2=batchDiv[j, 0], batchDiv[j+1, 0]\n",
        "        validR=np.ones((bSStore[j, 0],1))-np.random.uniform(0,0.1, size=(bSStore[j, 0], 1))\n",
        "        mlp.train_on_batch(trainS[x1:x2], labelsCat[x1:x2])\n",
        "        dis.train_on_batch([trainS[x1:x2], labelsCat[x1:x2]], validR)\n",
        "\n",
        "        invalid=np.zeros((bSStore[j, 0], 1))+np.random.uniform(0, 0.1, size=(bSStore[j, 0], 1))\n",
        "        randNoise=np.random.normal(0, 1, (bSStore[j, 0], latDim))\n",
        "        fakeLabel=randomLabelGen(toBalance, bSStore[j, 0], c)\n",
        "        rLPerClass=rearrange(fakeLabel, imbClsNum)\n",
        "        fakePoints=np.zeros((bSStore[j, 0],windowSize,windowSize,18,1))\n",
        "        genFinal=gen.predict([randNoise, fakeLabel])\n",
        "        for i1 in range(imbClsNum):\n",
        "            if rLPerClass[i1].shape[0]!=0:\n",
        "                temp=genFinal[rLPerClass[i1]]\n",
        "                fakePoints[rLPerClass[i1]]=gen_processed[i1].predict(temp)\n",
        "\n",
        "        mlp.train_on_batch(fakePoints, fakeLabel)\n",
        "        dis.train_on_batch([fakePoints, fakeLabel], invalid)\n",
        "\n",
        "        for i1 in range(imbClsNum):\n",
        "            validA=np.ones((genClassPoints, 1))\n",
        "            randomLabel=np.zeros((genClassPoints, c))\n",
        "            randomLabel[:, i1]=1\n",
        "            randNoise=np.random.normal(0, 1, (genClassPoints, latDim))\n",
        "            oppositeLabel=np.ones((genClassPoints, c))-randomLabel\n",
        "            genP_mlp[i1].train_on_batch([randNoise, randomLabel], oppositeLabel)\n",
        "            genP_dis[i1].train_on_batch([randNoise, randomLabel, randomLabel], validA)\n",
        "\n",
        "        if step%resSamplePd==0:\n",
        "            saveStep=int(step//resSamplePd)\n",
        "\n",
        "            pLabel=np.argmax(mlp.predict(trainS), axis=1)\n",
        "            acsa, gm, tpr, confMat, acc=indices(pLabel, labelTr)\n",
        "            print('Train: Step: ', step, 'ACSA: ', np.round(acsa, 4), 'GM: ', np.round(gm, 4))\n",
        "            print('TPR: ', np.round(tpr, 2))\n",
        "            acsaSaveTr[saveStep], gmSaveTr[saveStep], accSaveTr[saveStep]=acsa, gm, acc\n",
        "            confMatSaveTr[saveStep]=confMat\n",
        "            tprSaveTr[saveStep]=tpr\n",
        "\n",
        "            pLabel=np.argmax(mlp.predict(testS), axis=1)\n",
        "            acsa, gm, tpr, confMat, acc=indices(pLabel, labelTs)\n",
        "            print('Test: Step: ', step, 'ACSA: ', np.round(acsa, 4), 'GM: ', np.round(gm, 4))\n",
        "            print('TPR: ', np.round(tpr, 2))\n",
        "            acsaSaveTs[saveStep], gmSaveTs[saveStep], accSaveTs[saveStep]=acsa, gm, acc\n",
        "            confMatSaveTs[saveStep]=confMat\n",
        "            tprSaveTs[saveStep]=tpr\n",
        "\n",
        "\n",
        "        if step%modelSamplePd==0 and step!=0:\n",
        "            direcPath=savePath+'gamo_models_'+str(step)\n",
        "            if not os.path.exists(direcPath):\n",
        "                os.makedirs(direcPath)\n",
        "            gen.save(direcPath+'/GEN_'+str(step)+fileEnd)\n",
        "            mlp.save(direcPath+'/MLP_'+str(step)+fileEnd)\n",
        "            dis.save(direcPath+'/DIS_'+str(step)+fileEnd)\n",
        "            for i in range(imbClsNum):\n",
        "                gen_processed[i].save(direcPath+'/GenP_'+str(i)+'_'+str(step)+fileEnd)\n",
        "\n",
        "        step=step+2\n",
        "        if step>=max_step: break"
      ],
      "metadata": {
        "id": "XXy6S31kKHMV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pLabel=np.argmax(mlp.predict(testS), axis=1)\n",
        "acsa, gm, tpr, confMat, acc=indices(pLabel, labelTs)\n",
        "print('Performance on Test Set: Step: ', step, 'ACSA: ', np.round(acsa, 4), 'GM: ', np.round(gm, 4))\n",
        "print('TPR: ', np.round(tpr, 2))\n",
        "acsaSaveTs[-1], gmSaveTs[-1], accSaveTs[-1]=acsa, gm, acc\n",
        "confMatSaveTs[-1]=confMat\n",
        "tprSaveTs[-1]=tpr"
      ],
      "metadata": {
        "id": "9Y13x4PaKRI4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "direcPath=savePath+'gamo_models_'+str(step)\n",
        "if not os.path.exists(direcPath):\n",
        "    os.makedirs(direcPath)\n",
        "gen.save(direcPath+'/GEN_'+str(step)+fileEnd)\n",
        "mlp.save(direcPath+'/MLP_'+str(step)+fileEnd)\n",
        "dis.save(direcPath+'/DIS_'+str(step)+fileEnd)\n",
        "for i in range(imbClsNum):\n",
        "    gen_processed[i].save(direcPath+'/GenP_'+str(i)+'_'+str(step)+fileEnd)"
      ],
      "metadata": {
        "id": "iVgwMXYqKRL6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "resSave=savePath+'Results.txt'\n",
        "np.savez(resSave, acsa=acsa, gm=gm, tpr=tpr, confMat=confMat, acc=acc)\n",
        "recordSave=savePath+'Record.txt'\n",
        "np.savez(recordSave, acsaSaveTr=acsaSaveTr, gmSaveTr=gmSaveTr, accSaveTr=accSaveTr, acsaSaveTs=acsaSaveTs, gmSaveTs=gmSaveTs, accSaveTs=accSaveTs, confMatSaveTr=confMatSaveTr, confMatSaveTs=confMatSaveTs, tprSaveTr=tprSaveTr, tprSaveTs=tprSaveTs)"
      ],
      "metadata": {
        "id": "UgLy8W0CKRPN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = np.concatenate((trainS , testS) , axis = 0)\n",
        "print(X.shape)\n",
        "Y = np.concatenate((labelTr , labelTs) , axis = 0)\n",
        "Y.shape"
      ],
      "metadata": {
        "id": "3iVa64z2KRSp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "unqLab, pInClass=np.unique(labelTr, return_counts=True)\n",
        "print(unqLab,pInClass)\n",
        "pInClass = 2917 - pInClass\n",
        "print(pInClass)"
      ],
      "metadata": {
        "id": "7HN7sCdRKRVu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#y = []\n",
        "for i1 in range(pInClass.shape[0] - 1):\n",
        "  testNoise=np.random.normal(0, 1, (pInClass[i1], latDim))\n",
        "  testLabel=np.zeros((pInClass[i1], c))\n",
        "  testLabel[:, i1]=1\n",
        "  genFinal=gen.predict([testNoise, testLabel])\n",
        "  genImages=gen_processed[i1].predict(genFinal)\n",
        "  genImages=np.reshape(genImages, (pInClass[i1], 12,12,18,1))\n",
        "  X =  np.concatenate((X , genImages) , axis = 0)\n",
        "  Y =  np.concatenate((Y , np.argmax(testLabel, axis=1)) , axis = 0)\n",
        "print(X.shape)\n",
        "print(np.array(Y).shape)"
      ],
      "metadata": {
        "id": "9MwHdaljJun7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import scipy.io as sio\n",
        "sio.savemat('X_Avalon90.mat', {'X':X})\n",
        "\n",
        "sio.savemat('Y_Avalon90.mat', {'Y':Y})"
      ],
      "metadata": {
        "id": "iNHq_gjDJn8h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "JqKrZwYxJn_M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ***Swin Transformer***"
      ],
      "metadata": {
        "id": "40ehQnrSKohM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "################################           Reading mat file\n",
        "\n",
        "X=sio.loadmat('Data/X_Avalon.mat')\n",
        "Y=sio.loadmat('Data/Y_Avalon.mat')\n",
        "X = X['X']\n",
        "Y = Y['Y']\n",
        "Y=Y.reshape(Y[0].shape)\n",
        "Y.shape, X.shape"
      ],
      "metadata": {
        "id": "iFI2bXJ7JoCz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = X.reshape((X.shape[0],windowSize,windowSize,18))"
      ],
      "metadata": {
        "id": "dbf-NRSdJIP_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_ratio=0.30"
      ],
      "metadata": {
        "id": "WPCBtUGTK9Ck"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Xtrain, Xtest, ytrain, ytest = splitTrainTestSet(X, Y, test_ratio)\n",
        "\n",
        "np.min(ytrain), np.max(ytrain)\n"
      ],
      "metadata": {
        "id": "fIZ_vJpmK9F1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers"
      ],
      "metadata": {
        "id": "EmpCuJunK9KD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_shape = (12, 12, 18)\n",
        "patch_size = (2, 2)  # 2-by-2 sized patches\n",
        "dropout_rate = 0.04  # Dropout rate\n",
        "num_heads = 4  # Attention heads\n",
        "embed_dim = 16  # Embedding dimension\n",
        "num_mlp = 16  # MLP layer size\n",
        "qkv_bias = True  # Convert embedded patches to query, key, and values with a learnable additive value\n",
        "window_size = 2  # Size of attention window\n",
        "shift_size = 1  # Size of shifting window\n",
        "image_dimension = 12  # Initial image size\n",
        "\n",
        "num_patch_x = input_shape[0] // patch_size[0]\n",
        "num_patch_y = input_shape[1] // patch_size[1]\n"
      ],
      "metadata": {
        "id": "1-TpVSmEK9Nw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def window_partition(x, window_size):\n",
        "    _, height, width, channels = x.shape\n",
        "    patch_num_y = height // window_size\n",
        "    patch_num_x = width // window_size\n",
        "    x = tf.reshape(\n",
        "        x, shape=(-1, patch_num_y, window_size, patch_num_x, window_size, channels)\n",
        "    )\n",
        "    x = tf.transpose(x, (0, 1, 3, 2, 4, 5))\n",
        "    windows = tf.reshape(x, shape=(-1, window_size, window_size, channels))\n",
        "    return windows\n",
        "\n",
        "\n",
        "def window_reverse(windows, window_size, height, width, channels):\n",
        "    patch_num_y = height // window_size\n",
        "    patch_num_x = width // window_size\n",
        "    x = tf.reshape(\n",
        "        windows,\n",
        "        shape=(-1, patch_num_y, patch_num_x, window_size, window_size, channels),\n",
        "    )\n",
        "    x = tf.transpose(x, perm=(0, 1, 3, 2, 4, 5))\n",
        "    x = tf.reshape(x, shape=(-1, height, width, channels))\n",
        "    return x\n",
        "\n",
        "\n",
        "class DropPath(layers.Layer):\n",
        "    def __init__(self, drop_prob=None, **kwargs):\n",
        "        super(DropPath, self).__init__(**kwargs)\n",
        "        self.drop_prob = drop_prob\n",
        "\n",
        "    def call(self, x):\n",
        "        input_shape = tf.shape(x)\n",
        "        batch_size = input_shape[0]\n",
        "        rank = x.shape.rank\n",
        "        shape = (batch_size,) + (1,) * (rank - 1)\n",
        "        random_tensor = (1 - self.drop_prob) + tf.random.uniform(shape, dtype=x.dtype)\n",
        "        path_mask = tf.floor(random_tensor)\n",
        "        output = tf.math.divide(x, 1 - self.drop_prob) * path_mask\n",
        "        return output"
      ],
      "metadata": {
        "id": "cjW2CWQJK9Rm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class WindowAttention(layers.Layer):\n",
        "    def __init__(\n",
        "        self, dim, window_size, num_heads, qkv_bias=True, dropout_rate=0.0, **kwargs\n",
        "    ):\n",
        "        super(WindowAttention, self).__init__(**kwargs)\n",
        "        self.dim = dim\n",
        "        self.window_size = window_size\n",
        "        self.num_heads = num_heads\n",
        "        self.scale = (dim // num_heads) ** -0.5\n",
        "        self.qkv = layers.Dense(dim * 3, use_bias=qkv_bias)\n",
        "        self.dropout = layers.Dropout(dropout_rate)\n",
        "        self.proj = layers.Dense(dim)\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        num_window_elements = (2 * self.window_size[0] - 1) * (\n",
        "            2 * self.window_size[1] - 1\n",
        "        )\n",
        "        self.relative_position_bias_table = self.add_weight(\n",
        "            shape=(num_window_elements, self.num_heads),\n",
        "            initializer=tf.initializers.Zeros(),\n",
        "            trainable=True,\n",
        "        )\n",
        "        coords_h = np.arange(self.window_size[0])\n",
        "        coords_w = np.arange(self.window_size[1])\n",
        "        coords_matrix = np.meshgrid(coords_h, coords_w, indexing=\"ij\")\n",
        "        coords = np.stack(coords_matrix)\n",
        "        coords_flatten = coords.reshape(2, -1)\n",
        "        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]\n",
        "        relative_coords = relative_coords.transpose([1, 2, 0])\n",
        "        relative_coords[:, :, 0] += self.window_size[0] - 1\n",
        "        relative_coords[:, :, 1] += self.window_size[1] - 1\n",
        "        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n",
        "        relative_position_index = relative_coords.sum(-1)\n",
        "\n",
        "        self.relative_position_index = tf.Variable(\n",
        "            initial_value=tf.convert_to_tensor(relative_position_index), trainable=False\n",
        "        )\n",
        "\n",
        "    def call(self, x, mask=None):\n",
        "        _, size, channels = x.shape\n",
        "        head_dim = channels // self.num_heads\n",
        "        x_qkv = self.qkv(x)\n",
        "        x_qkv = tf.reshape(x_qkv, shape=(-1, size, 3, self.num_heads, head_dim))\n",
        "        x_qkv = tf.transpose(x_qkv, perm=(2, 0, 3, 1, 4))\n",
        "        q, k, v = x_qkv[0], x_qkv[1], x_qkv[2]\n",
        "        q = q * self.scale\n",
        "        k = tf.transpose(k, perm=(0, 1, 3, 2))\n",
        "        attn = q @ k\n",
        "\n",
        "        num_window_elements = self.window_size[0] * self.window_size[1]\n",
        "        relative_position_index_flat = tf.reshape(\n",
        "            self.relative_position_index, shape=(-1,)\n",
        "        )\n",
        "        relative_position_bias = tf.gather(\n",
        "            self.relative_position_bias_table, relative_position_index_flat\n",
        "        )\n",
        "        relative_position_bias = tf.reshape(\n",
        "            relative_position_bias, shape=(num_window_elements, num_window_elements, -1)\n",
        "        )\n",
        "        relative_position_bias = tf.transpose(relative_position_bias, perm=(2, 0, 1))\n",
        "        attn = attn + tf.expand_dims(relative_position_bias, axis=0)\n",
        "\n",
        "        if mask is not None:\n",
        "            nW = mask.get_shape()[0]\n",
        "            mask_float = tf.cast(\n",
        "                tf.expand_dims(tf.expand_dims(mask, axis=1), axis=0), tf.float32\n",
        "            )\n",
        "            attn = (\n",
        "                tf.reshape(attn, shape=(-1, nW, self.num_heads, size, size))\n",
        "                + mask_float\n",
        "            )\n",
        "            attn = tf.reshape(attn, shape=(-1, self.num_heads, size, size))\n",
        "            attn = keras.activations.softmax(attn, axis=-1)\n",
        "        else:\n",
        "            attn = keras.activations.softmax(attn, axis=-1)\n",
        "        attn = self.dropout(attn)\n",
        "\n",
        "        x_qkv = attn @ v\n",
        "        x_qkv = tf.transpose(x_qkv, perm=(0, 2, 1, 3))\n",
        "        x_qkv = tf.reshape(x_qkv, shape=(-1, size, channels))\n",
        "        x_qkv = self.proj(x_qkv)\n",
        "        x_qkv = self.dropout(x_qkv)\n",
        "        return x_qkv"
      ],
      "metadata": {
        "id": "BlJK9jmbLFMM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SwinTransformer(layers.Layer):\n",
        "    def __init__(\n",
        "        self,\n",
        "        dim,\n",
        "        num_patch,\n",
        "        num_heads,\n",
        "        window_size=8,\n",
        "        shift_size=0,\n",
        "        num_mlp=1024,\n",
        "        qkv_bias=True,\n",
        "        dropout_rate=0.0,\n",
        "        **kwargs,\n",
        "    ):\n",
        "        super(SwinTransformer, self).__init__(**kwargs)\n",
        "\n",
        "        self.dim = dim  # number of input dimensions\n",
        "        self.num_patch = num_patch  # number of embedded patches\n",
        "        self.num_heads = num_heads  # number of attention heads\n",
        "        self.window_size = window_size  # size of window\n",
        "        self.shift_size = shift_size  # size of window shift\n",
        "        self.num_mlp = num_mlp  # number of MLP nodes\n",
        "\n",
        "        self.norm1 = layers.LayerNormalization(epsilon=1e-5)\n",
        "        self.attn = WindowAttention(\n",
        "            dim,\n",
        "            window_size=(self.window_size, self.window_size),\n",
        "            num_heads=num_heads,\n",
        "            qkv_bias=qkv_bias,\n",
        "            dropout_rate=dropout_rate,\n",
        "        )\n",
        "        self.drop_path = DropPath(dropout_rate)\n",
        "        self.norm2 = layers.LayerNormalization(epsilon=1e-5)\n",
        "\n",
        "        self.mlp = keras.Sequential(\n",
        "            [\n",
        "                layers.Dense(num_mlp),\n",
        "                layers.Activation(keras.activations.gelu),\n",
        "                layers.Dropout(dropout_rate),\n",
        "                layers.Dense(dim),\n",
        "                layers.Dropout(dropout_rate),\n",
        "            ]\n",
        "        )\n",
        "\n",
        "        if min(self.num_patch) < self.window_size:\n",
        "            self.shift_size = 0\n",
        "            self.window_size = min(self.num_patch)\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        if self.shift_size == 0:\n",
        "            self.attn_mask = None\n",
        "        else:\n",
        "            height, width = self.num_patch\n",
        "            h_slices = (\n",
        "                slice(0, -self.window_size),\n",
        "                slice(-self.window_size, -self.shift_size),\n",
        "                slice(-self.shift_size, None),\n",
        "            )\n",
        "            w_slices = (\n",
        "                slice(0, -self.window_size),\n",
        "                slice(-self.window_size, -self.shift_size),\n",
        "                slice(-self.shift_size, None),\n",
        "            )\n",
        "            mask_array = np.zeros((1, height, width, 1))\n",
        "            count = 0\n",
        "            for h in h_slices:\n",
        "                for w in w_slices:\n",
        "                    mask_array[:, h, w, :] = count\n",
        "                    count += 1\n",
        "            mask_array = tf.convert_to_tensor(mask_array)\n",
        "\n",
        "            # mask array to windows\n",
        "            mask_windows = window_partition(mask_array, self.window_size)\n",
        "            mask_windows = tf.reshape(\n",
        "                mask_windows, shape=[-1, self.window_size * self.window_size]\n",
        "            )\n",
        "            attn_mask = tf.expand_dims(mask_windows, axis=1) - tf.expand_dims(\n",
        "                mask_windows, axis=2\n",
        "            )\n",
        "            attn_mask = tf.where(attn_mask != 0, -100.0, attn_mask)\n",
        "            attn_mask = tf.where(attn_mask == 0, 0.0, attn_mask)\n",
        "            self.attn_mask = tf.Variable(initial_value=attn_mask, trainable=False)\n",
        "\n",
        "    def call(self, x):\n",
        "        height, width = self.num_patch\n",
        "        _, num_patches_before, channels = x.shape\n",
        "        x_skip = x\n",
        "        x = self.norm1(x)\n",
        "        x = tf.reshape(x, shape=(-1, height, width, channels))\n",
        "        if self.shift_size > 0:\n",
        "            shifted_x = tf.roll(\n",
        "                x, shift=[-self.shift_size, -self.shift_size], axis=[1, 2]\n",
        "            )\n",
        "        else:\n",
        "            shifted_x = x\n",
        "\n",
        "        x_windows = window_partition(shifted_x, self.window_size)\n",
        "        x_windows = tf.reshape(\n",
        "            x_windows, shape=(-1, self.window_size * self.window_size, channels)\n",
        "        )\n",
        "        attn_windows = self.attn(x_windows, mask=self.attn_mask)\n",
        "\n",
        "        attn_windows = tf.reshape(\n",
        "            attn_windows, shape=(-1, self.window_size, self.window_size, channels)\n",
        "        )\n",
        "        shifted_x = window_reverse(\n",
        "            attn_windows, self.window_size, height, width, channels\n",
        "        )\n",
        "        if self.shift_size > 0:\n",
        "            x = tf.roll(\n",
        "                shifted_x, shift=[self.shift_size, self.shift_size], axis=[1, 2]\n",
        "            )\n",
        "        else:\n",
        "            x = shifted_x\n",
        "\n",
        "        x = tf.reshape(x, shape=(-1, height * width, channels))\n",
        "        x = self.drop_path(x)\n",
        "        x = x_skip + x\n",
        "        x_skip = x\n",
        "        x = self.norm2(x)\n",
        "        x = self.mlp(x)\n",
        "        x = self.drop_path(x)\n",
        "        x = x_skip + x\n",
        "        return x"
      ],
      "metadata": {
        "id": "HRPWxaaWLFPm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class PatchExtract(layers.Layer):\n",
        "    def __init__(self, patch_size, **kwargs):\n",
        "        super(PatchExtract, self).__init__(**kwargs)\n",
        "        self.patch_size_x = patch_size[0]\n",
        "        self.patch_size_y = patch_size[0]\n",
        "\n",
        "    def call(self, images):\n",
        "        batch_size = tf.shape(images)[0]\n",
        "        patches = tf.image.extract_patches(\n",
        "            images=images,\n",
        "            sizes=(1, self.patch_size_x, self.patch_size_y, 1),\n",
        "            strides=(1, self.patch_size_x, self.patch_size_y, 1),\n",
        "            rates=(1, 1, 1, 1),\n",
        "            padding=\"VALID\",\n",
        "        )\n",
        "        patch_dim = patches.shape[-1]\n",
        "        patch_num = patches.shape[1]\n",
        "        return tf.reshape(patches, (batch_size, patch_num * patch_num, patch_dim))\n",
        "\n",
        "\n",
        "class PatchEmbedding(layers.Layer):\n",
        "    def __init__(self, num_patch, embed_dim, **kwargs):\n",
        "        super(PatchEmbedding, self).__init__(**kwargs)\n",
        "        self.num_patch = num_patch\n",
        "        self.proj = layers.Dense(embed_dim)\n",
        "        self.pos_embed = layers.Embedding(input_dim=num_patch, output_dim=embed_dim)\n",
        "\n",
        "    def call(self, patch):\n",
        "        pos = tf.range(start=0, limit=self.num_patch, delta=1)\n",
        "        return self.proj(patch) + self.pos_embed(pos)\n",
        "\n",
        "\n",
        "class PatchMerging(tf.keras.layers.Layer):\n",
        "    def __init__(self, num_patch, embed_dim):\n",
        "        super(PatchMerging, self).__init__()\n",
        "        self.num_patch = num_patch\n",
        "        self.embed_dim = embed_dim\n",
        "        self.linear_trans = layers.Dense(2 * embed_dim, use_bias=False)\n",
        "\n",
        "    def call(self, x):\n",
        "        height, width = self.num_patch\n",
        "        _, _, C = x.get_shape().as_list()\n",
        "        x = tf.reshape(x, shape=(-1, height, width, C))\n",
        "        x0 = x[:, 0::2, 0::2, :]\n",
        "        x1 = x[:, 1::2, 0::2, :]\n",
        "        x2 = x[:, 0::2, 1::2, :]\n",
        "        x3 = x[:, 1::2, 1::2, :]\n",
        "        x = tf.concat((x0, x1, x2, x3), axis=-1)\n",
        "        x = tf.reshape(x, shape=(-1, (height // 2) * (width // 2), 4 * C))\n",
        "        return self.linear_trans(x)"
      ],
      "metadata": {
        "id": "oGhNik8vLFTP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_ST_model():\n",
        " \n",
        "    input_shape1 =  8, 8, 18\n",
        "    output_units=8\n",
        "    \n",
        " \n",
        "    input1_ = Input(shape=input_shape1)\n",
        "\n",
        "\n",
        "############################ Feature extractor\n",
        "    \n",
        "    conv_b1 = Conv2D(filters=128, kernel_size=(3, 3), padding='same',activation='relu', name='conv_b1')(input1_)\n",
        "    max_b_1 = MaxPool2D((1,1), strides=(1,1), padding='same')(conv_b1)\n",
        "    conv_b2 = Conv2D(filters=256, kernel_size=(3, 3), padding='same',activation='relu', name='conv_b2')(max_b_1)\n",
        "    conv_b3 = Conv2D(filters=256, kernel_size=(3, 3), padding='same',activation='relu', name='conv_b3')(conv_b2)\n",
        "    norm_b = BatchNormalization(name='norm_a')(conv_b3)\n",
        "    \n",
        "  \n",
        " \n",
        "    ######################################## Swin Transformers \n",
        "    \n",
        "    x = layers.RandomCrop(image_dimension, image_dimension)(norm_b)\n",
        "    x = layers.RandomFlip(\"horizontal\")(x)\n",
        "    x = PatchExtract(patch_size)(x)\n",
        "    x = PatchEmbedding(num_patch_x * num_patch_y, embed_dim)(x)\n",
        "    x = SwinTransformer(\n",
        "    dim=embed_dim,\n",
        "    num_patch=(num_patch_x, num_patch_y),\n",
        "    num_heads=num_heads,\n",
        "    window_size=window_size,\n",
        "    shift_size=0,\n",
        "    num_mlp=num_mlp,\n",
        "    qkv_bias=qkv_bias,\n",
        "    dropout_rate=dropout_rate,\n",
        "     )(x)\n",
        "    x = SwinTransformer(\n",
        "    dim=embed_dim,\n",
        "    num_patch=(num_patch_x, num_patch_y),\n",
        "    num_heads=num_heads,\n",
        "    window_size=window_size,\n",
        "    shift_size=shift_size,\n",
        "    num_mlp=num_mlp,\n",
        "    qkv_bias=qkv_bias,\n",
        "    dropout_rate=dropout_rate,\n",
        "     )(x)\n",
        "    x = PatchMerging((num_patch_x, num_patch_y), embed_dim=embed_dim)(x)\n",
        "    x = layers.GlobalAveragePooling1D()(x)\n",
        "    x = layers.Dense(50, activation=\"relu\")(x)\n",
        "    \n",
        "    #############################################\n",
        "    \n",
        "\n",
        "    output_layer = Dense(units=output_units, activation='softmax')(x)\n",
        " \n",
        "    model = Model(inputs=input1_, outputs=output_layer)\n",
        "    model.summary()\n",
        "    \n",
        "    \n",
        "   \n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "juPHdL4rLFXE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow.keras.utils import plot_model"
      ],
      "metadata": {
        "id": "-hIbB4yGLFbN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = get_ST_model()"
      ],
      "metadata": {
        "id": "Lk1AuLn0K9VQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow\n",
        "ytrain = tensorflow.keras.utils.to_categorical(ytrain)\n",
        "ytrain.shape"
      ],
      "metadata": {
        "id": "VWHzU2ChK9Yf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])\n",
        "model_checkpoint_callback = keras.callbacks.ModelCheckpoint(\"Data/ST.h5\",save_best_only=True)\n",
        "history = model.fit(x=Xtrain, y=ytrain, batch_size = 256, epochs=100,callbacks=model_checkpoint_callback)\n"
      ],
      "metadata": {
        "id": "ZVkU0icqK9cC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "VSpcvW_rK9g7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "Jyu7QNijJISs"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}