diff --git a/recognition/s4571796_ImprovedUNET/README.md b/recognition/s4571796_ImprovedUNET/README.md new file mode 100644 index 0000000000..36457ed25d --- /dev/null +++ b/recognition/s4571796_ImprovedUNET/README.md @@ -0,0 +1,109 @@ +# Image Segmentation for ISIC 2018 Melanoma Dataset +The goal of this project is to perform Image Segmentation on the ISIC (International Skin Imaging Collaboration) 2018 Melanoma Dataset. Melanoma, a form of skin cancer is a public health problem with extremely high mortality. To reduce this, Melanoma needs to be detected early. In this project, this is done using an Improved UNET Model based on the paper [1]. The model will be trained on the dataset to produce a mask for each image containing a skin lesion. Performance wise, the model should achieve a Dice Coefficient for each class (Background and Lesion) greater than 0.8. + +## ISIC 2018 Melanoma Dataset +*** +The ISIC 2018 Melanoma Dataset is a large dataset containing a total of 2,594 images of skin lesions and their respective labels. For this project, a preprocessed variant of the dataset was used. This was obtained from Course Help / Resources under COMP3710 Pattern Recognition, Semester 2, 2021. + +The images containing the skin lesions are RGB Images, and the respective masks are binary images with the black color representing the background, and the white background representing the masked skin lesion. A random image was chosen from the dataset to be displayed below. + +![Example_Image](./Resources/image.png) + +*Figure 1: Image of a Skin Lesion* + +![Example_Mask](./Resources/mask.png) + +*Figure 2: Mask for the Skin Lesion* + +## Data Preprocessing +*** +For data preprocessing, the directories containing the images and the masks are converted into a 4D numpy array. The arrays will have 4 dimensions, with each dimension corresponding to \[Number of Images, Width, Height, Number of Channels] +The images have a shape of \[2594, 128, 96, 3] and the masks have a shape of \[2594, 128, 96, 1]. + +Each of these arrays are divided using train_test_split to create a training set of size 0.7, validation set of 0.15, and test set of 0.15. + +The masks go through another step to convert all values for each pixel into either one of 2 values; 0 (Black) or 255 (White). The values are then encoded using a LabelEncoder to 0 or 1. Finally, the function to_categorical is used to change the number of channels of the masks from 1 to 2. + +## Improved UNET Architecture +*** +The Improved UNET Architecture used for this project is strongly based off the original UNET, with an additional segmentation layer and element wise sums performed after each contect module. The structure obtained from the paper [1] can be seen below. + +![Example_Mask](./Resources/ImprovedUNET.png) + +*Figure 3: Improved UNET Architecture* + +The entire Architecture is a connected series of convolution operations, concatentations and upscaling. Methods used to reduce overfitting have also been added to the model's implementations such as Batch Normalization and Dropout. + +A detailed explanation about the how the undefined modules have been implemented are included below. By default, all convolutional layers have identical padding to allow for the element wise sum and concatenation operations and use the same LeakyReLU Activation Function throughout. The number of output filters at for any module is shown in the image above contained within the green box. + +- Context Module: + > 1. Convolutional Layer with kernel size of 3 + > 2. Batch Normalization + > 3. Dropout with value of 0.1 + > 4. Convolutional Layer with kernel size of 3 + > 5. Element-wise Sum with initial input and calculated output + +- Upsampling Module: + > 1. Convolutional Transpose Layer with kernel size of 3, stride of 2 + +- Localization Module: + > 1. Convolutional Layer with kernel size of 3 + > 2. Convolutional Layer with kernel size of 1 + +- Segmentation Layer: + > 1. Convolutional Layer with kernel size of 1 + > 2. Upsampling with factor of 2 except for final layer + +## Results and Discussion +*** +From the results attained and the graph seen below, the loss graph shows a decreasing trend with a possible plateau around 0.22. The predicted plateau pattern means that the model is not overfitting, but since the difference between the training and validation loss is still somewhat significant, this implies that the model could be further improved by tuning the Hyperparameters. The accuracy graph also shows an upward trend with the latest value at 30 epochs around 0.94. + +![Loss_Graph](./Results/Loss.png) + +*Figure 4: Loss Graph* + +![Accuracy_Graph](./Results/Accuracy.png) + +*Figure 5: Accuracy Graph* + +As seen below, I included images for the real image containing a skin lesion, the real masks that have been used for testing, and my predicted results using the model. We can observe that the results are quite similar. + +![Loss_Graph](./Results/test_real_image.png) + +*Figure 6: Actual Image* + +![Loss_Graph](./Results/segmentation_real_masks.png) + +*Figure 7: Predicted Image* + +![Loss_Graph](./Results/segmentation_predicted_masks.png) + +*Figure 8: Predicted Masks* + +For the results based off DSC (Dice Score), I obtained these values. +- Overall: 0.9051 +- Label 0 (Background): 0.9605 +- Label 1 (Lesion): 0.8497 + +## Hyperparameter Tuning +*** +The initial hyperparameters chosen was 10 epochs with a batch size of 2, with the image widths and heights being 128 and 96 to preserve the aspect ratio. For the model downsampling step in the context module, this was done with 2 Convolutional Layers and a dropout in between with a value of 0.3. However, this resulted in many attempts with exploding gradients, which made the model unstable and the training step had to be restarted. To fix this issue, an additional batch normalization layer was added, and this improved the training significantly. + +However, there was another issue where the model would overfit, as the validation loss graph would start an upward trend after 15 epochs. This was resolved by choosing a lower value for the dropout layer, which was 0.1. This produced a better model that could train for 30 epochs instead and return a better result overall for the dice coefficient. + + +## Dependencies +*** +- Tensorflow +- Keras +- Operating System (os) +- Glob +- OpenCV (cv2) +- Numpy +- Matplotlib +- Sklearn +- ISIC 2018 Melanoma Dataset, Retrieved from https://challenge2018.isic-archive.com/ + +## References +*** +[1] F. Isensee, P. Kickingereder, W. Wick, M. Bendszus, and K. H. Maier-Hein, “Brain Tumor Segmentation and Radiomics Survival Prediction: Contribution to the BRATS 2017 Challenge,” Feb. 2018. Retrieved from: https://arxiv.org/pdf/1802.10508v1.pdf \ No newline at end of file diff --git a/recognition/s4571796_ImprovedUNET/Resources/ImprovedUNET.png b/recognition/s4571796_ImprovedUNET/Resources/ImprovedUNET.png new file mode 100644 index 0000000000..c39c3a3eba Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Resources/ImprovedUNET.png differ diff --git a/recognition/s4571796_ImprovedUNET/Resources/image.png b/recognition/s4571796_ImprovedUNET/Resources/image.png new file mode 100644 index 0000000000..1803162399 Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Resources/image.png differ diff --git a/recognition/s4571796_ImprovedUNET/Resources/mask.png b/recognition/s4571796_ImprovedUNET/Resources/mask.png new file mode 100644 index 0000000000..2f61f92127 Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Resources/mask.png differ diff --git a/recognition/s4571796_ImprovedUNET/Results/Accuracy.png b/recognition/s4571796_ImprovedUNET/Results/Accuracy.png new file mode 100644 index 0000000000..950b440607 Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Results/Accuracy.png differ diff --git a/recognition/s4571796_ImprovedUNET/Results/Loss.png b/recognition/s4571796_ImprovedUNET/Results/Loss.png new file mode 100644 index 0000000000..4d62d6358f Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Results/Loss.png differ diff --git a/recognition/s4571796_ImprovedUNET/Results/segmentation_predicted_masks.png b/recognition/s4571796_ImprovedUNET/Results/segmentation_predicted_masks.png new file mode 100644 index 0000000000..85d8c92dd4 Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Results/segmentation_predicted_masks.png differ diff --git a/recognition/s4571796_ImprovedUNET/Results/segmentation_real_masks.png b/recognition/s4571796_ImprovedUNET/Results/segmentation_real_masks.png new file mode 100644 index 0000000000..a27fce710f Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Results/segmentation_real_masks.png differ diff --git a/recognition/s4571796_ImprovedUNET/Results/test_real_image.png b/recognition/s4571796_ImprovedUNET/Results/test_real_image.png new file mode 100644 index 0000000000..17f3defa91 Binary files /dev/null and b/recognition/s4571796_ImprovedUNET/Results/test_real_image.png differ diff --git a/recognition/s4571796_ImprovedUNET/driver.py b/recognition/s4571796_ImprovedUNET/driver.py new file mode 100644 index 0000000000..b8744ab745 --- /dev/null +++ b/recognition/s4571796_ImprovedUNET/driver.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +__author__ = "Joshua Yu Xuan Soo" +__studentID__ = "s4571796" +__email__ = "s4571796@uqconnect.edu.au" + +""" +Driver Script for Image Segmentation on ISIC Melanoma Dataset by utilizing +an Improved UNET Model as defined in the imports. + +References: + https://arxiv.org/pdf/1802.10508v1.pdf + +""" + +# Imports +from model import unet_model +from tensorflow.keras.utils import to_categorical +import os +import glob +import cv2 +import numpy as np +from matplotlib import pyplot as plt +from sklearn.preprocessing import LabelEncoder +from sklearn.model_selection import train_test_split + +# Check for GPU -> Use Anaconda3\python.exe +# GPU Verified +""" +print("You are using TensorFlow version", tf.__version__) +if len(tf.config.list_physical_devices('GPU')) > 0: + print("You have a GPU enabled.") +else: + print("Please enable a GPU first.") +""" + +# Global Variables +classes = 2 # Number of features as Output +# Maintain Aspect Ratio of Image +height = 96 # 96 Pixels +width = 128 # 128 Pixels + +def format_images_png(directory=None): + """ Performs a formatting that returns an original folder containing + images into a numpy array in the form of (E, H, W) only if a directory + is provided. Requires images to all be in .png format. Does not + perform any operations otherwise. + + Arguments: + directory: The directory containing the images in png format // ("dir") + + Returns: + images: Numpy Array consisting of all images from a directory + """ + # Represent training image information as a list form + images = [] + + if directory != None: + # Open all files in specific folder in directory + for directory_path in glob.glob(directory): + for img_path in glob.glob(os.path.join(directory_path, "*.png")): + # Open the image + img = cv2.imread(img_path, 0) + # Resize image as per limitations of GPU + img = cv2.resize(img, (width, height)) + # Append images to list + images.append(img) + + # Convert training data list form to a numpy array + images = np.array(images) + + return images + +def format_images_jpg(directory=None): + """ Performs a formatting that returns an original folder containing + images into a numpy array in the form of (E, H, W) only if a directory + is provided. Requires images to all be in .png format. Does not + perform any operations otherwise. + + Arguments: + directory: The directory containing the images in png format // ("dir") + + Returns: + images: Numpy Array consisting of all images from a directory + """ + # Represent training image information as a list form + images = [] + + if directory != None: + # Open all files in specific folder in directory + for directory_path in glob.glob(directory): + for img_path in glob.glob(os.path.join(directory_path, "*.jpg")): + # Open the image + img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) + # Resize image as per limitations of GPU + img = cv2.resize(img, (width, height)) + # Append images to list + images.append(img) + + # Convert training data list form to a numpy array + images = np.array(images) + + return images + +def map_images(masks): + """ Takes in a set of images and masks, and performs a mapping + which converts the images and masks into a 4D Array. Also maps + each of the values in the array representing the masks into either + of 2 prominent values representing the feature set. [0, 255] + + Obtain 2 Prominent Feature Values by taking the 4 numbers with the + largest occurences in the array. Refer to Lines 133 to 143 below + 2 Feature Peak Values: 0, 255 + Midpoint Formula: (x + y) / 2 + Midpoints: 127.5 (128) + Define Ranges to convert values to + Range (0, 128) set to 0 + Range (129, 255) set to 255 + + Arguments: + masks: The set of masks represented in a 3D numpy array + + Requires: + images and masks have to come from the same set, i.e Training, Test + + Returns: + images: The set of images represented in a 4D numpy array + masks: The set of masks represented in a 4D numpy array with 2 values + representing the 2 features + """ + encoder = LabelEncoder() + + # Get the first mask as an example + # firstmask = masks[0] + # print(firstmask.shape) + + # # Change the 2D Array of size 64*64 into a single 1D Array of size 4096 + # firstmask = np.ravel(firstmask) + + # # Print the frequency of occurences for each of the values of range 0 to 255 + # unique, counts = np.unique(firstmask, return_counts=True) + # frequencies = np.asarray((unique, counts)).T + # print(frequencies) + + # Store Original Size of Array into 3 Variables: + # Number of Images (n), Height (h), Width (w) + n, h, w = masks.shape + + # Convert Numpy Array consisting of Training Masks to 1D + masks = masks.reshape(-1,1) + + # Iterate through all pixels in the entire training set and replace values + # while using Numpy's fast time complexity methods. + masks[(masks > 0)&(masks < 129)] = 0 + masks[(masks >= 129)&(masks <= 255)] = 255 + + # Transform (0, 255) into (0, 1) + masks = encoder.fit_transform(masks) + + # Validate that the array is supposed to only consist of 4 feature values, + # Correct Values: 0, 1 + # print(np.unique(masks)) + + # Transform the Train Masks back into the original n, h and w + masks = masks.reshape(n, h, w) + + # Perform Addition of Channel Dimension of Train Masks + masks = np.expand_dims(masks, axis=3) + + return masks + +def get_model(output_classes, height, width, input_classes): + """ Gets the imported model (Improved UNET) with specific parameters + + Arguments: + output_classes: Number of features as Output + height: The height of each image + width: The width of each image + input_classes: Number of features as Input + + Returns: + model: A UNET model with the specified parameters + """ + return unet_model(num_channels=output_classes, image_height=height, image_width=width, image_channels=input_classes) + +def plot_graphs(history, epoch, type=None): + """ Given the history of the model, plot a graph for viewing. + + Arguments: + history: The history of the model during training + epoch: Number of epochs + type: The graph to be plotted, i.e. Loss or Accuracy + + Returns: + plot: A plot of the type of graph specified displayed + """ + # Define x-axis as Epochs + epochs = range(1, epoch+1) + + # Plot loss graph + if type == "loss": + loss = history.history['loss'] + val_loss = history.history['val_loss'] + plt.plot(epochs, loss, 'g', label='Training loss') + plt.plot(epochs, val_loss, 'b', label='Validation loss') + plt.title('Training and Validation Loss') + plt.xlabel('Epochs') + plt.ylabel('Loss') + + # Plot accuracy graph + elif type == "accuracy": + accuracy = history.history['accuracy'] + val_accuracy = history.history['val_accuracy'] + epochs = range(1, len(accuracy) + 1) + plt.plot(epochs, accuracy, 'g', label='Training Accuracy') + plt.plot(epochs, val_accuracy, 'b', label='Validation Accuracy') + plt.title('Training and Validation Accuracy') + plt.xlabel('Number of Epochs') + plt.ylabel('Accuracy') + + plt.legend() + plt.grid() + plt.show() + +def dice_coef(y_true, y_pred): + """ Returns the Dice Score for binary image segmentation + + Arguments: + y_true: The true array + y_pred: The predicted array + + Returns: + score: The dice score according to the dice formula + """ + y_true_f = y_true.flatten() + y_pred_f = y_pred.flatten() + intersection = np.sum(y_true_f * y_pred_f) + smooth = 0.0001 + score = (2. * intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth) + return score + + +def dice_coef_multilabel(y_true, y_pred, labels): + """ Returns the Dice Score for multiclass image segmentation + + Arguments: + y_true: The true array + y_pred: The predicted array + labels: The number of classes for the output + + Returns: + score: The total Dice Score divided by the number of classes + """ + # Initialize Dice Score as 0 + dice = 0 + # Iterate through all classes + for index in range(labels): + coeff = dice_coef(y_true[:,:,:,index], y_pred[:,:,:,index]) + dice += coeff + print("The dice score for class " + str(index) + " is " + str(coeff)) + # Return the Dice Score + score = dice/labels + return score + +def main(): + # Lines 268 to 289 are Preprocessing of Data + + # Loading the Directories containing the Images + images = format_images_jpg("ISIC_Images") + masks = format_images_png("ISIC_Labels") + + # Create Training Set, Validation Set, Test Set, choose arbitrary random state value + # Training (0.7), Validation (0.15), Test (0.15) + train_images, val_images, train_masks, val_masks = train_test_split(images, masks, test_size=0.3, random_state=42) + val_images, test_images, val_masks, test_masks = train_test_split(val_images, val_masks, test_size=0.5, random_state=42) + + # Save the Real Image for Comparison (Do this beforehand, or to_categorical will produce errors) + # Taking the image with index 1 + seg_real_masks = test_masks[1] + cv2.imwrite("segmentation_real_masks.png", seg_real_masks) + + # Change the values of Masks to [0,1] + train_masks = map_images(train_masks) + val_masks = map_images(val_masks) + test_masks = map_images(test_masks) + + # Convert the Masks to 2 channels + train_masks = to_categorical(train_masks, num_classes=classes) + val_masks = to_categorical(val_masks, num_classes=classes) + test_masks = to_categorical(test_masks, num_classes=classes) + + # Input Channels = The number of the last dimension of the images = 3 + UNET_model = get_model(classes, height, width, 3) + UNET_model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=['accuracy']) + # Summary of the UNET Model + #print(UNET_model.summary()) + + #history = UNET_model.fit(train_images, train_masks, + # batch_size=2, + # verbose=1, + # epochs=30, + # validation_data=(val_images, val_masks), + # shuffle=False) + + # Save the model + #UNET_model.save('UNET_ISIC_30Epochs.hdf5') + + # Plot the Training and Validation Loss and Accuracy + #plot_graphs(history, epoch=30, type="loss") + #plot_graphs(history, epoch=30, type="accuracy") + + # Load the model + UNET_model.load_weights('UNET_ISIC_30Epochs.hdf5') + + # Make Predictions + predictions = UNET_model.predict(test_images) + + # Get Dice Score Coefficient + print("Average Dice Score", dice_coef_multilabel(predictions, test_masks, classes)) + + predictions_argmax = np.argmax(predictions, axis=3) + + # Save the Test Images and Predicted Images for Comparison + # Taking the image with index 1 + test_real_image = test_images[1] + cv2.imwrite("test_real_image.png", test_real_image) + + seg_pred_masks = predictions_argmax[1] + cv2.imwrite("segmentation_predicted_masks.png", seg_pred_masks) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/recognition/s4571796_ImprovedUNET/model.py b/recognition/s4571796_ImprovedUNET/model.py new file mode 100644 index 0000000000..c4c39cf6a2 --- /dev/null +++ b/recognition/s4571796_ImprovedUNET/model.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +__author__ = "Joshua Yu Xuan Soo" +__studentID__ = "s4571796" +__email__ = "s4571796@uqconnect.edu.au" + +""" +Improved UNET Model for Image Segmentation on ISIC Melanoma Dataset. + +References: + https://arxiv.org/pdf/1802.10508v1.pdf + +""" + +# Imports +from keras.models import Model +from keras.layers import Input, Conv2D, Conv2DTranspose, Dropout, Concatenate, LeakyReLU, UpSampling2D, BatchNormalization +import tensorflow as tf + +# Functions + +def context_module(input, num_filters): + """Performs the convolution in the context module + + Arguments: + input: The image to perform convolution on + num_filters: The number of filters used to perform convolution + """ + # Store the initial values of the image + temp = input + + # Used 3*3 filter size + # Activation Function is LeakyReLU + x = Conv2D(num_filters, 3, padding="same", activation=LeakyReLU(alpha=0.1))(input) + + # BatchNormalization to reduce Overfitting + x = BatchNormalization()(x) + + # Dropout to reduce Overfitting, value = 0.3 + x = Dropout(0.1)(x) + + # UNET Architecture performs Convolution2D two times for a single layer + x = Conv2D(num_filters, 3, padding="same", activation=LeakyReLU(alpha=0.1))(x) + + # Add the initial values to the convoluted result element wise + y = tf.math.add(temp, x) + + return y + +def downs(input, num_filters): + """Performs the full down convolution per layer of UNET. To half dimensions + of image, convolution with stride 2 is included. + + Arguments: + input: The image to perform convolution on + num_filters: The number of filters used to perform convolution + """ + x = context_module(input, num_filters) + + # Choosing Filter Size of 3*3, Stride set to Filter Size = 2*2. Note that + # the number of filters of this convolution is doubled. + p = Conv2D((num_filters * 2), 3, padding="same", strides=(2,2), activation=LeakyReLU(alpha=0.1))(x) + return x, p + +def localization_module(input, num_filters): + """Performs the convolution in the localization module + + Arguments: + input: The image to perform convolution on + num_filters: The number of filters used to perform convolution + """ + # Used 3*3 filter size + # Activation Function is LeakyReLU + x = Conv2D(num_filters, 3, padding="same", activation=LeakyReLU(alpha=0.1))(input) + + # UNET Architecture performs Convolution2D two times for a single layer + x = Conv2D(num_filters, 1, padding="same", activation=LeakyReLU(alpha=0.1))(x) + + return x + +def ups(input, skip_features, num_filters): + """Performs the full up convolution per layer of UNET. Similar to downs, + uses Conv2dTranspose to double dimensions of image and Concatenates + current image with skip connections, then performs a localization module. + + Arguments: + input: The image to perform convolution on + skip_features: Feature to skip, pass in input from downsampling + num_filters: The number of filters used to perform convolution + """ + x = Conv2DTranspose(num_filters, 3, strides=2, padding="same", activation=LeakyReLU(alpha=0.1))(input) + x = Concatenate()([x, skip_features]) + x = localization_module(x, num_filters) + + return x + +def ups_end(input, skip_features, num_filters): + """A clone of the ups function, but instead of performing the localization + module at the end, only perform a single convolution with kernel size 3. + + Arguments: + input: The image to perform convolution on + skip_features: Feature to skip, pass in input from downsampling + num_filters: The number of filters used to perform convolution + """ + + x = Conv2DTranspose(num_filters, 3, strides=2, padding="same", activation=LeakyReLU(alpha=0.1))(input) + x = Concatenate()([x, skip_features]) + x = Conv2D((num_filters * 2), 3, padding="same", activation=LeakyReLU(alpha=0.1))(x) + + return x + +def segmentation_module(input, output_filters): + """Performs the convolution in the localization module + + Arguments: + input: The image to perform convolution on + output_filters: The number of classes of the output image + """ + # Used 1*1 filter size + x = Conv2D(output_filters, 1, padding="same", activation=LeakyReLU(alpha=0.1))(input) + + return x + +def unet_model(num_channels, image_height, image_width, image_channels): + """Builds the UNET Model using defined variable values + + Arguments: + num_channels: The number of channels for the output image, i.e. 1 for Grayscale, 3 for RGB, etc. + image_height: The height of the original image + image_width: The width of the original image + image_channels: The number of channels on the original image + + Notes: + This project aims to classify Melanoma. The number of channels for the output is 2 because the + segmentations contain the background as black, and the melanoma as white. The original number + of channels is 3, since it is an RGB Image. + """ + inputs = Input((image_height, image_width, image_channels)) + + layer0 = Conv2D(16, 3, padding="same", activation=LeakyReLU(alpha=0.1))(inputs) + + # Inputting the original image and iterative downsampling + # Store Layers to transfer as skip connections to upsampling step + # Store Poolings to pass to next layers as inputs + layer1, pool1 = downs(layer0, 16) + layer2, pool2 = downs(pool1, 32) + layer3, pool3 = downs(pool2, 64) + layer4, pool4 = downs(pool3, 128) + + # Bridge: Only context module applied here, filters = 256 + bridge = context_module(pool4, 256) + + # Inputting the bridge and iterative upsampling + # Taking in the layers defined earlier as skip connections + # First upsampling does not need segmentation layer + layer5 = ups(bridge, layer4, 128) + layer6 = ups(layer5, layer3, 64) + layer7 = ups(layer6, layer2, 32) + layer8 = ups_end(layer7, layer1, 16) + + # First upsampling does not need segmentation layer, but layer6 + # onwards require segmentation layer + seg1 = segmentation_module(layer6, 2) + # Upsampling first segmentation layer + seg1 = UpSampling2D(size=(2,2), interpolation='nearest')(seg1) + seg2 = segmentation_module(layer7, 2) + # First element wise addition + seg_add1 = tf.math.add(seg1, seg2) + # Upsampling the sum of the first and second segmentation layer + seg_add1 = UpSampling2D(size=(2,2), interpolation='nearest')(seg1) + seg3 = segmentation_module(layer8, 2) + # Second element wise addition + segmentation = tf.math.add(seg_add1, seg3) + + # Utilizing Softmax Function + outputs = Conv2D(num_channels, (1,1), activation="softmax")(segmentation) + model = Model(inputs, outputs) + return model \ No newline at end of file