Merge pull request #138 from VikParuchuri/dev

Faster text extraction

marker/convert.py

@@ -64,6 +64,7 @@ def convert_single_pdf(
     doc = pdfium.PdfDocument(fname)
     pages, toc = get_text_blocks(
         doc,
+        fname,
         max_pages=max_pages,
     )
     out_meta.update({

marker/ocr/recognition.py

@@ -1,3 +1,4 @@
+import tempfile
 from itertools import repeat
 from typing import List, Optional, Dict
 
@@ -160,9 +161,12 @@ def _tesseract_recognition(in_pdf, langs: List[str]) -> Optional[Page]:
         tesseract_non_ocr_timeout=settings.TESSERACT_TIMEOUT,
     )
 
-    new_doc = pdfium.PdfDocument(out_pdf.getvalue())
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(out_pdf.getvalue())
+        f.seek(0)
+        new_doc = pdfium.PdfDocument(f.name)
+        blocks, _ = get_text_blocks(new_doc, f.name, max_pages=1)
 
-    blocks, _ = get_text_blocks(new_doc, max_pages=1)
     page = blocks[0]
     page.ocr_method = "tesseract"
     return page

marker/pdf/extract_text.py

@@ -74,15 +74,15 @@ def pdftext_format_to_blocks(page, pnum: int) -> Page:
     return out_page
 
 
-def get_text_blocks(doc, max_pages: Optional[int] = None) -> (List[Page], Dict):
+def get_text_blocks(doc, fname, max_pages: Optional[int] = None) -> (List[Page], Dict):
     toc = get_toc(doc)
 
     page_range = range(len(doc))
     if max_pages:
         range_end = min(max_pages, len(doc))
         page_range = range(range_end)
 
-    char_blocks = dictionary_output(doc, page_range=page_range, keep_chars=True)
+    char_blocks = dictionary_output(fname, page_range=page_range, keep_chars=True, workers=settings.PDFTEXT_CPU_WORKERS)
     marker_blocks = [pdftext_format_to_blocks(page, pnum) for pnum, page in enumerate(char_blocks)]
 
     return marker_blocks, toc

marker/settings.py

@@ -34,6 +34,9 @@ def TORCH_DEVICE_MODEL(self) -> str:
         "application/pdf": "pdf",
     }
 
+    # Text extraction
+    PDFTEXT_CPU_WORKERS: int = 4 # How many CPU workers to use for pdf text extraction
+
     # Text line Detection
     DETECTOR_BATCH_SIZE: Optional[int] = None # Defaults to 6 for CPU, 12 otherwise
     SURYA_DETECTOR_DPI: int = 96