import gradio as gr
from transformers import AutoProcessor, Pix2StructForConditionalGeneration
from PIL import Image

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

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

demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Chart to Text Converter",
    description="Upload a chart image to extract its data into table format"
)

demo.launch()