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

Automatically generated by Colaboratory.

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

# T1
"""

from tensorflow.keras.regularizers import l2
import pathlib
import tensorflow 
from tensorflow import keras
from tensorflow.keras.layers import  Conv2D, MaxPooling2D, Flatten, Dense,Dropout,BatchNormalization
import tensorflow.keras 
import pathlib
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow.keras.utils as utils
from tensorflow.keras.optimizers import Adam as adam
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.optimizers import Adagrad 
from tensorflow.keras.callbacks import  EarlyStopping ,ModelCheckpoint
import tensorflow as tf
from tensorflow.keras import Model
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras.layers import  Conv2D, MaxPooling2D, Flatten, Dense, GlobalAveragePooling2D, Dropout, Input
import gradio as gr
from tensorflow.keras.applications import VGG16

from tensorflow.keras.applications.resnet50 import preprocess_input
from matplotlib import pyplot
from PIL import Image
from numpy import asarray
from PIL import Image

import glob
import cv2
from tensorflow.keras import layers
from keras.models import load_model
from matplotlib import pyplot
from PIL import Image
from numpy import asarray
from mtcnn.mtcnn import MTCNN
import cv2
from mask_the_face import *
import numpy as np

def get_paths():
  classes = []
  for file in sorted(glob.iglob('./lfw-deepfunneled/*/')):
    classes.append(file)
  for i,d in enumerate(classes):
        paths=d+'*.jpg'
        class_=[]     
        for file in sorted(glob.iglob(paths)):
            class_.append(file)
        classes[i]=class_ 
  return classes

classLabels=np.load('classLabels.npy',)

def extract_face(photo, required_size=(224, 224)):
  # load image from file
  pixels = photo
  print(pixels.shape)
  maxH=(pixels.shape[0])
  maxW=(pixels.shape[1])
  if (pixels.shape[-1])>3 or (pixels.shape[-1])<3:
    image = Image.fromarray(pixels)
    return image
  
  # create the detector, using default weights
  detector = MTCNN()
  # detect faces in the image
  results = detector.detect_faces(pixels)
  if not results:
    image = Image.fromarray(pixels)
    image = image.resize(required_size)
    print('not cropped')
    return image
  # extract the bounding box from the first face
  print('cropped')
  x1, y1, width, height = results[0]['box']
  x2, y2 = x1 + width, y1 + height
  
  face = pixels[y1:int(y2), int(x1):int(x2)]
  # resize pixels to the model size
  image = Image.fromarray(face)
  image = image.resize(required_size)

  return image

class FaceNet():
    def __init__(self,Weights_loading_path,facenet_path):
        self.loading=Weights_loading_path
        self.modelPath=facenet_path
        self.data_augmentation = keras.Sequential([layers.Rescaling(scale=1./127.5, offset=-1),layers.Resizing(160, 160),],name="data_augmentation",)
        self.Facenet=tf.keras.models.load_model(self.modelPath)
        self.Facenet.load_weights(self.loading)
        
      

    def Transfer_FacenetModel_withNormlization(self):
    
        facenetmodel = tf.keras.models.load_model(self.modelPath)
        # facenetmodel.load_weights('/content/drive/MyDrive/FaceNet/facenet_keras_weights.h5')
        for layer in facenetmodel.layers[:-50]:
          layer.trainable = False
        inputs = layers.Input(shape=(224,224,3))
          # Augment data.
        augmented = self.data_augmentation(inputs)
        # This is 'bootstrapping' a new top_model onto the pretrained layers.
        top_model = facenetmodel(augmented)
        top_model = Dropout(0.5)(top_model)
        top_model = BatchNormalization()(top_model)
        top_model = Flatten(name="flatten")(top_model)
        output_layer = Dense(5750, activation='softmax')(top_model)

        # Group the convolutional base and new fully-connected layers into a Model object.
        model = Model(inputs=inputs, outputs=output_layer)

        return model
    def predict(self,testsSamples):
        predictionProbabilty=self.Facenet.predict(testsSamples)
        return predictionProbabilty

class PatchEncoder(layers.Layer):
    def __init__(self, num_patches, projection_dim):
        super(PatchEncoder, self).__init__()
        self.num_patches = num_patches
        self.projection = layers.Dense(units=projection_dim)
        self.position_embedding = layers.Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )

    def call(self, patch):
        positions = tf.range(start=0, limit=self.num_patches, delta=1)
        encoded = self.projection(patch) + self.position_embedding(positions)
        return encoded

class Patches(layers.Layer):
    def __init__(self, patch_size):
        super(Patches, self).__init__()
        self.patch_size = patch_size

    def call(self, images):
        batch_size = tf.shape(images)[0]
        patches = tf.image.extract_patches(
            images=images,
            sizes=[1, self.patch_size, self.patch_size, 1],
            strides=[1, self.patch_size, self.patch_size, 1],
            rates=[1, 1, 1, 1],
            padding="VALID",
        )
        patch_dims = patches.shape[-1]
        patches = tf.reshape(patches, [batch_size, -1, patch_dims])
        return patches

class Transforemer():
    def __init__(self,loading_path):
        self.learning_rate = 0.001
        self.weight_decay = 0.0001
        self.batch_size = 32
        self.num_epochs = 300
        self.image_size = 72  
        self.patch_size = 6  # Size of the patches to be extract from the input images
        self.num_patches = (self.image_size // self.patch_size) ** 2
        self.projection_dim = 64
        self.num_heads = 8
        self.transformer_units = [self.projection_dim * 2,self.projection_dim,]  # Size of the transformer layers
        self.transformer_layers = 10

        self.mlp_head_units = [2048, 1024]  # Size of the dense layers of the final classifier
        self.loading=loading_path
        self.data_augmentation = keras.Sequential([ layers.Rescaling(1./255), layers.Resizing(self.image_size, self.image_size), layers.RandomFlip("horizontal")],name="data_augmentation",)
        self.transformer = self.create_vit_classifier()
        self.trnaformer  = self.transformer.load_weights(self.loading)

    def mlp(self,x, hidden_units, dropout_rate):
        for units in hidden_units:
            x = layers.Dense(units, activation=tf.nn.gelu)(x)
            x = layers.Dropout(dropout_rate)(x)
        return x
    def create_vit_classifier(self):

        inputs = layers.Input(shape=(224,224,3))

        augmented = self.data_augmentation(inputs)
        patches = Patches(self.patch_size)(augmented)
        encoded_patches = PatchEncoder(self.num_patches, self.projection_dim)(patches)
        
        for _ in range(self.transformer_layers):
            x1 = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)
            attention_output = layers.MultiHeadAttention(num_heads=self.num_heads, key_dim=self.projection_dim, dropout=0.3)(x1, x1)
            x2 = layers.Add()([attention_output, encoded_patches])
            x3 = layers.LayerNormalization(epsilon=1e-6)(x2)
            x3 = self.mlp(x3, hidden_units=self.transformer_units, dropout_rate=0.3)
            encoded_patches = layers.Add()([x3, x2])

        representation = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)
        representation = layers.Flatten()(representation)
        representation = layers.Dropout(0.6)(representation)
        features = self.mlp(representation, hidden_units=self.mlp_head_units, dropout_rate=0.6)
        logits = layers.Dense(5750, activation='softmax')(features)
        model = keras.Model(inputs=inputs, outputs=logits)

        return model
 
    def predict(self,testsSamples):
        predictionProbabilty=self.transformer.predict(testsSamples)
        return predictionProbabilty

class EnsembleModel():
    def __init__(self,classLabels,model1,model2,model3,model4):
        self.labels=classLabels
        self.model1 =model1
        self.model2 =model2
        self.model3 =model3
        self.model4 =model4
    def predict(self,testSample,):
        pred_prob1=self.model1.predict(testSample)
        pred_prob2=self.model2.predict(testSample)
        pred_prob3=self.model3.predict(testSample)
        pred_prob4=self.model4.predict(testSample)
        pred_sum=pred_prob1+pred_prob2+pred_prob3+pred_prob4
        print(pred_sum.shape)
        preds_classes_sum = np.argmax(pred_sum, axis=-1)
        total=sum(pred_sum[0])
        print(total)
        percentages=[x/total for x in pred_sum[0]]
        percentages=np.asarray(percentages)
        idx = np.argsort(pred_sum, axis=1)[:,-5:]
        print(pred_sum[0][idx])
        print(percentages[idx])
        return self.labels[preds_classes_sum][0],np.flip(self.labels[idx]),np.flip(percentages[idx])

"""# Test

"""

faceModel1=FaceNet('MyEn3facenet.h5','facenetModel.h5')
faceModel2=FaceNet('MyEn4facenet.h5','facenetModel.h5')
transformerModel1=Transforemer('FirstTransformer3Ensamble1.h5')
transformerModel2=Transforemer('FirstTransformer3Ensamble2.h5')

Ensemble=EnsembleModel(classLabels,faceModel1,faceModel2,transformerModel1,transformerModel2)



def grid_display(list_of_images, list_of_titles=[], no_of_columns=2, figsize=(10,10)):

    fig = plt.figure(figsize=figsize)
    column = 0
    for i in range(len(list_of_images)):
        column += 1
        #  check for end of column and create a new figure
        if column == no_of_columns+1:
            fig = plt.figure(figsize=figsize)
            column = 1
        fig.add_subplot(1, no_of_columns, column)
        plt.imshow(list_of_images[i])
        plt.axis('off')
        if len(list_of_titles) >= len(list_of_images):
            plt.title(list_of_titles[i])

def reconitionPipline(img,mask):
    
    im = Image.fromarray(img.astype('uint8'), 'RGB')
    im=np.array(im) 
    im2= im[:,:,::-1].copy()

    if mask:
        im2=maskThisImages(im2)
        if len(im2)==0:
            im2=im.copy()
            im2= im2[:,:,::-1]
    
    im2= im2[:,:,::-1]
    temp=extract_face(im2)
    cropped = np.array(temp) 
    open_cv_image = cropped[:, :, ::-1].copy() 

    prediction,top5,percentage=Ensemble.predict(open_cv_image[None,...])
    return dict(zip(np.reshape(top5, -1), np.reshape(percentage, -1))),cropped

with gr.Blocks() as demo:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 650px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                
                <h1 style="font-weight: 900; margin-bottom: 7px;">
                  LFW-Masked Recognition
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
                An AI model developed using Ensamble learning method
                with transformer and facenet to recognize celebrties classes in LFW dataset (+5700 class)
                
              </p>
            </div>
        """
    )
    with gr.Row():
        with gr.Column():
            imagein = gr.Image(label='User-Input',interactive=True)
            
        # with gr.Column():
        #     gr.Examples(['1.jpg','2.jpg','3.jpg'],inputs=imagein)
    with gr.Row():
        checkbox=gr.Checkbox(label='Mask the face')
        image_button = gr.Button("Submit")
       
    with gr.Row():
        mOut = gr.Image(type='numpy',label=' (Model-input)')
        label = gr.Label(num_top_classes=5)
        
    
    gr.Markdown("<p style='text-align: center'>Made  with 🖤 by  Mohammed & Aseel </p>")
    image_button.click(fn=reconitionPipline,inputs=[imagein,checkbox],outputs=[label,mOut])
demo.launch()