Andrei Shadrikov commited on
Commit
77eac00
·
1 Parent(s): f875f63
Files changed (2) hide show
  1. app.py +92 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from time import time
3
+ from pathlib import Path
4
+ from gradio_pdf import PDF
5
+ from pdf2image import convert_from_path
6
+ import shutil
7
+ import tempfile
8
+ from transformers import pipeline
9
+
10
+ out_files = gr.State([])
11
+ FILE_TIMEOUT = 10 ** 0
12
+ MAX_FILES = 10
13
+
14
+ p = pipeline(
15
+ "document-question-answering",
16
+ model="impira/layoutlm-document-qa",
17
+ )
18
+
19
+ def handle_files(cur_files):
20
+ cur_time = cur_files[-1][0]
21
+ deleted_indices = set()
22
+ for other_idx, (other_time, other_file) in enumerate(cur_files[:-1]):
23
+ if abs(cur_time - other_time) > FILE_TIMEOUT:
24
+ shutil.rmtree(other_file.parent)
25
+ deleted_indices.add(other_idx)
26
+ cur_files = [cur_files[idx] for idx in range(len(cur_files)) if idx not in deleted_indices]
27
+
28
+ if len(cur_files) > MAX_FILES:
29
+ for _, other_file in cur_files[:-MAX_FILES]:
30
+ shutil.rmtree(other_file.parent)
31
+ cur_files = cur_files[-MAX_FILES:]
32
+ return cur_files
33
+
34
+ # Function to process PDF and generate ZIP file
35
+ def process_pdf(pdf_file, cur_files):
36
+
37
+ zip_output = Path(tempfile.mkdtemp()) / f'{Path(pdf_file).stem}'
38
+ # zip_output.parent.mkdir()
39
+
40
+ with tempfile.TemporaryDirectory() as path:
41
+ pdf_output = path
42
+ convert_from_path(pdf_file, output_folder=str(pdf_output))
43
+
44
+ # Create a BytesIO object to store zip file in memory
45
+ shutil.make_archive(zip_output, 'zip', pdf_output)
46
+
47
+ zip_output = zip_output.with_suffix('.zip')
48
+
49
+ cur_time = time()
50
+ cur_files.append((cur_time, zip_output))
51
+ cur_files = handle_files(cur_files)
52
+
53
+ return str(zip_output), cur_files
54
+
55
+ # Function to process image and text query
56
+ def process_image_and_text(doc, text_query):
57
+ images = convert_from_path(doc)
58
+ outputs = []
59
+ for img in images:
60
+ outputs += p(img, question)
61
+ return sorted(outputs, key=lambda x: x["score"], reverse=True)[0]['answer']
62
+
63
+ # Interface for the Gradio app
64
+ pdf_interface = gr.Interface(
65
+ fn=process_pdf,
66
+ inputs=[PDF(label="Upload PDF"), out_files],
67
+ outputs=[gr.File(label="Download ZIP"), out_files],
68
+ title="PDF to Image Converter",
69
+ description="Converts PDF pages to images and outputs a ZIP file."
70
+ )
71
+
72
+ image_interface = gr.Interface(
73
+ fn=process_image_and_text,
74
+ inputs=[
75
+ gr.Image(label="Upload Image"),
76
+ gr.Textbox(label="Text Query")
77
+ ],
78
+ outputs="text",
79
+ title="Image Text Search",
80
+ description="Searches for text in the uploaded image based on the provided query."
81
+ )
82
+
83
+ # Create a tabbed interface
84
+ tabbed_interface = gr.TabbedInterface(
85
+ [pdf_interface, image_interface],
86
+ title="PDF interaction",
87
+ tab_names=["Converter", "Interaction"],
88
+ # description="Choose a tab to perform the desired task."
89
+ )
90
+
91
+ # Launch the app
92
+ tabbed_interface.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ gradio
2
+ gradio_pdf
3
+ transformers