import streamlit as st
from transformers import ViltProcessor, ViltForQuestionAnswering
from PIL import Image

def load_model():
    processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
    model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
    return processor, model

def predict(image, text, processor, model):
    encoding = processor(image, text, return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits
    idx = logits.argmax(-1).item()
    return model.config.id2label[idx]

def main():
    st.title("VQA")
    st.write("Upload an image and input a question to get an answer.")

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

    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        question = st.text_input("Question about the image:")
        
        if question:
            processor, model = load_model()
            answer = predict(image, question, processor, model)
            
            col1, col2 = st.columns(2)
            with col1:
                st.image(image, caption='Uploaded Image.', use_column_width=True)
            with col2:
                st.write(f"**Question:** {question}")
                st.write(f"**Answer:** {answer}")

if __name__ == "__main__":
    main()