import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Charger les modèles (remplacez les chemins par vos fichiers de modèles)
model1 = tf.keras.models.load_model('model_v2_Q1.keras')
model2 = tf.keras.models.load_model('model_v2_Q2.keras')
model3 = tf.keras.models.load_model('model_Q3.keras')
#model4 = tf.keras.models.load_model('Model_4_Final.keras')


# Gère le formattage des images
def format_image_type1(image):
    image = image.resize((64, 64)) # Redimensionner selon la taille attendue
    image = np.array(image) / 255.0 # Normalisation
    return image_single_fix(image)

def format_image_type2(image):
    image = load_img(image, color_mode='grayscale') # Charge en noir et blanc mais avec une meilleure résolution
    image = img.resize((128, 128), Image.LANCZOS) # Redimensionner selon la taille attendue
    image = np.array(img) / 255.0 # Normalisation
    return image_single_fix(image)

def image_single_fix(image_array):
    return np.expand_dims(image_array, axis=0) # Ajouter batch dimension (demandé par tensorflow/keras qui demande à recevoir un groupe d'image plutôt que juste une seule)


# Configuration des infos des modèles
models = {
    "Modèle 1 (genre uniquement)": {
        "model": model1,
        "type": "gender_only",
        "image_format": format_image_type1
    },
    "Modèle 2 (âge uniquement)": {
        "model": model2,
        "type": "age_only",
        "image_format": format_image_type1
    },
    "Modèle 3 (genre et âge)": {
        "model": model3,
        "type": "gender_age",
        "image_format": format_image_type2
    }
}


# Fonction de prédiction
def predict(image, model_name):
    model_data = models[model_name]
    
    image_array = model_data["image_format"](image) # Formatter l'image de la façon qui correspond
    
    model = model_data["model"]
    model_type = model_data["type"]
    
    prediction = model.predict(image_array)
    
    return_text = ""
    if model_type == "gender_only":
        gender_value = prediction[0][0]
        gender = "Homme" if gender_value < 0.5 else "Femme" # Genre basé sur probabilité
        return_text = f"Genre: {gender} ({gender_value} - {get_gender_confidence(gender_value)} certitude)"
    elif model_type == "age_only":
        age_value = prediction[0][0]
        age = round(prediction[0][0]) # Arroundi à l'entier le plus proche "half away zero"
        return_text = f"Age: {age} ({age_value})"
    else: # "gender_age"
        gender_value = prediction[0][0]
        gender = "Homme" if gender_value < 0.5 else "Femme" # Genre basé sur probabilité
        age_value = prediction[0][1]
        age = round(prediction[0][0])
        return_text = f"Genre: {gender} ({gender_value} - {get_gender_confidence(gender_value)} certitude)\nAge: {age} ({age_value})"
    
    return f"{return_text}"


# Fonctions utiles
def get_gender_confidence(gender_value):
    return f"{round(abs(gender_value - 0.5) * 2 * 100)}%"


# Interface Gradio
iface = gr.Interface(
    fn=predict,
    inputs=[gr.Image(type="pil"), gr.Dropdown(choices=list(models.keys()), label="Choisir un modèle")],
    outputs=gr.Textbox(label="Prédictions des modèles")
)


# Lancer l'interface
if __name__ == "__main__":
    iface.launch()