Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import requests | |
| from PIL import Image | |
| from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor | |
| import spaces | |
| def infer_infographics(image, question): | |
| model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-ai2d-base").to("cuda") | |
| processor = Pix2StructProcessor.from_pretrained("google/pix2struct-ai2d-base") | |
| inputs = processor(images=image, text=question, return_tensors="pt").to("cuda") | |
| predictions = model.generate(**inputs) | |
| return processor.decode(predictions[0], skip_special_tokens=True) | |
| def infer_ui(image, question): | |
| model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-base").to("cuda") | |
| processor = Pix2StructProcessor.from_pretrained("google/pix2struct-screen2words-base") | |
| inputs = processor(images=image,text=question, return_tensors="pt").to("cuda") | |
| predictions = model.generate(**inputs) | |
| return processor.decode(predictions[0], skip_special_tokens=True) | |
| def infer_chart(image, question): | |
| model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-chartqa-base").to("cuda") | |
| processor = Pix2StructProcessor.from_pretrained("google/pix2struct-chartqa-base") | |
| inputs = processor(images=image, text=question, return_tensors="pt").to("cuda") | |
| predictions = model.generate(**inputs) | |
| return processor.decode(predictions[0], skip_special_tokens=True) | |
| def infer_doc(image, question): | |
| model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-base").to("cuda") | |
| processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-base") | |
| inputs = processor(images=image, text=question, return_tensors="pt").to("cuda") | |
| predictions = model.generate(**inputs) | |
| return processor.decode(predictions[0], skip_special_tokens=True) | |
| css = """ | |
| #mkd { | |
| height: 500px; | |
| overflow: auto; | |
| border: 1px solid #ccc; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| gr.HTML("<h1><center>Pix2Struct 📄<center><h1>") | |
| gr.HTML("<h3><center>Pix2Struct is a powerful backbone for visual question answering. ⚡</h3>") | |
| 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>") | |
| gr.HTML("<h3><center>This app has base versions of each model. For better performance, use large checkpoints.<h3>") | |
| with gr.Tab(label="Visual Question Answering over Documents"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input Document") | |
| question = gr.Text(label="Question") | |
| submit_btn = gr.Button(value="Submit") | |
| output = gr.Text(label="Answer") | |
| gr.Examples( | |
| [["docvqa_example.png", "How many items are sold?"]], | |
| inputs = [input_img, question], | |
| outputs = [output], | |
| fn=infer_doc, | |
| cache_examples=True, | |
| label='Click on any Examples below to get Document Question Answering results quickly 👇' | |
| ) | |
| submit_btn.click(infer_doc, [input_img, question], [output]) | |
| with gr.Tab(label="Visual Question Answering over Infographics"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input Image") | |
| question = gr.Text(label="Question") | |
| submit_btn = gr.Button(value="Submit") | |
| output = gr.Text(label="Answer") | |
| gr.Examples( | |
| [["infographics_example.jpeg", "What is this infographic about?"]], | |
| inputs = [input_img, question], | |
| outputs = [output], | |
| fn=infer_doc, | |
| cache_examples=True, | |
| label='Click on any Examples below to get Infographics QA results quickly 👇' | |
| ) | |
| submit_btn.click(infer_infographics, [input_img, question], [output]) | |
| with gr.Tab(label="Caption User Interfaces"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input UI Image") | |
| question = gr.Text(label="Question") | |
| submit_btn = gr.Button(value="Submit") | |
| output = gr.Text(label="Caption") | |
| submit_btn.click(infer_chart, [input_img, question], [output]) | |
| gr.Examples( | |
| [["screen2words_ui_example.png", "What is this UI about?"]], | |
| inputs = [input_img, question], | |
| outputs = [output], | |
| fn=infer_doc, | |
| cache_examples=True, | |
| label='Click on any Examples below to get UI question answering results quickly 👇' | |
| ) | |
| with gr.Tab(label="Ask about Charts"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input Chart") | |
| question = gr.Text(label="Question") | |
| submit_btn = gr.Button(value="Submit") | |
| output = gr.Text(label="Caption") | |
| submit_btn.click(infer_chart, [input_img, question], [output]) | |
| gr.Examples( | |
| [["chartqa_example.png", "How much percent is bicycle?"]], | |
| inputs = [input_img, question], | |
| outputs = [output], | |
| fn=infer_doc, | |
| cache_examples=True, | |
| label='Click on any Examples below to get Chart question answering results quickly 👇' | |
| ) | |
| demo.launch(debug=True) |