"""
## **Project Overview**: 

| <br/><font size="3"><b> Focus </b></font><br/><br/>  | <br/><font size="3"><b> Description </b></font><br/><br/> |
| :-- | :-- |
| <br/><font size="3"> <b>Project Title</b> </font>  <br/><br/>| <br/><font size="3"> CARLA Image Semantic Segmentation with DeepLabV3+</font> <br/><br/>|
| <br/><font size="3"> <b>Project Type</b> </font> <br/><br/>| <br/><font size="3">Image Segmentation (Semantic Segmenetation) </font> <br/><br/>|
| <br/><font size="3"><b>Project Objectives</b> </font>  <br/><br/>| <br/><font size="3">1. Create a Model to predict semantic segmentations of CARLA images<br/><br/> 2. Predict masks using the model and compare with ground-truth masks </font> <br/><br/>|
| <br/><font size="3"> <b>Dataset Overview </b> </font> <br/><br/>| <br/><font size="3"> This dataset provides data images and labeled semantic segmentations captured via CARLA self-driving car simulator. The data which was generated as part of the 2018 Lyft Udacity Perception Challenge consists of **5000** images and their semantic segmentations. </font> <br/><br/>|
| <br/><font size="3"> <b> Model Evaluation Metrics</b> </font> <br/><br/>| <br/><font size="3"> Accuracy:- Minimum of 92% <br/> Intersection over Union (IoU):- Minimum IoU: 50%, Max IoU: 80%, Mean IoU: 65% </font> <br/><br/>|
| <font size="3"> <b>Image Segmentation Model Type</b> </font> | <br/><font size="3"> <a href="https://arxiv.org/pdf/1802.02611.pdf"> <b> DeepLabV3+ Architecture</b></a> </font> <br/><br/>|
| <br/><font size="3"> <b>Major Libraries Used</b> </font> <br/><br/>| <br/><font size="3"> <a href="https://keras.io"><b>Keras</b></a>, <a href="https://imageio.readthedocs.io/"><b>ImageIO</b></a>, <a href="https://scikit-learn.org"><b> Scikit-Learn </b></a>, <a href="https://numpy.org"><b>Numpy</b></a>, <a href="https://matplpotlib.org"><b>Matplotlib</b></a></font> <br/><br/>|

<br/>
"""
"""
## **Table of Contents**: 
<font size="3">

- [**1 - Import Required Packages**](#1)

    
- [**2 - Data Preparation**](#2)
    - [2.1. Load the images and masks from their directories](#2-1)
    - [2.2.  Create a data pipeline to read and preprocess our data](#2-2)

    
- [**3 - Model Architecture and Training**](#3)
    - [3.1. - DeepLabV3+ Model design](#3-1)
    - [3.2. - Model training](#3-2)    

    
- [**4 - Model Evaluation**](#4)
    - [4.1. - Model Accuracy](#4-1)
    - [4.2. - Intersection-over-Union (IoU)](#4-2)

    
- [**5 - Predict image segmentations using the trained Model**](#5)
    - [5.1. - Create functions to preprocess selected images and display their true state, true mask and predicted mask](#5-1)
    - [5.2. - Prediction on the train set](#5-2)
    - [5.3. - Prediction on the validation set](#5-3)
    - [5.4. - Prediction on the test set](#5-4)    
    
<font>
"""
"""
<a name='1'></a>
## **1. Import Required Packages**
"""
import numpy as np
import imageio
import random
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 

import tensorflow as tf 
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import models
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.layers import Activation, Input, Conv2D, MaxPooling2D, BatchNormalization, Conv2DTranspose, concatenate
from tensorflow.keras.models import Model, load_model
from sklearn.model_selection import train_test_split

import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
"""
<a name='2'></a>
## **2. Data Preparation**

The Lyft Udacity Semantic Segmentation for Self-driving Cars Challenge data (images and masks) is splitted across five directories (dataA, dataB, dataC, dataD, and dataE). As part of the data preparation step, we will load images and masks from all the five directories and carry out the a few preprocessing steps to ensure we provide our model with quality dataset.
"""
"""
<a name='2-1'></a>
### **2.1. Load the images and masks from their directories**

In this data preparation step, we will:
1. Create 2 lists containing the paths of images and masks
2. Split the lists into training, validation and test sets
"""
"""
#### **2.1.1. Create lists containing the paths of images and masks**

In this step, we will

* Create a list that contains all the paths to all directories in the main directory (a list that contains the path to dataA, dataB, dataC, dataD, and dataE)
* Create a function to iterate over all the direcory paths where our data are located (list in 1.) and return the list of the image paths in those directories. 
* Create lists of image and mask paths by initializing the function above
* Preview some masked and unmasked images by reading them from their paths
"""
"""
**A. Create a list that contains all the paths to all directories in the main directory (a list that contains the path to dataA, dataB, dataC, dataD, and dataE)**
"""
image_path = ["../input/lyft-udacity-challenge/"+"data"+i+"/"+"data"+i+"/CameraRGB/" for i in ['A', 'B', 'C', 'D', 'E']]
mask_path = ["../input/lyft-udacity-challenge/"+"data"+i+"/"+"data"+i+"/CameraSeg/" for i in ['A', 'B', 'C', 'D', 'E']]
"""
**Create a function to iterate over all the direcory paths where our data are located (list in 2.1.1.) and return the list of the image paths in those directories**
"""
def list_image_paths(directory_paths):
    image_paths = []
    for directory in range(len(directory_paths)):
        image_filenames = os.listdir(directory_paths[directory])
        for image_filename in image_filenames:
            image_paths.append(directory_paths[directory] + image_filename)
    return image_paths
"""
**Create lists of image and mask paths by initializing the function above**
"""
image_paths = list_image_paths(image_path) 
mask_paths = list_image_paths(mask_path)
number_of_images, number_of_masks = len(image_paths), len(mask_paths)
print(f"1. There are {number_of_images} images and {number_of_masks} masks in our dataset")
print(f"2. An example of an image path is: \n {image_paths[0]}")
print(f"3. An example of a mask path is: \n {mask_paths[0]}")

"""
**Preview random masked and unmasked images by reading them from their paths**
"""
import random
number_of_samples = len(image_paths)

for i in range(3):
    N = random.randint(0, number_of_samples - 1)

    img = imageio.imread(image_paths[N])
    mask = imageio.imread(mask_paths[N])
    mask = np.array([max(mask[i, j]) for i in range(mask.shape[0]) for j in range(mask.shape[1])]).reshape(img.shape[0], img.shape[1])

    fig, arr = plt.subplots(1, 3, figsize=(20, 8))
    arr[0].imshow(img)
    arr[0].set_title('Image')
    arr[0].axis("off")
    arr[1].imshow(mask)
    arr[1].set_title('Segmentation')
    arr[1].axis("off")    
    arr[2].imshow(mask, cmap='Paired')
    arr[2].set_title('Segmentation')
    arr[2].axis("off")
"""
#### **2.1.2. Split the image and mask paths into training, validation, and test sets**
"""
# First split the image paths into training and validation sets
train_image_paths, val_image_paths, train_mask_paths, val_mask_paths = train_test_split(image_paths, mask_paths, train_size=0.8, random_state=0)

# Keep part of the validation set as test set
validation_image_paths, test_image_paths, validation_mask_paths, test_mask_paths = train_test_split(val_image_paths, val_mask_paths, train_size = 0.80, random_state=0)

print(f'There are {len(train_image_paths)} images in the Training Set')
print(f'There are {len(validation_image_paths)} images in the Validation Set')
print(f'There are {len(test_image_paths)} images in the Test Set')
"""
<a name='2-2'></a>
### **2.2. - Create a data pipeline to read and preprocess our data**

We will be using the tf.data.Dataset API to load our images and masks for our model to process. The Dataset API allows us to build an asynchronous, highly optimized data pipeline to prevent our GPU from data starvation. It loads data from the disk (images or text), applies optimized transformations, creates batches and sends it to the GPU. Unlike former data pipelines made the GPU, the Dataset API wait for the CPU to load the data, leading to performance issues.

To do this, we will 
1. Create a function to read image and mask paths and return equivalent arrays
2. Create a data generator function to read and load images and masks in batches
3. Create data pipelines for the training, validation and test sets using both functions
4. Preview sample images and their segmentations from the three dataset categories
"""
"""
#### **2.2.1. Create a function to read image and mask paths and return equivalent arrays**

The **read_image** function will
1. Read an image and its mask from their paths
2. Convert the digital image and its mask to image arrays 
3. Normalize the datasets
4. Resize the image and its masks to a desired dimension
"""
def read_image(image_path, mask_path):
    
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=3)
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, (256, 256), method='nearest')

    mask = tf.io.read_file(mask_path)
    mask = tf.image.decode_png(mask, channels=3)
    mask = tf.math.reduce_max(mask, axis=-1, keepdims=True)
    mask = tf.image.resize(mask, (256, 256), method='nearest')
    
    return image, mask
"""
#### **2.2.2. Create a data generator function to read and load images and masks in batches**


"""
def data_generator(image_paths, mask_paths, buffer_size, batch_size):
    
    image_list = tf.constant(image_paths) 
    mask_list = tf.constant(mask_paths)
    dataset = tf.data.Dataset.from_tensor_slices((image_list, mask_list))
    dataset = dataset.map(read_image, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.cache().shuffle(buffer_size).batch(batch_size)
    
    return dataset
"""
#### **2.2.3. Create data pipelines for the training, validation and test sets using both functions**
"""
batch_size = 32
buffer_size = 500

train_dataset = data_generator(train_image_paths, train_mask_paths, buffer_size, batch_size)
validation_dataset = data_generator(validation_image_paths, validation_mask_paths, buffer_size, batch_size)
test_dataset = data_generator(test_image_paths, test_mask_paths, buffer_size, batch_size)
"""
#### **2.2.4. Preview sample images and masks from the three dataset categories**
"""
# Take a batch (32 images and their labelled segmentations from each category of data)
for train_images, train_masks in train_dataset:
    break
for validation_images, validation_masks in validation_dataset:
    break
for test_images, test_masks in test_dataset:
    break
    

for i in range(3):
    N = random.randint(0, batch_size-1)
    
    images = [train_images[N], validation_images[N], test_images[N]]
    masks = [train_masks[N], validation_masks[N], test_masks[N]]
    title = ['Train Image', 'Validation Image', 'Test Image', 'Train Mask', 'Validation Mask', 'Test Mask']

    fig, arr = plt.subplots(1, 3, figsize=(20, 8))
    arr[0].imshow(images[i])
    arr[0].set_title(title[i])
    arr[0].axis("off")
    arr[1].imshow(masks[i])
    arr[1].set_title(title[i+3])
    arr[1].axis("off")
    arr[2].imshow(masks[i], cmap='Paired')
    arr[2].set_title(title[i+3])
    arr[2].axis("off")
"""
<a name='3'></a>
## **3. Model Architecture and Training**
We will using a the **DeepLabv3+ architecture** to train our semantic segmentation model. The DeepLabv3+ is a semantic segmentation architecture that improves upon DeepLabv3 with several improvements, such as adding a simple yet effective decoder module to achieve an encoder-decoder structure. The encoder module processes multiscale contextual information by applying dilated convolution at multiple scales, while the decoder module refines the segmentation results along object boundaries.

<center><img src="https://i.ibb.co/cXmRSr3/deeplabv3-plus-diagram.png" alt="deeplabv3-plus-diagram" border="0"></center>

"""
"""
<a name='3-1'></a>

### **3.1. DeepLabV3+ Model Design**
"""
def convolution_block(block_input, num_filters=256, kernel_size=3, dilation_rate=1, padding="same", use_bias=False):
    x = layers.Conv2D(
        num_filters,
        kernel_size=kernel_size,
        dilation_rate=dilation_rate,
        padding="same",
        use_bias=use_bias,
        kernel_initializer=keras.initializers.HeNormal(),
        )(block_input)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    
    return x
def DilatedSpatialPyramidPooling(dspp_input):
    dims = dspp_input.shape
    x = layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input)
    x = convolution_block(x, kernel_size=1, use_bias=True)
    out_pool = layers.UpSampling2D(
        size=(dims[-3] // x.shape[1], dims[-2] // x.shape[2]), interpolation="bilinear",
    )(x)

    out_1 = convolution_block(dspp_input, kernel_size=1, dilation_rate=1)
    out_6 = convolution_block(dspp_input, kernel_size=3, dilation_rate=6)
    out_12 = convolution_block(dspp_input, kernel_size=3, dilation_rate=12)
    out_18 = convolution_block(dspp_input, kernel_size=3, dilation_rate=18)

    x = layers.Concatenate(axis=-1)([out_pool, out_1, out_6, out_12, out_18])
    output = convolution_block(x, kernel_size=1)
    return output
def DeeplabV3(image_size, num_classes):
    model_input = keras.Input(shape=(image_size, image_size, 3))
    resnet50 = keras.applications.ResNet50(
        weights="imagenet", include_top=False, input_tensor=model_input
    )
    x = resnet50.get_layer("conv4_block6_2_relu").output
    x = DilatedSpatialPyramidPooling(x)

    input_a = layers.UpSampling2D(
        size=(image_size // 4 // x.shape[1], image_size // 4 // x.shape[2]),
        interpolation="bilinear",
    )(x)
    input_b = resnet50.get_layer("conv2_block3_2_relu").output
    input_b = convolution_block(input_b, num_filters=48, kernel_size=1)

    x = layers.Concatenate(axis=-1)([input_a, input_b])
    x = convolution_block(x)
    x = convolution_block(x)
    x = layers.UpSampling2D(
        size=(image_size // x.shape[1], image_size // x.shape[2]),
        interpolation="bilinear",
    )(x)
    model_output = layers.Conv2D(num_classes, kernel_size=(1, 1), padding="same")(x)
    model = tf.keras.Model(inputs=model_input, outputs=model_output)
    
    return model
img_height = 256
img_width = 256
num_channels = 3
filters = 32
n_classes = 23

model = DeeplabV3(img_height, num_classes=23)
model.summary()
"""
<a name='3-2'></a>
### **3.2. Model Training**
"""
model.compile(optimizer = 'adam', loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])
callback = EarlyStopping(monitor='val_accuracy', patience=20, restore_best_weights=True)
reduce_lr = ReduceLROnPlateau(monitor='val_accuracy',factor=1e-1, patience=5, verbose=1, min_lr = 2e-6)
batch_size = 32
epochs = 30
history = model.fit(train_dataset, 
                    validation_data = validation_dataset, 
                    epochs = epochs, 
                    verbose=1, 
                    callbacks = [callback, reduce_lr], 
                    batch_size = batch_size, 
                    shuffle = True)
acc = [0.] + history.history['accuracy']
val_acc = [0.] + history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

plt.figure(figsize=(8, 8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.ylabel('Accuracy')
plt.ylim([min(plt.ylim()),1])
plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.ylabel('Cross Entropy')
plt.ylim([0,1.0])
plt.title('Training and Validation Loss')
plt.xlabel('epoch')
plt.show()
model.save('carla-image-segmentation-model.h5')
"""
<a name='4'></a>
## **4. Model Evaluation**

We will be using Model Accuracy and Mean Intersection-over-Union (mIoU) to evaluate our model performance.
"""
"""
<a name='4-1'></a>
### **4.1. Model Accuracy**

"""
train_loss, train_accuracy = model.evaluate(train_dataset, batch_size = 32)
validation_loss, validation_accuracy = model.evaluate(validation_dataset, batch_size = 32)
test_loss, test_accuracy = model.evaluate(test_dataset, batch_size = 32)
print(f'Model Accuracy on the Training Dataset: {round(train_accuracy * 100, 2)}%')
print(f'Model Accuracy on the Validation Dataset: {round(validation_accuracy * 100, 2)}%')
print(f'Model Accuracy on the Test Dataset: {round(test_accuracy * 100, 2)}%')
"""
<a name='4-2'></a>
### **4.2. Intersection-over-Union (IoU)**

The Intersection-over-Union (IoU), also referred to as the Jaccard index is useful in quantifying the percent overlap between the target mask and the mask prediction output from our model. Recall that the task of semantic segmentation is simply to predict the class of each pixel in an image. So, the IoU aims to evaluate the similarities in the pixels of both the ground true mask and the predicted mask.

Since the size of our ground true masks and the predicted masks size is 256 * 256 * 23 (which is equivalent to 59,069,888 pixels). Hence, our Model's IoU Score for any image in our dataset is the proprortion of the 59,069,888 pixels in the predicted mask that matches with the ground true mask's 59,069,888 pixels. 


[Kindly Read more on IoU and other Image Segmentation Metrics here](https://www.jeremyjordan.me/evaluating-image-segmentation-models/)
"""
def iou_score(dataset):
    
    """
    Argument:
        dataset -- the dataset to calculate IoU on
    
    Returns:
        min_iou -- minimum IoU
        max_iou -- maximum IoU
        mean_iou -- mean IoU ()
        """
    # Create empty lists 
    intersections, unions, max_ious, min_ious = [], [], [], []
    
    for images, masks in dataset:
        pred_mask = model.predict(images)
        intersection = np.logical_and(masks, pred_mask)
        union = np.logical_or(masks, pred_mask)
        intersection_sum = np.array([np.sum(inter) for inter in intersection])
        union_sum = np.array([np.sum(un) for un in union])
        batch_iou_score = intersection_sum / union_sum
        batch_min_iou = np.amin(batch_iou_score)
        batch_max_iou = np.amax(batch_iou_score)
        
        intersections.append(np.sum(intersection))
        unions.append(np.sum(union))
        min_ious.append(batch_min_iou)
        max_ious.append(batch_max_iou)

    min_iou = np.amin(min_ious)
    max_iou = np.amax(max_ious)
    mean_iou = np.sum(intersections) / np.sum(unions)   
    
    return min_iou, max_iou, mean_iou
train_min_iou, train_max_iou, train_mean_iou = iou_score(train_dataset)
validation_min_iou, validation_max_iou, validation_mean_iou = iou_score(validation_dataset)
test_min_iou, test_max_iou, test_mean_iou = iou_score(test_dataset)
print(f'IoU on the Training Dataset: \n Minimum IoU Score: {round(train_min_iou*100, 2)}% \n Maximum IoU Score: {round(train_max_iou*100, 2)}% \n Mean IoU Score: {round(train_mean_iou*100, 2)}% \n')
print(f'IoU on the Validation Dataset: \n Minimum IoU Score: {round(validation_min_iou*100, 2)}% \n Maximum IoU Score: {round(validation_max_iou*100, 2)}% \n Mean Iou Score: {round(validation_mean_iou*100, 2)}% \n')
print(f'IoU on the Test Dataset: \n Minimum IoU Score: {round(test_min_iou*100, 2)}% \n Maximum IoU Score: {round(test_max_iou*100, 2)}% \n Mean IoU Score: {round(test_mean_iou*100, 2)}% \n')
"""
<a name='5'></a>
## **5. Predict image segmentations using the trained Model**

In this section, we will

1. Create a function to preprocess selected images and display their true state, true mask and predicted mask
2. Predict and compare masks of images in the training set
3. Predict and compare masks of images in the validation set
4. Predict and compare masks of images in the test set
"""
"""
<a name='5-1'></a>
### **5.1. Create functions to preprocess selected images and display their true state, true mask and predicted mask**

In this step, we will:
1. Load our model
2. Define a function to create new masks using our model
3. Define a function to display outputs of this process: an input image, its true mask, and its predicted mask.
4. Define a function to select images from a specified dataset and return the images, their true masks and their predicted masks.
"""
"""
##### **5.1.1. Load our model** 
"""
# Load model
from tensorflow.keras.models import Model, load_model
model = load_model('carla-image-segmentation-model.h5')
"""
##### **5.1.2. Define a function to create new masks using our model** 
"""
def create_mask(pred_mask):
    pred_mask = tf.argmax(pred_mask, axis=-1)
    pred_mask = pred_mask[..., tf.newaxis]
    return pred_mask[0]
"""
##### **5.1.3. Define a function to display outputs of this process: an input image, its true mask, and its predicted mask** 
"""
def display(display_list):
    plt.figure(figsize=(15, 15))

    title = ['Input Image', 'True Mask', 'Predicted Mask']

    for i in range(len(display_list)):
        plt.subplot(1, len(display_list), i+1)
        plt.title(title[i])
        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
        plt.axis('off')
    plt.show()
"""
##### **5.1.4. Define a function to select images from a specified dataset and return the images, their true masks and their predicted masks** 
"""
def show_predictions(dataset, num):
    """
    Displays the first image of each of the num batches
    """
    if dataset:
        for image, mask in dataset.take(num):
            pred_mask = model.predict(image)
            display([image[0], mask[0], create_mask(pred_mask)])
    else:
        display([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])
"""
<a name='5-2'></a>
### **5.2. Predict and compare masks of images in the training set**
"""
show_predictions(train_dataset, 6)
"""
<a name='5-3'></a>
### **5.3. Predict and compare masks of images in the validation set**
"""
show_predictions(validation_dataset, 6)
"""
<a name='5-4'></a>
### **5.4. Predict and compare masks of images in the test set**
"""
show_predictions(test_dataset, 6)