import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array, load_img
import matplotlib.pyplot as plt
from tensorflow.keras.losses import MeanSquaredError

def apply_aging_effect(model, image_path):
    img = load_img(image_path, target_size=(128, 128)) 
    img_array = img_to_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)
    
    predicted_img_bgr = cv2.cvtColor(predicted_img, cv2.COLOR_RGB2BGR)
    
    return predicted_img_bgr

def main():
    model_path = "face_aging_autoencoder.h5"
    image_path = "visage.jpg"
    output_path = "visage_aged.jpg"
    
    model = load_model(model_path, custom_objects={'mse': MeanSquaredError()})
    print("Modèle chargé avec succès !")
    
    aged_image = apply_aging_effect(model, image_path)
    
    cv2.imwrite(output_path, aged_image)
    print(f"Image vieillie sauvegardée sous {output_path}")

if __name__ == "__main__":
    main()