CBAM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "093f3581",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary libraries\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Input, Conv2D, GlobalAveragePooling2D, Dense, Activation, Add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7d12d39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Channel(input_shape, r):\n",
    "    # Calculate average and maximum values along height and width dimensions\n",
    "    avg_p = tf.reduce_mean(input_shape, axis=[1, 2], keepdims=True)\n",
    "    max_p = tf.reduce_max(input_shape, axis=[1, 2], keepdims=True)\n",
    "    \n",
    "    reduced_features = int(input_shape.shape[-1] // r) # Calculate the reduced number of features\n",
    "    \n",
    "    # Define dense layers\n",
    "    dense1 = Dense(reduced_features, activation=\"relu\")\n",
    "    dense2 = Dense(input_shape.shape[-1], activation=\"linear\")\n",
    "    \n",
    "    # Pass average and maximum values through the first dense layer\n",
    "    Dense1_avg = dense1(avg_p)\n",
    "    Dense1_max = dense1(max_p)\n",
    "\n",
    "    # Pass the outputs of the first dense layer through the second dense layer\n",
    "    Dense2_avg = dense2(Dense1_avg)\n",
    "    Dense2_max = dense2(Dense1_max)\n",
    "\n",
    "    # Apply sigmoid activation to the sum of outputs of the second dense layer\n",
    "    out = Activation('sigmoid')(Dense2_avg + Dense2_max)\n",
    "\n",
    "    return out * input_shape # Scale the input_shape by the sigmoid output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "39a94c41",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Spatial(input_shape, ks=7):\n",
    "    # Calculate average and maximum values along the channel dimension\n",
    "    avg_p = tf.reduce_mean(input_shape, axis=[-1], keepdims=True)\n",
    "    max_p = tf.reduce_max(input_shape, axis=[-1], keepdims=True)\n",
    "\n",
    "    # Concatenate average and maximum values along the channel dimension\n",
    "    concat_pool = tf.concat([avg_p, max_p], axis=-1)\n",
    "    \n",
    "    # Apply convolutional layer with kernel size (ks, ks) and sigmoid activation\n",
    "    out = Conv2D(1, (ks, ks), padding='same', activation='sigmoid')(concat_pool)\n",
    "\n",
    "    # Scale the input_shape by the sigmoid output\n",
    "    return out * input_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5d030cb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def CBAM(input_shape, ks=7, r=2):\n",
    "    # Apply channel attention mechanism\n",
    "    channel_out = Channel(input_shape, r)\n",
    "    \n",
    "    # Apply spatial attention mechanism on the output of channel attention\n",
    "    Spatial_out = Spatial(channel_out, ks)\n",
    "    \n",
    "    # Add the output of channel attention and spatial attention\n",
    "    out = Add()([channel_out, Spatial_out])\n",
    "    \n",
    "    return Activation('relu')(out) # Apply ReLU activation to the combined output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d924af8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_CBAM_model(input_shape, num_classes):\n",
    "    \n",
    "    input_tensor = Input(shape=input_shape) # Define input tensor\n",
    "    \n",
    "    # First convolutional layer with 64 filters and ReLU activation, followed by CBAM\n",
    "    x = tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu')(input_tensor)\n",
    "    x = CBAM(x)  # Apply CBAM\n",
    "    \n",
    "    # Second convolutional layer with 128 filters and ReLU activation, followed by CBAM\n",
    "    x = Conv2D(128, (3, 3), padding='same', activation='relu')(x)\n",
    "    x = CBAM(x)  # Apply CBAM\n",
    "    \n",
    "    x = GlobalAveragePooling2D()(x) # Global average pooling layer\n",
    "    \n",
    "    # Fully connected layer with num_classes nodes and softmax activation\n",
    "    x = Dense(num_classes, activation='softmax')(x)\n",
    "    \n",
    "    model = tf.keras.Model(inputs=input_tensor, outputs=x) # Create the model by specifying input and output tensors\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "94bba95f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_1 (InputLayer)            [(None, 224, 224, 12 0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d (Conv2D)                 (None, 224, 224, 64) 6976        input_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_mean (TFOpLambda (None, 1, 1, 64)     0           conv2d[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_max (TFOpLambda) (None, 1, 1, 64)     0           conv2d[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "dense (Dense)                   (None, 1, 1, 32)     2080        tf.math.reduce_mean[0][0]        \n",
      "                                                                 tf.math.reduce_max[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "dense_1 (Dense)                 (None, 1, 1, 64)     2112        dense[0][0]                      \n",
      "                                                                 dense[1][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "tf.__operators__.add (TFOpLambd (None, 1, 1, 64)     0           dense_1[0][0]                    \n",
      "                                                                 dense_1[1][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "activation (Activation)         (None, 1, 1, 64)     0           tf.__operators__.add[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.multiply (TFOpLambda)   (None, 224, 224, 64) 0           activation[0][0]                 \n",
      "                                                                 conv2d[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_mean_1 (TFOpLamb (None, 224, 224, 1)  0           tf.math.multiply[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_max_1 (TFOpLambd (None, 224, 224, 1)  0           tf.math.multiply[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "tf.concat (TFOpLambda)          (None, 224, 224, 2)  0           tf.math.reduce_mean_1[0][0]      \n",
      "                                                                 tf.math.reduce_max_1[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_1 (Conv2D)               (None, 224, 224, 1)  99          tf.concat[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.multiply_1 (TFOpLambda) (None, 224, 224, 64) 0           conv2d_1[0][0]                   \n",
      "                                                                 tf.math.multiply[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "add (Add)                       (None, 224, 224, 64) 0           tf.math.multiply[0][0]           \n",
      "                                                                 tf.math.multiply_1[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "activation_1 (Activation)       (None, 224, 224, 64) 0           add[0][0]                        \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_2 (Conv2D)               (None, 224, 224, 128 73856       activation_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_mean_2 (TFOpLamb (None, 1, 1, 128)    0           conv2d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_max_2 (TFOpLambd (None, 1, 1, 128)    0           conv2d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (None, 1, 1, 64)     8256        tf.math.reduce_mean_2[0][0]      \n",
      "                                                                 tf.math.reduce_max_2[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "dense_3 (Dense)                 (None, 1, 1, 128)    8320        dense_2[0][0]                    \n",
      "                                                                 dense_2[1][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf.__operators__.add_1 (TFOpLam (None, 1, 1, 128)    0           dense_3[0][0]                    \n",
      "                                                                 dense_3[1][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "activation_2 (Activation)       (None, 1, 1, 128)    0           tf.__operators__.add_1[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.multiply_2 (TFOpLambda) (None, 224, 224, 128 0           activation_2[0][0]               \n",
      "                                                                 conv2d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_mean_3 (TFOpLamb (None, 224, 224, 1)  0           tf.math.multiply_2[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.reduce_max_3 (TFOpLambd (None, 224, 224, 1)  0           tf.math.multiply_2[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "tf.concat_1 (TFOpLambda)        (None, 224, 224, 2)  0           tf.math.reduce_mean_3[0][0]      \n",
      "                                                                 tf.math.reduce_max_3[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_3 (Conv2D)               (None, 224, 224, 1)  99          tf.concat_1[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "tf.math.multiply_3 (TFOpLambda) (None, 224, 224, 128 0           conv2d_3[0][0]                   \n",
      "                                                                 tf.math.multiply_2[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "add_1 (Add)                     (None, 224, 224, 128 0           tf.math.multiply_2[0][0]         \n",
      "                                                                 tf.math.multiply_3[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "activation_3 (Activation)       (None, 224, 224, 128 0           add_1[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "global_average_pooling2d (Globa (None, 128)          0           activation_3[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "dense_4 (Dense)                 (None, 10)           1290        global_average_pooling2d[0][0]   \n",
      "==================================================================================================\n",
      "Total params: 103,088\n",
      "Trainable params: 103,088\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# Example usage\n",
    "input_shape = (224, 224, 12)  # Adjust based on your input image size and channels\n",
    "num_classes = 10  # Adjust based on your classification task\n",
    "CBAM_model = create_CBAM_model(input_shape, num_classes)\n",
    "\n",
    "# Display model summary\n",
    "CBAM_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c6dfac9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "language": "python",
   "name": "tensorflow"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}