# Copyright 2021 The Layout Parser team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import List, Union, Optional, Dict, Tuple
import pdfplumber
import pandas as pd
from ..elements import Layout
from .basic import load_dataframe
DEFAULT_PDF_DPI = 72
def extract_words_for_page(
page: pdfplumber.page.Page,
x_tolerance=1.5,
y_tolerance=2,
keep_blank_chars=False,
use_text_flow=True,
horizontal_ltr=True,
vertical_ttb=True,
extra_attrs=None,
) -> Layout:
"""The helper function used for extracting words from a pdfplumber page
object.
Returns:
Layout: a layout object representing all extracted pdf tokens on this page.
"""
if extra_attrs is None:
extra_attrs = ["fontname", "size"]
tokens = page.extract_words(
x_tolerance=x_tolerance,
y_tolerance=y_tolerance,
keep_blank_chars=keep_blank_chars,
use_text_flow=use_text_flow,
horizontal_ltr=horizontal_ltr,
vertical_ttb=vertical_ttb,
extra_attrs=extra_attrs,
)
df = pd.DataFrame(tokens)
if len(df) == 0:
return Layout()
df[["x0", "x1"]] = (
df[["x0", "x1"]].clip(lower=0, upper=int(page.width)).astype("float")
)
df[["top", "bottom"]] = (
df[["top", "bottom"]].clip(lower=0, upper=int(page.height)).astype("float")
)
page_tokens = load_dataframe(
df.reset_index().rename(
columns={
"x0": "x_1",
"x1": "x_2",
"top": "y_1",
"bottom": "y_2",
"index": "id",
"fontname": "type", # also loading fontname as "type"
}
),
block_type="rectangle",
)
return page_tokens
[docs]def load_pdf(
filename: str,
load_images: bool = False,
x_tolerance: int = 1.5,
y_tolerance: int = 2,
keep_blank_chars: bool = False,
use_text_flow: bool = True,
horizontal_ltr: bool = True,
vertical_ttb: bool = True,
extra_attrs: Optional[List[str]] = None,
dpi: int = DEFAULT_PDF_DPI,
) -> Union[List[Layout], Tuple[List[Layout], List["Image.Image"]]]:
"""Load all tokens for each page from a PDF file, and save them
in a list of Layout objects with the original page order.
Args:
filename (str): The path to the PDF file.
load_images (bool, optional):
Whether load screenshot for each page of the PDF file.
When set to true, the function will return both the layout and
screenshot image for each page.
Defaults to False.
x_tolerance (int, optional):
The threshold used for extracting "word tokens" from the pdf file.
It will merge the pdf characters into a word token if the difference
between the x_2 of one character and the x_1 of the next is less than
or equal to x_tolerance. See details in `pdf2plumber's documentation
<https://github.com/jsvine/pdfplumber#the-pdfplumberpage-class>`_.
Defaults to 1.5.
y_tolerance (int, optional):
The threshold used for extracting "word tokens" from the pdf file.
It will merge the pdf characters into a word token if the difference
between the y_2 of one character and the y_1 of the next is less than
or equal to y_tolerance. See details in `pdf2plumber's documentation
<https://github.com/jsvine/pdfplumber#the-pdfplumberpage-class>`_.
Defaults to 2.
keep_blank_chars (bool, optional):
When keep_blank_chars is set to True, it will treat blank characters
are treated as part of a word, not as a space between words. See
details in `pdf2plumber's documentation
<https://github.com/jsvine/pdfplumber#the-pdfplumberpage-class>`_.
Defaults to False.
use_text_flow (bool, optional):
When use_text_flow is set to True, it will use the PDF's underlying
flow of characters as a guide for ordering and segmenting the words,
rather than presorting the characters by x/y position. (This mimics
how dragging a cursor highlights text in a PDF; as with that, the
order does not always appear to be logical.) See details in
`pdf2plumber's documentation
<https://github.com/jsvine/pdfplumber#the-pdfplumberpage-class>`_.
Defaults to True.
horizontal_ltr (bool, optional):
When horizontal_ltr is set to True, it means the doc should read
text from left to right, vice versa.
Defaults to True.
vertical_ttb (bool, optional):
When vertical_ttb is set to True, it means the doc should read
text from top to bottom, vice versa.
Defaults to True.
extra_attrs (Optional[List[str]], optional):
Passing a list of extra_attrs (e.g., ["fontname", "size"]) will
restrict each words to characters that share exactly the same
value for each of those `attributes extracted by pdfplumber
<https://github.com/jsvine/pdfplumber/blob/develop/README.md#char-properties>`_,
and the resulting word dicts will indicate those attributes.
See details in `pdf2plumber's documentation
<https://github.com/jsvine/pdfplumber#the-pdfplumberpage-class>`_.
Defaults to `["fontname", "size"]`.
dpi (int, optional):
When loading images of the pdf, you can also specify the resolution
(or `DPI, dots per inch <https://en.wikipedia.org/wiki/Dots_per_inch>`_)
for rendering the images. Higher DPI values mean clearer images (also
larger file sizes).
Setting dpi will also automatically resizes the extracted pdf_layout
to match the sizes of the images. Therefore, when visualizing the
pdf_layouts, it can be rendered appropriately.
Defaults to `DEFAULT_PDF_DPI=72`, which is also the default rendering dpi
from the pdfplumber PDF parser.
Returns:
List[Layout]:
When `load_images=False`, it will only load the pdf_tokens from
the PDF file. Each element of the list denotes all the tokens appeared
on a single page, and the list is ordered the same as the original PDF
page order.
Tuple[List[Layout], List["Image.Image"]]:
When `load_images=True`, besides the `all_page_layout`, it will also
return a list of page images.
Examples::
>>> import layoutparser as lp
>>> pdf_layout = lp.load_pdf("path/to/pdf")
>>> pdf_layout[0] # the layout for page 0
>>> pdf_layout, pdf_images = lp.load_pdf("path/to/pdf", load_images=True)
>>> lp.draw_box(pdf_images[0], pdf_layout[0])
"""
plumber_pdf_object = pdfplumber.open(filename)
all_page_layout = []
for page_id in range(len(plumber_pdf_object.pages)):
cur_page = plumber_pdf_object.pages[page_id]
page_tokens = extract_words_for_page(
cur_page,
x_tolerance=x_tolerance,
y_tolerance=y_tolerance,
keep_blank_chars=keep_blank_chars,
use_text_flow=use_text_flow,
horizontal_ltr=horizontal_ltr,
vertical_ttb=vertical_ttb,
extra_attrs=extra_attrs,
)
# Adding metadata for the current page
page_tokens.page_data["width"] = float(cur_page.width)
page_tokens.page_data["height"] = float(cur_page.height)
page_tokens.page_data["index"] = page_id
all_page_layout.append(page_tokens)
if not load_images:
return all_page_layout
else:
import pdf2image
pdf_images = pdf2image.convert_from_path(filename, dpi=dpi)
for page_id, page_image in enumerate(pdf_images):
image_width, image_height = page_image.size
page_layout = all_page_layout[page_id]
layout_width = page_layout.page_data["width"]
layout_height = page_layout.page_data["height"]
if image_width != layout_width or image_height != layout_height:
scale_x = image_width / layout_width
scale_y = image_height / layout_height
page_layout = page_layout.scale((scale_x, scale_y))
page_layout.page_data["width"] = image_width
page_layout.page_data["height"] = image_height
all_page_layout[page_id] = page_layout
return all_page_layout, pdf_images