# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from PIL import Image # used for loading images
import numpy as np
import os # used for navigating to image path
import imageio # used for writing images
import random
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import matplotlib.pyplot as plt
from tensorflow.keras.optimizers import SGD
import pydot
from timeit import default_timer as timer
from tensorflow.keras import backend as K

"ResNet 50 dependencies"
from tensorflow.keras.applications.resnet50 import ResNet50 
from tensorflow.keras.applications import resnet50

"GoogLeNet dependencies"
from tensorflow.keras import regularizers
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Input, Flatten, Dense, Dropout, BatchNormalization
from tensorflow.keras.layers import Conv2D, MaxPooling2D, AveragePooling2D, ZeroPadding2D
from tensorflow.keras.layers import Concatenate
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# -

"""Load numpy output files"""
pr_im_64 = np.load('../../data/tidy/preprocessed_images/size64_exp5_Pr_Im.npy', allow_pickle=True)
pr_po_im_64 = np.load('../../data/tidy/preprocessed_images/size64_exp5_Pr_Po_Im.npy', allow_pickle=True)
pr_poim_64 = np.load('../../data/tidy/preprocessed_images/size64_exp5_Pr_PoIm.npy', allow_pickle=True)
prpo_im_64 = np.load('../../data/tidy/preprocessed_images/size64_exp5_PrPo_Im.npy', allow_pickle=True)


def getImageOneHotVector(image_file_name, classification_scenario = "B"):
    """Returns one-hot vector encoding for each image based on specified classification scenario:
    Classification Scenario A (3 classes): {probable, possible, improbable}
    Classification Scenario B (2 classes): {probable, improbable}
    Classification Scenario C (2 classes): {{probable, possible}, improbable}
    Classification Scenario D (2 classes): {probable, {possible, improbable}}
    """
    word_label = image_file_name.split('-')[0]
    if classification_scenario == "A":
        if word_label == 'probable' : 
            return np.array([1, 0, 0])
        elif word_label == 'possible' : 
            return np.array([0, 1, 0])    
        elif word_label == 'improbable':
            return np.array([0, 0, 1])
        else :
            return np.array([0, 0, 0]) # if label is not present for current image
    elif classification_scenario == "B":
        if word_label == 'probable' : 
            return np.array([1, 0])
        elif word_label == 'improbable' : 
            return np.array([0, 1])
        else :
            return np.array([0, 0]) # if label is not present for current image
    elif classification_scenario == "C":
        if word_label in ['probable', 'possible'] : 
            return np.array([1, 0])
        elif word_label == 'improbable' : 
            return np.array([0, 1])
        else :
            return np.array([0, 0]) # if label is not present for current image        
    elif classification_scenario == "D":
        if word_label == 'probable' : 
            return np.array([1, 0])
        elif word_label in ['possible', 'improbable'] : 
            return np.array([0, 1])
        else :
            return np.array([0, 0]) # if label is not present for current image        


#IMG_SIZE = 300
NUM_CLASS = 3
NUM_CHANNEL = 1
CLASSIFICATION_SCENARIO = "A"

#IMG_SIZE = 300
NUM_CLASS = 2
NUM_CHANNEL = 1
CLASSIFICATION_SCENARIO = "B"
DIR = '../../data/tidy/labeled_images'
def processImageData(img_size, channels=1, l=400,t=0,r=3424,b=3024):    
    data = []
    image_list = os.listdir(DIR)
    for img in image_list:
        label = getImageOneHotVector(img, CLASSIFICATION_SCENARIO)
        if label.sum() == 0:
            continue
        path = os.path.join(DIR, img)
        img = Image.open(path)
        if channels == 1:
            img = img.convert('L') # convert image to monochrome 
            img = img.crop((l, t, r, b)) # after cropping, image size is 3024 x 3024 pixels
            #img_size_w, img_size_h = img.size
            img = img.resize((img_size, img_size), Image.BICUBIC)
            data.append([(np.array(img)/255.).T, label])#scale to 0-1 and transpose
#             flip_img = np.fliplr((np.array(img)/255.).T)# Basic Data Augmentation - Horizontal Flipping
#             data.append([flip_img, label])#scale to 0-1 and transpose
        elif channels == 3:
            img = img.crop((l, t, r, b)) # after cropping, image size is 3024 x 3024 pixels  
            img = img.resize((img_size, img_size), Image.BICUBIC)
            data.append([(np.array(img)/255.).T, label])#scale to 0-1 and transpose            
    return (data)


def splitData(image_array, prop = 0.80, seed_num = 111):
    """Returns training and test arrays of images with specified proportion - prop:1-prop"""
    random.Random(seed_num).shuffle(image_array)
    train_size = int(prop*np.shape(image_array)[0])
    train = image_array[:train_size]
    test = image_array[train_size:]
    return(train, test)


processed_image_data = processImageData(108, channels = NUM_CHANNEL)
train_data, test_data = splitData(processed_image_data, seed_num = 111) 

tr_data, te_data = splitData(pr_im_64, seed_num = 111)

processed_image_data[202][1]

plt.imshow(processed_image_data[202][0], cmap = 'gist_gray')
#plt.savefig( "../../figures/image0.png", dpi=100)

def getImageShape(image_array):
    if NUM_CHANNEL==1:
        image_shape = np.array([np.expand_dims(x[0],axis=2) for x in image_array]).shape[1:4]
    elif NUM_CHANNEL==3:
        image_shape = np.array([x[0] for x in image_array]).shape[1:4][::-1]
    print(image_shape)
    return image_shape


input_image_shape = getImageShape(tr_data)

# +
# Train the model.
if NUM_CHANNEL == 1:
    train_array = np.array([np.expand_dims(x[0],axis=2) for x in tr_data])
    validation_array = np.array([np.expand_dims(x[0],axis=2) for x in te_data])
elif NUM_CHANNEL == 3:
    train_array = np.array([x[0] for x in tr_data]) 
    train_array = np.moveaxis(train_array, 1, -1)
    validation_array = np.array([x[0] for x in te_data])
    validation_array = np.moveaxis(validation_array, 1, -1)

train_labels = np.array([x[1] for x in tr_data])
validation_labels = np.array([x[1] for x in te_data])


# +
#https://github.com/keras-team/keras/issues/5400#issuecomment-408743570
def check_units(y_true, y_pred):
    if y_pred.shape[1] != 1:
      y_pred = y_pred[:,1:2]
      y_true = y_true[:,1:2]
    return y_true, y_pred

def precision(y_true, y_pred):
    y_true, y_pred = check_units(y_true, y_pred)
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision = true_positives / (predicted_positives + K.epsilon())
    return precision

def recall(y_true, y_pred):
    y_true, y_pred = check_units(y_true, y_pred)
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall = true_positives / (possible_positives + K.epsilon())
    return recall

def f1(y_true, y_pred):
    def recall(y_true, y_pred):
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
        recall = true_positives / (possible_positives + K.epsilon())
        return recall

    def precision(y_true, y_pred):
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
        precision = true_positives / (predicted_positives + K.epsilon())
        return precision
    y_true, y_pred = check_units(y_true, y_pred)
    precision = precision(y_true, y_pred)
    recall = recall(y_true, y_pred)
    return 2*((precision*recall)/(precision+recall+K.epsilon()))


# +
def plot_model_accuracy(hist):
    plt.plot(hist.history["accuracy"])
    plt.plot(hist.history["val_accuracy"])
    plt.title("Model Accuracy")
    plt.ylabel("Accuracy")
    plt.xlabel("Epoch")
    plt.legend(["Train", "Validation"], loc="lower right")
    plt.show()

def plot_model_loss(hist):
    plt.plot(hist.history["loss"])
    plt.plot(hist.history["val_loss"])
    plt.title("Model Loss")
    plt.ylabel("Loss")
    plt.xlabel("Epoch")
    plt.legend(["Train", "Validation"], loc="upper right")
    plt.show()


# +
# opt = keras.optimizers.Adam(learning_rate=0.01)
# model.compile(loss='categorical_crossentropy', optimizer=opt)

model = models.Sequential([
    layers.Conv2D(filters = 64, kernel_size = 7, strides = 2, activation="relu", padding="same", input_shape = input_image_shape),
    layers.MaxPooling2D(2),
    layers.Conv2D(128, 3, activation="relu", padding="same"),
    layers.Conv2D(128, 3, activation="relu", padding="same"),
    layers.MaxPooling2D(2),
    layers.Conv2D(256, 3, activation="relu", padding="same"),
    layers.Conv2D(256, 3, activation="relu", padding="same"),
    layers.MaxPooling2D(2),
    layers.Flatten(),
    layers.Dense(128, activation="relu"),
    layers.Dropout(0.5), # randomly drop out 50% of the neuorns at each training step
    layers.Dense(64, activation="relu"), # flatten all outputs
    layers.Dropout(0.5),
    layers.Dense(NUM_CLASS, activation="softmax")
])

# +
# Compile the model.
opt = SGD(lr = 0.001) #default learning rate (lr) = 0.1
model.compile(loss='categorical_crossentropy',  optimizer = "adam",
              metrics=[precision,recall, f1, 'accuracy'])

start = timer()
hist_seq = model.fit(
    train_array, train_labels, 
    batch_size = 32,
    epochs = 5,
    validation_data=(validation_array, validation_labels)
)
end = timer()
print(end - start) # Time in seconds, e.g. 5.380919524002
# -

plot_model_accuracy(hist_seq)

plot_model_loss(hist_seq)

# + jupyter={"outputs_hidden": true}
#loss, acc = model.evaluate(testImages, testLabels, verbose = 0)
#print(acc * 100)
y_pred = model.predict(validation_array, batch_size=32, verbose=1)
y_pred_bool = np.argmax(y_pred, axis=1)

#print(classification_report(validation_labels, y_pred_bool))

# + jupyter={"outputs_hidden": true}
tf.keras.utils.plot_model(model, "../../figures/cnn-model1.png", expand_nested = False, rankdir = "TB", show_shapes=True, dpi=192)


# -

# ## ResNet 34

class ResidualUnit(tf.keras.layers.Layer):
    def __init__(self, filters, strides=1, activation="relu", **kwargs):
        super().__init__(**kwargs)
        self.activation = tf.keras.activations.get(activation)
        self.main_layers = [
            tf.keras.layers.Conv2D(filters, 3, strides=strides, padding="same", use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Conv2D(filters, 3, strides=1, padding="same", use_bias=False),
            tf.keras.layers.BatchNormalization()
        ]
        self.skip_layers = []
        if strides > 1:
            self.skip_layers = [
                tf.keras.layers.Conv2D(filters, 1, strides=strides, padding="same", use_bias=False),    
                tf.keras.layers.BatchNormalization()
            ]
    
    def call(self, inputs):
        Z = inputs
        for layer in self.main_layers:
            Z = layer(Z)
        skip_Z = inputs
        for layer in self.skip_layers:
            skip_Z = layer(skip_Z)
        return self.activation(Z + skip_Z)


resnet34mod = tf.keras.models.Sequential()
resnet34mod.add(tf.keras.layers.Conv2D(64, 7, strides=2, input_shape=input_image_shape, padding="same", use_bias=False))
resnet34mod.add(tf.keras.layers.BatchNormalization())
resnet34mod.add(tf.keras.layers.Activation("relu"))
resnet34mod.add(tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding="same"))
prev_filters = 64
for filters in [64] * 3 + [128] * 4 + [256] * 6 + [512] * 3:
    strides = 1 if filters == prev_filters else 2
    resnet34mod.add(ResidualUnit(filters, strides=strides))
    prev_filters = filters
resnet34mod.add(tf.keras.layers.GlobalAvgPool2D())
resnet34mod.add(tf.keras.layers.Flatten())
resnet34mod.add(tf.keras.layers.Flatten())
resnet34mod.add(tf.keras.layers.Dense(NUM_CLASS, activation="softmax"))

resnet34mod.compile(loss='binary_crossentropy',  optimizer = "adam",# metrics=[ 'accuracy']) #tf.keras.metrics.SpecificityAtSensitivity(0.5), tf.keras.metrics.SensitivityAtSpecificity(0.5), 
              metrics=[precision,recall, f1, 'accuracy']) #metrics=['accuracy'])

# + jupyter={"outputs_hidden": true}
start = timer()
resnet34mod.fit(
    train_array, train_labels, 
    batch_size = 32,
    epochs = 4,
    validation_data=(validation_array, validation_labels)
)
end = timer()
print(end - start) # Time in seconds, e.g. 5.380919524002
# -

# ## ResNet 50

# + jupyter={"outputs_hidden": true}
"Define ResNet 50 model instance (Keras built-in)"
rn50 = resnet50.ResNet50(include_top=True, 
                           weights=None, 
                           input_tensor=None, 
                           input_shape=input_image_shape, 
                           pooling= 'max', 
                           classes=2)

rn50.summary()

# +
"Configure the model with losses and metrics"
rn50.compile(loss='categorical_crossentropy',  optimizer = "adam",
              metrics=[precision,recall, f1, 'accuracy']) 

start = timer()

"Fit ResNet 50 to data"
hist_rn50 = rn50.fit(
    train_array, train_labels, 
    batch_size = 32,
    epochs = 5,
    validation_data=(validation_array, validation_labels)
)

end = timer()
print(end - start)
# -

plot_model_accuracy(hist_rn50)

plot_model_loss(hist_rn50)

# ## GoogLeNet
# ### InceptionV3

"""Instantiate the Inception v3 architecture"""
iv3 = tf.keras.applications.InceptionV3(
    include_top=True,
    weights=None,
    input_tensor=None,
    input_shape=input_image_shape,
    pooling='avg',
    classes=2,
    classifier_activation="softmax",
)

# +
"Configure the model with losses and metrics"
iv3.compile(loss='categorical_crossentropy',  optimizer = "adam",
              metrics=[precision,recall, f1, 'accuracy']) 

start = timer()

"Fit Inception v3 to data"
hist_iv3 = iv3.fit(
    train_array, train_labels, 
    batch_size = 32,
    epochs = 10,
    validation_data=(validation_array, validation_labels)
)

end = timer()
print(end - start)
# -

plot_model_accuracy(hist_iv3)

plot_model_loss(hist_iv3)

StarcoderdataJupyterScriptsDedupFiltered