ECE471 - Landcover Classification

# -*- coding: utf-8 -*-
"""Homework 2.ipynb

Automatically generated by Colaboratory.

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

# Mark Koszykowski

## ECE471 - Landcover Classification


### For this assignment, you will be training a landcover classification model from the [National Landcover Database](https://www.usgs.gov/centers/eros/science/national-land-cover-database?qt-science_center_objects=0#).

The dataset consists of 1340 total samples. Every image is a 256x256 Landsat 8 image taken from spring to fall in 2016. You have been provided the following bands in each image file: 

* [coastal-aerosol, blue, green, red, nir, swir1, swir2, tirs1, cloud-mask, alpha]

An attempt was made to render each scene mostly cloud free, but a cloud mask is provided along with each image chip.

The corresponding target image to each image is a single-band image with NLCD classification labels taken from 2016.

Your over-arching goal is to build a pixel-wise segmentation model that outputs the following classes:

* [water, snow/ice, built area, bare, forest, shrub/scrub, grass, crops, wetlands]

Note that the number of classes your model needs to output is less than the number of classes in NLCD.

### Dataset details

The data exists in google cloud storage, in a public facing read-only storage bucket. You will access all the file links to access the data via a CSV file which also exists in the bucket, that contains the entire dataset that is available to be processed. See the example code for how to access these data. Rasterio will be able to load in a GeoTIFF file from a URL that corresponds to each file in cloud storage.

**Do not make an attempt to modify this storage bucket in any way, or I will get in trouble.**

### Task 1: Analyze the distribution of the dataset

Your first task is to analyze the distribution of the dataset. This will inform properties about your model (e.g. how well it will generalize). In particular, you should look at both the geographic spread of the where the dataset has been sampled, and the number of pixels that belong to each class. 

Note that since you will need to re-map the input NLCD labels to the specified output labels, you will want to look at the re-mapped distribution of pixels across the dataset.

Any other analysis that you do to better understand the underlying data (and which may affect your model parameters, architecture, etc.) falls under this section.

### Task 2: Train a landcover classification model

Your second task is to train a landcover classification model. I have given you some tensorflow code to get started, which should give you a sense of how to build a data loader and input pipeline. When implemented properly, tensorflow will manage the "lazy-loading" of the dataset for you (e.g. only loading as much data as it needs to at once to keep the GPU fed). 

A neural network architecture such as a UNet is a good choice, but you are free to use whatever model you choose. You are not required to use tensorflow (e.g. you can use Pytorch or just scikit-learn) but it may be the most straightforward choice given the code to get you started.

I expect you to think about ML concepts (and implement them if you decide they are applicable) such as:
* Train/val split
* Data normalization
* Data augmentation
* Regularization

Once your model is trained, I expect you to discuss the results (e.g. accuracy, validation loss) and show plots and analysis as appropriate.

### Evauation

Every model will be scored against a test set that I have that you do not have access to. In order to facilitate evaluation, you will be expected to provide a function that reads in your model, accepts takes an input file (in the same format as every training image) and outputs a single band classification raster that follows the classification scheme above.

## Code examples to get you started
"""

# we will use rasterio to read GeoTIFF files from cloud storage
!pip install rasterio

# Commented out IPython magic to ensure Python compatibility.
# base imports
from glob import glob
import itertools
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline
import numpy as np
import os
from osgeo import gdal
import pandas as pd
import rasterio
from random import shuffle
import sklearn.metrics as skmetrics
# %tensorflow_version 2.x
import tensorflow as tf

# check that a GPU is enabled
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))

"""### Loading in file links"""

# read in the CSV file that contains file links to the entire dataset
dataset_file = 'super_secret_csv_link'
df = pd.read_csv(dataset_file)
print(len(df))
df.head(5)

# use these as base directories for reading the dataset from cloud storage
image_base = 'super_secret_images_link'
target_base = 'super_secret_targets_link'

# load in file links from cloud storage
samples = list()
for idx, row in df.iterrows():
  image_file = os.path.join(image_base, row['images'])
  target_file = os.path.join(target_base, row['targets'])

  samples.append((image_file, target_file))

with rasterio.open(samples[0][0]) as src:
  img = src.read()

with rasterio.open(samples[0][1]) as src:
  tgt = src.read()

print(img.shape)
print(tgt.shape)
print(len(samples))

"""### Visualization"""

# NLCD color scheme
nlcd_color_mapping = {
    11: (70, 107, 159),
    12: (209, 222, 248),
    21: (222, 197, 255),
    22: (217, 146, 130),
    23: (235, 0, 0),
    24: (171, 0, 0),
    31: (179, 172, 159),
    41: (104, 171, 95),
    42: (28, 95, 44),
    43: (181, 197, 143),
    52: (204, 184, 121),
    71: (223, 223, 194),
    81: (220, 217, 57),
    82: (171, 108, 40),
    90: (184, 217, 235),
    95: (108, 159, 184),
}

# NLCD class mapping
nlcd_name_mapping = {
    11: 'Open Water',
    12: 'Perennial Ice/Snow',
    21: 'Developed, Open Space',
    22: 'Developed, Low Intensity',
    23: 'Developed, Medium Intensity',
    24: 'Developed, High Intensity',
    31: 'Barren Land',
    41: 'Deciduous Forest',
    42: 'Evergreen Forest',
    43: 'Mixed Forest',
    52: 'Shrub/Scrub',
    71: 'Grassland/Herbaceous',
    81: 'Pasture/Hay',
    82: 'Cultivated Crops',
    90: 'Woody Wetlands',
    95: 'Emergent Herbaceous Wetlands',
}

# build a matplotlib colormap so we can visualize this data
colors = list()
for idx, class_val in enumerate(nlcd_color_mapping.keys()):
  red, green, blue = nlcd_color_mapping[class_val]

  color_vec = np.array([red/255, green/255, blue/255])
  colors.append(color_vec)

colors = np.stack(colors)
cmap = matplotlib.colors.ListedColormap(colors=colors, N=len(colors))

bounds = list(range(len(colors)))
norm = matplotlib.colors.BoundaryNorm(bounds, len(colors))

# write a function so that we can display image/target/predictions
def display_image_target(display_list):
  plt.figure(dpi=200)
  title = ['Image', 'Target', 'Prediction']

  for idx, disp in enumerate(display_list):
    plt.subplot(1, len(display_list), idx+1)
    plt.title(title[idx], fontsize=6)
    plt.axis('off')

    if title[idx] == 'Image':
      arr = disp.numpy()
      rgb = np.stack([arr[:, :, 3], arr[:, :, 2], arr[:, :, 1]], axis=-1) / 3000.0
      plt.imshow(rgb)

    elif title[idx] == 'Target':
      tgt = disp.numpy().squeeze()
      plt.imshow(tgt, interpolation='none', norm=norm, cmap=cmap)

    elif title[idx] == 'Prediction':
      pred = np.argmax(disp, axis=-1) # argmax across probabilities to get class outputs
      plt.imshow(pred, interpolation='none', norm=norm, cmap=cmap)

  plt.show()
  plt.close()

"""
### Tensorflow dataset loader"""

def read_sample(data_path: str) -> tuple:
  path = data_path.numpy()
  image_path, target_path = path[0].decode('utf-8'), path[1].decode('utf-8')

  with rasterio.open(image_path) as src:
    img = np.transpose(src.read(), axes=(1, 2, 0)).astype(np.uint16)

  with rasterio.open(target_path) as src:
    tgt = np.transpose(src.read(), axes=(1, 2, 0)).astype(np.uint8)

  tgt_out = np.zeros(tgt.shape, dtype=np.uint8)
  for class_val, nlcd_val in enumerate(nlcd_color_mapping.keys()):
    tgt_out[tgt == nlcd_val] = class_val

  return (img, tgt_out)

@tf.function
def tf_read_sample(data_path: str) -> dict:
  # wrap custom dataloader into tensorflow
  [image, target] = tf.py_function(read_sample, [data_path], [tf.uint16, tf.uint8])

  # explicitly set tensor shapes
  image.set_shape((256, 256, 10))
  target.set_shape((256, 256, 1))

  return {'image': image, 'target': target}

@tf.function
def load_sample(sample: dict) -> tuple:
  # convert to tf image
  image = tf.image.resize(sample['image'], (256, 256))
  target = tf.image.resize(sample['target'], (256, 256))

  # cast to proper data types
  image = tf.cast(image, tf.float32)
  target = tf.cast(target, tf.uint8)

  return image, target

# create tensorflow dataset from file links
ds = tf.data.Dataset.from_tensor_slices(samples)

# read in image/target pairs
ds = ds.map(tf_read_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE)

# read in as tensors
ds = ds.map(load_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE)

for image, target in ds.take(5):
  display_image_target([image, target])

"""### Task 1"""

import rasterio.features
import rasterio.warp

long_lat = []

for ind, (x, y) in enumerate(samples):
  # https://rasterio.readthedocs.io/en/latest/
  with rasterio.open(x) as src:
    mask = src.dataset_mask()
    longs = []
    lats = []
    for geom, val in rasterio.features.shapes(mask, transform=src.transform):
      geom = rasterio.warp.transform_geom(src.crs, 'EPSG:4326', geom, precision=6)
      # get the mean (long, lat) which should be the center assuming parallelogram
      # avoiding the last set of coordinates, since the same as the first
      for i in range(4):
        longs.append(geom['coordinates'][0][i][0])
        lats.append(geom['coordinates'][0][i][1])
    long_lat.append((sum(longs) / len(longs), sum(lats) / len(lats)))
  #print(ind)

import folium

# coordinates for geographic center of contiguous USA
map = folium.Map(location=[39.8283, -98.5795], zoom_start=5)

for long, lat in long_lat:
  folium.CircleMarker(location=[lat, long], radius=3, weight=3).add_to(map)

map

# https://www.mrlc.gov/data/legends/national-land-cover-database-2016-nlcd2016-legend
# mapping from provided classes to desired classes
nlcd_to_new = {
    0: 0,
    1: 1,
    2: 2,
    3: 2,
    4: 2,
    5: 2,
    6: 3,
    7: 4,
    8: 4,
    9: 4,
    10: 5,
    11: 6,
    12: 7,
    13: 7,
    14: 8,
    15: 8
}

# shorter set of RGB values for classes
new_color_mapping = [
    [70, 107, 159],
    [209, 222, 248],
    [222, 197, 255],
    [179, 172, 159],
    [104, 171, 95],
    [204, 184, 121],
    [223, 223, 194],
    [220, 217, 57],
    [184, 217, 235],
]

for ind, (r, g, b) in enumerate(new_color_mapping):
  new_color_mapping[ind] = [r/255, g/255, b/255]

new_name_mapping = {
    0: 'Water',
    1: 'Snow/Ice',
    2: 'Built Area',
    3: 'Bare',
    4: 'Forest',
    5: 'Shrub/Scrub',
    6: 'Grass',
    7: 'Crops',
    8: 'Wetlands'
}

# total amount of pixels per class
tot_counts = {
    0: 0,
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0,
    8: 0
}

# total amount of pixels from valid images per class
new_counts = {
    0: 0,
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0,
    8: 0
}

# mean value for each band across dataset
means = {
    0: 0,
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0
}

# standard deviation value for each band across dataset
stds = {
    0: 0,
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0
}

invalid_pixels = {
    'alpha': 0,
    'cloud': 0
}

invalid_images = {
    'alpha': 0,
    'cloud': 0,
    'both': 0,
}

# stores the percentage of invalid pixels for each image with invalid pixels
percentages_of_invalid = {}

# stores the classes which are in each image
# useful for training/validation split
classes_in_image = {}

# build a new matplotlib colormap so we can visualize this data
colors = [np.array(rgb) for rgb in new_color_mapping]
colors = np.stack(colors)
cmap = matplotlib.colors.ListedColormap(colors=colors, N=len(colors))

bounds = list(range(len(colors)))
norm = matplotlib.colors.BoundaryNorm(bounds, len(colors))

for ind, (x, y) in enumerate(ds):
  # gets information on invalid pixels
  alpha_pixels = np.count_nonzero(x[:, :, 9] == 0)
  cloud_pixels = np.count_nonzero(x[:, :, 8] > 0)
  invalid_pixels['alpha'] += alpha_pixels
  invalid_pixels['cloud'] += cloud_pixels
  if alpha_pixels > 0 and cloud_pixels > 0:
    invalid_images['both'] += 1
    percentages_of_invalid[ind] = (alpha_pixels + cloud_pixels) / (256*256)
  elif alpha_pixels > 0:
    invalid_images['alpha'] += 1
    percentages_of_invalid[ind] = (alpha_pixels) / (256*256)
  elif cloud_pixels > 0:
    invalid_images['cloud'] += 1
    percentages_of_invalid[ind] = (cloud_pixels) / (256*256)

  # map output classes from standard NLCD to new classes
  y = np.vectorize(nlcd_to_new.get)(y)
  # get unique counts of each class
  counts = dict(zip(*np.unique(y, return_counts=True)))
  classes_in_image[ind] = []
  for i in range(9):
    if i in counts:
      tot_counts[i] += counts[i]
      classes_in_image[ind].append(i)
  #print(ind)

# gives idea of how much information will need to be tossed
invalid_pixels

# gives idea of how many images contain invalid information
invalid_images

percentages_of_invalid

classes_in_image

# make sure all pixels in dataset are accounted for in total counts
assert sum(tot_counts.values()) == 256*256*1340

# wanted to minimize this value without defecting the data so that as minimal 
# invalid data was fed into the model

# originally tested 1% threshold however even though this only reduced the data
# size by 95, this biasedly affected the Snow/Ice class based on observation of
# the graphs below
invalid_percent_thres = .05

# count the number of invalid pixels based on threshold
num_of_invalid = 0
for key, value in percentages_of_invalid.items():
  if value >= invalid_percent_thres:
    num_of_invalid += 1

num_of_invalid

# get information on valid images based on threshold
# making sure an individual class isnt too strongly effected by tossing data
# used to help determine threshold
for ind, (x, y) in enumerate(ds):
  if ind in percentages_of_invalid and percentages_of_invalid[ind] >= invalid_percent_thres:
    continue
  # map output classes from standard NLCD to new classes
  y = np.vectorize(nlcd_to_new.get)(y)
  # get unique counts of each class
  counts = dict(zip(*np.unique(y, return_counts=True)))
  for i in range(9):
    if i in counts:
      new_counts[i] += counts[i]
  # store a sum of all the values in a band for computation later
  for i in range(8):
    means[i] += np.sum(x[:, :, i], dtype=np.int64)
  #print(ind)

# calculate the actual means from sums
for band, sum in means.items():
  means[band] = sum / ((len(samples) - num_of_invalid) * 256 * 256)

# iterate through data again to get the values necessary for standard deviation
# calculation
for x, y in ds:
  if ind in percentages_of_invalid and percentages_of_invalid[ind] >= invalid_percent_thres:
    continue
  for i in range(8):
    stds[i] += np.sum(np.square(x[:, :, i] - means[i]), dtype=np.int64)

# calculate the actual (sample) standard deviations from sums
for band, sum in stds.items():
  stds[band] = np.sqrt(sum / ((len(samples) - num_of_invalid) * 256 * 256 - 1))

# look at mean values of bands across dataset
means

# look at standard deviation values of bands across dataset
stds

# plot distribution of classes with all pixels
graph_colors = ['blue',
         'snow',
         'black',
         'lightgray',
         'forestgreen',
         'lime',
         'green',
         'yellow',
         'royalblue'
]

plt.figure(figsize=(20,12))
plt.bar(new_name_mapping.values(), tot_counts.values(), color=graph_colors)
for ind, count in enumerate(tot_counts.values()):
  plt.annotate(str(count),
               xy=(ind, count),
               xytext=(0, 3),
               textcoords='offset points',
               ha='center', va='bottom')
plt.title("Total Pixel Distribution")
plt.show()

# plot distribution of classes with pixels from valid images
graph_colors = ['blue',
         'snow',
         'black',
         'lightgray',
         'forestgreen',
         'lime',
         'green',
         'yellow',
         'royalblue'
]

plt.figure(figsize=(20,12))
plt.bar(new_name_mapping.values(), new_counts.values(), color=graph_colors)
for ind, count in enumerate(new_counts.values()):
  plt.annotate(str(count),
               xy=(ind, count),
               xytext=(0, 3),
               textcoords='offset points',
               ha='center', va='bottom')
plt.title("Valid Pixel Distribution")
plt.show()

"""### Task 2"""

# create a new list to create a new TF Dataset, eliminating invalid images
new_samples = []
for ind, (image_file, target_file) in enumerate(samples):
  # skip invalid images
  if ind in percentages_of_invalid and percentages_of_invalid[ind] >= invalid_percent_thres:
    continue
  new_samples.append((image_file, target_file))

# tried reversing order of dataset since training accuracy started higher than
# validation accuracy, this helped for the first couple epochs but in the end led
# to a validation accuracy lower than training accuracy
# new_samples.reverse()

import sys

# function to normalize pixel values
# ideally want all pixel values to be on the same order of magnitude
# originally used a simple min/max scaling to normalize between -1 and 1
# however switched to z-scores
# using the nonlinear z-score scale, outliers are more strongly represented
def normalize_band(img_band, band):
  return (img_band - means[band]) / (stds[band] + sys.float_info.epsilon)

# new reading function that adds necessary band information
def read_and_preprocess(data_path: str) -> tuple:
  path = data_path.numpy()
  image_path, target_path = path[0].decode('utf-8'), path[1].decode('utf-8')

  with rasterio.open(image_path) as src:
    # switched from uint16 to int32 since cant produce negatives with uint
    # necessary for the normalized indices
    img = np.transpose(src.read(), axes=(1, 2, 0)).astype(np.int32)

  # add additional band information to help model
  new_img = np.zeros((256, 256, 12))
  # https://www.usna.edu/Users/oceano/pguth/md_help/html/norm_sat.htm
  # NDWI
  new_img[:, :, 8] = (img[:, :, 2] - img[:, :, 4]) / (img[:, :, 2] + img[:, :, 4] + sys.float_info.epsilon)
  # NDVI
  new_img[:, :, 9] = (img[:, :, 4] - img[:, :, 3]) / (img[:, :, 4] + img[:, :, 3] + sys.float_info.epsilon)
  # NDSI (Snow)
  new_img[:, :, 10] = (img[:, :, 2] - img[:, :, 5]) / (img[:, :, 2] + img[:, :, 5] + sys.float_info.epsilon)
  # NDSI (Soil) / NDBI
  # https://pro.arcgis.com/en/pro-app/latest/arcpy/spatial-analyst/ndbi.htm
  new_img[:, :, 11] = (img[:, :, 5] - img[:, :, 4]) / (img[:, :, 5] + img[:, :, 4] + sys.float_info.epsilon)
  # normalize all the bands except for cloud and alpha since not used in the model
  for i in range(8):
    new_img[:, :, i] = normalize_band(img[:, :, i], i)

  with rasterio.open(target_path) as src:
    tgt = np.transpose(src.read(), axes=(1, 2, 0)).astype(np.uint8)

  tgt_out = np.zeros(tgt.shape, dtype=np.uint8)
  for class_val, nlcd_val in enumerate(nlcd_color_mapping.keys()):
    tgt_out[tgt == nlcd_val] = class_val

  # map output classes from standard NLCD to new classes
  tgt_out = np.vectorize(nlcd_to_new.get)(tgt_out)

  return (new_img, tgt_out)

# new TF function that calls new reading and preprocessing function
@tf.function
def tf_read_and_preprocess(data_path: str) -> dict:
  # wrap custom dataloader into tensorflow
  [image, target] = tf.py_function(read_and_preprocess, [data_path], [tf.float32, tf.uint8])

  # explicitly set tensor shapes
  image.set_shape((256, 256, 12))
  target.set_shape((256, 256, 1))

  return {'image': image, 'target': target}

# set variables required for parameters in model
DATASET_SIZE = len(new_samples)
print(DATASET_SIZE)

# create tensorflow dataset from new file links
new_ds = tf.data.Dataset.from_tensor_slices(new_samples)

# read in image/target pairs and preprocess them
new_ds = new_ds.map(tf_read_and_preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)

# read in as tensors
new_ds = new_ds.map(load_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE)

# divide up provided data between training and validation
train_size = int(0.8 * DATASET_SIZE)
val_size = DATASET_SIZE - train_size

train_ds = new_ds.take(train_size)
val_ds = new_ds.skip(train_size)

# new TF function to augment the data
# augmentation includes flipping and rotation since model should be invariant
# in both of these respects
@tf.function
def augment(img, tgt) -> tuple:
  # 2/3 chance data will be augmented
  rand1 = tf.random.uniform((), maxval=3, dtype=tf.int64)
  if rand1 == 1:
    # 50/50 chance of different flips
    rand2 = tf.random.uniform((), maxval=2, dtype=tf.int64)
    if rand2 == 0:
      img = tf.image.flip_left_right(img)
      tgt = tf.image.flip_left_right(tgt)
    else:
      img = tf.image.flip_up_down(img)
      tgt = tf.image.flip_up_down(tgt)
  elif rand1 == 2:
    # 33/33/33 chance of different rotations
    rand2 = tf.random.uniform((), maxval=3, dtype=tf.int64)
    for i in range(rand2 + 1):
      img = tf.image.rot90(img)
      tgt = tf.image.rot90(tgt)
  return img, tgt

# apply the augmentation to the training set but NOT the validation set
train_ds = train_ds.map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)

# print one set of tensors to make sure data is as desired
for img, tgt in train_ds.take(1):
  print(img)
  print(tgt)

# https://www.tensorflow.org/tutorials/images/segmentation
TRAIN_LENGTH = train_size
BATCH_SIZE = 8
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE
OUTPUT_CHANNELS = 9

# shuffle and repeat data 
train_dataset = train_ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
val_dataset = val_ds.batch(BATCH_SIZE)

# load in predefined keras model
# weights set to none because of the large number of input channels
base_model = tf.keras.applications.MobileNetV2(input_shape=[256, 256, 12], include_top=False, weights=None)

# 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

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

from tensorflow_examples.models.pix2pix import pix2pix

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
]

# function to return the Unet
def unet_model(output_channels):
  inputs = tf.keras.layers.Input(shape=[256, 256, 12])

  # 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])
    # added BatchNorm (commented on later)
    x = tf.keras.layers.BatchNormalization()(x)

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

  x = last(x)

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

# look at pixel distribution
new_counts

# scale down this distribution so it can be added to loss function without
# blowing up loss, potentially causing overflow
tot_valid = np.sum(np.stack(new_counts.values()))
for k, v in new_counts.items():
  # add factor of 100 to prevent underflow
  new_counts[k] = np.divide(v, tot_valid) * 100

# look at scaled pixel distribution
new_counts

model = unet_model(OUTPUT_CHANNELS)
# compile model with pixel distribution since data is heavily unbalanced
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              loss_weights=new_counts.values(),
              metrics=['accuracy'])

EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = val_size//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_dataset, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=val_dataset)

# final validation accuracies range from ~52% - 59%
# reasons for this range:
#   - probabilistic augmentation (rotations and flips)
#   - random shuffling of data switches the sequence in which data is fed into the model
#     between model fittings
#   - weight initializtion scheming
#     - tried to use BatchNorm to circumvent initialization issues and prevent
#       overfitting
#
# by intruducing these factors of randomness, each time the model is trained
# it is different

"""### Average Run
##### dont mind the times, this was before i purchased Colab Pro );

        Epoch 1/20
        129/129 [==============================] - 523s 2s/step - loss: 3.5611 - accuracy: 0.4199 - val_loss: 3.6305 - val_accuracy: 0.2061
        Epoch 2/20
        129/129 [==============================] - 479s 2s/step - loss: 2.7567 - accuracy: 0.5176 - val_loss: 3.5130 - val_accuracy: 0.2599
        Epoch 3/20
        129/129 [==============================] - 478s 2s/step - loss: 2.7177 - accuracy: 0.5256 - val_loss: 3.1926 - val_accuracy: 0.4801
        Epoch 4/20
        129/129 [==============================] - 466s 2s/step - loss: 2.7027 - accuracy: 0.5234 - val_loss: 2.7654 - val_accuracy: 0.5447
        Epoch 5/20
        129/129 [==============================] - 468s 2s/step - loss: 2.7022 - accuracy: 0.5216 - val_loss: 2.5573 - val_accuracy: 0.5504
        Epoch 6/20
        129/129 [==============================] - 472s 2s/step - loss: 2.6413 - accuracy: 0.5303 - val_loss: 2.6824 - val_accuracy: 0.5309
        Epoch 7/20
        129/129 [==============================] - 461s 2s/step - loss: 2.6736 - accuracy: 0.5325 - val_loss: 2.6552 - val_accuracy: 0.5392
        Epoch 8/20
        129/129 [==============================] - 452s 2s/step - loss: 2.6653 - accuracy: 0.5328 - val_loss: 2.5632 - val_accuracy: 0.5568
        Epoch 9/20
        129/129 [==============================] - 463s 2s/step - loss: 2.6684 - accuracy: 0.5270 - val_loss: 2.5211 - val_accuracy: 0.5582
        Epoch 10/20
        129/129 [==============================] - 485s 2s/step - loss: 2.5533 - accuracy: 0.5596 - val_loss: 2.4924 - val_accuracy: 0.5808
        Epoch 11/20
        129/129 [==============================] - 488s 2s/step - loss: 2.6024 - accuracy: 0.5410 - val_loss: 2.5934 - val_accuracy: 0.5643
        Epoch 12/20
        129/129 [==============================] - 505s 2s/step - loss: 2.6123 - accuracy: 0.5369 - val_loss: 2.6272 - val_accuracy: 0.5543
        Epoch 13/20
        129/129 [==============================] - 492s 2s/step - loss: 2.5837 - accuracy: 0.5543 - val_loss: 2.7456 - val_accuracy: 0.5279
        Epoch 14/20
        129/129 [==============================] - 473s 2s/step - loss: 2.5288 - accuracy: 0.5610 - val_loss: 2.6420 - val_accuracy: 0.5493
        Epoch 15/20
        129/129 [==============================] - 493s 2s/step - loss: 2.6039 - accuracy: 0.5509 - val_loss: 2.6434 - val_accuracy: 0.5307
        Epoch 16/20
        129/129 [==============================] - 471s 2s/step - loss: 2.5648 - accuracy: 0.5527 - val_loss: 2.6222 - val_accuracy: 0.5551
        Epoch 17/20
        129/129 [==============================] - 451s 2s/step - loss: 2.5880 - accuracy: 0.5421 - val_loss: 2.6397 - val_accuracy: 0.5415
        Epoch 18/20
        129/129 [==============================] - 461s 2s/step - loss: 2.5728 - accuracy: 0.5505 - val_loss: 2.5334 - val_accuracy: 0.5605
        Epoch 19/20
        129/129 [==============================] - 460s 2s/step - loss: 2.5507 - accuracy: 0.5495 - val_loss: 2.7932 - val_accuracy: 0.5366
        Epoch 20/20
        129/129 [==============================] - 501s 2s/step - loss: 2.6420 - accuracy: 0.5343 - val_loss: 2.4754 - val_accuracy: 0.5714
"""

# plot the losses
plt.figure(figsize=(20, 12))
plt.plot(range(EPOCHS), model_history.history['loss'], 'r', label='Training loss')
plt.plot(range(EPOCHS), model_history.history['val_loss'], 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.legend()
plt.show()

# plot the accuracies
plt.figure(figsize=(20, 12))
plt.plot(range(EPOCHS), model_history.history['accuracy'], 'r', label='Training accuracy')
plt.plot(range(EPOCHS), model_history.history['val_accuracy'], 'bo', label='Validation accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy Value')
plt.legend()
plt.show()

# save the model just in case if runtime crash or to just prevent extensive retraining
model.save("/content/drive/MyDrive/ECE471/homework2/unetWbatchnorm")

# function to load the model if dont want to retrain
model = tf.keras.models.load_model("/content/drive/MyDrive/ECE471/homework2/unetWbatchnorm")

# function to test the model
# assuming test data is in the same format as the data provided
def test_model(samples, model):
  sample_size = len(samples)
  
  data = tf.data.Dataset.from_tensor_slices(samples)

  data = data.map(tf_read_and_preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  data = data.map(load_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  dataset = data.batch(BATCH_SIZE).take(sample_size)

  cur = 0
  targs = None
  preds = None
  for ind, (img, tgt) in enumerate(dataset):
    guesses = model.predict(img)
    if cur == 0:
      targs = tgt
      preds = guesses
    else:
      targs = tf.concat([targs, tgt], 0)
      preds = tf.concat([preds, guesses], 0)
    for i in range(BATCH_SIZE):
      cur += 1
      if cur == (sample_size + 1):
        break
      # denormalize data for displaying
      viz = np.zeros((256, 256, 4))
      for j in range(1, 4):
        viz[:, :, j] = img[i, :, :, j] * stds[j] + means[j]
      display_image_target([tf.convert_to_tensor(viz), tgt[i], guesses[i]])

  return targs, preds

# test the testing function on the data provided
targs, preds = test_model(new_samples, model)

# compute a confusion matrix for data analysis
conf = skmetrics.confusion_matrix(tf.reshape(targs, [-1]),
                                tf.reshape(np.argmax(preds, axis=-1), [-1]),
                                normalize='true')

# put it all on a percentage scale so the color mapping is more uniform
# especially since classes were heavily unbalanced
conf = conf * 100
conf = np.rint(np.nextafter(conf, conf+1))

import seaborn as sn

# plot the confusion matrix
pd_cm = pd.DataFrame(conf, index=new_name_mapping.values(),
                  columns=new_name_mapping.values())

plt.figure(figsize=(17, 15))
plt.title("Confusion Matrix as Rounded Percentages")
sn.heatmap(pd_cm, annot=True, fmt='g')
plt.xlabel("True Labels")
plt.ylabel("Predicted Labels")
plt.show()