Upload 4 files

app.py

@@ -0,0 +1,182 @@
+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")  # 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,
+        )
+
+    # define pdf bytes as None
+    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(
+            "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="file",
+                value=_here / "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="file",
+                interactive=False,
+            )
+
+        convert_button.click(
+            fn=convert_PDF,
+            inputs=[uploaded_file],
+            outputs=[OCR_text, out_placeholder, text_file],
+        )
+    demo.launch(enable_queue=True)

pdf2text.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """
+
 easyocr.py - A wrapper for easyocr to convert pdf to images to text
 """
 
@@ -40,6 +41,7 @@ def simple_rename(filepath, target_ext=".txt"):
 def rm_local_text_files(name_contains="RESULT_"):
     """
     rm_local_text_files - remove local text files
+
     Args:
         name_contains (str, optional): [description]. Defaults to "OCR_".
     """
@@ -60,10 +62,12 @@ def corr(
     exceptions=["e.g.", "i.e.", "etc.", "cf.", "vs.", "p."],
 ) -> str:
     """corrects spacing in a string
+
     Args:
         s (str): the string to correct
         add_space_when_numerics (bool, optional): [add a space when a period is between two numbers, example 5.73]. Defaults to False.
         exceptions (list, optional): [do not change these substrings]. Defaults to ['e.g.', 'i.e.', 'etc.', 'cf.', 'vs.', 'p.'].
+
     Returns:
         str: the corrected string
     """
@@ -90,9 +94,11 @@ def corr(
 def fix_punct_spaces(string):
     """
     fix_punct_spaces - replace spaces around punctuation with punctuation. For example, "hello , there" -> "hello, there"
+
     Parameters
     ----------
     string : str, required, input string to be corrected
+
     Returns
     -------
     str, corrected string
@@ -108,9 +114,11 @@ def fix_punct_spaces(string):
 def clean_OCR(ugly_text: str):
     """
     clean_OCR - clean the OCR text files.
+
     Parameters
     ----------
     ugly_text : str, required, input string to be cleaned
+
     Returns
     -------
     str, cleaned string
@@ -154,6 +162,7 @@ def move2completed(from_dir, filename, new_folder="completed", verbose=False):
 
 
 """## pdf2text functions
+
 """
 
 
@@ -179,8 +188,10 @@ spell = SpellChecker()
 def check_word_spelling(word: str) -> bool:
     """
     check_word_spelling - check the spelling of a word
+
     Args:
         word (str): word to check
+
     Returns:
         bool: True if word is spelled correctly, False if not
     """
@@ -193,9 +204,11 @@ def check_word_spelling(word: str) -> bool:
 def eval_and_replace(text: str, match_token: str = "- ") -> str:
     """
     eval_and_replace  - conditionally replace all instances of a substring in a string based on whether the eliminated substring results in a valid word
+
     Args:
         text (str): text to evaluate
         match_token (str, optional): token to replace. Defaults to "- ".
+
     Returns:
         str:  text with replaced tokens
     """
@@ -228,10 +241,12 @@ def eval_and_replace(text: str, match_token: str = "- ") -> str:
 def cleantxt_ocr(ugly_text, lower=False, lang: str = "en") -> str:
     """
     cleantxt_ocr - clean text from OCR
+
     Args:
         ugly_text (str): text to clean
         lower (bool, optional): _description_. Defaults to False.
         lang (str, optional): _description_. Defaults to "en".
+
     Returns:
         str: cleaned text
     """
@@ -363,6 +378,7 @@ def translate_text(text, source_l, target_l="en"):
 
 def translate_doc(filepath, lang_start, lang_end="en", verbose=False):
     """translate a document from lang_start to lang_end
+
         {'code': 'en', 'name': 'English'},
     {'code': 'fr', 'name': 'French'},
     {'code': 'de', 'name': 'German'},
@@ -387,4 +403,4 @@ def translate_doc(filepath, lang_start, lang_end="en", verbose=False):
         f_o.writelines(translated_doc)
     if verbose:
         print("finished translating the document! - ", datetime.now())
-    return out_path
+    return out_path

requirements.txt

@@ -1,6 +1,6 @@
 clean-text[gpl]
 python-doctr[torch]
-gradio
+gradio==3.4.0
 libretranslatepy
 natsort
 nltk