segmentationAI

# -*- coding: utf-8 -*-
"""segmentation.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1tBLA0L4ktZwyZe2433AjRoMcyXw9ZLKI

##### Copyright 2019 The TensorFlow Authors.

Licensed under the Apache License, Version 2.0 (the "License");
"""

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""## Load Packages"""

!pip install git+https://github.com/tensorflow/examples.git

import tensorflow as tf

import tensorflow_datasets as tfds

from tensorflow_examples.models.pix2pix import pix2pix

from IPython.display import clear_output
import matplotlib.pyplot as plt

import os
import glob #for loading images from a directory
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import cv2
import random
import numpy as np
import pandas as pd
from PIL import Image
from urllib.request import urlopen
from PIL.ExifTags import TAGS 

from scipy import ndimage

import numpy as np
from sklearn.model_selection import train_test_split

from google.colab import drive
drive.mount('/content/drive')

#%cd {'gdrive/My Drive/ImageSegmentation'}

"""**bold text**## Download the Oxford-IIIT Pets dataset

The dataset is [available from TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/oxford_iiit_pet). The segmentation masks are included in version 3+.

## **Upload Snow Images (train data) and Mask Data (test data)**
"""

inputImages = []
inputImagesNames = []
for file in glob.glob("drive/My Drive/ImageSegmentation/InputImage/*"):
  if os.path.splitext(file)[1].lower() in ('.jpg', '.jpeg'):
    filename = file.split('/')[-1]
    image = cv2.imread(file)
    inputImages.append(image)
    inputImagesNames.append(filename)

maskedImages = []
maskedImagesNames = []
for file in glob.glob("drive/My Drive/ImageSegmentation/TrueMask/*"):
  if os.path.splitext(file)[1].lower() in ('.jpg', '.jpeg'):
    image = cv2.imread(file)
    maskedImages.append(image)
    filename = file.split('_',1)[1]
    name = (filename.split('.')[0])
    ending = (filename.split('.')[1]).upper()
    filename = name + '.' + ending
    maskedImagesNames.append(filename)

## glob doesn't always work in the same order so we want to be sure that everything is in the same order, we will sort by name (ascending) and then make sure everything lines up 
def sortMasks(inputImages, inputImagesNames, maskedImages, maskedImagesNames):

  df1 = pd.DataFrame(inputImagesNames)
  df2 = pd.DataFrame(maskedImagesNames)

  sortedImages = []#[None] * len(df1) 
  sortedImagesNames = []#[None] * len(df1)
  sortedMasks = []#[None] * len(df1)
  sortedMasksNames = []#[None] * len(df1)
  for i in range(0, len(df1)):
    if sum(df2[0] == df1[0][i]) > 0:
     sortedImages.append(inputImages[i])
     sortedImagesNames.append(inputImagesNames[i])    
     index = (df2[df2[0] == df1[0][i]]).index[0]
     mask = maskedImages[index]
     name = maskedImagesNames[index]
     sortedMasks.append(mask) 
     sortedMasksNames.append(name) 
     # index = (df2[df2[0] == df1[0][i]]).index[0]
      #mask = maskedImages[index]
     # name = maskedImagesNames[index]
      #sortedMasks[i]= mask
      #sortedMasksNames[i] = name

  return sortedImages, sortedImagesNames, sortedMasks, sortedMasksNames

inputImages, inputImagesNames, maskedImages, maskedImagesNames = sortMasks(inputImages, inputImagesNames, maskedImages, maskedImagesNames)

r = 20
display([inputImages[r], maskedImages[r]])

"""### Generic function to load images

We will resize the image and set to float32. In the loop we will separate the binary testing data into two separate binary images, which are technically just inverted masks of each other. Then, we will put them into one tensor using stack so that the tensor is 128 x 128 x 2 channels
"""

### could add cropping here because we dont need the top and bottom of the image 

def load_image(image, mask):
  input_image = tf.image.resize(image, (128, 128))
  input_mask = tf.image.resize(mask, (128, 128))
  input_image = tf.cast(input_image, tf.float32) / 255.0
  return input_image, input_mask

def Preprocessing(inputImages, maskedImages):
  x_data = []
  #y_data1 = []
  #y_data2 = []
  y_data = []
  for i in range(0, len(inputImages)):
    x, y = load_image(inputImages[i], maskedImages[i])
    y = y[:,:,0]//255
    y1 = tf.cast(y, dtype = tf.uint16)
    y2 = tf.bitwise.invert(y1)
    y2= y2[:,:]//65535 ## I don't know why it's doing that but it works if you divide it by 65535. We may need to convert back to float 32
    #(print(x.shape))
    t1 = y1 ### snow
    t2 = y2 ### no snow
    stack = tf.stack([y2, y1], 2) ## channel 0 = no snow; channel 2 = snow
    x_data.append(x)
    #y_data1.append(y1)
    #y_data2.append(y2)
    y_data.append(stack)
  return x_data, y_data

x_data, y_data = Preprocessing(inputImages, maskedImages)

"""Here is an example of how the dataset is two channels of inverted masks of each other.

## Split data for model
"""

X_train, X_test, y_train, y_test = train_test_split(x_data, y_data, test_size=0.33, random_state=42)

train_images = {'image':X_train, 'segmentation_mask':y_train}
test_imags = {'image':X_test, 'segmentation_mask':y_test}

"""## Visualize Data"""

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(display_list[i])
    plt.axis('off')
  plt.show()

for i in range(0, len(train_images)):
  sample_image, sample_mask = train_images['image'][i], train_images['segmentation_mask'][i]
  display([sample_image, sample_mask[:,:,1]])

"""## Define the model
The model being used here is a modified [U-Net](https://arxiv.org/abs/1505.04597). A U-Net consists of an encoder (downsampler) and decoder (upsampler). In-order to learn robust features and reduce the number of trainable parameters, you will use a pretrained model - MobileNetV2 - as the encoder. For the decoder, you will use the upsample block, which is already implemented in the [pix2pix](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py) example in the TensorFlow Examples repo. (Check out the [pix2pix: Image-to-image translation with a conditional GAN](../generative/pix2pix.ipynb) tutorial in a notebook.)

As mentioned, the encoder will be a pretrained MobileNetV2 model which is prepared and ready to use in `tf.keras.applications`. The encoder consists of specific outputs from intermediate layers in the model. Note that the encoder will not be trained during the training process.
"""

#TRAIN_LENGTH = info.splits['train'].num_examples
TRAIN_LENGTH = len(X_train)
BATCH_SIZE =1 #64 ### lower because our dataset is so small
BUFFER_SIZE = 5#1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE

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

# Use the activations of these layers
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)

down_stack.trainable = False

"""The decoder/upsampler is simply a series of upsample blocks implemented in TensorFlow examples."""

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

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

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

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    concat = tf.keras.layers.Concatenate()
    x = concat([x, skip])

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

  x = last(x)

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

"""Note that the number of filters on the last layer is set to the number of `output_channels`. This will be one output channel per class.

## Train the model

Now, all that is left to do is to compile and train the model. 

Since this is a multiclass classification problem, use the `tf.keras.losses.CategoricalCrossentropy` loss function with the `from_logits` argument set to `True`, since the labels are scalar integers instead of vectors of scores for each pixel of every class. 

When running inference, the label assigned to the pixel is the channel with the highest value. This is what the `create_mask` function is doing.
"""

OUTPUT_CLASSES = 2

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

"""Have a quick look at the resulting model architecture:"""

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

"""## Now we will add convert our lists to tensors, and add a dimension to the test data in the process. 

We have to add a dimension to the test data using .expand_dims() becaus the binary information isn't automatically included in the channel. 

We also need our tensors to match the shape of the model for inputs and outputs (None, X, Y, channel) and (none, X, Y, 1)
"""

x_train = tf.convert_to_tensor(X_train) #tf.expand_dims(X_train, 0)
y_train = tf.convert_to_tensor(y_train)#tf.expand_dims(y_train, 3)
x_test =  tf.convert_to_tensor(X_test)#tf.expand_dims(X_test, 0)
y_test = tf.convert_to_tensor(y_test)#tf.expand_dims(y_test,3)

print(x_train.shape)
print(y_train.shape)

"""## Predict model before training

Try out the model to check what it predicts before training.
"""

def show_predictions(image):
  image = tf.expand_dims(image, axis=0) 
  pred_mask = model.predict(image)
  vis1 = pred_mask[0,:,:,0]
  pred_mask1 = tf.math.argmax(pred_mask, axis = -1)
  print(pred_mask1.shape)
  vis2 = pred_mask1[0,:,:]
  fig, (ax1, ax2) = plt.subplots(1,2)
  ax1.imshow(vis1)
  ax2.imshow(vis2)

show_predictions(x_train[0])

"""## Train Model"""

EPOCHS = 10
VAL_SUBSPLITS = 5
STEPS_PER_EPOCH = 2
VALIDATION_STEPS = 2 

model_history = model.fit(x_train, y_train, 
                          epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=(x_test, y_test))

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

plt.figure()
plt.plot(model_history.epoch, loss, 'r', label='Training loss')
plt.plot(model_history.epoch, 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()

"""## Make predictions

Now, make some predictions. In the interest of saving time, the number of epochs was kept small, but you may set this higher to achieve more accurate results.
"""

def make_predictions(image):
  image = tf.expand_dims(image, axis=0) #x_train[0:1,:]
  pred_mask = model.predict(image)
  #print(pred_mask.shape) ## probabilities

  pred_mask1 = tf.math.argmax(pred_mask, axis = -1)
  #print(pred_mask1.shape)
  #plt.imshow(pred_mask1[0,:,:]) ## final mask
  prediction = pred_mask1[0,:,:]
  return prediction

prediction = []
for i in range(0, len(x_train)):
  output = make_predictions(x_train[i])
  prediction.append(output)

k = 1
display([x_train[1], y_train[1][:,:,1], prediction[1]])

"""## Model Performance Evaluation
F1 score, accuracy and precision between true mask and predicted mask

Precision: the percentage of snow classifications predicted by our model that are also snow classifications in the true mask.
 precision = (True Positives / True Positives + False Positives) 

Recall: the percentage of true snow classifications predicted by our model that are also true snow classifications in the compared dataset
  recall = (True Positives) / (True Positives + False Negatives) 

FScore: The harmonic mean of precision and recall. 
  FScore = 2 * (Precision X Recall) / (Precision + Recall) 


"""

p = tf.keras.metrics.Precision()
p.update_state([y_train[1][:,:,1]], [prediction[1]])
precision = p.result().numpy()
precision

m = tf.keras.metrics.Recall()
m.update_state([y_train[1][:,:,1]], [prediction[1]])
recall = m.result().numpy()

#Fscore 
Fscore = 2 * (precision * recall)/(precision + recall)

def modelEvaluation(trueMask, predictedMask):
  p = tf.keras.metrics.Precision()
  p.update_state([trueMask], [predictedMask])
  precision = p.result().numpy()
  m = tf.keras.metrics.Recall()
  m.update_state([trueMask], [predictedMask])
  recall = m.result().numpy()
  Fscore = 2 * (precision * recall)/(precision + recall)
  if recall == 0:
    Fscore = 0
  return  precision, recall, Fscore

modelEvaluation(trueMask = y_train[1][:,:,1], predictedMask = prediction[1])

precision, recall, Fscore = [], [], []
for i in range(0, len(x_train)):
  precision1, recall1, Fscore1 = modelEvaluation(trueMask = y_train[i][:,:,1], predictedMask = prediction[i])
  precision.append(precision1)
  recall.append(recall1)
  Fscore.append(Fscore1)

print(np.mean(precision))
print(np.mean(recall))
print(np.mean(Fscore))

Fscore

modelEvaluation(trueMask = y_train[11][:,:,1], predictedMask = prediction[11])

display([y_train[11][:,:,1], prediction[11]])

"""## White Pixel Fraction"""

image = prediction[1]
height, width = image.shape
print(image.shape)
bottom_half = round(height/2)

  #plt.imshow(image[(round(height/2)):height,:]) just bottom half 
    # errors include (trees, animal in image, etc)
cropped = image[(round(height/2)):(height-10),:] ## above 10 pixels of the bottom ribbon
plt.imshow(cropped)

width, height = cropped.shape
white_pixel_fraction = (np.sum(cropped))/(width *height)
white_pixel_fraction

width *height

#image = tf.expand_dims(x_train[5], axis=0) #x_train[0:1,:]
#pred_mask = model.predict(image)
#print(pred_mask.shape) ## probabilities

#pred_mask1 = tf.math.argmax(pred_mask, axis = -1)
#print(pred_mask1.shape)
#plt.imshow(pred_mask1[0,:,:]) ## final mask

#show predictions
#display([x_train[5], y_train[5][:,:,1], pred_mask1[0,:,:]])

"""# Predict on unseen data 


"""

## load in the Labeled Images 
path = 'drive/My Drive/LabeledImages/'
originalData = []
unseenData = []
filenames = []
for file in glob.glob(os.path.join(path + "*")):
  #print(file)
  if os.path.splitext(file)[1].lower() in ('.jpg', '.jpeg'):
    image = cv2.imread(file)
    input_image = tf.image.resize(image, (128, 128))
    input_image = tf.cast(input_image, tf.float32) / 255.0
    originalData.append(input_image)
    output = make_predictions(input_image)
    unseenData.append(output)

path = 'drive/My Drive/LabeledImages/'
filenames = []
for file in glob.glob(os.path.join(path + "*")):
  #print(file)
  if os.path.splitext(file)[1].lower() in ('.jpg', '.jpeg'):
    image = cv2.imread(file)
    filename = file.split('/')[-1]
    filenames.append(filename)

filenames

f = 8000
display([originalData[f], unseenData[f]])

## count pixels
fractions = []
for i in range(0, len(unseenData)):
  image = unseenData[i]
  height, width = image.shape
  #print(image.shape)
  bottom_half = round(height/2)
  #plt.imshow(image[(round(height/2)):height,:]) just bottom half 
    # errors include (trees, animal in image, etc)
  cropped = image[(round(height/2)):(height-10),:]
  width, height = cropped.shape
  white_pixel_fraction = (np.sum(cropped))/(width *height)
  fractions.append(white_pixel_fraction)

df.to_csv("AI_fractions.csv", index=False)

df = pd.DataFrame({"filename" : filenames, "whitePixelIndexAI" : fractions})
df.to_csv("drive/My Drive/AI_fractions.csv", index=False)

df.head()

os.getcwd()

"""## Optional: Imbalanced classes and class weights

Semantic segmentation datasets can be highly imbalanced meaning that particular class pixels can be present more inside images than that of other classes. Since segmentation problems can be treated as per-pixel classification problems, you can deal with the imbalance problem by weighing the loss function to account for this. It's a simple and elegant way to deal with this problem. Refer to the [Classification on imbalanced data](../structured_data/imbalanced_data.ipynb) tutorial to learn more.

To [avoid ambiguity](https://github.com/keras-team/keras/issues/3653#issuecomment-243939748), `Model.fit` does not support the `class_weight` argument for inputs with 3+ dimensions.
"""

try:
  model_history = model.fit(train_batches, epochs=EPOCHS,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            class_weight = {0:2.0, 1:2.0, 2:1.0})
  assert False
except Exception as e:
  print(f"Expected {type(e).__name__}: {e}")

"""So, in this case you need to implement the weighting yourself. You'll do this using sample weights: In addition to `(data, label)` pairs, `Model.fit` also accepts `(data, label, sample_weight)` triples.

`Model.fit` propagates the `sample_weight` to the losses and metrics, which also accept a `sample_weight` argument. The sample weight is multiplied by the sample's value before the reduction step. For example:
"""

label = [0,0]
prediction = [[-3., 0], [-3, 0]] 
sample_weight = [1, 10] 

loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True,
                                               reduction=tf.losses.Reduction.NONE)
loss(label, prediction, sample_weight).numpy()

"""So to make sample weights for this tutorial you need a function that takes a `(data, label)` pair and returns a `(data, label, sample_weight)` triple. Where the `sample_weight` is a 1-channel image containing the class weight for each pixel. 

The simplest possible implementation is to use the label as an index into a `class_weight` list:
"""

def add_sample_weights(image, label):
  # The weights for each class, with the constraint that:
  #     sum(class_weights) == 1.0
  class_weights = tf.constant([2.0, 2.0, 1.0])
  class_weights = class_weights/tf.reduce_sum(class_weights)

  # Create an image of `sample_weights` by using the label at each pixel as an 
  # index into the `class weights` .
  sample_weights = tf.gather(class_weights, indices=tf.cast(label, tf.int32))

  return image, label, sample_weights

"""The resulting dataset elements contain 3 images each:"""

train_batches.map(add_sample_weights).element_spec

"""Now you can train a model on this weighted dataset:"""

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

weighted_model.fit(
    train_batches.map(add_sample_weights),
    epochs=1,
    steps_per_epoch=10)

"""## Next steps
Now that you have an understanding of what image segmentation is and how it works, you can try this tutorial out with different intermediate layer outputs, or even different pretrained models. You may also challenge yourself by trying out the [Carvana](https://www.kaggle.com/c/carvana-image-masking-challenge/overview) image masking challenge hosted on Kaggle.

You may also want to see the [Tensorflow Object Detection API](https://github.com/tensorflow/models/blob/master/research/object_detection/README.md) for another model you can retrain on your own data. Pretrained models are available on [TensorFlow Hub](https://www.tensorflow.org/hub/tutorials/tf2_object_detection#optional)
"""