visualizingjp/pdf-ocr

1import logging2import time3from pathlib import Path4import contextlib5 6logging.basicConfig(7    level=logging.INFO,8    format="%(asctime)s - %(levelname)s - %(message)s",9)10 11 12import gradio as gr13import nltk14import torch15 16from pdf2text import *17 18_here = Path(__file__).parent19 20nltk.download("stopwords")  # TODO=find where this requirement originates from21 22 23def load_uploaded_file(file_obj, temp_dir: Path = None):24    """25    load_uploaded_file - process an uploaded file26 27    Args:28        file_obj (POTENTIALLY list): Gradio file object inside a list29 30    Returns:31        str, the uploaded file contents32    """33 34    # check if mysterious file object is a list35    if isinstance(file_obj, list):36        file_obj = file_obj[0]37    file_path = Path(file_obj.name)38 39    if temp_dir is None:40        _temp_dir = _here / "temp"41    _temp_dir.mkdir(exist_ok=True)42 43    try:44        pdf_bytes_obj = open(file_path, "rb").read()45        temp_path = temp_dir / file_path.name if temp_dir else file_path46        # save to PDF file47        with open(temp_path, "wb") as f:48            f.write(pdf_bytes_obj)49        logging.info(f"Saved uploaded file to {temp_path}")50        return str(temp_path.resolve())51 52    except Exception as e:53        logging.error(f"Trying to load file with path {file_path}, error: {e}")54        print(f"Trying to load file with path {file_path}, error: {e}")55        return None56 57 58def convert_PDF(59    pdf_obj,60    language: str = "en",61    max_pages=20,62):63    """64    convert_PDF - convert a PDF file to text65 66    Args:67        pdf_bytes_obj (bytes): PDF file contents68        language (str, optional): Language to use for OCR. Defaults to "en".69 70    Returns:71        str, the PDF file contents as text72    """73    # clear local text cache74    rm_local_text_files()75    global ocr_model76    st = time.perf_counter()77    if isinstance(pdf_obj, list):78        pdf_obj = pdf_obj[0]79    file_path = Path(pdf_obj.name)80    if not file_path.suffix == ".pdf":81        logging.error(f"File {file_path} is not a PDF file")82 83        html_error = f"""84        <div style="color: red; font-size: 20px; font-weight: bold;">85        File {file_path} is not a PDF file. Please upload a PDF file.86        </div>87        """88        return "File is not a PDF file", html_error, None89 90    conversion_stats = convert_PDF_to_Text(91        file_path,92        ocr_model=ocr_model,93        max_pages=max_pages,94    )95    converted_txt = conversion_stats["converted_text"]96    num_pages = conversion_stats["num_pages"]97    was_truncated = conversion_stats["truncated"]98    # if alt_lang: # TODO: fix this99 100    rt = round((time.perf_counter() - st) / 60, 2)101    print(f"Runtime: {rt} minutes")102    html = ""103    if was_truncated:104        html += f"<p>WARNING - PDF was truncated to {max_pages} pages</p>"105    html += f"<p>Runtime: {rt} minutes on CPU for {num_pages} pages</p>"106 107    _output_name = f"RESULT_{file_path.stem}_OCR.txt"108    with open(_output_name, "w", encoding="utf-8", errors="ignore") as f:109        f.write(converted_txt)110 111    return converted_txt, html, _output_name112 113 114if __name__ == "__main__":115    logging.info("Starting app")116 117    use_GPU = torch.cuda.is_available()118    logging.info(f"Using GPU status: {use_GPU}")119    logging.info("Loading OCR model")120    with contextlib.redirect_stdout(None):121        ocr_model = ocr_predictor(122            "db_resnet50",123            "crnn_mobilenet_v3_large",124            pretrained=True,125            assume_straight_pages=True,126        )127 128    # define pdf bytes as None129    pdf_obj = _here / "example_file.pdf"130    pdf_obj = str(pdf_obj.resolve())131    _temp_dir = _here / "temp"132    _temp_dir.mkdir(exist_ok=True)133 134    logging.info("starting demo")135    demo = gr.Blocks()136 137    with demo:138 139        gr.Markdown("# PDF to Text")140        gr.Markdown(141            "A basic demo of pdf-to-text conversion using OCR from the [doctr](https://mindee.github.io/doctr/index.html) package"142        )143        gr.Markdown("---")144 145        with gr.Column():146 147            gr.Markdown("## Load Inputs")148            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.")149            gr.Markdown(150                "_If no file is uploaded, a sample PDF will be used. PDFs are truncated to 20 pages._"151            )152 153            uploaded_file = gr.File(154                label="Upload a PDF file",155                file_count="single",156                type="file",157                value=_here / "example_file.pdf",158            )159 160            gr.Markdown("---")161 162        with gr.Column():163            gr.Markdown("## Convert PDF to Text")164            convert_button = gr.Button("Convert PDF!", variant="primary")165            out_placeholder = gr.HTML("<p><em>Output will appear below:</em></p>")166            gr.Markdown("### Output")167            OCR_text = gr.Textbox(168                label="OCR Result", placeholder="The OCR text will appear here"169            )170            text_file = gr.File(171                label="Download Text File",172                file_count="single",173                type="file",174                interactive=False,175            )176 177        convert_button.click(178            fn=convert_PDF,179            inputs=[uploaded_file],180            outputs=[OCR_text, out_placeholder, text_file],181        )182    demo.launch(enable_queue=True)183