# -*- coding: utf-8 -*-

"""

Created on Tue Jul 21 16:11:54 2020

@author: Billy

"""

import tensorflow as tf

from PIL import Image

import numpy as np

import math 

import matplotlib.pyplot as plt

import gc

import os

import time

import random

Image.MAX_IMAGE_PIXELS = 933120000

class uNet_segmentor:

    #Generates UNet Model using the model weights (supply the file location in the initialising line)

    def __init__(self, checkpoint_loc, first_conv = 8, window = 224, outputs=11):

        self.n = n = first_conv

        self.input_size = input_size = (window,window,3)

        self.outputs = outputs

        inputs = tf.keras.layers.Input(input_size)

        conv1 = tf.keras.layers.Conv2D(n, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)

        conv1 = tf.keras.layers.Conv2D(n, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)

        pool1 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv1)

        conv2 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)

        conv2 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)

        pool2 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv2)

        conv3 = tf.keras.layers.Conv2D(n*(2**2), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)

        conv3 = tf.keras.layers.Conv2D(n*(2**2), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)

        drop3 = tf.keras.layers.Dropout(0.5)(conv3)

        pool3 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(drop3)

        conv4 = tf.keras.layers.Conv2D(n*(2**3), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)

        conv4 = tf.keras.layers.Conv2D(n*(2**3), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)

        drop4 = tf.keras.layers.Dropout(0.5)(conv4)

        pool4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(drop4)

        conv_e = tf.keras.layers.Conv2D(n*(2**4), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)

        conv_e = tf.keras.layers.Conv2D(n*(2**4), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv_e)

        drop_e = tf.keras.layers.Dropout(0.5)(conv_e)

        pool_e = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(drop_e)

        conv5 = tf.keras.layers.Conv2D(n*(2**5), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool_e)

        conv5 = tf.keras.layers.Conv2D(n*(2**5), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)

        drop5 = tf.keras.layers.Dropout(0.5)(conv5)

        up6_e = tf.keras.layers.Conv2D(n*(2**4), 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(drop5))

        merge6_e = tf.keras.layers.concatenate([drop_e,up6_e], axis = 3)

        conv6_e = tf.keras.layers.Conv2D(n*(2**4), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6_e)

        conv6_e = tf.keras.layers.Conv2D(n*(2**4), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6_e)

        up6 = tf.keras.layers.Conv2D(n*(2**3), 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv6_e))

        merge6 = tf.keras.layers.concatenate([drop4,up6], axis = 3)

        conv6 = tf.keras.layers.Conv2D(n*(2**3), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)

        conv6 = tf.keras.layers.Conv2D(n*(2**3), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)

        up7 = tf.keras.layers.Conv2D(n*(2**2), 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv6))

        merge7 = tf.keras.layers.concatenate([conv3,up7], axis = 3)

        conv7 = tf.keras.layers.Conv2D(n*(2**2), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)

        conv7 = tf.keras.layers.Conv2D(n*(2**2), 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)

        up8 = tf.keras.layers.Conv2D(n*2, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv7))

        merge8 = tf.keras.layers.concatenate([conv2,up8], axis = 3)

        conv8 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)

        conv8 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)

        up9 = tf.keras.layers.Conv2D(n*2, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv8))

        merge9 = tf.keras.layers.concatenate([conv1,up9], axis = 3)

        conv9 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)

        conv9 = tf.keras.layers.Conv2D(n*2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)

        conv9 = tf.keras.layers.Conv2D(self.outputs, 3, activation = 'sigmoid', padding = 'same', kernel_initializer = 'he_normal')(conv9)

        model = tf.keras.Model(inputs=inputs, outputs=conv9)

        model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001), loss = 'binary_crossentropy', metrics = ['accuracy'])

        #If checkpoint path exists, this section loads the checkpoint, updating the model to the latest trained set of weights

        parentdir,file = os.path.split(checkpoint_loc)

        cpkt_save = os.path.join(parentdir,'checkpoint')

        if os.path.exists(cpkt_save):

            model.load_weights(checkpoint_loc)

        self.model = model

    #Converts an input image into a set of tiles, according to the Neural Networks input size.

    #

    #For example, if the neural network input had an input space of 4x4x3 (4x4 RGB), and the input image to the tile function was of size 7x7x3,

    #the tile function would insert the 7x7x3 input image into an 8x8x3 template (because 8x8 is the nearest set of dimension that is divisible by 4x4 tiles)

    #and then divide that modified template (the 8x8x3 section with the insert image) into 4 4x4x3 sections. 

    #

    #This is useful, because those 4 4x4x3 sections can then be fed directly into the neural network

    def tile(self, im_loc):

        height,width,channels = self.input_size

        im = np.asarray(Image.open(im_loc))[:,:,0:3]

        newy = math.ceil(im.shape[0]/self.input_size[0])

        newx = math.ceil(im.shape[1]/self.input_size[1])

        template = np.ones((newy*height,newx*width,channels))*255

        template[0:im.shape[0],0:im.shape[1]] = im

        template = template.reshape(newy, height,newx, width,3).swapaxes(1,2).reshape(-1,height,width,3)

        return template, im.shape[0], im.shape[1]

    #predict method receives an input image file location, it opens that image and then allows the AI to analyse it,

    #before returning the AI output.

    #

    #predict method uses the tile function to split the input image into a set of analysable windows, before analysing

    #them, stitching the outputs of the AI together and saving them.

    #

    #the default save location for this function is in te same folder as the input image, where the output file name is identical to the input,

    # except for it being prefixed by 'predicted_'

    def predict(self, im_loc, save_loc=None, show = False):

        if save_loc==None:

            parent_dir, file = os.path.split(im_loc)

            save_loc = os.path.join(parent_dir,'predicted_'+file)

        tiles, og_width, og_height = self.tile(im_loc)

        print(np.shape(tiles))

        pred_tiles = self.model.predict(np.array(tiles),verbose=1)

        print(np.shape(pred_tiles))

        splitx = math.ceil(og_width/self.input_size[0])

        splity = math.ceil(og_height/self.input_size[1])

        pred_tiles = pred_tiles.reshape(splitx,splity,self.input_size[0],self.input_size[1],self.outputs).swapaxes(1,2).reshape(splitx*self.input_size[0],splity*self.input_size[1],self.outputs)

        pred_tiles = np.argmax(pred_tiles,axis=2)

        if show:

            plt.imshow(pred_tiles,cmap='gray', vmax=self.outputs-1, vmin=0).write_png(save_loc)

        else:

            plt.imshow(pred_tiles,cmap='gray', vmax=self.outputs-1, vmin=0).write_png(save_loc)

            plt.close()

    # This method is designed to train the Neural network based on a set of MATLAB app images, but this is not necessary.

    #

    # The MATLAB app is called 'Image Labeller'

    #

    # The MATLAB app allows one to assign semantic labels to an RGB/Greyscale image, by annotating pixel regions with certain colours. Each colour

    # corresponds to a number, which corresponds to a pixel label. Therefore, for each input image, one can generate a 'ground truth' pixel map, where 

    # every pixel is labelled according to the semantic labels set up in the MATLAB App.

    #

    # In reality, this app is not needed for this function. Just have two folders of images; one with with a set of RGB input images and another

    # with a set of ground truths, where the pixel's semantic labels correspond to distinct integers, starting at 0 and incrementing.

    # I.e. In the ground truth, 0 = background, 1 = background, 2 = Stroma etc... The ground truth and corresponding image should be of identical

    # size and shape. Training images and their corresponding ground truths do not need to be the same size as the Neural network input.

    #

    # This function receives the folder lcoations of the training images and ground truths, and chooses random windows to learn from within each corresponding set of images.

    # one can vary how many windows are processed per batch (batch = 32 by default), the number of photos per training image (photos_per_mask = 100 by default) and how many times each

    # training image is revisted (epochs = 500 by default)

    def learn_from_matlab(self,image_loc, label_loc,photo_per_mask=100,batch = 32, epochs=500,colour=True):

        checkpoint_save = "N:\\"+str(self.n)+"_"+str(self.input_size)+"\\cp.ckpt"

        assert os.path.exists(image_loc)

        assert os.path.exists(label_loc)

        def gen(photo_per_mask = photo_per_mask): 

            image_files = os.listdir(image_loc)

            label_files = os.listdir(label_loc)

            gc.collect()

            assert len(label_files)== len(image_files)

            for i in range(len(image_files)):

                in_ = np.asarray(Image.open(os.path.join(image_loc,image_files[i])).convert('RGB'))

                label = np.asarray(Image.open(os.path.join(label_loc,label_files[i])).convert('L'))

                if colour:

                    in_ = np.asarray(Image.open(os.path.join(image_loc,image_files[i])).convert('RGB'))

                    label = np.asarray(Image.open(os.path.join(label_loc,label_files[i])).convert('L'))

                    padded_image_data = [0]*3

                    for j in range(3):

                        padded_image_data[j] = np.pad(in_[:,:,j], (int(self.input_size[0]/2),int(self.input_size[1]/2)), 'symmetric')

                    padded_image_data = np.dstack((padded_image_data[0],padded_image_data[1],padded_image_data[2]))

                    padded_label_data = np.pad(label, (int(self.input_size[0]/2),int(self.input_size[1]/2)), 'symmetric')

                else:

                    in_ = np.asarray(Image.open(os.path.join(image_loc,image_files[i])).convert('L'))

                    label = np.asarray(Image.open(os.path.join(label_loc,label_files[i])).convert('L'))

                    padded_image_data = np.pad(in_,(int(self.input_size[0]/2),int(self.input_size[1]/2)), 'symmetric')

                    padded_label_data = np.pad(label, (int(self.input_size[0]/2),int(self.input_size[1]/2)), 'symmetric')

                for j in range(photo_per_mask):

                    height, width = label.shape

                    row = random.randint(0,height-1)

                    col = random.randint(0,width-1)

                    rotate = random.randint(0,2)

                    window = np.rot90(padded_image_data[row:row+self.input_size[0],col:col+self.input_size[1]],rotate)

                    mask = np.asarray(np.rot90(padded_label_data[row:row+self.input_size[0],col:col+self.input_size[1]],rotate))

                    if random.uniform(0,1)>0.5:

                        window = np.flip(window,axis=0)

                        mask = np.flip(mask, axis=0)

                    yield (window, (np.arange(self.outputs) == mask[...,None]).astype(int))

        shuffle = photo_per_mask*len(os.listdir(image_loc))

        dataset = tf.data.Dataset.from_generator(gen, (tf.int64, tf.int64), (self.input_size, (self.input_size[0],self.input_size[1],self.outputs)))

        dataset = dataset.shuffle(shuffle).batch(batch)

        cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save,

                                                         save_weights_only=True,

                                                         verbose=1)

        self.model.fit(dataset,epochs=epochs,callbacks=[cp_callback])

if __name__ == '__main__':

    #initialise the unet with the file location of the model checkpoint

    unet = uNet_segmentor('N:\\8_(384, 384, 3)\\cp.ckpt',window = 384, first_conv=8)

    #insert the location of input images and files here

    image_loc = "N:\\Bill_Mcgough\\Correct Labelling\\Labels\\ImageData"

    label_loc = "N:\\Bill_Mcgough\\Correct Labelling\\Labels\\MatlabLabelData"

    #unet training here based on location of images and files

    unet.learn_from_matlab(image_loc, label_loc, photo_per_mask=65, batch = 10, epochs = 5000, colour = True)

    predictions = "C:\\Users\\Billy\\Downloads\\Prepared_SVS"

    images = os.listdir(predictions)

    print(images)

    for image in images:

        complete_loc = os.path.join(predictions,image)

        print(complete_loc)

        unet.predict(complete_loc)

    before = time.time()

    unet.predict("C:\\Users\\Billy\\Downloads\\Output tester\\14-2302 (1)_0_mag20.png")

    end = time.time()

    print(end-before)