Chetan Khadke commited on
Commit
20fdcac
1 Parent(s): bc65916

added all files

Browse files
Files changed (4) hide show
  1. app.py +42 -0
  2. example_1.png +0 -0
  3. example_2.jpeg +0 -0
  4. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import gradio as gr
3
+
4
+ import torch
5
+ from functools import partial
6
+ from PIL import Image
7
+ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
8
+
9
+
10
+ model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-large")
11
+ processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-large")
12
+
13
+
14
+ device = "cuda" if torch.cuda.is_available() else "cpu"
15
+ model.to(device)
16
+
17
+ def generate(img, questions):
18
+ global model, processor
19
+ inputs = processor(images=[img for _ in range(len(questions))], text=questions, return_tensors="pt").to(device)
20
+ predictions = model.generate(**inputs, max_new_tokens=256)
21
+ return zip(questions, processor.batch_decode(predictions, skip_special_tokens=True))
22
+
23
+
24
+ def process_document(image, question):
25
+ return generate(image, [question])
26
+
27
+
28
+ description = "Gradio Demo for Pix2Struct, an instance of `VisionEncoderDecoderModel` fine-tuned on DocVQA (document visual question answering). To use it, simply upload your image and type a question and click 'submit', or click one of the examples to load them. Read more at the links below."
29
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.15664' target='_blank'>Pix2Struct for DocVQA</a> | <a href='https://arxiv.org/pdf/2210.03347.pdf' target='_blank'>Paper link</a></p>"
30
+
31
+ demo = gr.Interface(
32
+ fn=process_document,
33
+ inputs=["image", "text"],
34
+ outputs="json",
35
+ title="Demo: Pix2Struct for DocVQA",
36
+ description=description,
37
+ article=article,
38
+ enable_queue=True,
39
+ examples=[["example_1.png", "When is the coffee break?"], ["example_2.jpeg", "What's the population of Stoddard?"]],
40
+ cache_examples=False)
41
+
42
+ demo.launch()
example_1.png ADDED
example_2.jpeg ADDED
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch
2
+ git+https://github.com/huggingface/transformers.git
3
+ sentencepiece