app.py

import streamlit as st
from beluga import load_model, process_emotions, generate_prompt
from emodeepface import check_image_rotation, process_photo

@st.cache_resource
def load_cached_model():
    return load_model()

if 'model' not in st.session_state:
    loading_message = st.empty()
    loading_message.text("Loading model... Please wait.")
    
    st.session_state.model, st.session_state.tokenizer = load_cached_model()
    loading_message.empty()

st.title("Affective Journaling Assistant")

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!
""")

file_name = st.file_uploader("Please upload your photo.")

if file_name is not None:
    image = check_image_rotation(file_name)
    
    processing_message = st.empty()
    processing_message.text("Analyzing your image... Please wait.")
    
    emotion_predictions = process_photo(file_name) 
    result = process_emotions(st.session_state.model, st.session_state.tokenizer, emotion_predictions)
    processing_message.empty()
    
    prompt = generate_prompt(result)

    # Create columns to place the image and the prompt side by side
    col1, col2 = st.columns(2)

    # Show image in the left column
    col1.image(image, width=300)  

    # Show generated prompt in the right column
    col2.write(prompt)