'''
    Build Project
'''

# imported necessary libraries
import gradio as gr
import tensorflow as tf
import numpy as np
import os
import yaml
import cv2
from PIL import Image
import io
import json

# read config file
def read_config():
    config = {}
    print(os.path.curdir)
    with open('models_config.yaml', 'r') as cf:
        config = yaml.safe_load(cf)

    for var in config:
        config[var] = config[var].replace(';', os.sep)

    return config

config = read_config()

# loading the models
age_model = tf.keras.models.load_model(config['A_M_PATH'])
gender_model = tf.keras.models.load_model(config['G_M_PATH'])
face_cascade = cv2.CascadeClassifier(config['FD_M_PATH'])
 

def main_pipeline(image):
    
    if image is None:
        image = Image.open(config['BANNER_IMG'])
        return image, "NO FACE DETECTED", "NO FACE DETECTED"


    original_image = image
    image = np.asarray(image)

    gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    faces = face_cascade.detectMultiScale(gray_img, 1.1, 4)
    
    if len(faces) == 0:
        return original_image , "NO FACE DETECTED", "NO FACE DETECTED"

    x, y, w, h =  faces[0]


    image = Image.fromarray(image)
    image = image.crop((x, y, x + w, y + h))

    image = image.resize((224,224))
    #image = tf.image.resize(image, [224,224]) 
    image = tf.keras.preprocessing.image.img_to_array(image)
    image = image / 255.0      
    image = tf.expand_dims(image, axis=0)

    age_prds = age_model.predict(image)
    gender_prds = gender_model.predict(image)

    age_prds = np.around(age_prds)
    gender_prds = np.around(gender_prds)
    gender = ""

    if gender_prds[0][0] == 0:
        gender = 'male'
    else:
        gender = 'female'

    data = {}
    data['age'] = int(age_prds[0][0])
    data['gender'] = gender

    image_b_box = cv2.rectangle(np.asarray(original_image), (x, y), (x+w, y+h), (255, 0, 0), 2)

    result_image = Image.fromarray(image_b_box)

    return result_image, f"~ {(data['age'])}", data['gender'] 


# initializing the input component
image = gr.inputs.Image()

# initializzing examples
test_imgs = [os.path.join(config['TEST_IMG_PATH'],img) for img in os.listdir(config['TEST_IMG_PATH'])]
examples = test_imgs
 
# launching the interface
gr.Interface(fn = main_pipeline, title= 'Age & Gender predictions', inputs = image,
             outputs = [gr.Image(label= 'Result', shape= (250,250)),
                        gr.Text(label= 'Age'), 
                        gr.Text(label= 'Gender')], flagging_options=["correct", "incorrect", "other"],
                        examples= examples, capture_session = True).launch()