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('model3.h5')
#model4 = tf.keras.models.load_model('Model_4_Final.keras')

models = {
    "Modèle 1 (genre uniquement)": {
        "model": model1, "type": "gender_only"
    },
    "Modèle 2 (âge uniquement)": {
        "model": model2, "type": "age_only"
    }
}
#models = []

IMG_SIZE = (64, 64)

# Fonction de prédiction
def predict(image, model_name):
    image = image.resize(IMG_SIZE)  # Redimensionner selon la taille attendue
    image_array = np.array(image) / 255.0  # Normalisation
    image_array = 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)
    
    model_data = models[model_name]
    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])  # Exemple : âge comme valeur continue
        return_text = f"Age: {age} ({age_value})"
    
    return f"{return_text}"

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()