import numpy as np
import cv2
from tensorflow.keras.models import load_model
import matplotlib.pyplot as plt

# Load the trained generator model
generator = load_model("aging_generator_model.h5")

# Function to load and preprocess an image
def load_image(image_path, img_size=(200, 200)):
    img = cv2.imread(image_path)
    img = cv2.resize(img, img_size)
    img = img / 255.0  # Normalize
    return np.expand_dims(img, axis=0)  # Add batch dimension

# Function to save and display the original and aged images
def save_and_display_images(original_img_path, aged_img, output_path):
    original_img = cv2.imread(original_img_path)
    aged_img = (aged_img[0] * 255).astype(np.uint8)  # Denormalize
    aged_img = cv2.cvtColor(aged_img, cv2.COLOR_RGB2BGR)
    
    # Save the aged image
    cv2.imwrite(output_path, aged_img)
    
    # Display the original and aged images
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
    axes[0].imshow(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
    axes[0].set_title("Original Image")
    axes[0].axis("off")
    
    axes[1].imshow(cv2.cvtColor(aged_img, cv2.COLOR_BGR2RGB))
    axes[1].set_title("Aged Image")
    axes[1].axis("off")
    
    plt.tight_layout()
    plt.show()

# Path to the input image
input_image_path = "visage.jpg"  # Update with your input image path
output_image_path = "aged_face_result.jpg"

# Load and preprocess the input image
input_image = load_image(input_image_path)

# Generate the aged face
aged_face = generator.predict(input_image)

# Save and display the original and aged images
save_and_display_images(input_image_path, aged_face, output_image_path)