import streamlit as st
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np

# Load the model
model = load_model('model.h5', compile=False)

# Image preprocessing function
def process_image(img):
    # Resize, normalize, and add batch dimension
    img = img.resize((128, 128))
    img = np.array(img)
    img = img / 255.0  # Normalize pixel values
    img = np.expand_dims(img, axis=0)  # Add batch dimension (1, 128, 128, 3)
    return img

# App title
st.title('Age Prediction from Image')
st.write("Upload a photo to predict the subject's age.")

# Sidebar instructions
st.sidebar.header("Instructions")
st.sidebar.write("""
1. Upload a clear face photo
2. Model will analyze facial features
3. Predicted age appears below image
4. This dataset contains images of people generated by Stable Diffusion.
For each image, you are to determine how old the person is.
Ages will range between 20 and 90 years old.
""")

# File uploader widget
file = st.file_uploader('Choose an image (JPG, JPEG, PNG)', type=['jpg', 'jpeg', 'png'])

if file is not None:
    # Convert to RGB if image has alpha channel (critical fix)
    img = Image.open(file).convert('RGB')
    
    # Display uploaded image
    st.image(img, caption='Uploaded Image', use_column_width=True)
    
    # Process and predict
    processed_image = process_image(img)
    prediction = model.predict(processed_image)
    prediction = np.round(prediction).astype(int)  # Round to nearest integer

    # Show prediction
    st.subheader("Prediction Result:")
    st.markdown(f"**Estimated Age:** {prediction[0][0]} years")
    
    # Disclaimer
    st.info("""
    Note: Predictions are based on patterns learned during training. 
    Results may vary with image quality and lighting conditions.
    """)
    
    st.divider()
else:
    st.write("Please upload an image to begin.")