rishabh062 commited on
Commit
0c825a3
1 Parent(s): 91ce4be

Provided download button

Browse files
Files changed (1) hide show
  1. app.py +19 -4
app.py CHANGED
@@ -1,5 +1,6 @@
1
- import re
2
  import gradio as gr
 
 
3
 
4
  import torch
5
  from transformers import DonutProcessor, VisionEncoderDecoderModel
@@ -36,12 +37,22 @@ def process_document(image):
36
  sequence = processor.batch_decode(outputs.sequences)[0]
37
  sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
38
  sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() # remove first task start token
 
 
 
39
 
40
- return processor.token2json(sequence)
 
 
 
 
 
41
 
42
  description = "To use it, simply upload your image and click 'submit', or click one of the examples to load them. Read more at the links below."
43
  article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.15664' target='_blank'>Donut: OCR-free Document Understanding Transformer</a> | <a href='https://github.com/clovaai/donut' target='_blank'>Github Repo</a></p>"
44
 
 
 
45
  demo = gr.Interface(
46
  fn=process_document,
47
  inputs="image",
@@ -51,6 +62,10 @@ demo = gr.Interface(
51
  article=article,
52
  enable_queue=True,
53
  examples=[["example.png"], ["example_2.png"], ["example_3.png"]],
54
- cache_examples=False)
 
 
 
 
55
 
56
- demo.launch()
 
 
1
  import gradio as gr
2
+ import csv
3
+ import json
4
 
5
  import torch
6
  from transformers import DonutProcessor, VisionEncoderDecoderModel
 
37
  sequence = processor.batch_decode(outputs.sequences)[0]
38
  sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
39
  sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() # remove first task start token
40
+ with open('output.csv', 'a', newline='') as file:
41
+ writer = csv.writer(file)
42
+ writer.writerow([image_filename, sequence])
43
 
44
+ return {"json": processor.token2json(sequence)}
45
+
46
+ def download_csv(output_csv):
47
+ with open(output_csv) as f:
48
+ response = f.read()
49
+ return response, {"Content-Type": "text/csv"}
50
 
51
  description = "To use it, simply upload your image and click 'submit', or click one of the examples to load them. Read more at the links below."
52
  article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.15664' target='_blank'>Donut: OCR-free Document Understanding Transformer</a> | <a href='https://github.com/clovaai/donut' target='_blank'>Github Repo</a></p>"
53
 
54
+ output_csv = "output.csv"
55
+
56
  demo = gr.Interface(
57
  fn=process_document,
58
  inputs="image",
 
62
  article=article,
63
  enable_queue=True,
64
  examples=[["example.png"], ["example_2.png"], ["example_3.png"]],
65
+ cache_examples=False,
66
+ allow_download=True,
67
+ download_name="output.csv",
68
+ download=download_csv
69
+ )
70
 
71
+ demo.launch()