|
import logging |
|
import time |
|
from pathlib import Path |
|
import contextlib |
|
|
|
logging.basicConfig( |
|
level=logging.INFO, |
|
format="%(asctime)s - %(levelname)s - %(message)s", |
|
) |
|
|
|
|
|
import gradio as gr |
|
import nltk |
|
import torch |
|
|
|
from pdf2text import * |
|
|
|
_here = Path(__file__).parent |
|
|
|
nltk.download("stopwords") |
|
|
|
|
|
def load_uploaded_file(file_obj, temp_dir: Path = None): |
|
""" |
|
load_uploaded_file - process an uploaded file |
|
|
|
Args: |
|
file_obj (POTENTIALLY list): Gradio file object inside a list |
|
|
|
Returns: |
|
str, the uploaded file contents |
|
""" |
|
|
|
|
|
|
|
if isinstance(file_obj, list): |
|
file_obj = file_obj[0] |
|
file_path = Path(file_obj.name) |
|
|
|
if temp_dir is None: |
|
_temp_dir = _here / "temp" |
|
_temp_dir.mkdir(exist_ok=True) |
|
|
|
try: |
|
pdf_bytes_obj = open(file_path, "rb").read() |
|
temp_path = temp_dir / file_path.name if temp_dir else file_path |
|
|
|
with open(temp_path, "wb") as f: |
|
f.write(pdf_bytes_obj) |
|
logging.info(f"Saved uploaded file to {temp_path}") |
|
return str(temp_path.resolve()) |
|
|
|
except Exception as e: |
|
logging.error(f"Trying to load file with path {file_path}, error: {e}") |
|
print(f"Trying to load file with path {file_path}, error: {e}") |
|
return None |
|
|
|
|
|
def convert_PDF(pdf_obj, language: str = "en"): |
|
""" |
|
convert_PDF - convert a PDF file to text |
|
|
|
Args: |
|
pdf_bytes_obj (bytes): PDF file contents |
|
language (str, optional): Language to use for OCR. Defaults to "en". |
|
|
|
Returns: |
|
str, the PDF file contents as text |
|
""" |
|
|
|
global ocr_model |
|
st = time.perf_counter() |
|
|
|
conversion_stats = convert_PDF_to_Text( |
|
pdf_obj, |
|
ocr_model=ocr_model, |
|
max_pages=20, |
|
) |
|
converted_txt = conversion_stats["converted_text"] |
|
num_pages = conversion_stats["num_pages"] |
|
|
|
|
|
rt = round((time.perf_counter() - st) / 60, 2) |
|
print(f"Runtime: {rt} minutes") |
|
html = "" |
|
html += f"<p>Runtime: {rt} minutes on CPU for {num_pages} pages</p>" |
|
|
|
return converted_txt, html |
|
|
|
|
|
if __name__ == "__main__": |
|
logging.info("Starting app") |
|
|
|
use_GPU = torch.cuda.is_available() |
|
logging.info(f"Using GPU status: {use_GPU}") |
|
logging.info("Loading OCR model") |
|
with contextlib.redirect_stdout(None): |
|
ocr_model = ocr_predictor('db_resnet50', 'crnn_mobilenet_v3_large', pretrained=True, assume_straight_pages=True) |
|
|
|
|
|
pdf_obj = _here / "example_file.pdf" |
|
pdf_obj = str(pdf_obj.resolve()) |
|
_temp_dir = _here / "temp" |
|
_temp_dir.mkdir(exist_ok=True) |
|
|
|
logging.info("starting demo") |
|
demo = gr.Blocks() |
|
|
|
with demo: |
|
|
|
gr.Markdown("# PDF to Text") |
|
gr.Markdown("**Upload a PDF file to convert to text**") |
|
gr.Markdown("If no file is uploaded, a sample PDF will be used") |
|
|
|
with gr.Column(): |
|
|
|
gr.Markdown("## Load Inputs") |
|
gr.Markdown("Upload your own file:") |
|
pdf_obj = gr.Textbox( |
|
lines=1, |
|
label="VM file path", |
|
placeholder="When the file is uploaded, the path will appear here", |
|
value=pdf_obj, |
|
) |
|
with gr.Row(): |
|
uploaded_file = gr.File( |
|
label="Upload a PDF file", |
|
file_count="single", |
|
type="file", |
|
) |
|
load_file_button = gr.Button("Load Uploaded File") |
|
|
|
gr.Markdown("---") |
|
|
|
with gr.Column(): |
|
gr.Markdown("## Convert PDF to Text") |
|
convert_button = gr.Button("Convert PDF!") |
|
out_placeholder = gr.HTML("<p><em>Output will appear below:</em></p>") |
|
gr.Markdown("### Output") |
|
OCR_text = gr.Textbox( |
|
label="OCR Result", placeholder="The OCR text will appear here" |
|
) |
|
|
|
load_file_button.click( |
|
fn=load_uploaded_file, inputs=uploaded_file, outputs=[pdf_obj] |
|
) |
|
|
|
convert_button.click( |
|
fn=convert_PDF, inputs=[pdf_obj], outputs=[OCR_text, out_placeholder] |
|
) |
|
demo.launch(enable_queue=True) |
|
|