import streamlit as st
from PIL import Image
from logic import *

model, processor= load_Model()

def load_image(img):
    im = Image.open(img).convert("RGB")
    return im

st.title('Ask Image Questions UDOP')
uploaded_file = st.file_uploader("Input Image",type=['png', 'jpg'] )

if uploaded_file is not None:
    img = load_image(uploaded_file)
    st.image(img)
    st.write("Image Uploaded Successfully")
    with st.form("my_form"):
        user_query = st.text_input("Ask the image")
              
        new_submitted = st.form_submit_button("Submit")
        if new_submitted:
            encoding = processor(images=img, text=user_query, return_tensors="pt")
            outputs = model.generate(**encoding, max_new_tokens=20)
            generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
            st.write(generated_text)