import streamlit as st
import cv2
import numpy as np
from keras.models import load_model

def getAge(distr):
    distr = distr * 4
    if 0.65 <= distr <= 1.4:
        return "0-18"
    elif 1.65 <= distr <= 2.4:
        return "19-30"
    elif 2.65 <= distr <= 3.4:
        return "31-80"
    elif 3.65 <= distr <= 4.4:
        return "80 +"
    return "Unknown"

def getGender(prob):
    return "Male" if prob < 0.5 else "Female"

def getAgeGender(image_path):
    # Loading the uploaded Image:
    image = cv2.imread(image_path, 0)
    image = cv2.resize(image, dsize=(64, 64))
    image = image.reshape((image.shape[0], image.shape[1], 1))

    # Loading the trained model:
    model = load_model('data.h5')

    # Getting the predictions:
    image = image / 255
    val = model.predict(np.array([image]))
    age = getAge(val[0])
    gender = getGender(val[1])
    return age, gender

def main():
    st.title("Age and Gender Prediction with Streamlit")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # Save the uploaded file to a temporary location
        temp_image_path = "temp_image.jpg"
        with open(temp_image_path, "wb") as f:
            f.write(uploaded_file.getvalue())

        # Get age and gender predictions
        age, gender = getAgeGender(temp_image_path)

        # Display the uploaded image
        st.image(temp_image_path, caption="Uploaded Image", use_column_width=True)

        # Display the predictions
        st.write(f"Predicted Age: {age}")
        st.write(f"Predicted Gender: {gender}")

if __name__ == "__main__":
    main()