import streamlit as st
from transformers import AutoProcessor, UdopForConditionalGeneration
from datasets import load_dataset

processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=True)
model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large")

st.title("Image Selection and Upload")
st.write("Select or upload an image to test the model.")

# Image selection
image_option = st.selectbox("Select an image:", ("Image 1", "Image 2", "Image 3"))
if image_option == "Image 1":
    st.image("image1.jpg", caption="Image 1")
elif image_option == "Image 2":
    st.image("image2.jpg", caption="Image 2")
elif image_option == "Image 3":
    st.image("image3.jpg", caption="Image 3")

# Image upload
uploaded_file = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    st.image(uploaded_file, caption="Uploaded Image")

# Model testing button
if st.button("Test Model"):
    if image_option == "Image 1":
        # Test the model with Image 1
        st.write("Testing the model with Image 1...")
    elif image_option == "Image 2":
        # Test the model with Image 2
        st.write("Testing the model with Image 2...")
    elif image_option == "Image 3":
        # Test the model with Image 3
        st.write("Testing the model with Image 3...")
    elif uploaded_file is not None:
        # Test the model with the uploaded image
        st.write("Testing the model with the uploaded image...")
    else:
        st.write("Please select or upload an image.")

# encoding = processor(image, question, words, boxes=boxes, return_tensors="pt")

# predicted_ids = model.generate(**encoding)
# print(processor.batch_decode(predicted_ids, skip_special_tokens=True)[0])