from PIL import Image
import requests
from io import BytesIO
from transformers import pipeline
import streamlit as st

def predict(image):
    type_food = oracle(image, "What type of food is this?")
    cal_est = oracle(image, "About how many calories are in this meal?")
    guess1, guess2 = int(cal_est[0]['answer']), int(cal_est[1]['answer'])
    return f"This is {type_food[0]['answer']}. I estimate this to contain {min(guess1, guess2)}-{max(guess1, guess2)} calories"
oracle = pipeline(model="dandelin/vilt-b32-finetuned-vqa")

def main():
    st.title("Image Question Answering App")
    st.write("Upload an image and ask a question to get answers!")

    # File uploader for image
    uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
    
    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        response = predict(image)
        st.write(response)
            
if __name__ == "__main__":
    main()