from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
import requests
from PIL import Image
import streamlit as st

processor = Pix2StructProcessor.from_pretrained('google/deplot')
model = Pix2StructForConditionalGeneration.from_pretrained('google/deplot')

document = st.file_uploader(label="Upload the document you want to explore",type=["png",'jpg', "jpeg","pdf"])

if document == None:
    st.write("Please upload the document in the box above")
else:
    image = Image.open(document)
    st.image(image,"Document uploaded")

    inputs = processor(images=image, text="Generate underlying data table of the figure below:", return_tensors="pt")
    predictions = model.generate(**inputs, max_new_tokens=512)
    st.write(processor.decode(predictions[0], skip_special_tokens=True))