Spaces:
Runtime error
Runtime error
File size: 5,348 Bytes
90e70a5 7992cef 90e70a5 9b3a39c 90e70a5 7992cef 90e70a5 7992cef 90e70a5 3b3978c |
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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
import logging
import time
from pathlib import Path
import contextlib
import gradio as gr
import nltk
import torch
from pdf2text import *
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
)
_here = Path(__file__).parent
nltk.download("stopwords") # TODO=find where this requirement originates from
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
"""
# check if mysterious file object is a list
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
# save to PDF file
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",
max_pages=20,
):
"""
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
"""
# clear local text cache
rm_local_text_files()
global ocr_model
st = time.perf_counter()
if isinstance(pdf_obj, list):
pdf_obj = pdf_obj[0]
file_path = Path(pdf_obj.name)
if not file_path.suffix == ".pdf":
logging.error(f"File {file_path} is not a PDF file")
html_error = f"""
<div style="color: red; font-size: 20px; font-weight: bold;">
File {file_path} is not a PDF file. Please upload a PDF file.
</div>
"""
return "File is not a PDF file", html_error, None
conversion_stats = convert_PDF_to_Text(
file_path,
ocr_model=ocr_model,
max_pages=max_pages,
)
converted_txt = conversion_stats["converted_text"]
num_pages = conversion_stats["num_pages"]
was_truncated = conversion_stats["truncated"]
# if alt_lang: # TODO: fix this
rt = round((time.perf_counter() - st) / 60, 2)
print(f"Runtime: {rt} minutes")
html = ""
if was_truncated:
html += f"<p>WARNING - PDF was truncated to {max_pages} pages</p>"
html += f"<p>Runtime: {rt} minutes on CPU for {num_pages} pages</p>"
_output_name = f"RESULT_{file_path.stem}_OCR.txt"
with open(_output_name, "w", encoding="utf-8", errors="ignore") as f:
f.write(converted_txt)
return converted_txt, html, _output_name
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 / "try_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(
"A basic demo of pdf-to-text conversion using OCR from the [doctr](https://mindee.github.io/doctr/index.html) package"
)
gr.Markdown("---")
with gr.Column():
gr.Markdown("## Load Inputs")
gr.Markdown("Upload your own file & replace the default. Files should be < 10MB to avoid upload issues - search for a PDF compressor online as needed.")
gr.Markdown(
"_If no file is uploaded, a sample PDF will be used. PDFs are truncated to 20 pages._"
)
uploaded_file = gr.File(
label="Upload a PDF file",
file_count="single",
type="filepath",
value=str(_here / "try_example_file.pdf"),
)
gr.Markdown("---")
with gr.Column():
gr.Markdown("## Convert PDF to Text")
convert_button = gr.Button("Convert PDF!", variant="primary")
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"
)
text_file = gr.File(
label="Download Text File",
file_count="single",
type="filepath",
interactive=False,
)
convert_button.click(
fn=convert_PDF,
inputs=[uploaded_file],
outputs=[OCR_text, out_placeholder, text_file],
)
demo.launch()
|