import tensorflow as tf
print(tf.version.VERSION)
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))

"""
Explanation: Enable GPU and set up helper functions
This notebook and pretty much every other notebook in this repository
will run faster if you are using a GPU.
On Colab:
- Navigate to Edit→Notebook Settings
- Select GPU from the Hardware Accelerator drop-down
On Cloud AI Platform Notebooks:
- Navigate to https://console.cloud.google.com/ai-platform/notebooks
- Create an instance with a GPU or select your instance and add a GPU
Next, we'll confirm that we can connect to the GPU with tensorflow:
End of explanation
"""


import tensorflow_datasets as tfds
dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)

"""
Explanation: Read Oxford-IIIT Pets dataset
The dataset is part of TensorFlow datasets.
Version 3 and higher of the dataset has ground truth
segmentation masks.
The dataset already contains test and train splits.
End of explanation
"""


def read_and_preprocess(data):
  input_image = tf.image.resize(data['image'], (128, 128))
  input_mask = tf.image.resize(data['segmentation_mask'], (128, 128))
  
  input_image = tf.image.convert_image_dtype(input_image, tf.float32) # [0,1]
  input_mask -= 1 # {1,2,3} to {0,1,2}

  return input_image, input_mask

train = dataset['train'].map(read_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
test = dataset['test'].map(read_and_preprocess)

"""
Explanation: When loading an image, we do the following:
* Resize the image and mask to 128x128
* Scale the image to [0,1]
* Create input_mask labels. The original segmentation mask consists of {1,2,3} for background, outline, and interior respectively. We subtract 1 to get class indexes of {0,1,2}.
End of explanation
"""


import matplotlib.pylab as plt
import numpy as np
f, ax = plt.subplots(2, 5, figsize=(16,5))
for idx, (img, mask) in enumerate(train.take(5)):
  ax[0, idx].imshow(tf.keras.preprocessing.image.array_to_img(img))
  ax[0, idx].axis('off')
  mask = tf.reshape(mask, [128, 128])
  ax[1, idx].imshow(mask.numpy())
  ax[1, idx].axis('off')

"""
Explanation: Display some example images and their labels
End of explanation
"""


BATCH_SIZE = 64
BUFFER_SIZE = 1000

def augment(img, mask):
  if tf.random.uniform(()) > 0.5:
    img = tf.image.flip_left_right(img)
    mask = tf.image.flip_left_right(mask)
  return img, mask

train_dataset = train.cache().map(augment).shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
train_dataset = train_dataset.prefetch(buffer_size=tf.data.AUTOTUNE)

test_dataset = test.batch(BATCH_SIZE)

TRAIN_LENGTH = info.splits['train'].num_examples
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE
OUTPUT_CHANNELS = 3

base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)

# Use the activations of these layers to form skip connections
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs,
                            name='pretrained_mobilenet')

down_stack.trainable = False

"""
Explanation: Modified U-Net
The model being used here is a modified U-Net.
A U-Net typically consists of an encoder 
which downsamples an image to an encoding,
and a mirrored decoder which upsamples the encoding
back to the desired mask. The decoder blocks have a number of skip connections
that directly connect the encoder blocks to the decoder.
Doing this from scratch requires a lot of trainable parameters and is better done using Generative
Adversarial Networks (GANs) which are covered in
Chapter 12. 
In this notebook, we will use a pretrained MobileNetV2 to
create the encoding and a set of upsampling layers to
get back to the desired mask. 
When doing so, however, we will pull out layers with the
desired sizes so that the upsampling skip layers use
corresponding weights from the pretrained models.
To make the dataset go a little bit farther,
we also augment the training dataset to flip the images.
(data augmentation is covered in Chapter 6)
Note that the output consists of 3 channels. This is because we have 3 possible labels for each pixel (background, outline, interior). Think of this as multi-classification problem with three possible classes.
End of explanation
"""


def upsample(filters, size, name):
  return tf.keras.Sequential([
     tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same'),
     tf.keras.layers.BatchNormalization(),
     tf.keras.layers.ReLU()
  ], name=name)

up_stack = [
    upsample(512, 3, 'upsample_4x4_to_8x8'),
    upsample(256, 3, 'upsample_8x8_to_16x16'),
    upsample(128, 3, 'upsample_16x16_to_32x32'),
    upsample(64, 3,  'upsample_32x32_to_64x64')
]

import re

def unet_model(output_channels):
  inputs = tf.keras.layers.Input(shape=[128, 128, 3], name='input_image')

  # Downsampling through the model
  skips = down_stack(inputs)
  x = skips[-1]
  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for idx, (up, skip) in enumerate(zip(up_stack, skips)):
    x = up(x)
    concat = tf.keras.layers.Concatenate(name='expand_{}'.format(idx))
    x = concat([x, skip])

  # This is the last layer of the model
  last = tf.keras.layers.Conv2DTranspose(
      output_channels, 3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

model = unet_model(OUTPUT_CHANNELS)
tf.keras.utils.plot_model(model, show_shapes=True)

"""
Explanation: The decoder/upsampler is a series of upsample blocks.
Conv2DTranspose is what does the upsampling.
We also add batchnorm and relu.
End of explanation
"""


model = unet_model(OUTPUT_CHANNELS)
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# Assign to the pixel the label with the highest probability
def create_mask(pred_mask):
  pred_mask = tf.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tf.newaxis]
  return pred_mask[0]

"""
Explanation: Train the model
We treat this as a multi-class prediction problem.
So, the loss function used is SparseCategoricalCrossentropy
and the metric is accuracy.
End of explanation
"""


# display helper functions
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()

def show_predictions(dataset, num):
  for image, mask in dataset.take(num):
    pred_mask = model.predict(image)
    display([image[0], mask[0], create_mask(pred_mask)])

# before the model is trained
show_predictions(train.batch(1), 1)

"""
Explanation: Helper functions to see model output as it is training
End of explanation
"""


# from IPython.display import clear_output
class DisplayCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    if epoch%5 == 0:
      # clear_output(wait=True) # if you want replace the images each time, uncomment this
      show_predictions(train.batch(1), 1)
      print ('\nSample Prediction after epoch {}\n'.format(epoch+1))

EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_dataset, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_dataset,
                          callbacks=[DisplayCallback()])

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

epochs = range(EPOCHS)

plt.figure()
plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()

"""
Explanation: To observe how the model improves while it is training we define a callback function to show the predictions when the epoch finishes.
End of explanation
"""


