{

 "cells": [

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "nDXW7A7pwZ7a"

   },

   "source": [

    "# UNET SEGMENTATION (SUJAL)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "M5oF6qqZwZ70"

   },

   "outputs": [],

   "source": [

    "## Imports\n",

    "import os\n",

    "import sys\n",

    "import random\n",

    "\n",

    "import numpy as np\n",

    "import cv2\n",

    "import matplotlib.pyplot as plt\n",

    "\n",

    "import tensorflow as tf\n",

    "from tensorflow import keras\n",

    "\n",

    "## Seeding\n",

    "seed = 42\n",

    "random.seed = seed\n",

    "np.random.seed = seed\n",

    "tf.seed = seed"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "70jPovBLwZ7-"

   },

   "source": [

    "## Data "

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "SvRbQp8LwZ8B"

   },

   "outputs": [],

   "source": [

    "    class DataGen(keras.utils.Sequence):\n",

    "        def __init__(self, ids, path, batch_size=8, image_size=128):\n",

    "            self.ids = ids\n",

    "            self.path = path\n",

    "            self.batch_size = batch_size\n",

    "            self.image_size = image_size\n",

    "            self.on_epoch_end()\n",

    "\n",

    "        def __load__(self, id_name):\n",

    "            # Construct image path\n",

    "            image_path = os.path.join(self.path, \"images\", id_name)\n",

    "\n",

    "            # Construct mask path\n",

    "            mask_path = os.path.join(self.path, \"masks\", id_name)\n",

    "\n",

    "            # Read and validate image\n",

    "            image = cv2.imread(image_path, 1)\n",

    "            if image is None:\n",

    "                print(f\"Warning: Failed to load image: {image_path}\")\n",

    "                return None, None\n",

    "\n",

    "            image = cv2.resize(image, (self.image_size, self.image_size))\n",

    "\n",

    "            # Read and validate mask\n",

    "            mask = cv2.imread(mask_path, 0)  # mask often is a single-channel image\n",

    "            if mask is None:\n",

    "                print(f\"Warning: Failed to load mask: {mask_path}\")\n",

    "                return None, None\n",

    "\n",

    "            mask = cv2.resize(mask, (self.image_size, self.image_size))\n",

    "            mask = np.expand_dims(mask, axis=-1)\n",

    "\n",

    "            # Normalize\n",

    "            image = image / 255.0\n",

    "            mask = mask / 255.0\n",

    "\n",

    "            return image, mask\n",

    "\n",

    "        def __getitem__(self, index):\n",

    "            if (index + 1) * self.batch_size > len(self.ids):\n",

    "                self.batch_size = len(self.ids) - index * self.batch_size\n",

    "\n",

    "            files_batch = self.ids[index * self.batch_size: (index + 1) * self.batch_size]\n",

    "\n",

    "            image = []\n",

    "            mask = []\n",

    "\n",

    "            for id_name in files_batch:\n",

    "                _img, _mask = self.__load__(id_name)\n",

    "                if _img is not None and _mask is not None:  # Only add valid samples\n",

    "                    image.append(_img)\n",

    "                    mask.append(_mask)\n",

    "\n",

    "            if len(image) == 0:  # Handle case where all files in the batch fail\n",

    "                raise ValueError(\"No valid images or masks found in the current batch.\")\n",

    "\n",

    "            image = np.array(image)\n",

    "            mask = np.array(mask)\n",

    "\n",

    "            return image, mask\n",

    "\n",

    "        def on_epoch_end(self):\n",

    "            pass\n",

    "\n",

    "        def __len__(self):\n",

    "            return int(np.ceil(len(self.ids) / float(self.batch_size)))"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "rUG6cDK1wZ8I"

   },

   "source": [

    "## Hyperparameters"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "2QwbWknPwZ8M"

   },

   "outputs": [],

   "source": [

    "image_size = 128\n",

    "train_path = \"/Users/sandhyakilari/Desktop/Fall Semester 2024/Computer Vision/Segmentation Project\"\n",

    "epochs = 30\n",

    "batch_size = 8\n",

    "\n",

    "## Training Ids\n",

    "train_ids = os.listdir(os.path.join(train_path, \"images\"))\n",

    "\n",

    "\n",

    "## Validation Data Size\n",

    "val_data_size = 10\n",

    "\n",

    "valid_ids = train_ids[:val_data_size]\n",

    "train_ids = train_ids[val_data_size:]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/"

    },

    "id": "6_cMh29vwZ8R",

    "outputId": "534e3fce-2174-4f8d-99d3-c19550f24f70"

   },

   "outputs": [],

   "source": [

    "gen = DataGen(train_ids, train_path, batch_size=batch_size, image_size=image_size)\n",

    "x, y = gen.__getitem__(0)\n",

    "print(x.shape, y.shape)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/",

     "height": 309

    },

    "id": "NCG2jWo9wZ8c",

    "outputId": "f158147f-b3b7-4c7f-e0a0-13219d136aa7",

    "scrolled": false

   },

   "outputs": [],

   "source": [

    "r = random.randint(0, len(x)-1)\n",

    "\n",

    "fig = plt.figure()\n",

    "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",

    "ax = fig.add_subplot(1, 2, 1)\n",

    "ax.imshow(x[r])\n",

    "ax = fig.add_subplot(1, 2, 2)\n",

    "ax.imshow(np.reshape(y[r], (image_size, image_size)), cmap=\"gray\")"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "1xJMvRqswZ8o"

   },

   "source": [

    "## Different Convolutional Blocks"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "xtA9Rr75wZ8p"

   },

   "outputs": [],

   "source": [

    "def down_block(x, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(x)\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",

    "    p = keras.layers.MaxPool2D((2, 2), (2, 2))(c)\n",

    "    return c, p\n",

    "\n",

    "def up_block(x, skip, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",

    "    us = keras.layers.UpSampling2D((2, 2))(x)\n",

    "    concat = keras.layers.Concatenate()([us, skip])\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(concat)\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",

    "    return c\n",

    "\n",

    "def bottleneck(x, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(x)\n",

    "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",

    "    return c"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "QMdNQ_HVwZ82"

   },

   "source": [

    "## UNet Model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "YAa1RbB8wZ89"

   },

   "outputs": [],

   "source": [

    "def UNet():\n",

    "    f = [16, 32, 64, 128, 256]\n",

    "    inputs = keras.layers.Input((image_size, image_size, 3))\n",

    "\n",

    "    p0 = inputs\n",

    "    c1, p1 = down_block(p0, f[0])  # 128 --> 64\n",

    "    c2, p2 = down_block(p1, f[1])  # 64  --> 32\n",

    "    c3, p3 = down_block(p2, f[2])  # 32  --> 16\n",

    "    c4, p4 = down_block(p3, f[3])  # 16  --> 8\n",

    "\n",

    "    bn = bottleneck(p4, f[4])\n",

    "\n",

    "    u1 = up_block(bn, c4, f[3])  # 8  --> 16\n",

    "    u2 = up_block(u1, c3, f[2])  # 16 --> 32\n",

    "    u3 = up_block(u2, c2, f[1])  # 32 --> 64\n",

    "    u4 = up_block(u3, c1, f[0])  # 64 --> 128\n",

    "\n",

    "    outputs = keras.layers.Conv2D(1, (1, 1), padding=\"same\", activation=\"sigmoid\")(u4)\n",

    "    model = keras.models.Model(inputs, outputs)\n",

    "    return model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/",

     "height": 1000

    },

    "id": "ntBPY06awZ9E",

    "outputId": "bc8241cf-1ca4-46a5-ddf1-afcb43ddc395",

    "scrolled": true

   },

   "outputs": [],

   "source": [

    "model = UNet()\n",

    "model.compile(optimizer=\"adam\", loss=\"binary_crossentropy\", metrics=[\"acc\"])\n",

    "model.summary()"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "Nmkxq4-BwZ9Q"

   },

   "source": [

    "## Training the model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/"

    },

    "id": "6I34e4tcwZ9T",

    "outputId": "427bd73a-0c95-4b03-ce95-6584de960233",

    "scrolled": false

   },

   "outputs": [],

   "source": [

    "train_gen = DataGen(train_ids, train_path, image_size=image_size, batch_size=batch_size)\n",

    "valid_gen = DataGen(valid_ids, train_path, image_size=image_size, batch_size=batch_size)\n",

    "\n",

    "train_steps = len(train_ids)//batch_size\n",

    "valid_steps = len(valid_ids)//batch_size\n",

    "\n",

    "history = model.fit(train_gen, validation_data=valid_gen, steps_per_epoch=train_steps,\n",

    "          validation_steps=valid_steps, epochs=epochs)\n",

    "\n",

    "# Save the weights\n",

    "model.save_weights(\"UNetW.weights.h5\")\n",

    "\n",

    "# Plot training & validation accuracy and loss\n",

    "acc = history.history['acc']\n",

    "val_acc = history.history['val_acc']\n",

    "loss = history.history['loss']\n",

    "val_loss = history.history['val_loss']\n",

    "num_epochs = min(len(acc), len(val_acc))\n",

    "epochs_range = range(1, num_epochs + 1)\n",

    "\n",

    "plt.figure(figsize=(16, 6))\n",

    "\n",

    "# Accuracy plot\n",

    "plt.subplot(1, 2, 1)\n",

    "plt.plot(epochs_range, acc[:num_epochs], label='Training Accuracy')\n",

    "plt.plot(epochs_range, val_acc[:num_epochs], label='Validation Accuracy')\n",

    "plt.title('Training & Validation Accuracy')\n",

    "plt.xlabel('Epoch')\n",

    "plt.ylabel('Accuracy')\n",

    "plt.legend()\n",

    "\n",

    "# Loss plot\n",

    "plt.subplot(1, 2, 2)\n",

    "plt.plot(epochs_range, loss[:num_epochs], label='Training Loss')\n",

    "plt.plot(epochs_range, val_loss[:num_epochs], label='Validation Loss')\n",

    "plt.title('Training & Validation Loss')\n",

    "plt.xlabel('Epoch')\n",

    "plt.ylabel('Loss')\n",

    "plt.legend()\n",

    "\n",

    "plt.show()"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "id": "iDUykT_awZ9h"

   },

   "source": [

    "## Testing the model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/"

    },

    "id": "i_oY4H3LwZ9l",

    "outputId": "371b3f40-25d0-416c-fd0f-e3eb6b0df943"

   },

   "outputs": [],

   "source": [

    "## Save the Weights\n",

    "model.save_weights(\"UNetW.weights.h5\")\n",

    "\n",

    "## Dataset for prediction\n",

    "x, y = valid_gen.__getitem__(2)\n",

    "result = model.predict(x)\n",

    "\n",

    "result = result > 0.5"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "id": "DLdgSxZKpQcQ"

   },

   "outputs": [],

   "source": [

    "\n",

    "# Load the weights\n",

    "model.load_weights(\"UNetW.weights.h5\")\n",

    "\n",

    "# Now you can use the model for predictions or further training\n",

    "x, y = valid_gen.__getitem__(2)\n",

    "result = model.predict(x)\n",

    "result = result > 0.5"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/",

     "height": 309

    },

    "id": "uSdH42qVwZ9q",

    "outputId": "95c65b4e-5f31-4a9e-88a1-d416166ec239"

   },

   "outputs": [],

   "source": [

    "fig = plt.figure()\n",

    "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",

    "\n",

    "ax = fig.add_subplot(1, 2, 1)\n",

    "ax.imshow(np.reshape(y[0]*255, (image_size, image_size)), cmap=\"gray\")\n",

    "\n",

    "ax = fig.add_subplot(1, 2, 2)\n",

    "ax.imshow(np.reshape(result[0]*255, (image_size, image_size)), cmap=\"gray\")"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/",

     "height": 309

    },

    "id": "1OKwJGbVwZ-A",

    "outputId": "8ef32166-359f-4023-f3a4-f12aeb895db0"

   },

   "outputs": [],

   "source": [

    "fig = plt.figure()\n",

    "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",

    "\n",

    "ax = fig.add_subplot(1, 2, 1)\n",

    "ax.imshow(np.reshape(y[1]*255, (image_size, image_size)), cmap=\"gray\")\n",

    "\n",

    "ax = fig.add_subplot(1, 2, 2)\n",

    "ax.imshow(np.reshape(result[1]*255, (image_size, image_size)), cmap=\"gray\")"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# CNN Model (Madhurya)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "colab": {

     "base_uri": "https://localhost:8080/",

     "height": 1000

    },

    "id": "xBk2k9vzKRD5",

    "outputId": "e9880f4e-a1a6-4c3b-aac7-bacca57154e5"

   },

   "outputs": [],

   "source": [

    "import os\n",

    "import cv2\n",

    "import numpy as np\n",

    "import matplotlib.pyplot as plt\n",

    "\n",

    "import keras\n",

    "from keras.models import Sequential\n",

    "from keras.layers import Dense, Activation, Flatten, Reshape\n",

    "from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D\n",

    "from keras import regularizers\n",

    "\n",

    "# Set your data directory\n",

    "data_dir = '/Users/sandhyakilari/Desktop/Fall Semester 2024/Computer Vision/Segmentation Project'\n",

    "image_dir = os.path.join(data_dir, \"images\")\n",

    "mask_dir = os.path.join(data_dir, \"masks\")\n",

    "\n",

    "# List files\n",

    "image_names = sorted(os.listdir(image_dir))\n",

    "mask_names = sorted(os.listdir(mask_dir))\n",

    "\n",

    "image_names = [name for name in image_names if name in mask_names]\n",

    "\n",

    "# Make sure that image_names and mask_names correspond one-to-one\n",

    "# If they differ, ensure file naming consistency before running.\n",

    "assert len(image_names) == len(mask_names), \"Number of images and masks do not match.\"\n",

    "\n",

    "# Desired shapes\n",

    "input_image_shape = (64, 64)  # For the CNN input\n",

    "mask_shape = (32, 32)         # For the CNN output\n",

    "\n",

    "X = []\n",

    "Y = []\n",

    "\n",

    "for img_name, msk_name in zip(image_names, mask_names):\n",

    "    img_path = os.path.join(image_dir, img_name)\n",

    "    msk_path = os.path.join(mask_dir, msk_name)\n",

    "\n",

    "    # Read images in grayscale\n",

    "    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n",

    "    msk = cv2.imread(msk_path, cv2.IMREAD_GRAYSCALE)\n",

    "\n",

    "    # Resize to desired shapes\n",

    "    img = cv2.resize(img, input_image_shape)\n",

    "    msk = cv2.resize(msk, mask_shape)\n",

    "\n",

    "    # Normalize (if you want to normalize)\n",

    "    img = img.astype(np.float32) / 255.0\n",

    "    msk = msk.astype(np.float32) / 255.0\n",

    "\n",

    "    # Add channel dimension\n",

    "    img = np.expand_dims(img, axis=-1)  # (64,64,1)\n",

    "    msk = np.expand_dims(msk, axis=-1)  # (32,32,1)\n",

    "\n",

    "    X.append(img)\n",

    "    Y.append(msk)\n",

    "\n",

    "X = np.array(X)\n",

    "Y = np.array(Y)\n",

    "\n",

    "print('Dataset shape :', X.shape, Y.shape)  # X: (N,64,64,1), Y: (N,32,32,1)\n",

    "\n",

    "# Create the CNN model\n",

    "def create_model(input_shape=(64, 64, 1)):\n",

    "    \"\"\"\n",

    "    Simple convnet model: one convolution, one average pooling and one fully connected layer\n",

    "    ending with a reshape to (32,32,1).\n",

    "    \"\"\"\n",

    "    model = Sequential()\n",

    "    # Conv layer\n",

    "    model.add(Conv2D(100, (11,11), padding='valid', strides=(1, 1), input_shape=input_shape))\n",

    "    # Average Pooling\n",

    "    model.add(AveragePooling2D((6,6)))\n",

    "    # Flatten/Reshape step\n",

    "    # After Conv+Pool:\n",

    "    # Input: (64x64x1) -> Conv(11x11): (54x54x100) -> AvgPool(6x6): (9x9x100) = 8100 features\n",

    "    model.add(Reshape((8100,)))  # Flatten to (8100,)\n",

    "    model.add(Dense(1024, activation='sigmoid', kernel_regularizer=regularizers.l2(0.0001)))\n",

    "    # Now we have (1024,). We want (32,32,1):\n",

    "    # 32*32 = 1024, so we reshape to (32,32,1)\n",

    "    model.add(Reshape((32,32,1)))\n",

    "    return model\n",

    "\n",

    "m = create_model()\n",

    "m.compile(loss='mean_squared_error', optimizer='adam', metrics=['accuracy'])\n",

    "print('Model Summary:')\n",

    "m.summary()\n",

    "\n",

    "# Train the model\n",

    "epochs = 20\n",

    "batch_size = 16\n",

    "history = m.fit(X, Y, batch_size=batch_size, epochs=epochs, validation_split=0.2)\n",

    "\n",

    "# Plot training and validation loss\n",

    "plt.figure(figsize=(10,5))\n",

    "plt.plot(history.history['loss'], label='Training Loss')\n",

    "plt.plot(history.history['val_loss'], label='Validation Loss', linestyle='--')\n",

    "plt.title(\"Learning Curve\")\n",

    "plt.xlabel(\"Epochs\")\n",

    "plt.ylabel(\"Loss\")\n",

    "plt.legend()\n",

    "plt.show()\n",

    "\n",

    "# Save the weights\n",

    "m.save_weights(\"UNetW.weights.h5\")\n",

    "\n",

    "# Example prediction\n",

    "y_pred = m.predict(X, batch_size=batch_size)\n",

    "print(\"y_pred shape:\", y_pred.shape)\n",

    "\n",

    "# Visualize a sample\n",

    "idx = 0\n",

    "fig, ax = plt.subplots(1, 2, figsize=(8,4))\n",

    "ax[0].imshow(X[idx].reshape(64,64), cmap='gray')\n",

    "ax[0].set_title('Input Image')\n",

    "ax[1].imshow(y_pred[idx].reshape(32,32), cmap='gray')\n",

    "ax[1].set_title('Predicted Mask')\n",

    "plt.show()"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# AttentionUNET (Vishal)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "import tensorflow as tf\n",

    "from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, concatenate, Activation, BatchNormalization, Add, Multiply\n",

    "from tensorflow.keras.models import Model\n",

    "import os\n",

    "import numpy as np\n",

    "from tensorflow.keras.preprocessing.image import load_img, img_to_array\n",

    "\n",

    "def attention_block(x, g, inter_channel):\n",

    "    \"\"\"\n",

    "    Attention Block: Refines encoder features based on decoder signals.\n",

    "    x: Input tensor from the encoder (skip connection)\n",

    "    g: Gating signal from the decoder (upsampled tensor)\n",

    "    inter_channel: Number of intermediate channels (reduces computation)\n",

    "    \"\"\"\n",

    "    # 1x1 Convolution on input tensor\n",

    "    theta_x = Conv2D(inter_channel, kernel_size=(1, 1), strides=(1, 1), padding='same')(x)\n",

    "    # 1x1 Convolution on gating tensor\n",

    "    phi_g = Conv2D(inter_channel, kernel_size=(1, 1), strides=(1, 1), padding='same')(g)\n",

    "    \n",

    "    # Add the transformed inputs and apply ReLU\n",

    "    add_xg = Add()([theta_x, phi_g])\n",

    "    relu_xg = Activation('relu')(add_xg)\n",

    "    \n",

    "    # Another 1x1 Convolution to generate attention coefficients\n",

    "    psi = Conv2D(1, kernel_size=(1, 1), strides=(1, 1), padding='same')(relu_xg)\n",

    "    # Sigmoid activation to normalize attention weights\n",

    "    sigmoid_psi = Activation('sigmoid')(psi)\n",

    "    \n",

    "    # Multiply the input tensor with the attention weights\n",

    "    return Multiply()([x, sigmoid_psi])\n",

    "\n",

    "def conv_block(x, filters):\n",

    "    \"\"\"\n",

    "    Convolutional Block: Apply two 3x3 convolutions followed by BatchNorm and ReLU.\n",

    "    x: Input tensor\n",

    "    filters: Number of output filters for the convolutions\n",

    "    \"\"\"\n",

    "    x = Conv2D(filters, kernel_size=(3, 3), padding='same')(x)\n",

    "    x = BatchNormalization()(x)\n",

    "    x = Activation('relu')(x)\n",

    "    x = Conv2D(filters, kernel_size=(3, 3), padding='same')(x)\n",

    "    x = BatchNormalization()(x)\n",

    "    x = Activation('relu')(x)\n",

    "    return x\n",

    "\n",

    "def attention_unet(input_shape, num_classes):\n",

    "    \"\"\"\n",

    "    Attention U-Net model architecture.\n",

    "    input_shape: Shape of input images (H, W, C)\n",

    "    num_classes: Number of output segmentation classes\n",

    "    \"\"\"\n",

    "    # Input layer for the images\n",

    "    inputs = Input(input_shape)\n",

    "    \n",

    "    # Encoder (Downsampling path)\n",

    "    c1 = conv_block(inputs, 64)              # First Conv Block\n",

    "    p1 = MaxPooling2D((2, 2))(c1)            # Downsample by 2\n",

    "    \n",

    "    c2 = conv_block(p1, 128)                 # Second Conv Block\n",

    "    p2 = MaxPooling2D((2, 2))(c2)            # Downsample by 2\n",

    "    \n",

    "    c3 = conv_block(p2, 256)                 # Third Conv Block\n",

    "    p3 = MaxPooling2D((2, 2))(c3)            # Downsample by 2\n",

    "    \n",

    "    c4 = conv_block(p3, 512)                 # Fourth Conv Block\n",

    "    p4 = MaxPooling2D((2, 2))(c4)            # Downsample by 2\n",

    "    \n",

    "    # Bottleneck (lowest level of the U-Net)\n",

    "    c5 = conv_block(p4, 1024)\n",

    "    \n",

    "    # Decoder (Upsampling path)\n",

    "    up6 = UpSampling2D((2, 2))(c5)           # Upsample\n",

    "    att6 = attention_block(c4, up6, 512)     # Attention Block\n",

    "    merge6 = concatenate([up6, att6], axis=-1)  # Concatenate features\n",

    "    c6 = conv_block(merge6, 512)             # Conv Block after concatenation\n",

    "    \n",

    "    up7 = UpSampling2D((2, 2))(c6)\n",

    "    att7 = attention_block(c3, up7, 256)\n",

    "    merge7 = concatenate([up7, att7], axis=-1)\n",

    "    c7 = conv_block(merge7, 256)\n",

    "    \n",

    "    up8 = UpSampling2D((2, 2))(c7)\n",

    "    att8 = attention_block(c2, up8, 128)\n",

    "    merge8 = concatenate([up8, att8], axis=-1)\n",

    "    c8 = conv_block(merge8, 128)\n",

    "    \n",

    "    up9 = UpSampling2D((2, 2))(c8)\n",

    "    att9 = attention_block(c1, up9, 64)\n",

    "    merge9 = concatenate([up9, att9], axis=-1)\n",

    "    c9 = conv_block(merge9, 64)\n",

    "    \n",

    "    # Output layer for segmentation\n",

    "    outputs = Conv2D(num_classes, (1, 1), activation='softmax' if num_classes > 1 else 'sigmoid')(c9)\n",

    "    \n",

    "    # Define the model\n",

    "    model = Model(inputs=inputs, outputs=outputs)\n",

    "    return model\n",

    "\n",

    "# Function to load and preprocess images and masks\n",

    "def load_data(image_dir, mask_dir, image_size):\n",

    "    \"\"\"\n",

    "    Load and preprocess images and masks for training.\n",

    "    image_dir: Path to the directory containing input images\n",

    "    mask_dir: Path to the directory containing segmentation masks\n",

    "    image_size: Tuple specifying the size (height, width) to resize the images and masks\n",

    "    \"\"\"\n",

    "    images = []\n",

    "    masks = []\n",

    "    image_files = sorted(os.listdir(image_dir))\n",

    "    mask_files = sorted(os.listdir(mask_dir))\n",

    "    \n",

    "    for img_file, mask_file in zip(image_files, mask_files):\n",

    "        try:\n",

    "            # Load and preprocess images\n",

    "            img_path = os.path.join(image_dir, img_file)\n",

    "            mask_path = os.path.join(mask_dir, mask_file)\n",

    "            \n",

    "            img = load_img(img_path, target_size=image_size)  # Resize image\n",

    "            mask = load_img(mask_path, target_size=image_size, color_mode='grayscale')  # Resize mask\n",

    "            \n",

    "            # Convert to numpy arrays and normalize\n",

    "            img = img_to_array(img) / 255.0\n",

    "            mask = img_to_array(mask) / 255.0\n",

    "            mask = np.round(mask)  # Ensure masks are binary\n",

    "            \n",

    "            images.append(img)\n",

    "            masks.append(mask)\n",

    "        except Exception as e:\n",

    "            print(f\"Error loading {img_file} or {mask_file}: {e}. Skipping...\")\n",

    "    \n",

    "    return np.array(images), np.array(masks)\n",

    "\n",

    "# Example usage\n",

    "if __name__ == \"__main__\":\n",

    "    # Load data\n",

    "    image_dir = \"./images/\"  # Replace with your image directory\n",

    "    mask_dir = \"./masks/\"    # Replace with your mask directory\n",

    "    image_size = (128, 128)       # Resize all images to 128x128\n",

    "    images, masks = load_data(image_dir, mask_dir, image_size)\n",

    "    \n",

    "    # Define the model\n",

    "    model = attention_unet(input_shape=(128, 128, 3), num_classes=1)\n",

    "    \n",

    "    # Compile the model\n",

    "    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])\n",

    "    \n",

    "    # Train the model\n",

    "    model.fit(images, masks, batch_size=8, epochs=20, validation_split=0.1)"

   ]

  }

 ],

 "metadata": {

  "accelerator": "GPU",

  "colab": {

   "provenance": []

  },

  "kernelspec": {

   "display_name": "base",

   "language": "python",

   "name": "python3"

  },

  "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.12.2"

  }

 },

 "nbformat": 4,

 "nbformat_minor": 0

}