merve/pix2struct

1import gradio as gr2import requests3from PIL import Image4from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor5import spaces6 7@spaces.GPU8def infer_infographics(image, question):9    model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-ai2d-base").to("cuda")10    processor = Pix2StructProcessor.from_pretrained("google/pix2struct-ai2d-base")11 12    inputs = processor(images=image, text=question, return_tensors="pt").to("cuda")13 14    predictions = model.generate(**inputs)15    return processor.decode(predictions[0], skip_special_tokens=True)16 17@spaces.GPU18def infer_ui(image, question):19    model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-base").to("cuda")20    processor = Pix2StructProcessor.from_pretrained("google/pix2struct-screen2words-base")21 22    inputs = processor(images=image,text=question, return_tensors="pt").to("cuda")23 24    predictions = model.generate(**inputs)25    return processor.decode(predictions[0], skip_special_tokens=True)26 27@spaces.GPU28def infer_chart(image, question):29  model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-chartqa-base").to("cuda")30  processor = Pix2StructProcessor.from_pretrained("google/pix2struct-chartqa-base")31 32  inputs = processor(images=image, text=question, return_tensors="pt").to("cuda")33 34  predictions = model.generate(**inputs)35  return processor.decode(predictions[0], skip_special_tokens=True)36 37@spaces.GPU38def infer_doc(image, question):39  model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-base").to("cuda")40  processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-base")41  inputs = processor(images=image, text=question, return_tensors="pt").to("cuda")42  predictions = model.generate(**inputs)43  return processor.decode(predictions[0], skip_special_tokens=True)44 45css = """46  #mkd {47    height: 500px; 48    overflow: auto; 49    border: 1px solid #ccc; 50  }51"""52 53with gr.Blocks(css=css) as demo:54  gr.HTML("<h1><center>Pix2Struct 📄<center><h1>")55  gr.HTML("<h3><center>Pix2Struct is a powerful backbone for visual question answering. ⚡</h3>")56  gr.HTML("<h3><center>Each tab in this app demonstrates Pix2Struct models fine-tuned on document question answering, infographics question answering, question answering on user interfaces, and charts. 📄📱📊<h3>")57  gr.HTML("<h3><center>This app has base versions of each model. For better performance, use large checkpoints.<h3>")58 59  with gr.Tab(label="Visual Question Answering over Documents"):60    with gr.Row():61      with gr.Column():62        input_img = gr.Image(label="Input Document")63        question = gr.Text(label="Question")64        submit_btn = gr.Button(value="Submit")65      output = gr.Text(label="Answer")66    gr.Examples(67    [["docvqa_example.png", "How many items are sold?"]],68    inputs = [input_img, question],69    outputs = [output],70    fn=infer_doc,71    cache_examples=True,72    label='Click on any Examples below to get Document Question Answering results quickly 👇'73    )74 75    submit_btn.click(infer_doc, [input_img, question], [output])76 77  with gr.Tab(label="Visual Question Answering over Infographics"):78    with gr.Row():79      with gr.Column():80        input_img = gr.Image(label="Input Image")81        question = gr.Text(label="Question")82        submit_btn = gr.Button(value="Submit")83      output = gr.Text(label="Answer")84    gr.Examples(85    [["infographics_example.jpeg", "What is this infographic about?"]],86    inputs = [input_img, question],87    outputs = [output],88    fn=infer_doc,89    cache_examples=True,90    label='Click on any Examples below to get Infographics QA results quickly 👇'91    )92 93    submit_btn.click(infer_infographics, [input_img, question], [output])94  with gr.Tab(label="Caption User Interfaces"):95    with gr.Row():96      with gr.Column():97        input_img = gr.Image(label="Input UI Image")98        question = gr.Text(label="Question")99        submit_btn = gr.Button(value="Submit")100      output = gr.Text(label="Caption")101    submit_btn.click(infer_chart, [input_img, question], [output])102    gr.Examples(103    [["screen2words_ui_example.png", "What is this UI about?"]],104    inputs = [input_img, question],105    outputs = [output],106    fn=infer_doc,107    cache_examples=True,108    label='Click on any Examples below to get UI question answering results quickly 👇'109    )110 111  with gr.Tab(label="Ask about Charts"):112    with gr.Row():113      with gr.Column():114        input_img = gr.Image(label="Input Chart")115        question = gr.Text(label="Question")116        submit_btn = gr.Button(value="Submit")117      output = gr.Text(label="Caption")118 119    submit_btn.click(infer_chart, [input_img, question], [output])120    gr.Examples(121    [["chartqa_example.png", "How much percent is bicycle?"]],122    inputs = [input_img, question],123    outputs = [output],124    fn=infer_doc,125    cache_examples=True,126    label='Click on any Examples below to get Chart question answering results quickly 👇'127    )128 129demo.launch(debug=True)