pdftext 0.3.0__tar.gz → 0.3.2__tar.gz

{pdftext-0.3.0 → pdftext-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pdftext
-Version: 0.3.0
+Version: 0.3.2
 Summary: Extract structured text from pdfs quickly
 Home-page: https://github.com/VikParuchuri/pdftext
 License: Apache-2.0
@@ -49,25 +49,27 @@ pdftext PDF_PATH --out_path output.txt
 - `--out_path` path to the output txt file.  If not specified, will write to stdout.
 - `--sort` will attempt to sort in reading order if specified.
 - `--keep_hyphens` will keep hyphens in the output (they will be stripped and words joined otherwise)
+- `--pages` will specify pages (comma separated) to extract
 
 ## JSON
 
 This command outputs structured blocks and lines with font and other information.
 
 ```shell
-pdftext PDF_PATH --out_path output.txt --output_type json
+pdftext PDF_PATH --out_path output.txt --json
 ```
 
 - `PDF_PATH` must be a single pdf file.
 - `--out_path` path to the output txt file.  If not specified, will write to stdout.
-- `--output_type` specifies whether to write out plain text (default) or json
+- `--json` specifies json output
 - `--sort` will attempt to sort in reading order if specified.
+- `--pages` will specify pages (comma separated) to extract
 
 The output will be a json list, with each item in the list corresponding to a single page in the input pdf (in order).  Each page will include the following keys:
 
 - `bbox` - the page bbox, in `[x1, y1, x2, y2]` format
 - `rotation` - how much the page is rotated, in degrees (`0`, `90`, `180`, or `270`)
-- `page_idx` - the index of the page
+- `page` - the index of the page
 - `blocks` - the blocks that make up the text in the pdf.  Approximately equal to a paragraph.
   - `bbox` - the block bbox, in `[x1, y1, x2, y2]` format
   - `lines` - the lines inside the block
@@ -91,17 +93,21 @@ If the pdf is rotated, the bboxes will be relative to the rotated page (they're
 Extract plain text:
 
 ```python
+import pypdfium2 as pdfium
 from pdftext.extraction import plain_text_output
 
-text = plain_text_output(PDF_PATH, sort=False)
+pdf = pdfium.PdfDocument(PDF_PATH)
+text = plain_text_output(pdf, sort=False, hyphens=False, page_range=[1,2,3]) # Optional arguments explained above
 ```
 
 Extract structured blocks and lines:
 
 ```python
+import pypdfium2 as pdfium
 from pdftext.extraction import dictionary_output
 
-text = dictionary_output(PDF_PATH)
+pdf = pdfium.PdfDocument(PDF_PATH)
+text = dictionary_output(pdf, sort=False, page_range=[1,2,3]) # Optional arguments explained above
 ```
 
 If you want more customization, check out the `pdftext.extraction._get_pages` function for a starting point to dig deeper.  pdftext is a pretty thin wrapper around [pypdfium2](https://pypdfium2.readthedocs.io/en/stable/), so you might want to look at the documentation for that as well.

{pdftext-0.3.0 → pdftext-0.3.2}/README.md

@@ -26,25 +26,27 @@ pdftext PDF_PATH --out_path output.txt
 - `--out_path` path to the output txt file.  If not specified, will write to stdout.
 - `--sort` will attempt to sort in reading order if specified.
 - `--keep_hyphens` will keep hyphens in the output (they will be stripped and words joined otherwise)
+- `--pages` will specify pages (comma separated) to extract
 
 ## JSON
 
 This command outputs structured blocks and lines with font and other information.
 
 ```shell
-pdftext PDF_PATH --out_path output.txt --output_type json
+pdftext PDF_PATH --out_path output.txt --json
 ```
 
 - `PDF_PATH` must be a single pdf file.
 - `--out_path` path to the output txt file.  If not specified, will write to stdout.
-- `--output_type` specifies whether to write out plain text (default) or json
+- `--json` specifies json output
 - `--sort` will attempt to sort in reading order if specified.
+- `--pages` will specify pages (comma separated) to extract
 
 The output will be a json list, with each item in the list corresponding to a single page in the input pdf (in order).  Each page will include the following keys:
 
 - `bbox` - the page bbox, in `[x1, y1, x2, y2]` format
 - `rotation` - how much the page is rotated, in degrees (`0`, `90`, `180`, or `270`)
-- `page_idx` - the index of the page
+- `page` - the index of the page
 - `blocks` - the blocks that make up the text in the pdf.  Approximately equal to a paragraph.
   - `bbox` - the block bbox, in `[x1, y1, x2, y2]` format
   - `lines` - the lines inside the block
@@ -68,17 +70,21 @@ If the pdf is rotated, the bboxes will be relative to the rotated page (they're
 Extract plain text:
 
 ```python
+import pypdfium2 as pdfium
 from pdftext.extraction import plain_text_output
 
-text = plain_text_output(PDF_PATH, sort=False)
+pdf = pdfium.PdfDocument(PDF_PATH)
+text = plain_text_output(pdf, sort=False, hyphens=False, page_range=[1,2,3]) # Optional arguments explained above
 ```
 
 Extract structured blocks and lines:
 
 ```python
+import pypdfium2 as pdfium
 from pdftext.extraction import dictionary_output
 
-text = dictionary_output(PDF_PATH)
+pdf = pdfium.PdfDocument(PDF_PATH)
+text = dictionary_output(pdf, sort=False, page_range=[1,2,3]) # Optional arguments explained above
 ```
 
 If you want more customization, check out the `pdftext.extraction._get_pages` function for a starting point to dig deeper.  pdftext is a pretty thin wrapper around [pypdfium2](https://pypdfium2.readthedocs.io/en/stable/), so you might want to look at the documentation for that as well.

{pdftext-0.3.0 → pdftext-0.3.2}/extract_text.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+import pypdfium2 as pdfium
 
 from pdftext.extraction import plain_text_output, dictionary_output
 
@@ -8,18 +9,23 @@ def main():
     parser = argparse.ArgumentParser(description="Extract plain text from PDF.  Not guaranteed to be in order.")
     parser.add_argument("pdf_path", type=str, help="Path to the PDF file")
     parser.add_argument("--out_path", type=str, help="Path to the output text file, defaults to stdout", default=None)
-    parser.add_argument("--output_type", type=str, help="Type of output to generate", default="plain_text")
+    parser.add_argument("--json", action="store_true", help="Output json instead of plain text", default=False)
     parser.add_argument("--sort", action="store_true", help="Attempt to sort the text by reading order", default=False)
     parser.add_argument("--keep_hyphens", action="store_true", help="Keep hyphens in words", default=False)
+    parser.add_argument("--pages", type=str, help="Comma separated pages to extract, like 1,2,3", default=None)
     args = parser.parse_args()
 
-    assert args.output_type in ["plain_text", "json"], "Invalid output type, must be 'plain_text' or 'json'"
+    pdf_doc = pdfium.PdfDocument(args.pdf_path)
+    pages = None
+    if args.pages is not None:
+        pages = [int(p) for p in args.pages.split(",")]
+        assert all(p <= len(pdf_doc) for p in pages), "Invalid page number(s) provided"
 
-    if args.output_type == "plain_text":
-        text = plain_text_output(args.pdf_path, sort=args.sort, hyphens=args.keep_hyphens)
-    elif args.output_type == "json":
-        text = dictionary_output(args.pdf_path, sort=args.sort)
+    if args.json:
+        text = dictionary_output(pdf_doc, sort=args.sort, page_range=pages)
         text = json.dumps(text)
+    else:
+        text = plain_text_output(pdf_doc, sort=args.sort, hyphens=args.keep_hyphens, page_range=pages)
 
     if args.out_path is None:
         print(text)

{pdftext-0.3.0 → pdftext-0.3.2}/pdftext/extraction.py

@@ -7,29 +7,29 @@ from pdftext.pdf.utils import unnormalize_bbox
 from pdftext.postprocessing import merge_text, sort_blocks, postprocess_text, handle_hyphens
 
 
-def _get_pages(pdf_path, model=None):
+def _get_pages(pdf_doc, model=None, page_range=None):
     if model is None:
         model = get_model()
-    text_chars = get_pdfium_chars(pdf_path)
+    text_chars = get_pdfium_chars(pdf_doc, page_range=page_range)
     pages = inference(text_chars, model)
     return pages
 
 
-def plain_text_output(pdf_path, sort=False, model=None, hyphens=False) -> str:
-    text = paginated_plain_text_output(pdf_path, sort=sort, model=model, hyphens=hyphens)
+def plain_text_output(pdf_doc, sort=False, model=None, hyphens=False, page_range=None) -> str:
+    text = paginated_plain_text_output(pdf_doc, sort=sort, model=model, hyphens=hyphens, page_range=page_range)
     return "\n".join(text)
 
 
-def paginated_plain_text_output(pdf_path, sort=False, model=None, hyphens=False) -> List[str]:
-    pages = _get_pages(pdf_path, model)
+def paginated_plain_text_output(pdf_doc, sort=False, model=None, hyphens=False, page_range=None) -> List[str]:
+    pages = _get_pages(pdf_doc, model, page_range)
     text = []
     for page in pages:
         text.append(merge_text(page, sort=sort, hyphens=hyphens).strip())
     return text
 
 
-def dictionary_output(pdf_path, sort=False, model=None):
-    pages = _get_pages(pdf_path, model)
+def dictionary_output(pdf_doc, sort=False, model=None, page_range=None):
+    pages = _get_pages(pdf_doc, model, page_range)
     for page in pages:
         for block in page["blocks"]:
             bad_keys = [key for key in block.keys() if key not in ["lines", "bbox"]]

{pdftext-0.3.0 → pdftext-0.3.2}/pdftext/inference.py

@@ -3,6 +3,7 @@ from itertools import chain
 import sklearn
 
 from pdftext.pdf.utils import LINE_BREAKS, TABS, SPACES
+from pdftext.settings import settings
 
 
 def update_current(current, new_char):
@@ -98,7 +99,7 @@ def update_block(blocks, block):
     return block
 
 
-def infer_single_page(text_chars):
+def infer_single_page(text_chars, block_threshold=settings.BLOCK_THRESHOLD):
     prev_char = None
     prev_font_info = None
 
@@ -120,14 +121,15 @@ def infer_single_page(text_chars):
             sorted_keys = sorted(training_row.keys())
             training_row = [training_row[key] for key in sorted_keys]
 
-            prediction = yield training_row
-            if prediction == 0:
+            prediction_probs = yield training_row
+            # First item is probability of same line/block, second is probability of new line, third is probability of new block
+            if prediction_probs[0] >= .5:
                 pass
-            elif prediction == 2 and prev_char["char"] in LINE_BREAKS:
+            elif prediction_probs[2] > block_threshold:
                 span = update_span(line, span)
                 line = update_line(block, line)
                 block = update_block(blocks, block)
-            elif prev_char["char"] in LINE_BREAKS and prediction == 1: # Look for newline character as a forcing signal for a new line
+            elif prev_char["char"] in LINE_BREAKS and prediction_probs[1] >= .5: # Look for newline character as a forcing signal for a new line
                 span = update_span(line, span)
                 line = update_line(block, line)
             elif prev_font_info != font_info:
@@ -180,7 +182,7 @@ def inference(text_chars, model):
 
         # Disable nan, etc, validation for a small speedup
         with sklearn.config_context(assume_finite=True):
-            predictions = model.predict(training_rows)
+            predictions = model.predict_proba(training_rows)
         for pred, page_idx in zip(predictions, training_idxs):
             next_prediction[page_idx] = pred
     sorted_keys = sorted(page_blocks.keys())

{pdftext-0.3.0 → pdftext-0.3.2}/pdftext/pdf/chars.py

@@ -23,11 +23,12 @@ def update_previous_fonts(text_chars: Dict, i: int, fontname: str, fontflags: in
         text_chars["chars"][j]["font"]["flags"] = fontflags
 
 
-def get_pdfium_chars(pdf_path, fontname_sample_freq=settings.FONTNAME_SAMPLE_FREQ):
-    pdf = pdfium.PdfDocument(pdf_path)
+def get_pdfium_chars(pdf, fontname_sample_freq=settings.FONTNAME_SAMPLE_FREQ, page_range=None):
     blocks = []
+    if page_range is None:
+        page_range = range(len(pdf))
 
-    for page_idx in range(len(pdf)):
+    for page_idx in page_range:
         page = pdf.get_page(page_idx)
         text_page = page.get_textpage()
         mediabox = page.get_mediabox()

{pdftext-0.3.0 → pdftext-0.3.2}/pdftext/postprocessing.py

@@ -29,14 +29,14 @@ def postprocess_text(text: str) -> str:
 
 def handle_hyphens(text: str, keep_hyphens=False) -> str:
     if keep_hyphens:
-        text = text.replace(HYPHEN_CHAR, "-")
+        text = text.replace(HYPHEN_CHAR, "-\n")
     elif len(text) == 0:
         pass
     else:
         new_text = ""
         found_hyphen = False
         for i in range(len(text) - 1):
-            if text[i] == HYPHEN_CHAR and text[i+1] in LINE_BREAKS:
+            if text[i] == HYPHEN_CHAR:
                 found_hyphen = True
             elif found_hyphen:
                 if text[i] in LINE_BREAKS:

{pdftext-0.3.0 → pdftext-0.3.2}/pdftext/settings.py

@@ -7,10 +7,13 @@ class Settings(BaseSettings):
     BASE_PATH: str = os.path.dirname(os.path.dirname(__file__))
     MODEL_PATH: str = os.path.join(BASE_PATH, "models", "dt.joblib")
 
-    # How many characters to buffer when reading a font name
-    FONT_BUFFER_SIZE: int = 1024
+    # Fonts
+    FONT_BUFFER_SIZE: int = 1024 # How many characters to buffer when reading a font name
     FONTNAME_SAMPLE_FREQ: int = 10
 
+    # Inference
+    BLOCK_THRESHOLD: float = 0.8 # Confidence threshold for block detection
+
     # Benchmark
     RESULTS_FOLDER: str = "results"
     BENCH_DATASET_NAME: str = "vikp/pdf_bench"

{pdftext-0.3.0 → pdftext-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pdftext"
-version = "0.3.0"
+version = "0.3.2"
 description = "Extract structured text from pdfs quickly"
 authors = ["Vik Paruchuri <vik.paruchuri@gmail.com>"]
 license = "Apache-2.0"