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

Automatically generated by Colaboratory.

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

class DeepCNNClassifier:
  def __init__(self , dataset_link , model_name):
    self.dataset_link = dataset_link
    self.model_name = model_name
    self.number_of_classes = 2 #Default


  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def extract_dataset(self):
    import zipfile
    import os
    destination_path = '/content'
    zip_ref = zipfile.ZipFile(self.dataset_link, "r")
    extracted = zip_ref.namelist()
    zip_ref.extractall(destination_path)
    zip_ref.close()
    extracted_file = os.path.join(destination_path, extracted[0])
    print("=========== Dataset Extracted to - {}===========".format(extracted_file))
    self.extracted_file_path = extracted_file
    #return extracted_file
  
  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def set_details(self , train_dir , test_dir , transfer_model_url = None , image_height= 224 , image_width = 224 , isGrayScale = False):
    self.train_dir = train_dir
    self.test_dir  = test_dir
    self.transfer_model_url = transfer_model_url
    self.image_height =image_height
    self.image_width = image_width
    self.isGrayScale = isGrayScale

    if(self.isGrayScale):
      self.num_color_channels = 1
    else:
      self.num_color_channels = 3

    print("=========== Train/Test Directories Configured===========")

  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def show_dataset_summary(self):
    import os
    for dirpath , dirnames , filenames in os.walk(self.extracted_file_path):
      print(f"-->{len(dirnames)} directories AND  {len(filenames)} images in '{dirpath}'.")

  #----------------------------------------------------WILL_USE_BY_ANOTHER_FUNCTION-----------------------------------------------------------------------------
  def get_class_names(self):
    import pathlib
    import numpy as np

    data_dir = pathlib.Path(self.train_dir)
    class_names = np.array(sorted([item.name for item in data_dir.glob('*')]))
    print("=========== Class Names Obtained  - {} ".format(class_names))
    self.class_names = class_names
    self.number_of_classes = len(class_names)
  
  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def view_random_image(self):
    self.get_class_names()
    import matplotlib.pyplot as plt
    import matplotlib.image as mpimg
    import random
    import os

    target_dir = self.train_dir
    target_class = random.choice(self.class_names)

    target_folder = target_dir+"/"+target_class
    random_image = random.sample(os.listdir(target_folder) ,1)

    img = mpimg.imread(target_folder+"/"+random_image[0])
    plt.imshow(img)
    plt.title(target_class)
    plt.axis("off")

    print(f"Image Shape:{img.shape}") 

  
  #----------------------------------------------------WILL_USE_BY_ANOTHER_FUNCTION-----------------------------------------------------------------------------
  def preprocess_dataset(self):
    print("=========== Data Pre-Processing Started ===========")
    from tensorflow.keras.preprocessing.image import ImageDataGenerator
    import tensorflow as tf

    PRE_RPOCESS_FUNC = None
    CLASS_MODE = 'binary'

    def to_grayscale_then_rgb(image):
      image = tf.image.rgb_to_grayscale(image)
      image = tf.image.grayscale_to_rgb(image)
      return image

    if(self.isGrayScale):
      print("=========== Training will use GrayScale Images !!!")
      PRE_RPOCESS_FUNC = to_grayscale_then_rgb

    if(self.number_of_classes >2):
      CLASS_MODE = 'categorical'

    train_data_gen  = ImageDataGenerator(rescale=1/255.,
                                     preprocessing_function= PRE_RPOCESS_FUNC,
                                     rotation_range = 0.2,
                                     shear_range  = 0.2,
                                     zoom_range = 0.2,
                                     width_shift_range=0.3,
                                     height_shift_range = 0.3,
                                     horizontal_flip= True)
    test_data_gen = ImageDataGenerator(rescale=1/255. ,
                                   preprocessing_function=PRE_RPOCESS_FUNC)
    
    
    #Import and Transform/pre process the data
    train_data_ = train_data_gen.flow_from_directory(self.train_dir,
                                                          target_size = (self.image_height,self.image_width),
                                                          batch_size = 32,
                                                          class_mode = CLASS_MODE,
                                                          shuffle = True)
    test_data_ = test_data_gen.flow_from_directory(self.test_dir,
                                                        target_size = (self.image_height, self.image_width),
                                                        batch_size =  32,
                                                        class_mode = CLASS_MODE, #This will gives us one-hot encoded data
                                                        shuffle = True)
    print("=========== Data Pre-Processing Completed ===========")
    return (train_data_ , test_data_)
  
   #----------------------------------------------------WILL_USE_BY_ANOTHER_FUNCTION-----------------------------------------------------------------------------
  def create_model(self):
    import tensorflow as tf
    import tensorflow_hub as hub
    from tensorflow.keras import layers

    #Download the pretrained model and save it as a keras layer
    feature_extractor_layer = hub.KerasLayer(self.transfer_model_url,
                                            trainable = False, #freeze the already learned patterns 
                                            name = "feature_extraction_layer",
                                            input_shape = (self.image_height, self.image_width, self.num_color_channels)) 
    self.activation_func = 'sigmoid'
    if(self.number_of_classes > 2):
      self.activation_func = 'softmax'


    model  =  tf.keras.Sequential([
      feature_extractor_layer,
      layers.Dense(self.number_of_classes , activation= self.activation_func , name="output_layer")
    ])
    print("=========== Model is Created ==========")
    return model
  
  
  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def start_training(self , num_epochs = 5):
    import tensorflow as tf
    train_data_ , test_data_ = self.preprocess_dataset()
    cnn_model =  self.create_model()

    self.loss_function = 'binary_crossentropy'
    if(self.number_of_classes > 2):
      self.loss_function = 'categorical_crossentropy'

    cnn_model.compile(loss=self.loss_function,
                     optimizer = tf.keras.optimizers.Adam(),
                     metrics=["accuracy"])
    print("=========== Model Compiled SuccessFully ==========")

    model_hist = cnn_model.fit(train_data_,
                               epochs=num_epochs,
                               steps_per_epoch=len(train_data_),
                               validation_data = test_data_,
                               validation_steps = len(test_data_))
    print("=========== Model Trained SuccessFully ==========")
    self.trained_model = cnn_model
    self.model_train_history = model_hist
    return cnn_model
  
  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def save_model(self, save_to_path = '/content'):
    model_name = str(self.model_name)+".h5"
    model = self.trained_model
    model.save(model_name)
    print("=========== Model Saved to -->  {}".format(model_name))

  #---------------------------------------------------SHOULD_CALL_BY_USER / Function------------------------------------------------------------------------------
  def plot_loss_curves(self, ):
    import matplotlib.pyplot as plt

    history = self.model_train_history
    loss = history.history["loss"]
    val_loss = history.history["val_loss"]

    accuracy = history.history["accuracy"]
    val_accuracy = history.history["val_accuracy"]

    #get the number of epochs that we run for
    epochs =  range(len(history.history["loss"]))

    #Plot the lost
    plt.plot(epochs , loss , label="Training Loss")
    plt.plot(epochs , val_loss , label="Validation Loss")
    plt.title("Loss")
    plt.xlabel("Epochs")
    plt.legend()

    #Plot the accuracy
    plt.figure()
    plt.plot(epochs , accuracy , label="Training accuracy")
    plt.plot(epochs , val_accuracy , label="Validation accuracy")
    plt.title("accuracy")
    plt.xlabel("Epochs")
    plt.legend()
  
  #----------------------------------------------------WILL_USE_BY_ANOTHER_FUNCTION-----------------------------------------------------------------------------
  def load_and_prep_image(self, filename):
    import tensorflow as tf

    img = tf.io.read_file(filename)
    img = tf.image.decode_image(img)
    #Resiz the Image
    img = tf.image.resize(img , size=[self.image_height , self.image_width])
    #Re scale the image - get all values between 0 and 255
    img = img/255.
    return img

  #---------------------------------------------------SHOULD_CALL_BY_USER------------------------------------------------------------------------------
  def pred_and_plot(self, filename):
    import tensorflow as tf
    import matplotlib.pyplot as plt
    import numpy as np

    model = self.trained_model
    class_names = self.class_names
  
    img = self.load_and_prep_image(filename)
    
    pred = model.predict(tf.expand_dims(img , axis=0))

    pred_class =  class_names[np.argmax(tf.round(pred))]

    #Plot the image and the predicted class
    plt.imshow(img)
    plt.title(f"Prediction :{pred_class}")
    plt.axis(False)


  
 

