import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.losses import MeanSquaredError

def generate_aged_image(model, noise_dim):
    # Générer un vecteur de bruit
    noise = np.random.normal(0, 1, (1, noise_dim))
    
    # Générer l'image vieillie
    generated_img = model.predict(noise)
    
    print(generated_img)
    print(type(generated_img[0][0]))
    print(f"Shape of generated image: {generated_img.shape}")
    print(f"Generated image values (min, max): {generated_img.min()}, {generated_img.max()}")
    
    generated_img = np.clip(generated_img[0], 0, 1)  # S'assurer que les valeurs sont entre 0 et 1
    generated_img = (generated_img * 255).astype(np.uint8)
    
    # Convertir en BGR pour OpenCV
    generated_img_bgr = cv2.cvtColor(generated_img, cv2.COLOR_RGB2BGR)
    
    return generated_img_bgr

def main():
    # Chemin vers le modèle et l'image
    model_path = "models/generator_epoch_7400.h5"
    output_path = "visage_aged.jpg"
    
    # Charger le modèle
    model = load_model(model_path, custom_objects={'mse': MeanSquaredError()})
    print("Modèle chargé avec succès !")
    
    # Générer l'image vieillie
    aged_image = generate_aged_image(model, noise_dim=100)
    
    # Sauvegarder l'image résultante
    cv2.imwrite(output_path, aged_image)
    print(f"Image vieillie sauvegardée sous {output_path}")

if __name__ == "__main__":
    main()