import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import tensorflow as tf
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.utils import get_custom_objects


def mse(y_true, y_pred):
    return MeanSquaredError()(y_true, y_pred)

get_custom_objects().update({'mse': mse})

# Chargement des modèles
face_aging_model = load_model("face_aging_model.h5", custom_objects={"mse": mse})
face_aging_autoencoder = load_model("face_aging_autoencoder.h5", custom_objects={"mse": mse})

def load_image():
    file_path = filedialog.askopenfilename(title="Sélectionner une image", filetypes=[("Image Files", "*.jpg;*.png;*.jpeg")])
    if file_path:
        img = Image.open(file_path)
        img = img.resize((256, 256))  # Redimensionner l'image
        img_tk = ImageTk.PhotoImage(img)
        original_panel.configure(image=img_tk)
        original_panel.image = img_tk
        global image_data, image_path
        image_path = file_path
        image_data = np.array(img) / 255.0  # Normalisation de l'image
        image_data = np.expand_dims(image_data, axis=0) 

def predict_age():
    if 'image_path' in globals():
        predicted_age = predict_age_from_model(face_aging_model, image_path)
        messagebox.showinfo("Prédiction d'âge", f"L'âge prédit du visage est : {predicted_age:.2f} ans")
    else:
        messagebox.showerror("Erreur", "Veuillez d'abord charger une image.")

def predict_age_from_model(model, image_path):
    img = Image.open(image_path).resize((128, 128))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    prediction = model.predict(img_array)
    return prediction[0][0]

def apply_aging_effect(model, image_path):
    img = Image.open(image_path).resize((128, 128))
    img_array = np.array(img) / 255.0
    predicted_img = model.predict(np.expand_dims(img_array, axis=0))
    predicted_img = np.clip(predicted_img[0], 0, 1) 
    predicted_img = (predicted_img * 255).astype(np.uint8)
    return Image.fromarray(predicted_img)

def show_aged_image():
    if 'image_path' in globals():
        aged_image = apply_aging_effect(face_aging_autoencoder, image_path)
        aged_image_tk = ImageTk.PhotoImage(aged_image.resize((256, 256)))
        aged_panel.configure(image=aged_image_tk)
        aged_panel.image = aged_image_tk
    else:
        messagebox.showerror("Erreur", "Veuillez d'abord charger une image.")

# Création de la fenêtre principale Tkinter
root = tk.Tk()
root.title("Face Age Prediction and Aging")

# Affichage de l'image originale
original_panel = tk.Label(root)
original_panel.pack(pady=10)

# Boutons et widgets
load_button = tk.Button(root, text="Charger une image", command=load_image)
load_button.pack(pady=5)

predict_button = tk.Button(root, text="Prédire l'âge", command=predict_age)
predict_button.pack(pady=5)

age_button = tk.Button(root, text="Afficher l'image vieillie", command=show_aged_image)
age_button.pack(pady=5)

# Affichage de l'image vieillie
aged_panel = tk.Label(root)
aged_panel.pack(pady=10)

# Lancer l'interface
root.mainloop()