|
import os |
|
os.system('git clone https://github.com/facebookresearch/detectron2.git') |
|
os.system('pip install -e detectron2') |
|
os.system("git clone https://github.com/microsoft/unilm.git") |
|
os.system("sed -i 's/from collections import Iterable/from collections.abc import Iterable/' unilm/dit/object_detection/ditod/table_evaluation/data_structure.py") |
|
os.system("curl -LJ -o publaynet_dit-b_cascade.pth 'https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_cascade.pth?sv=2022-11-02&ss=b&srt=o&sp=r&se=2033-06-08T16:48:15Z&st=2023-06-08T08:48:15Z&spr=https&sig=a9VXrihTzbWyVfaIDlIT1Z0FoR1073VB0RLQUMuudD4%3D'") |
|
|
|
import sys |
|
sys.path.append("unilm") |
|
sys.path.append("detectron2") |
|
|
|
import uuid |
|
import torch |
|
import gradio as gr |
|
from pdfextract_fun import * |
|
from pdfsummary_fun import * |
|
from imagesummary_fun import * |
|
|
|
|
|
|
|
|
|
def process_pdf(pdf_file,state): |
|
|
|
unique_id = str(uuid.uuid4()) |
|
|
|
output_folder = os.path.join("processed_files", unique_id) |
|
|
|
if not os.path.exists(output_folder): |
|
os.makedirs(output_folder) |
|
|
|
|
|
convert_pdf_to_jpg(pdf_file.name, output_folder) |
|
|
|
process_jpeg_images(output_folder) |
|
|
|
rename_files_sequentially(output_folder) |
|
ocr_folder(output_folder) |
|
image_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) |
|
if f.endswith('.jpg') and ('figure' in f or 'table' in f)] |
|
|
|
|
|
images = [(Image.open(f), os.path.basename(f).split('.')[0]) for f in image_files] |
|
|
|
|
|
|
|
return images, output_folder |
|
|
|
def call_pdf_summary(state): |
|
ocr_results_folder = os.path.join(state, "ocr_results") |
|
summary = pdf_summary(ocr_results_folder) |
|
return summary |
|
|
|
|
|
def handle_summary_button_click(selected_images): |
|
|
|
summary = get_image_summary(selected_images) |
|
return summary |
|
|
|
with gr.Blocks(theme=gr.themes.Monochrome()) as app: |
|
gr.Markdown("# ChatPaper!") |
|
state = gr.State() |
|
|
|
with gr.Row(): |
|
file_input = gr.File(type="filepath", label="Upload a PDF") |
|
|
|
with gr.Row(): |
|
gallery_output = gr.Gallery(label="Extracted Figures and Tables", show_label=True,columns=[3], rows=[1], object_fit="contain", height="auto") |
|
with gr.Column(): |
|
summary_output = gr.Textbox(label="PDF Summary") |
|
summary_button = gr.Button("Generate Summary") |
|
with gr.Row(): |
|
|
|
image_input = gr.Image(label="Select an Figure or Table for analysis",type='filepath',show_label=True, height="auto") |
|
with gr.Column(): |
|
image_summary_output = gr.Textbox(label="Figure or Table analysis") |
|
image_summary_button = gr.Button("Generate Figure or Table analysis") |
|
|
|
|
|
|
|
file_input.change(process_pdf, inputs=[file_input, state], outputs=[gallery_output, state]) |
|
summary_button.click(call_pdf_summary, inputs=[state], outputs=[summary_output]) |
|
image_summary_button.click(handle_summary_button_click, inputs=image_input, outputs=image_summary_output) |
|
|
|
|
|
app.launch() |