import streamlit as st

from beluga import load_model, process_emotions, generate_prompt
from emodeepface import check_image_rotation, process_photo

# begin loading beluga model and tokenizer
model, tokenizer = load_model()

# title webpage
st.title("Affective Journaling Assistant")

# provide user instructions
st.write("""
Welcome to the Affective Journaling Assistant!

For a tailored journaling experience, we analyze your facial expressions to gauge your emotions.
To proceed:

1. Ensure the image is well-lit and of high quality.
2. Your face should be fully visible without obstructions (e.g., no sunglasses or hats).
3. By uploading, you acknowledge and consent to our data processing.

Let's get started!
""")

# request user image upload
file_name = st.file_uploader("Please upload your photo.")

# once an image has been uploaded
if file_name is not None:
    # capture image with intended rotation
    image = check_image_rotation(file_name)
    # display the image directly
    st.image(image, use_column_width=True)
    # process uploaded image
    emotion_predictions = process_photo(file_name) 
    # process emotion predictions
    result = process_emotions(model, tokenizer, emotion_predictions)
    # generate affective journaling prompt based on emotion predictions
    prompt = generate_prompt(result)
    # display journal prompt
    st.write(prompt)