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