FHC_Unet_ac_copy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "BvCJ2P_MBtmp"
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'imgaug'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[1], line 7\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mrandom\u001b[39;00m\n\u001b[1;32m----> 7\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mimgaug\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maugmenters\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01miaa\u001b[39;00m\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mtf\u001b[39;00m\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlosses\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m binary_crossentropy\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'imgaug'"
     ]
    }
   ],
   "source": [
    "#utils.py\n",
    "\n",
    "import os\n",
    "from PIL import Image,ImageFilter\n",
    "import numpy as np\n",
    "import random\n",
    "import imgaug.augmenters as iaa\n",
    "import tensorflow as tf\n",
    "from keras.losses import binary_crossentropy\n",
    "\n",
    "def preprocess_images(image, annotation,dim,resize_only=False):\n",
    "\n",
    "        # Convert to grayscale\n",
    "        image_gray = image.convert('L')\n",
    "        ann_gray = annotation.convert('L')\n",
    "\n",
    "        if resize_only==False:\n",
    "\n",
    "          # Apply blur\n",
    "          image_blurred = image_gray.filter(ImageFilter.BLUR)\n",
    "          ann_blurred = ann_gray.filter(ImageFilter.BLUR)\n",
    "\n",
    "          # Augmentation\n",
    "          aug = iaa.Sequential([\n",
    "              iaa.GaussianBlur(sigma=(0, 3.0)),\n",
    "              iaa.AdditiveGaussianNoise(scale=(0, 0.05 * 255)),\n",
    "              iaa.Affine(rotate=(-45, 45)),\n",
    "              iaa.Multiply((0.8, 1.2))\n",
    "          ])\n",
    "\n",
    "          #Crop/resize\n",
    "          image_augmented = aug.augment_image(np.array(image_blurred))\n",
    "          ann_augmented = aug.augment_image(np.array(ann_blurred))\n",
    "\n",
    "          image_augmented = tf.convert_to_tensor(image_augmented)\n",
    "          image_augmented = tf.expand_dims(image_augmented,-1)\n",
    "          image_augmented = tf.image.resize_with_crop_or_pad(image_augmented,*dim)\n",
    "\n",
    "          ann_augmented = tf.convert_to_tensor(ann_augmented)\n",
    "          ann_augmented = tf.expand_dims(ann_augmented,-1)\n",
    "          ann_augmented = tf.image.resize_with_crop_or_pad(ann_augmented,*dim)\n",
    "\n",
    "          ann_augmented = tf.cast(ann_augmented,tf.float32)/255.0\n",
    "          ann_augmented=tf.cast(ann_augmented,tf.int32)\n",
    "\n",
    "        else:\n",
    "\n",
    "          #Crop/resize\n",
    "          image_augmented = np.array(image_gray)\n",
    "          ann_augmented = np.array(ann_gray)\n",
    "\n",
    "          image_augmented = tf.convert_to_tensor(image_augmented)\n",
    "          image_augmented = tf.expand_dims(image_augmented,-1)\n",
    "          image_augmented = tf.image.resize_with_crop_or_pad(image_augmented,*dim)\n",
    "\n",
    "          ann_augmented = tf.convert_to_tensor(ann_augmented)\n",
    "          ann_augmented = tf.expand_dims(ann_augmented,-1)\n",
    "          ann_augmented = tf.image.resize_with_crop_or_pad(ann_augmented,*dim)\n",
    "\n",
    "          ann_augmented = tf.cast(ann_augmented,tf.float32)/255.0\n",
    "          ann_augmented=tf.cast(ann_augmented,tf.int32)\n",
    "\n",
    "        return image_augmented,ann_augmented\n",
    "\n",
    "\n",
    "def dice_coeff(y_true, y_pred):\n",
    "    smooth = 1.\n",
    "    # Flatten\n",
    "    y_true = tf.cast(y_true, dtype=tf.float32)  # Convert labels to float32\n",
    "    y_pred = tf.cast(y_pred, dtype=tf.float32)\n",
    "    y_true_f = tf.reshape(y_true, [-1])\n",
    "    y_pred_f = tf.reshape(y_pred, [-1])\n",
    "    intersection = tf.reduce_sum(y_true_f * y_pred_f)\n",
    "    score = (2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)\n",
    "    return score\n",
    "\n",
    "def dice_loss(y_true, y_pred):\n",
    "    loss = 1 - dice_coeff(y_true, y_pred)\n",
    "    return loss\n",
    "\n",
    "def bce_dice_loss(y_true, y_pred):\n",
    "    loss = binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-nsf85-DCnR2"
   },
   "outputs": [],
   "source": [
    "#customdatagenerator.py\n",
    "\n",
    "import numpy as np\n",
    "import keras\n",
    "#from utils import preprocess_images\n",
    "from PIL import Image\n",
    "import os\n",
    "\n",
    "\n",
    "\n",
    "class CustomDataGenerator(keras.utils.Sequence):\n",
    "    'Generates data for Keras'\n",
    "    def __init__(self, path, batch_size=32, dim=(800,540), n_channels=1, shuffle=True, augmentation=True):\n",
    "        'Initialization'\n",
    "        self.dim = dim # dimensions of images (make sure they are all the same dimension)\n",
    "        self.batch_size = batch_size # choose batch size\n",
    "        self.n_channels = n_channels # = 1 for grayscale\n",
    "        self.shuffle = shuffle\n",
    "        self.augmentation = augmentation # whether or not you want to perform augmentation\n",
    "        self.path = path    #**whats the diff bw this and data_path?\n",
    "        self.img_folder = self.path + 'Image/'\n",
    "        self.mask_folder = self.path + 'Annotation/'\n",
    "        self.list_IDs = os.listdir(self.img_folder) # make sure the names of the corresponding files in the images and annotations folders are the same **not exactly.?\n",
    "        self.on_epoch_end()\n",
    "\n",
    "    def __len__(self):\n",
    "        'Denotes the number of batches per epoch'\n",
    "        return int(np.floor(len(self.list_IDs) / self.batch_size))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        'Generate one batch of data'\n",
    "        # Generate indexes of the batch\n",
    "        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]\n",
    "\n",
    "        # Find list of IDs\n",
    "        list_IDs_temp = [self.list_IDs[k] for k in indexes]\n",
    "\n",
    "        # Generate data\n",
    "        X, y = self.__data_generation(list_IDs_temp)\n",
    "\n",
    "        return X, y\n",
    "\n",
    "    def on_epoch_end(self):\n",
    "        'Updates indexes after each epoch'\n",
    "        self.indexes = np.arange(len(self.list_IDs))\n",
    "        if self.shuffle == True:\n",
    "            np.random.shuffle(self.indexes)\n",
    "\n",
    "    def __data_generation(self, list_IDs_temp):\n",
    "        'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)\n",
    "        # Initialization\n",
    "        X = np.empty((self.batch_size, *self.dim, self.n_channels))\n",
    "        y = np.empty((self.batch_size, *self.dim, self.n_channels), dtype=int)\n",
    "\n",
    "        # Generate data\n",
    "        for i, ID in enumerate(list_IDs_temp):\n",
    "\n",
    "            filename,ext = os.path.splitext(ID)\n",
    "            # Store sample\n",
    "            #img = Image.open(self.img_folder + ID + '.png')\n",
    "            #label = Image.open(self.mask_folder + ID + '.png')  #**adjusted preprocess function to match\n",
    "            img = Image.open(self.img_folder + ID)\n",
    "            label = Image.open(self.mask_folder + filename + '_Annotation.png')\n",
    "\n",
    "\n",
    "            if self.augmentation==True:\n",
    "                augmented_img, augmented_label = preprocess_images(img, label, self.dim) # ensure this \"augment\" function encompasses all transformations\n",
    "                X[i,] = augmented_img\n",
    "                y[i,] = augmented_label\n",
    "            else:\n",
    "                img, label = preprocess_images(img,label,self.dim,resize_only=True)\n",
    "                X[i,]=img\n",
    "                y[i,]=label\n",
    "\n",
    "        return X, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T1lkF1h0DNLP"
   },
   "outputs": [],
   "source": [
    "#data.py\n",
    "\n",
    "from __future__ import print_function\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "import numpy as np\n",
    "import os\n",
    "import glob\n",
    "import skimage.io as io\n",
    "import skimage.transform as trans\n",
    "\n",
    "#...\n",
    "\"\"\"\n",
    "Sky = [128,128,128]\n",
    "Building = [128,0,0]\n",
    "Pole = [192,192,128]\n",
    "Road = [128,64,128]\n",
    "Pavement = [60,40,222]\n",
    "Tree = [128,128,0]\n",
    "SignSymbol = [192,128,128]\n",
    "Fence = [64,64,128]\n",
    "Car = [64,0,128]\n",
    "Pedestrian = [64,64,0]\n",
    "Bicyclist = [0,128,192]\n",
    "Unlabelled = [0,0,0]\n",
    "\"\"\"\n",
    "AF = [255,255,255]\n",
    "Rest = [0,0,0]\n",
    "\n",
    "#...\n",
    "#COLOR_DICT = np.array([Sky, Building, Pole, Road, Pavement,Tree, SignSymbol, Fence, Car, Pedestrian, Bicyclist, Unlabelled])\n",
    "COLOR_DICT = np.array([AF, Rest])\n",
    "\n",
    "#...already using ImageDataGenerator - https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/image/ImageDataGenerator\n",
    "# M/S aug: rotate, add noise, crop, flip(?) // IDG flips/rotates. does not crop (?) or add noise\n",
    "# clarify: preprocessing should also done here/is part of dataloader right? point was to avoid having to upload separately\n",
    "# just import prepr function and run through that first right?\n",
    "# are annotation/label/mask/ground truth the same thing?\n",
    "# so: load data -> run prepr function -> run aug function [trainGenerator?] -> etc. ?\n",
    "# what's geneTrainNpy() doing?\n",
    "\n",
    "def adjustData(img,mask,flag_multi_class,num_class):\n",
    "    if(flag_multi_class):\n",
    "        img = img / 255\n",
    "        mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]\n",
    "        new_mask = np.zeros(mask.shape + (num_class,))\n",
    "        for i in range(num_class):\n",
    "            #for one pixel in the image, find the class in mask and convert it into one-hot vector\n",
    "            #index = np.where(mask == i)\n",
    "            #index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)\n",
    "            #new_mask[index_mask] = 1\n",
    "            new_mask[mask == i,i] = 1\n",
    "        new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))\n",
    "        mask = new_mask\n",
    "    elif(np.max(img) > 1):\n",
    "        img = img / 255\n",
    "        mask = mask /255\n",
    "        mask[mask > 0.5] = 1\n",
    "        mask[mask <= 0.5] = 0\n",
    "    return (img,mask)\n",
    "\n",
    "#change prefixes?\n",
    "def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = \"grayscale\",\n",
    "                    mask_color_mode = \"grayscale\",image_save_prefix  = \"image\",mask_save_prefix  = \"mask\",\n",
    "                    flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):\n",
    "    '''\n",
    "    can generate image and mask at the same time\n",
    "    use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same\n",
    "    if you want to visualize the results of generator, set save_to_dir = \"your path\"\n",
    "    '''\n",
    "    image_datagen = ImageDataGenerator(**aug_dict)\n",
    "    mask_datagen = ImageDataGenerator(**aug_dict)\n",
    "    image_generator = image_datagen.flow_from_directory(\n",
    "        train_path,\n",
    "        classes = [image_folder],\n",
    "        class_mode = None,\n",
    "        color_mode = image_color_mode,\n",
    "        target_size = target_size,\n",
    "        batch_size = batch_size,\n",
    "        save_to_dir = save_to_dir,\n",
    "        save_prefix  = image_save_prefix,\n",
    "        seed = seed)\n",
    "    mask_generator = mask_datagen.flow_from_directory(\n",
    "        train_path,\n",
    "        classes = [mask_folder],\n",
    "        class_mode = None,\n",
    "        color_mode = mask_color_mode,\n",
    "        target_size = target_size,\n",
    "        batch_size = batch_size,\n",
    "        save_to_dir = save_to_dir,\n",
    "        save_prefix  = mask_save_prefix,\n",
    "        seed = seed)\n",
    "    train_generator = zip(image_generator, mask_generator)\n",
    "    for (img,mask) in train_generator:\n",
    "        img,mask = adjustData(img,mask,flag_multi_class,num_class)\n",
    "        yield (img,mask)\n",
    "\n",
    "def testGenerator(test_path,num_image = 30,target_size = (256,256),flag_multi_class = False,as_gray = True):\n",
    "    for i in range(num_image):\n",
    "        files = os.listdir(test_path)\n",
    "        img = io.imread(os.path.join(test_path, files[i]),as_gray = as_gray)\n",
    "        img = img / 255\n",
    "        img = trans.resize(img,target_size)\n",
    "        img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img\n",
    "        img = np.reshape(img,(1,)+img.shape)\n",
    "        yield img\n",
    "\n",
    "\n",
    "def geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = \"image\",mask_prefix = \"mask\",image_as_gray = True,mask_as_gray = True):\n",
    "    image_name_arr = glob.glob(os.path.join(image_path,\"%s*.png\"%image_prefix))\n",
    "    image_arr = []\n",
    "    mask_arr = []\n",
    "    for index,item in enumerate(image_name_arr):\n",
    "        img = io.imread(item,as_gray = image_as_gray)\n",
    "        img = np.reshape(img,img.shape + (1,)) if image_as_gray else img\n",
    "        mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)\n",
    "        mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask\n",
    "        img,mask = adjustData(img,mask,flag_multi_class,num_class)\n",
    "        image_arr.append(img)\n",
    "        mask_arr.append(mask)\n",
    "    image_arr = np.array(image_arr)\n",
    "    mask_arr = np.array(mask_arr)\n",
    "    return image_arr,mask_arr\n",
    "\n",
    "\n",
    "def labelVisualize(num_class,color_dict,img):\n",
    "    img = img[:,:,0] if len(img.shape) == 3 else img\n",
    "    img_out = np.zeros(img.shape + (3,))\n",
    "    for i in range(num_class):\n",
    "        img_out[img == i,:] = color_dict[i]\n",
    "    return img_out / 255\n",
    "\n",
    "\n",
    "\n",
    "def saveResult(save_path,inp_path,npyfile,flag_multi_class = False,num_class = 2):\n",
    "    filenames = os.listdir(inp_path)\n",
    "    for i,item in enumerate(npyfile):\n",
    "        img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]\n",
    "        io.imsave(os.path.join(save_path,filenames[i]),img)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hJE0-suUCuBF"
   },
   "outputs": [],
   "source": [
    "#model.py\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import skimage.io as io\n",
    "import skimage.transform as trans\n",
    "import numpy as np\n",
    "from keras.models import *\n",
    "from keras.layers import *\n",
    "from keras.optimizers import *\n",
    "from keras.callbacks import ModelCheckpoint, LearningRateScheduler\n",
    "from keras import backend as keras\n",
    "#from utils import dice_coeff, bce_dice_loss\n",
    "\n",
    "\n",
    "def unet(pretrained_weights = None,input_size = (256,256,1)):\n",
    "    inputs = Input(input_size)\n",
    "    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)\n",
    "    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)\n",
    "    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)\n",
    "    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)\n",
    "    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)\n",
    "    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)\n",
    "    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)\n",
    "    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)\n",
    "    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)\n",
    "    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)\n",
    "    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)\n",
    "    drop4 = Dropout(0.5)(conv4)\n",
    "    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)\n",
    "\n",
    "    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)\n",
    "    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)\n",
    "    drop5 = Dropout(0.5)(conv5)\n",
    "\n",
    "    up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))\n",
    "    merge6 = concatenate([drop4,up6], axis = 3)\n",
    "    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)\n",
    "    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)\n",
    "\n",
    "    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))\n",
    "    merge7 = concatenate([conv3,up7], axis = 3)\n",
    "    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)\n",
    "    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)\n",
    "\n",
    "    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))\n",
    "    merge8 = concatenate([conv2,up8], axis = 3)\n",
    "    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)\n",
    "    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)\n",
    "\n",
    "    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))\n",
    "    merge9 = concatenate([conv1,up9], axis = 3)\n",
    "    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)\n",
    "    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n",
    "    conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n",
    "    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)\n",
    "\n",
    "    model = Model(inputs = inputs, outputs = conv10)\n",
    "\n",
    "    model.compile(optimizer = Adam(lr = 1e-4), loss = bce_dice_loss, metrics = [tf.keras.metrics.binary_accuracy,dice_coeff])\n",
    "\n",
    "    #model.summary()\n",
    "\n",
    "    if(pretrained_weights):\n",
    "      model.load_weights(pretrained_weights)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XqRN7b6N86Rg",
    "outputId": "6b6a2c51-3f2d-485a-f5d5-4a60bf3c5e8c"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FB_4C5Fm9XaU",
    "outputId": "0176198c-6f29-4692-b193-0a122f49168e"
   },
   "outputs": [],
   "source": [
    "train_path = '/content/drive/MyDrive/2023 AI4Good Lab Trainee Drive/4. Project Teams/T1 Project/Data/Training-Splitted/train/'\n",
    "os.path.exists(train_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "osEbSjjL9gsP",
    "outputId": "4c564e7c-0142-4d12-bf3a-708916016d7a"
   },
   "outputs": [],
   "source": [
    "val_path = '/content/drive/MyDrive/2023 AI4Good Lab Trainee Drive/4. Project Teams/T1 Project/Data/Training-Splitted/val/'\n",
    "os.path.exists(val_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3fHtcuuGCxYZ",
    "outputId": "2941a00e-0e9c-47bd-e57b-3658e4279d9f"
   },
   "outputs": [],
   "source": [
    "#main.py\n",
    "\n",
    "#from model import *\n",
    "#from data import *\n",
    "#from customdatagenerator import *\n",
    "\n",
    "\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "# these augmentation params not needed if \"augment\" function is hard-coded to do random transformations\n",
    "# data_gen_args = dict(rotation_range=0.2,\n",
    "                    # width_shift_range=0.05,\n",
    "                    # height_shift_range=0.05,\n",
    "                    # shear_range=0.05,\n",
    "                    # zoom_range=0.05,\n",
    "                    # horizontal_flip=True,\n",
    "                    # fill_mode='nearest')\n",
    "\n",
    "train_path = '/content/drive/MyDrive/2023 AI4Good Lab Trainee Drive/4. Project Teams/T1 Project/Data/Training-Splitted/train/'\n",
    "\n",
    "val_path = '/content/drive/MyDrive/2023 AI4Good Lab Trainee Drive/4. Project Teams/T1 Project/Data/Training-Splitted/val/'\n",
    "batch_size = 5\n",
    "dim = (640,640) # dimensions of images  **actual dim or 32,32,32 like customdatagen.py?\n",
    "n_channels = 1\n",
    "\n",
    "# using custom datagenerator:\n",
    "train_gen = CustomDataGenerator(path = train_path, batch_size = batch_size, dim = dim, n_channels = n_channels, shuffle = True, augmentation = True)\n",
    "val_gen = CustomDataGenerator(path = val_path, batch_size = batch_size, dim = dim, n_channels = n_channels, shuffle = True, augmentation = True)\n",
    "\n",
    "model = unet(input_size=(*dim,1))\n",
    "model_checkpoint = ModelCheckpoint('unet_membrane.hdf5', monitor='loss',verbose=1, save_best_only=True)\n",
    "#model.fit(train_gen, val_gen ,steps_per_epoch=300,epochs=1,callbacks=[model_checkpoint])\n",
    "model.fit(x=train_gen,epochs=20,callbacks=[model_checkpoint])\n",
    "#https://www.tensorflow.org/api_docs/python/tf/keras/Model\n",
    "\n",
    "\n",
    "#*****\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rsSNaCQtfR2x",
    "outputId": "343f5d10-1a3e-4582-96b6-4efe5493a612"
   },
   "outputs": [],
   "source": [
    "checkpoint_dir = '/content/drive/MyDrive/Colab Notebooks/Checkpoints'\n",
    "\n",
    "os.path.exists(checkpoint_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Np313SipgV9I"
   },
   "outputs": [],
   "source": [
    "# Define the checkpoint directory and file path\n",
    "checkpoint_path = f'{checkpoint_dir}/model.ckpt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4RWaciNlghd5"
   },
   "outputs": [],
   "source": [
    "#model = unet(input_size=(*dim,1)) #don above\n",
    "# Create a checkpoint instance\n",
    "checkpoint = tf.train.Checkpoint(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qPok21Eygp9i"
   },
   "outputs": [],
   "source": [
    "# Create a checkpoint manager to handle saving checkpoints\n",
    "checkpoint_manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "kKAuC6xTguvF",
    "outputId": "1fce801d-b358-4556-a0bb-a59aecd2bcfe"
   },
   "outputs": [],
   "source": [
    "# Save the checkpoint\n",
    "checkpoint_manager.save(checkpoint_number=1)\n",
    "\n",
    "# /content/drive/MyDrive/Colab Notebooks/Checkpoints/ckpt-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bNAk5bl5hTQp",
    "outputId": "ad5da1ea-912a-4398-e30b-aa2afb21d3d8"
   },
   "outputs": [],
   "source": [
    "model_path ='/content/drive/MyDrive/Colab Notebooks/Modelfile'\n",
    "os.path.exists(model_path)\n",
    "model.save(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TqERW7T3ilpC",
    "outputId": "e167c661-2c30-4de2-feb8-99d38fd61c34"
   },
   "outputs": [],
   "source": [
    "!pip install gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "B3LmWUewnB5i",
    "outputId": "1f4eafe1-3707-4c34-9732-6c40175e6fc5"
   },
   "outputs": [],
   "source": [
    "model_path ='/content/drive/MyDrive/Colab Notebooks/Modelfile/unet_custom.hdf5'\n",
    "\n",
    "os.path.exists(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Pzd6cVUrnk5-",
    "outputId": "ed37442b-be39-4574-f687-3bcc193dc25c"
   },
   "outputs": [],
   "source": [
    "!cd drive\n",
    "!ls\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 372
    },
    "id": "vWDRDmeIifeP",
    "outputId": "366d9d8e-f00d-4cce-8a97-b9f33e917071"
   },
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "import tensorflow as tf\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "# Load the UNet model\n",
    "model_path ='/content/drive/MyDrive/Colab Notebooks/Modelfile/unet_custom.hdf5'\n",
    "model = tf.keras.models.load_model(\"model_path\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1dHpFWMFj-x5"
   },
   "outputs": [],
   "source": [
    "val_gen = CustomDataGenerator(path = val_path, batch_size = batch_size, dim = dim, n_channels = n_channels, shuffle = False, augmentation = False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gh4Gs4uGj_og",
    "outputId": "8b31714b-3150-4ae4-ab8a-3a81de35095c"
   },
   "outputs": [],
   "source": [
    "#testGene = testGenerator(val_path,target_size=dim) # I don't think this needs to be changed (the function does not use ImageDataGenerator nor augmentations)\n",
    "results = model.predict_generator(val_gen,verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 745
    },
    "id": "yePJpSJgiPHS",
    "outputId": "7a740018-ac46-49f2-a463-929130eb0ea0"
   },
   "outputs": [],
   "source": [
    "\n",
    "# Define the image segmentation function\n",
    "def segment_head_circumference(image):\n",
    "    # Preprocess the input image\n",
    "    processed_image = cv2.imread(image, cv2.IMREAD_GRAYSCALE)\n",
    "    processed_image = processed_image / 255.0\n",
    "    processed_image = np.expand_dims(processed_image, axis=0)\n",
    "    processed_image = np.expand_dims(processed_image, axis=-1)\n",
    "\n",
    "    # Perform the segmentation\n",
    "    prediction = model.predict(processed_image)\n",
    "    prediction = np.squeeze(prediction)\n",
    "    binary_image = np.where(prediction > 0.5, 255, 0).astype(np.uint8)\n",
    "\n",
    "    return binary_image\n",
    "\n",
    "# Define the input and output interfaces for Gradio\n",
    "inputs = gr.inputs.Image(label=\"Ultrasound Image (PNG format)\")\n",
    "output = gr.outputs.Image(label=\"Segmented Head Circumference\", type=\"pil\")\n",
    "\n",
    "\n",
    "# Create the Gradio interface\n",
    "gr_interface = gr.Interface(\n",
    "    fn=segment_head_circumference,\n",
    "    inputs=inputs,\n",
    "    outputs=output,\n",
    "    title=\"Ultrasound Image Segmentation\",\n",
    "    description=\"Segment the head circumference in ultrasound images.\"\n",
    ")\n",
    "\n",
    "# Run the Gradio app on the Hugging Face Spaces platform\n",
    "gr_interface.launch()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WNZzUtTLootR"
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_TZ5Jw9vksY"
   },
   "outputs": [],
   "source": [
    "saveResult(\"/content/predictions\",val_path+\"Image\",results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EAaZYSfXxDSi"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.imshow(results[0,])\n",
    "np.unique(results[0,])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}