修改为东南天坐标系
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,341 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from typing import IO, Any, Optional, Sequence
|
||||
|
||||
import requests
|
||||
from unstructured_client import UnstructuredClient
|
||||
from unstructured_client.models import operations, shared
|
||||
from unstructured_client.utils import retries
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.staging.base import elements_from_dicts, elements_from_json
|
||||
|
||||
# Default retry configuration taken from the client code
|
||||
DEFAULT_RETRIES_INITIAL_INTERVAL_SEC = 3000
|
||||
DEFAULT_RETRIES_MAX_INTERVAL_SEC = 720000
|
||||
DEFAULT_RETRIES_EXPONENT = 1.5
|
||||
DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC = 1800000
|
||||
DEFAULT_RETRIES_CONNECTION_ERRORS = True
|
||||
|
||||
|
||||
def partition_via_api(
|
||||
filename: Optional[str] = None,
|
||||
content_type: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
file_filename: Optional[str] = None,
|
||||
api_url: str = "https://api.unstructured.io/general/v0/general",
|
||||
api_key: str = "",
|
||||
metadata_filename: Optional[str] = None,
|
||||
retries_initial_interval: [int] = None,
|
||||
retries_max_interval: Optional[int] = None,
|
||||
retries_exponent: Optional[float] = None,
|
||||
retries_max_elapsed_time: Optional[int] = None,
|
||||
retries_connection_errors: Optional[bool] = None,
|
||||
**request_kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions a document using the Unstructured REST API. This is equivalent to
|
||||
running the document through partition.
|
||||
|
||||
See https://api.unstructured.io/general/docs for the hosted API documentation or
|
||||
https://github.com/Unstructured-IO/unstructured-api for instructions on how to run
|
||||
the API locally as a container.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
content_type
|
||||
A string defining the file content in MIME type
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_filename
|
||||
When file is not None, the filename (string) to store in element metadata. E.g. "foo.txt"
|
||||
api_url
|
||||
The URL for the Unstructured API. Defaults to the hosted Unstructured API.
|
||||
api_key
|
||||
The API key to pass to the Unstructured API.
|
||||
retries_initial_interval
|
||||
Defines the time interval (in seconds) to wait before the first retry in case of a request
|
||||
failure. Defaults to 3000. If set should be > 0.
|
||||
retries_max_interval
|
||||
Defines the maximum time interval (in seconds) to wait between retries (the interval
|
||||
between retries is increased as using exponential increase algorithm
|
||||
- this setting limits it). Defaults to 720000. If set should be > 0.
|
||||
retries_exponent
|
||||
Defines the exponential factor to increase the interval between retries. Defaults to 1.5.
|
||||
If set should be > 0.0.
|
||||
retries_max_elapsed_time
|
||||
Defines the maximum time (in seconds) to wait for retries. If exceeded, the original
|
||||
exception is raised. Defaults to 1800000. If set should be > 0.
|
||||
retries_connection_errors
|
||||
Defines whether to retry on connection errors. Defaults to True.
|
||||
request_kwargs
|
||||
Additional parameters to pass to the data field of the request to the Unstructured API.
|
||||
For example the `strategy` parameter.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
if metadata_filename and file_filename:
|
||||
raise ValueError(
|
||||
"Only one of metadata_filename and file_filename is specified. "
|
||||
"metadata_filename is preferred. file_filename is marked for deprecation.",
|
||||
)
|
||||
|
||||
if file_filename is not None:
|
||||
metadata_filename = file_filename
|
||||
logger.warning(
|
||||
"The file_filename kwarg will be deprecated in a future version of unstructured. "
|
||||
"Please use metadata_filename instead.",
|
||||
)
|
||||
|
||||
# Note(austin) - the sdk takes the base url, but we have the full api_url
|
||||
# For consistency, just strip off the path when it's given
|
||||
base_url = api_url[:-19] if "/general/v0/general" in api_url else api_url
|
||||
sdk = UnstructuredClient(api_key_auth=api_key, server_url=base_url)
|
||||
|
||||
if filename is not None:
|
||||
with open(filename, "rb") as f:
|
||||
files = shared.Files(
|
||||
content=f.read(),
|
||||
file_name=filename,
|
||||
)
|
||||
|
||||
elif file is not None:
|
||||
if metadata_filename is None:
|
||||
raise ValueError(
|
||||
"If file is specified in partition_via_api, "
|
||||
"metadata_filename must be specified as well.",
|
||||
)
|
||||
files = shared.Files(content=file, file_name=metadata_filename)
|
||||
|
||||
req = operations.PartitionRequest(
|
||||
partition_parameters=shared.PartitionParameters(files=files, **request_kwargs)
|
||||
)
|
||||
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=retries_connection_errors,
|
||||
retries_exponent=retries_exponent,
|
||||
retries_initial_interval=retries_initial_interval,
|
||||
retries_max_elapsed_time=retries_max_elapsed_time,
|
||||
retries_max_interval=retries_max_interval,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
response = sdk.general.partition(
|
||||
request=req,
|
||||
retries=retries_config,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
return elements_from_json(text=response.raw_response.text)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Receive unexpected status code {response.status_code} from the API.",
|
||||
)
|
||||
|
||||
|
||||
def get_retries_config(
|
||||
retries_connection_errors: Optional[bool],
|
||||
retries_exponent: Optional[float],
|
||||
retries_initial_interval: Optional[int],
|
||||
retries_max_elapsed_time: Optional[int],
|
||||
retries_max_interval: Optional[int],
|
||||
sdk: UnstructuredClient,
|
||||
) -> Optional[retries.RetryConfig]:
|
||||
"""Constructs a RetryConfig object from the provided parameters. If any of the parameters
|
||||
are None, the default values are taken from the SDK configuration or the default constants.
|
||||
|
||||
If all parameters are None, returns None (and the SDK-managed defaults are used within the
|
||||
client)
|
||||
|
||||
The solution is not perfect as the RetryConfig object does not include the defaults by
|
||||
itself so we might need to construct it basing on our defaults.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
retries_connection_errors
|
||||
Defines whether to retry on connection errors. If not set the
|
||||
DEFAULT_RETRIES_CONNECTION_ERRORS constant is used.
|
||||
retries_exponent
|
||||
Defines the exponential factor to increase the interval between retries.
|
||||
If set, should be > 0.0 (otherwise the DEFAULT_RETRIES_EXPONENT constant is used)
|
||||
retries_initial_interval
|
||||
Defines the time interval to wait before the first retry in case of a request failure.
|
||||
If set, should be > 0 (otherwise the DEFAULT_RETRIES_INITIAL_INTERVAL_SEC constant is used)
|
||||
retries_max_elapsed_time
|
||||
Defines the maximum time to wait for retries. If exceeded, the original exception is raised.
|
||||
If set, should be > 0 (otherwise the DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC constant is used)
|
||||
retries_max_interval
|
||||
Defines the maximum time interval to wait between retries. If set, should be > 0
|
||||
(otherwise the DEFAULT_RETRIES_MAX_INTERVAL_SEC constant is used)
|
||||
sdk
|
||||
The UnstructuredClient object to take the default values from.
|
||||
"""
|
||||
retries_config = None
|
||||
sdk_default_retries_config = sdk.sdk_configuration.retry_config
|
||||
if any(
|
||||
setting is not None
|
||||
for setting in (
|
||||
retries_initial_interval,
|
||||
retries_max_interval,
|
||||
retries_exponent,
|
||||
retries_max_elapsed_time,
|
||||
retries_connection_errors,
|
||||
)
|
||||
):
|
||||
|
||||
def get_backoff_default(setting_name: str, default_value: Any) -> Any:
|
||||
if sdk_default_retries_config: # noqa: SIM102
|
||||
if setting_value := getattr(sdk_default_retries_config.backoff, setting_name):
|
||||
return setting_value
|
||||
return default_value
|
||||
|
||||
default_retries_connneciton_errors = (
|
||||
sdk_default_retries_config.retry_connection_errors
|
||||
if sdk_default_retries_config
|
||||
and sdk_default_retries_config.retry_connection_errors is not None
|
||||
else DEFAULT_RETRIES_CONNECTION_ERRORS
|
||||
)
|
||||
|
||||
backoff_strategy = retries.BackoffStrategy(
|
||||
initial_interval=(
|
||||
retries_initial_interval
|
||||
or get_backoff_default("initial_interval", DEFAULT_RETRIES_INITIAL_INTERVAL_SEC)
|
||||
),
|
||||
max_interval=(
|
||||
retries_max_interval
|
||||
or get_backoff_default("max_interval", DEFAULT_RETRIES_MAX_INTERVAL_SEC)
|
||||
),
|
||||
exponent=(
|
||||
retries_exponent or get_backoff_default("exponent", DEFAULT_RETRIES_EXPONENT)
|
||||
),
|
||||
max_elapsed_time=(
|
||||
retries_max_elapsed_time
|
||||
or get_backoff_default("max_elapsed_time", DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC)
|
||||
),
|
||||
)
|
||||
retries_config = retries.RetryConfig(
|
||||
strategy="backoff",
|
||||
backoff=backoff_strategy,
|
||||
retry_connection_errors=(
|
||||
retries_connection_errors
|
||||
if retries_connection_errors is not None
|
||||
else default_retries_connneciton_errors
|
||||
),
|
||||
)
|
||||
return retries_config
|
||||
|
||||
|
||||
def partition_multiple_via_api(
|
||||
filenames: Optional[list[str]] = None,
|
||||
content_types: Optional[list[str]] = None,
|
||||
files: Optional[Sequence[IO[bytes]]] = None,
|
||||
file_filenames: Optional[list[str]] = None,
|
||||
api_url: str = "https://api.unstructured.io/general/v0/general",
|
||||
api_key: str = "",
|
||||
metadata_filenames: Optional[list[str]] = None,
|
||||
**request_kwargs: Any,
|
||||
) -> list[list[Element]]:
|
||||
"""Partitions multiple documents using the Unstructured REST API by batching
|
||||
the documents into a single HTTP request.
|
||||
|
||||
See https://api.unstructured.io/general/docs for the hosted API documentation or
|
||||
https://github.com/Unstructured-IO/unstructured-api for instructions on how to run
|
||||
the API locally as a container.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filenames
|
||||
A list of strings defining the target filename paths.
|
||||
content_types
|
||||
A list of strings defining the file contents in MIME types.
|
||||
files
|
||||
A list of file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_filename
|
||||
When file is not None, the filename (string) to store in element metadata. E.g. "foo.txt"
|
||||
api_url
|
||||
The URL for the Unstructured API. Defaults to the hosted Unstructured API.
|
||||
api_key
|
||||
The API key to pass to the Unstructured API.
|
||||
request_kwargs
|
||||
Additional parameters to pass to the data field of the request to the Unstructured API.
|
||||
For example the `strategy` parameter.
|
||||
"""
|
||||
headers = {
|
||||
"ACCEPT": "application/json",
|
||||
"UNSTRUCTURED-API-KEY": api_key,
|
||||
}
|
||||
|
||||
if metadata_filenames and file_filenames:
|
||||
raise ValueError(
|
||||
"Only one of metadata_filenames and file_filenames is specified. "
|
||||
"metadata_filenames is preferred. file_filenames is marked for deprecation.",
|
||||
)
|
||||
|
||||
if file_filenames is not None:
|
||||
metadata_filenames = file_filenames
|
||||
logger.warning(
|
||||
"The file_filenames kwarg will be deprecated in a future version of unstructured. "
|
||||
"Please use metadata_filenames instead.",
|
||||
)
|
||||
|
||||
if filenames is not None:
|
||||
if content_types and len(content_types) != len(filenames):
|
||||
raise ValueError("content_types and filenames must have the same length.")
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(f, "rb")) for f in filenames] # type: ignore
|
||||
|
||||
_files = []
|
||||
for i, file in enumerate(files):
|
||||
filename = filenames[i]
|
||||
content_type = content_types[i] if content_types is not None else None
|
||||
_files.append(("files", (filename, file, content_type)))
|
||||
|
||||
response = requests.post(
|
||||
api_url,
|
||||
headers=headers,
|
||||
data=request_kwargs,
|
||||
files=_files, # type: ignore
|
||||
)
|
||||
|
||||
elif files is not None:
|
||||
if content_types and len(content_types) != len(files):
|
||||
raise ValueError("content_types and files must have the same length.")
|
||||
|
||||
if not metadata_filenames:
|
||||
raise ValueError("metadata_filenames must be specified if files are passed")
|
||||
elif len(metadata_filenames) != len(files):
|
||||
raise ValueError("metadata_filenames and files must have the same length.")
|
||||
|
||||
_files = []
|
||||
for i, _file in enumerate(files): # type: ignore
|
||||
content_type = content_types[i] if content_types is not None else None
|
||||
filename = metadata_filenames[i]
|
||||
_files.append(("files", (filename, _file, content_type)))
|
||||
|
||||
response = requests.post(
|
||||
api_url,
|
||||
headers=headers,
|
||||
data=request_kwargs,
|
||||
files=_files, # type: ignore
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
documents = []
|
||||
response_list = response.json()
|
||||
# NOTE(robinson) - this check is because if only one filename is passed, the return
|
||||
# type from the API is a list of objects instead of a list of lists
|
||||
if not isinstance(response_list[0], list):
|
||||
response_list = [response_list]
|
||||
|
||||
for document in response_list:
|
||||
documents.append(elements_from_dicts(document))
|
||||
return documents
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Receive unexpected status code {response.status_code} from the API.",
|
||||
)
|
||||
@@ -0,0 +1,383 @@
|
||||
"""Provides partitioning with automatic file-type detection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import importlib
|
||||
import io
|
||||
from typing import IO, Any, Callable, Optional
|
||||
|
||||
import requests
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
from unstructured.documents.elements import DataSourceMetadata, Element
|
||||
from unstructured.file_utils.filetype import (
|
||||
detect_filetype,
|
||||
is_json_processable,
|
||||
is_ndjson_processable,
|
||||
)
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.common import UnsupportedFileFormatError
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.lang import check_language_args
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
from unstructured.utils import dependency_exists
|
||||
|
||||
Partitioner: TypeAlias = Callable[..., list[Element]]
|
||||
|
||||
|
||||
def partition(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
encoding: Optional[str] = None,
|
||||
content_type: Optional[str] = None,
|
||||
url: Optional[str] = None,
|
||||
headers: dict[str, str] = {},
|
||||
ssl_verify: bool = True,
|
||||
request_timeout: Optional[int] = None,
|
||||
strategy: str = PartitionStrategy.AUTO,
|
||||
skip_infer_table_types: list[str] = ["pdf", "jpg", "png", "heic"],
|
||||
ocr_languages: Optional[str] = None, # changing to optional for deprecation
|
||||
languages: Optional[list[str]] = None,
|
||||
detect_language_per_element: bool = False,
|
||||
pdf_infer_table_structure: bool = False,
|
||||
extract_images_in_pdf: bool = False,
|
||||
extract_image_block_types: Optional[list[str]] = None,
|
||||
extract_image_block_output_dir: Optional[str] = None,
|
||||
extract_image_block_to_payload: bool = False,
|
||||
data_source_metadata: Optional[DataSourceMetadata] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
hi_res_model_name: Optional[str] = None,
|
||||
model_name: Optional[str] = None, # to be deprecated
|
||||
starting_page_number: int = 1,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions a document into its constituent elements.
|
||||
|
||||
Uses libmagic to determine the file's type and route it to the appropriate partitioning
|
||||
function. Applies the default parameters for each partitioning function. Use the document-type
|
||||
specific partitioning functions if you need access to additional kwarg options.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
encoding
|
||||
The character-encoding used to decode the input bytes when drawn from `filename` or `file`.
|
||||
Defaults to "utf-8".
|
||||
url
|
||||
The url for a remote document. Pass in content_type if you want partition to treat
|
||||
the document as a specific content_type.
|
||||
headers
|
||||
The headers to be used in conjunction with the HTTP request if URL is set.
|
||||
ssl_verify
|
||||
If the URL parameter is set, determines whether or not partition uses SSL verification
|
||||
in the HTTP request.
|
||||
request_timeout
|
||||
The timeout for the HTTP request if URL is set. Defaults to None meaning no timeout and
|
||||
requests will block indefinitely.
|
||||
content_type
|
||||
A string defining the file content in MIME type
|
||||
metadata_filename
|
||||
When file is not None, the filename (string) to store in element metadata. E.g. "foo.txt"
|
||||
strategy
|
||||
The strategy to use for partitioning PDF/image. Uses a layout detection model if set
|
||||
to 'hi_res', otherwise partition simply extracts the text from the document
|
||||
and processes it.
|
||||
skip_infer_table_types
|
||||
The document types that you want to skip table extraction with.
|
||||
languages
|
||||
The languages present in the document, for use in partitioning and/or OCR. For partitioning
|
||||
image or pdf documents with Tesseract, you'll first need to install the appropriate
|
||||
Tesseract language pack. For other partitions, language is detected using naive Bayesian
|
||||
filter via `langdetect`. Multiple languages indicates text could be in either language.
|
||||
Additional Parameters:
|
||||
detect_language_per_element
|
||||
Detect language per element instead of at the document level.
|
||||
pdf_infer_table_structure
|
||||
Deprecated! Use `skip_infer_table_types` to opt out of table extraction for any document
|
||||
type.
|
||||
If True and strategy=hi_res, any Table Elements extracted from a PDF will include an
|
||||
additional metadata field, "text_as_html," where the value (string) is a just a
|
||||
transformation of the data into an HTML <table>.
|
||||
The "text" field for a partitioned Table Element is always present, whether True or False.
|
||||
extract_images_in_pdf
|
||||
Only applicable if `strategy=hi_res`.
|
||||
If True, any detected images will be saved in the path specified by
|
||||
'extract_image_block_output_dir' or stored as base64 encoded data within metadata fields.
|
||||
Deprecation Note: This parameter is marked for deprecation. Future versions will use
|
||||
'extract_image_block_types' for broader extraction capabilities.
|
||||
extract_image_block_types
|
||||
Only applicable if `strategy=hi_res`.
|
||||
Images of the element type(s) specified in this list (e.g., ["Image", "Table"]) will be
|
||||
saved in the path specified by 'extract_image_block_output_dir' or stored as base64
|
||||
encoded data within metadata fields.
|
||||
extract_image_block_to_payload
|
||||
Only applicable if `strategy=hi_res`.
|
||||
If True, images of the element type(s) defined in 'extract_image_block_types' will be
|
||||
encoded as base64 data and stored in two metadata fields: 'image_base64' and
|
||||
'image_mime_type'.
|
||||
This parameter facilitates the inclusion of element data directly within the payload,
|
||||
especially for web-based applications or APIs.
|
||||
extract_image_block_output_dir
|
||||
Only applicable if `strategy=hi_res` and `extract_image_block_to_payload=False`.
|
||||
The filesystem path for saving images of the element type(s)
|
||||
specified in 'extract_image_block_types'.
|
||||
hi_res_model_name
|
||||
The layout detection model used when partitioning strategy is set to `hi_res`.
|
||||
model_name
|
||||
The layout detection model used when partitioning strategy is set to `hi_res`. To be
|
||||
deprecated in favor of `hi_res_model_name`.
|
||||
starting_page_number
|
||||
Indicates what page number should be assigned to the first page in the document.
|
||||
This information will be reflected in elements' metadata and can be be especially
|
||||
useful when partitioning a document that is part of a larger document.
|
||||
"""
|
||||
exactly_one(file=file, filename=filename, url=url)
|
||||
|
||||
kwargs.setdefault("metadata_filename", metadata_filename)
|
||||
|
||||
if pdf_infer_table_structure:
|
||||
logger.warning(
|
||||
"The pdf_infer_table_structure kwarg is deprecated. Please use skip_infer_table_types "
|
||||
"instead."
|
||||
)
|
||||
|
||||
languages = check_language_args(languages or [], ocr_languages)
|
||||
|
||||
if url is not None:
|
||||
file, file_type = file_and_type_from_url(
|
||||
url=url,
|
||||
content_type=content_type,
|
||||
headers=headers,
|
||||
ssl_verify=ssl_verify,
|
||||
request_timeout=request_timeout,
|
||||
)
|
||||
else:
|
||||
if headers != {}:
|
||||
logger.warning(
|
||||
"The headers kwarg is set but the url kwarg is not. "
|
||||
"The headers kwarg will be ignored.",
|
||||
)
|
||||
file_type = detect_filetype(
|
||||
file_path=filename,
|
||||
file=file,
|
||||
encoding=encoding,
|
||||
content_type=content_type,
|
||||
metadata_file_path=metadata_filename,
|
||||
)
|
||||
|
||||
if file is not None:
|
||||
file.seek(0)
|
||||
|
||||
# avoid double specification of infer_table_structure; this can happen when the kwarg passed
|
||||
# into a partition function, e.g., partition_email is reused to partition sub-elements, e.g.,
|
||||
# partition an image attachment buy calling partition with the kwargs. In that case here kwargs
|
||||
# would have a infer_table_structure already
|
||||
kwargs_infer_table_structure = kwargs.pop("infer_table_structure", None)
|
||||
infer_table_structure = (
|
||||
kwargs_infer_table_structure
|
||||
if kwargs_infer_table_structure is not None
|
||||
else decide_table_extraction(
|
||||
file_type,
|
||||
skip_infer_table_types,
|
||||
pdf_infer_table_structure,
|
||||
)
|
||||
)
|
||||
|
||||
partitioner_loader = _PartitionerLoader()
|
||||
|
||||
# -- extracting this post-processing to allow multiple exit-points from function --
|
||||
def augment_metadata(elements: list[Element]) -> list[Element]:
|
||||
"""Add some metadata fields to each element."""
|
||||
for element in elements:
|
||||
element.metadata.url = url
|
||||
element.metadata.data_source = data_source_metadata
|
||||
if content_type is not None:
|
||||
out_filetype = FileType.from_mime_type(content_type)
|
||||
element.metadata.filetype = out_filetype.mime_type if out_filetype else None
|
||||
else:
|
||||
element.metadata.filetype = file_type.mime_type
|
||||
|
||||
return elements
|
||||
|
||||
# -- handle PDF/Image partitioning separately because they have a lot of special-case
|
||||
# -- parameters. We'll come back to this after sorting out the other file types.
|
||||
if file_type == FileType.PDF:
|
||||
partition_pdf = partitioner_loader.get(file_type)
|
||||
elements = partition_pdf(
|
||||
filename=filename,
|
||||
file=file,
|
||||
url=None,
|
||||
infer_table_structure=infer_table_structure,
|
||||
strategy=strategy,
|
||||
languages=languages,
|
||||
detect_language_per_element=detect_language_per_element,
|
||||
hi_res_model_name=hi_res_model_name or model_name,
|
||||
extract_images_in_pdf=extract_images_in_pdf,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_output_dir=extract_image_block_output_dir,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
starting_page_number=starting_page_number,
|
||||
**kwargs,
|
||||
)
|
||||
return augment_metadata(elements)
|
||||
|
||||
if file_type.partitioner_shortname and file_type.partitioner_shortname == "image":
|
||||
partition_image = partitioner_loader.get(file_type)
|
||||
elements = partition_image(
|
||||
filename=filename,
|
||||
file=file,
|
||||
url=None,
|
||||
infer_table_structure=infer_table_structure,
|
||||
strategy=strategy,
|
||||
languages=languages,
|
||||
detect_language_per_element=detect_language_per_element,
|
||||
hi_res_model_name=hi_res_model_name or model_name,
|
||||
extract_images_in_pdf=extract_images_in_pdf,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_output_dir=extract_image_block_output_dir,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
starting_page_number=starting_page_number,
|
||||
**kwargs,
|
||||
)
|
||||
return augment_metadata(elements)
|
||||
|
||||
# -- JSON is a special case because it's not a document format per se and is insensitive to
|
||||
# -- most of the parameters that apply to other file types.
|
||||
if file_type == FileType.JSON:
|
||||
if not is_json_processable(filename=filename, file=file):
|
||||
raise ValueError(
|
||||
"Detected a JSON file that does not conform to the Unstructured schema. "
|
||||
"partition_json currently only processes serialized Unstructured output.",
|
||||
)
|
||||
partition_json = partitioner_loader.get(file_type)
|
||||
elements = partition_json(filename=filename, file=file, **kwargs)
|
||||
return augment_metadata(elements)
|
||||
|
||||
if file_type == FileType.NDJSON:
|
||||
if not is_ndjson_processable(filename=filename, file=file):
|
||||
raise ValueError(
|
||||
"Detected an NDJSON file that does not conform to the Unstructured schema. "
|
||||
"partition_json currently only processes serialized Unstructured output.",
|
||||
)
|
||||
partition_ndjson = partitioner_loader.get(file_type)
|
||||
elements = partition_ndjson(filename=filename, file=file, **kwargs)
|
||||
return augment_metadata(elements)
|
||||
|
||||
# -- EMPTY is also a special case because while we can't determine the file type, we can be
|
||||
# -- sure it doesn't contain any elements.
|
||||
if file_type == FileType.EMPTY:
|
||||
return []
|
||||
|
||||
# ============================================================================================
|
||||
# ALL OTHER FILE TYPES
|
||||
# ============================================================================================
|
||||
|
||||
partitioning_kwargs = copy.deepcopy(kwargs)
|
||||
partitioning_kwargs["detect_language_per_element"] = detect_language_per_element
|
||||
partitioning_kwargs["encoding"] = encoding
|
||||
partitioning_kwargs["infer_table_structure"] = infer_table_structure
|
||||
partitioning_kwargs["languages"] = languages
|
||||
partitioning_kwargs["starting_page_number"] = starting_page_number
|
||||
partitioning_kwargs["strategy"] = strategy
|
||||
partitioning_kwargs["extract_image_block_types"] = extract_image_block_types
|
||||
partitioning_kwargs["extract_image_block_to_payload"] = extract_image_block_to_payload
|
||||
|
||||
partition = partitioner_loader.get(file_type)
|
||||
elements = partition(filename=filename, file=file, **partitioning_kwargs)
|
||||
return augment_metadata(elements)
|
||||
|
||||
|
||||
def file_and_type_from_url(
|
||||
url: str,
|
||||
content_type: Optional[str] = None,
|
||||
headers: dict[str, str] = {},
|
||||
ssl_verify: bool = True,
|
||||
request_timeout: Optional[int] = None,
|
||||
) -> tuple[io.BytesIO, FileType]:
|
||||
response = requests.get(url, headers=headers, verify=ssl_verify, timeout=request_timeout)
|
||||
file = io.BytesIO(response.content)
|
||||
|
||||
if content_type := content_type or response.headers.get("Content-Type", None):
|
||||
content_type = content_type.split(";")[0].strip().lower()
|
||||
|
||||
# -- non-None when response is textual --
|
||||
encoding = response.encoding
|
||||
|
||||
filetype = detect_filetype(file=file, encoding=encoding, content_type=content_type)
|
||||
return file, filetype
|
||||
|
||||
|
||||
def decide_table_extraction(
|
||||
filetype: Optional[FileType],
|
||||
skip_infer_table_types: list[str],
|
||||
pdf_infer_table_structure: bool,
|
||||
) -> bool:
|
||||
doc_type = filetype.name.lower() if filetype else None
|
||||
|
||||
if doc_type == "pdf":
|
||||
# For backwards compatibility. Ultimately we want to remove pdf_infer_table_structure
|
||||
# completely and rely exclusively on `skip_infer_table_types` for all file types.
|
||||
# Until then for pdf files we first check pdf_infer_table_structure and then update
|
||||
# based on skip_infer_tables.
|
||||
return pdf_infer_table_structure or doc_type not in skip_infer_table_types
|
||||
|
||||
return doc_type not in skip_infer_table_types
|
||||
|
||||
|
||||
class _PartitionerLoader:
|
||||
"""Provides uniform helpful error when a partitioner dependency is not installed.
|
||||
|
||||
Used by `partition()` to encapsulate coping with the possibility the Python environment it is
|
||||
executing in may not have all dependencies installed for a particular partitioner.
|
||||
|
||||
Provides `.get()` to access partitioners by file-type, which raises when one or more
|
||||
dependencies for that partitioner are not installed.
|
||||
|
||||
The error message indicates what extra needs to be installed to enable that partitioner. This
|
||||
avoids an inconsistent variety of possibly puzzling exceptions arising from much deeper in the
|
||||
partitioner when access to the missing dependency is first attempted.
|
||||
"""
|
||||
|
||||
# -- module-lifetime cache for partitioners once loaded --
|
||||
_partitioners: dict[FileType, Partitioner] = {}
|
||||
|
||||
def get(self, file_type: FileType) -> Partitioner:
|
||||
"""Return partitioner for `file_type`.
|
||||
|
||||
Raises when one or more package dependencies for that file-type have not been
|
||||
installed. Also raises when the file-type is not partitionable.
|
||||
"""
|
||||
if not file_type.is_partitionable:
|
||||
raise UnsupportedFileFormatError(
|
||||
f"Partitioning is not supported for the {file_type} file type."
|
||||
)
|
||||
|
||||
# -- if the partitioner is not in the cache, load it; note this raises if one or more of
|
||||
# -- the partitioner's dependencies is not installed.
|
||||
if file_type not in self._partitioners:
|
||||
self._partitioners[file_type] = self._load_partitioner(file_type)
|
||||
|
||||
return self._partitioners[file_type]
|
||||
|
||||
def _load_partitioner(self, file_type: FileType) -> Partitioner:
|
||||
"""Load the partitioner for `file_type` after verifying dependencies."""
|
||||
# -- verify all package dependencies are installed --
|
||||
for pkg_name in file_type.importable_package_dependencies:
|
||||
if not dependency_exists(pkg_name):
|
||||
raise ImportError(
|
||||
f"{file_type.partitioner_function_name}() is not available because one or"
|
||||
f" more dependencies are not installed. Use:"
|
||||
f' pip install "unstructured[{file_type.extra_name}]" (including quotes)'
|
||||
f" to install the required dependencies",
|
||||
)
|
||||
|
||||
# -- load the partitioner and return it --
|
||||
assert file_type.is_partitionable # -- would be a programming error if this failed --
|
||||
partitioner_module = importlib.import_module(file_type.partitioner_module_qname)
|
||||
return getattr(partitioner_module, file_type.partitioner_function_name)
|
||||
@@ -0,0 +1,6 @@
|
||||
class UnsupportedFileFormatError(Exception):
|
||||
"""File-type is not supported for this operation.
|
||||
|
||||
For example, when receiving a file for auto-partitioning where its file-formatt cannot be
|
||||
identified or there is no partitioner available for that file-format.
|
||||
"""
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,460 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numbers
|
||||
import subprocess
|
||||
from enum import Enum
|
||||
from io import BufferedReader, BytesIO, TextIOWrapper
|
||||
from tempfile import SpooledTemporaryFile
|
||||
from time import sleep
|
||||
from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast
|
||||
|
||||
import emoji
|
||||
import psutil
|
||||
|
||||
from unstructured.documents.coordinates import CoordinateSystem, PixelSpace
|
||||
from unstructured.documents.elements import (
|
||||
TYPE_TO_TEXT_ELEMENT_MAP,
|
||||
CheckBox,
|
||||
CoordinatesMetadata,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
ElementType,
|
||||
ListItem,
|
||||
PageBreak,
|
||||
Text,
|
||||
)
|
||||
from unstructured.logger import logger
|
||||
from unstructured.nlp.patterns import ENUMERATED_BULLETS_RE, UNICODE_BULLETS_RE
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.layout import PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement
|
||||
|
||||
|
||||
def normalize_layout_element(
|
||||
layout_element: LayoutElement | Element | dict[str, Any],
|
||||
coordinate_system: Optional[CoordinateSystem] = None,
|
||||
infer_list_items: bool = True,
|
||||
source_format: Optional[str] = "html",
|
||||
) -> Element | list[Element]:
|
||||
"""Converts an unstructured_inference LayoutElement object to an unstructured Element."""
|
||||
|
||||
if isinstance(layout_element, Element) and source_format == "html":
|
||||
return layout_element
|
||||
|
||||
# NOTE(alan): Won't the lines above ensure this never runs (PageBreak is a subclass of Element)?
|
||||
if isinstance(layout_element, PageBreak):
|
||||
return PageBreak(text="")
|
||||
|
||||
if not isinstance(layout_element, dict):
|
||||
layout_dict = layout_element.to_dict()
|
||||
else:
|
||||
layout_dict = layout_element
|
||||
|
||||
text = layout_dict.get("text", "")
|
||||
# Both `coordinates` and `coordinate_system` must be present
|
||||
# in order to add coordinates metadata to the element.
|
||||
coordinates = layout_dict.get("coordinates") if coordinate_system else None
|
||||
element_type = layout_dict.get("type")
|
||||
prob = layout_dict.get("prob")
|
||||
aux_origin = layout_dict.get("source", None)
|
||||
origin = None
|
||||
if isinstance(layout_dict.get("is_extracted"), Enum):
|
||||
is_extracted = layout_dict["is_extracted"].value
|
||||
else:
|
||||
is_extracted = None
|
||||
if aux_origin:
|
||||
origin = aux_origin.value
|
||||
if prob and isinstance(prob, (int, str, float, numbers.Number)):
|
||||
class_prob_metadata = ElementMetadata(detection_class_prob=float(prob)) # type: ignore
|
||||
else:
|
||||
class_prob_metadata = ElementMetadata()
|
||||
class_prob_metadata.is_extracted = is_extracted
|
||||
common_kwargs = {
|
||||
"coordinates": coordinates,
|
||||
"coordinate_system": coordinate_system,
|
||||
"metadata": class_prob_metadata,
|
||||
"detection_origin": origin,
|
||||
}
|
||||
if element_type == ElementType.LIST:
|
||||
if infer_list_items:
|
||||
return layout_list_to_list_items(
|
||||
text,
|
||||
**common_kwargs,
|
||||
)
|
||||
else:
|
||||
return ListItem(
|
||||
text=text,
|
||||
**common_kwargs,
|
||||
)
|
||||
|
||||
elif element_type in TYPE_TO_TEXT_ELEMENT_MAP:
|
||||
assert isinstance(element_type, str) # Added to resolve type-error
|
||||
_element_class = TYPE_TO_TEXT_ELEMENT_MAP[element_type]
|
||||
_element_class = _element_class(
|
||||
text=text,
|
||||
**common_kwargs,
|
||||
)
|
||||
if element_type == ElementType.HEADLINE:
|
||||
_element_class.metadata.category_depth = 1
|
||||
elif element_type == ElementType.SUB_HEADLINE:
|
||||
_element_class.metadata.category_depth = 2
|
||||
return _element_class
|
||||
elif element_type in [
|
||||
ElementType.CHECK_BOX_CHECKED,
|
||||
ElementType.CHECK_BOX_UNCHECKED,
|
||||
ElementType.RADIO_BUTTON_CHECKED,
|
||||
ElementType.RADIO_BUTTON_UNCHECKED,
|
||||
ElementType.CHECKED,
|
||||
ElementType.UNCHECKED,
|
||||
]:
|
||||
checked = element_type in [
|
||||
ElementType.CHECK_BOX_CHECKED,
|
||||
ElementType.RADIO_BUTTON_CHECKED,
|
||||
ElementType.CHECKED,
|
||||
]
|
||||
return CheckBox(
|
||||
checked=checked,
|
||||
**common_kwargs,
|
||||
)
|
||||
else:
|
||||
return Text(
|
||||
text=text,
|
||||
**common_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def layout_list_to_list_items(
|
||||
text: Optional[str],
|
||||
coordinates: Optional[tuple[tuple[float, float], ...]],
|
||||
coordinate_system: Optional[CoordinateSystem],
|
||||
metadata: Optional[ElementMetadata],
|
||||
detection_origin: Optional[str],
|
||||
) -> list[Element]:
|
||||
"""Converts a list LayoutElement to a list of ListItem elements."""
|
||||
split_items = ENUMERATED_BULLETS_RE.split(text) if text else []
|
||||
# NOTE(robinson) - this means there wasn't a match for the enumerated bullets
|
||||
if len(split_items) == 1:
|
||||
split_items = UNICODE_BULLETS_RE.split(text) if text else []
|
||||
|
||||
list_items: list[Element] = []
|
||||
for text_segment in split_items:
|
||||
if len(text_segment.strip()) > 0:
|
||||
# Both `coordinates` and `coordinate_system` must be present
|
||||
# in order to add coordinates metadata to the element.
|
||||
item = ListItem(
|
||||
text=text_segment.strip(),
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
metadata=metadata,
|
||||
detection_origin=detection_origin,
|
||||
)
|
||||
list_items.append(item)
|
||||
|
||||
return list_items
|
||||
|
||||
|
||||
def add_element_metadata(
|
||||
element: Element,
|
||||
filename: Optional[str] = None,
|
||||
filetype: Optional[str] = None,
|
||||
page_number: Optional[int] = None,
|
||||
url: Optional[str] = None,
|
||||
text_as_html: Optional[str] = None,
|
||||
coordinates: Optional[tuple[tuple[float, float], ...]] = None,
|
||||
coordinate_system: Optional[CoordinateSystem] = None,
|
||||
image_path: Optional[str] = None,
|
||||
detection_origin: Optional[str] = None,
|
||||
languages: Optional[list[str]] = None,
|
||||
**kwargs: Any,
|
||||
) -> Element:
|
||||
"""Adds document metadata to the document element.
|
||||
|
||||
Document metadata includes information like the filename, source url, and page number.
|
||||
"""
|
||||
|
||||
coordinates_metadata = (
|
||||
CoordinatesMetadata(
|
||||
points=coordinates,
|
||||
system=coordinate_system,
|
||||
)
|
||||
if coordinates is not None and coordinate_system is not None
|
||||
else None
|
||||
)
|
||||
links = element.links if hasattr(element, "links") and len(element.links) > 0 else None
|
||||
link_urls = [link.get("url") for link in links] if links else None
|
||||
link_texts = [link.get("text") for link in links] if links else None
|
||||
link_start_indexes = [link.get("start_index") for link in links] if links else None
|
||||
emphasized_texts = (
|
||||
element.emphasized_texts
|
||||
if hasattr(element, "emphasized_texts") and len(element.emphasized_texts) > 0
|
||||
else None
|
||||
)
|
||||
emphasized_text_contents = (
|
||||
[emphasized_text.get("text") for emphasized_text in emphasized_texts]
|
||||
if emphasized_texts
|
||||
else None
|
||||
)
|
||||
emphasized_text_tags = (
|
||||
[emphasized_text.get("tag") for emphasized_text in emphasized_texts]
|
||||
if emphasized_texts
|
||||
else None
|
||||
)
|
||||
depth = element.metadata.category_depth if element.metadata.category_depth else None
|
||||
|
||||
metadata = ElementMetadata(
|
||||
coordinates=coordinates_metadata,
|
||||
filename=filename,
|
||||
filetype=filetype,
|
||||
page_number=page_number,
|
||||
url=url,
|
||||
text_as_html=text_as_html,
|
||||
link_urls=link_urls,
|
||||
link_texts=link_texts,
|
||||
link_start_indexes=link_start_indexes,
|
||||
emphasized_text_contents=emphasized_text_contents,
|
||||
emphasized_text_tags=emphasized_text_tags,
|
||||
category_depth=depth,
|
||||
image_path=image_path,
|
||||
languages=languages,
|
||||
)
|
||||
element.metadata.update(metadata)
|
||||
if detection_origin is not None:
|
||||
element.metadata.detection_origin = detection_origin
|
||||
return element
|
||||
|
||||
|
||||
def remove_element_metadata(layout_elements: list[Element]) -> list[Element]:
|
||||
"""Removes document metadata from the document element.
|
||||
|
||||
Document metadata includes information like the filename, source url, and page number.
|
||||
"""
|
||||
elements: list[Element] = []
|
||||
metadata = ElementMetadata()
|
||||
for layout_element in layout_elements:
|
||||
element = normalize_layout_element(layout_element)
|
||||
if isinstance(element, list):
|
||||
for _element in element:
|
||||
_element.metadata = metadata
|
||||
elements.extend(element)
|
||||
else:
|
||||
element.metadata = metadata
|
||||
elements.append(element)
|
||||
return elements
|
||||
|
||||
|
||||
def _is_soffice_running():
|
||||
for proc in psutil.process_iter():
|
||||
try:
|
||||
if "soffice" in proc.name().lower():
|
||||
return True
|
||||
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def convert_office_doc(
|
||||
input_filename: str,
|
||||
output_directory: str,
|
||||
target_format: str = "docx",
|
||||
target_filter: Optional[str] = None,
|
||||
wait_for_soffice_ready_time_out: int = 10,
|
||||
):
|
||||
"""Converts a .doc/.ppt file to a .docx/.pptx file using the libreoffice CLI.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_filename: str
|
||||
The name of the .doc file to convert to .docx
|
||||
output_directory: str
|
||||
The output directory for the convert .docx file
|
||||
target_format: str
|
||||
The desired output format
|
||||
target_filter: str
|
||||
The output filter name to use when converting. See references below
|
||||
for details.
|
||||
wait_for_soffice_ready_time_out: int
|
||||
The max wait time in seconds for soffice to become available to run
|
||||
|
||||
References
|
||||
----------
|
||||
https://stackoverflow.com/questions/52277264/convert-doc-to-docx-using-soffice-not-working
|
||||
https://git.libreoffice.org/core/+/refs/heads/master/filter/source/config/fragments/filters
|
||||
|
||||
"""
|
||||
if target_filter is not None:
|
||||
target_format = f"{target_format}:{target_filter}"
|
||||
# NOTE(robinson) - In the future can also include win32com client as a fallback for windows
|
||||
# users who do not have LibreOffice installed
|
||||
# ref: https://stackoverflow.com/questions/38468442/
|
||||
# multiple-doc-to-docx-file-conversion-using-python
|
||||
command = [
|
||||
"soffice",
|
||||
"--headless",
|
||||
"--convert-to",
|
||||
target_format,
|
||||
"--outdir",
|
||||
output_directory,
|
||||
input_filename,
|
||||
]
|
||||
try:
|
||||
# only one soffice process can be ran
|
||||
wait_time = 0
|
||||
sleep_time = 0.1
|
||||
output = subprocess.run(command, capture_output=True)
|
||||
message = output.stdout.decode().strip()
|
||||
# we can't rely on returncode unfortunately because on macOS it would return 0 even when the
|
||||
# command failed to run; instead we have to rely on the stdout being empty as a sign of the
|
||||
# process failed
|
||||
while (wait_time < wait_for_soffice_ready_time_out) and (message == ""):
|
||||
wait_time += sleep_time
|
||||
if _is_soffice_running():
|
||||
sleep(sleep_time)
|
||||
else:
|
||||
output = subprocess.run(command, capture_output=True)
|
||||
message = output.stdout.decode().strip()
|
||||
except FileNotFoundError:
|
||||
raise FileNotFoundError(
|
||||
"""soffice command was not found. Please install libreoffice
|
||||
on your system and try again.
|
||||
|
||||
- Install instructions: https://www.libreoffice.org/get-help/install-howto/
|
||||
- Mac: https://formulae.brew.sh/cask/libreoffice
|
||||
- Debian: https://wiki.debian.org/LibreOffice""",
|
||||
)
|
||||
|
||||
logger.info(message)
|
||||
if output.returncode != 0 or message == "":
|
||||
logger.error(
|
||||
"soffice failed to convert to format %s with code %i", target_format, output.returncode
|
||||
)
|
||||
logger.error(output.stderr.decode().strip())
|
||||
|
||||
|
||||
def exactly_one(**kwargs: Any) -> None:
|
||||
"""
|
||||
Verify arguments; exactly one of all keyword arguments must not be None.
|
||||
|
||||
Example:
|
||||
>>> exactly_one(filename=filename, file=file, text=text, url=url)
|
||||
"""
|
||||
if sum([(arg is not None and arg != "") for arg in kwargs.values()]) != 1:
|
||||
names = list(kwargs.keys())
|
||||
if len(names) > 1:
|
||||
message = f"Exactly one of {', '.join(names[:-1])} and {names[-1]} must be specified."
|
||||
else:
|
||||
message = f"{names[0]} must be specified."
|
||||
raise ValueError(message)
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def spooled_to_bytes_io_if_needed(file: _T | SpooledTemporaryFile[bytes]) -> _T | BytesIO:
|
||||
"""Convert `file` to `BytesIO` when it is a `SpooledTemporaryFile`.
|
||||
|
||||
Note that `file` does not need to be IO[bytes]. It can be `None` or `bytes` and this function
|
||||
will not complain.
|
||||
|
||||
In Python <3.11, `SpooledTemporaryFile` does not implement `.readable()` or `.seekable()` which
|
||||
triggers an exception when the file is loaded by certain packages. In particular, the stdlib
|
||||
`zipfile.Zipfile` raises on opening a `SpooledTemporaryFile` as does `Pandas.read_csv()`.
|
||||
"""
|
||||
if isinstance(file, SpooledTemporaryFile):
|
||||
file.seek(0)
|
||||
return BytesIO(cast(bytes, file.read()))
|
||||
|
||||
# -- return `file` unchanged otherwise --
|
||||
return file
|
||||
|
||||
|
||||
def convert_to_bytes(file: bytes | IO[bytes]) -> bytes:
|
||||
"""Extract the bytes from `file` without preventing it from being read again later.
|
||||
|
||||
As a convenience to simplify client code, also returns `file` unchanged if it is already bytes.
|
||||
"""
|
||||
if isinstance(file, bytes):
|
||||
return file
|
||||
|
||||
if isinstance(file, SpooledTemporaryFile):
|
||||
file.seek(0)
|
||||
f_bytes = file.read()
|
||||
file.seek(0)
|
||||
return f_bytes
|
||||
|
||||
if isinstance(file, BytesIO):
|
||||
return file.getvalue()
|
||||
|
||||
if isinstance(file, (TextIOWrapper, BufferedReader)):
|
||||
with open(file.name, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
raise ValueError("Invalid file-like object type")
|
||||
|
||||
|
||||
def contains_emoji(s: str) -> bool:
|
||||
"""
|
||||
Check if the input string contains any emoji characters.
|
||||
|
||||
Parameters:
|
||||
- s (str): The input string to check.
|
||||
|
||||
Returns:
|
||||
- bool: True if the string contains any emoji, False otherwise.
|
||||
"""
|
||||
|
||||
return bool(emoji.emoji_count(s))
|
||||
|
||||
|
||||
def get_page_image_metadata(page: PageLayout) -> dict[str, Any]:
|
||||
"""Retrieve image metadata and coordinate system from a page."""
|
||||
|
||||
image = getattr(page, "image", None)
|
||||
image_metadata = getattr(page, "image_metadata", None)
|
||||
|
||||
if image:
|
||||
image_format = image.format
|
||||
image_width = image.width
|
||||
image_height = image.height
|
||||
elif image_metadata:
|
||||
image_format = image_metadata.get("format")
|
||||
image_width = image_metadata.get("width")
|
||||
image_height = image_metadata.get("height")
|
||||
else:
|
||||
image_format = None
|
||||
image_width = None
|
||||
image_height = None
|
||||
|
||||
return {
|
||||
"format": image_format,
|
||||
"width": image_width,
|
||||
"height": image_height,
|
||||
}
|
||||
|
||||
|
||||
def ocr_data_to_elements(
|
||||
ocr_data: list["LayoutElement"],
|
||||
image_size: tuple[int | float, int | float],
|
||||
common_metadata: Optional[ElementMetadata] = None,
|
||||
infer_list_items: bool = True,
|
||||
source_format: Optional[str] = None,
|
||||
) -> list[Element]:
|
||||
"""Convert OCR layout data into `unstructured` elements with associated metadata."""
|
||||
|
||||
image_width, image_height = image_size
|
||||
coordinate_system = PixelSpace(width=image_width, height=image_height)
|
||||
elements: list[Element] = []
|
||||
for layout_element in ocr_data:
|
||||
element = normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
infer_list_items=infer_list_items,
|
||||
source_format=source_format if source_format else "html",
|
||||
)
|
||||
|
||||
if common_metadata:
|
||||
element.metadata.update(common_metadata)
|
||||
|
||||
elements.append(element)
|
||||
|
||||
return elements
|
||||
@@ -0,0 +1,538 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from functools import lru_cache
|
||||
from typing import Iterable, Iterator, Optional
|
||||
|
||||
import iso639 # pyright: ignore[reportMissingTypeStubs]
|
||||
from langdetect import ( # pyright: ignore[reportMissingTypeStubs]
|
||||
DetectorFactory,
|
||||
detect_langs, # pyright: ignore[reportUnknownVariableType]
|
||||
lang_detect_exception,
|
||||
)
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.utils.constants import (
|
||||
TESSERACT_LANGUAGES_AND_CODES,
|
||||
TESSERACT_LANGUAGES_SPLITTER,
|
||||
)
|
||||
|
||||
_ASCII_RE = re.compile(r"^[\x00-\x7F]+$")
|
||||
|
||||
# pytesseract.get_languages(config="") only shows user installed language packs,
|
||||
# so manually include the list of all currently supported Tesseract languages
|
||||
PYTESSERACT_LANG_CODES = [
|
||||
"afr",
|
||||
"amh",
|
||||
"ara",
|
||||
"asm",
|
||||
"aze",
|
||||
"aze_cyrl",
|
||||
"bel",
|
||||
"ben",
|
||||
"bod",
|
||||
"bos",
|
||||
"bre",
|
||||
"bul",
|
||||
"cat",
|
||||
"ceb",
|
||||
"ces",
|
||||
"chi_sim",
|
||||
"chi_sim_vert",
|
||||
"chi_tra",
|
||||
"chi_tra_vert",
|
||||
"chr",
|
||||
"cos",
|
||||
"cym",
|
||||
"dan",
|
||||
"deu",
|
||||
"div",
|
||||
"dzo",
|
||||
"ell",
|
||||
"eng",
|
||||
"enm",
|
||||
"epo",
|
||||
"equ",
|
||||
"est",
|
||||
"eus",
|
||||
"fao",
|
||||
"fas",
|
||||
"fil",
|
||||
"fin",
|
||||
"fra",
|
||||
"frk",
|
||||
"frm",
|
||||
"fry",
|
||||
"gla",
|
||||
"gle",
|
||||
"glg",
|
||||
"grc",
|
||||
"guj",
|
||||
"hat",
|
||||
"heb",
|
||||
"hin",
|
||||
"hrv",
|
||||
"hun",
|
||||
"hye",
|
||||
"iku",
|
||||
"ind",
|
||||
"isl",
|
||||
"ita",
|
||||
"ita_old",
|
||||
"jav",
|
||||
"jpn",
|
||||
"jpn_vert",
|
||||
"kan",
|
||||
"kat",
|
||||
"kat_old",
|
||||
"kaz",
|
||||
"khm",
|
||||
"kir",
|
||||
"kmr",
|
||||
"kor",
|
||||
"kor_vert",
|
||||
"lao",
|
||||
"lat",
|
||||
"lav",
|
||||
"lit",
|
||||
"ltz",
|
||||
"mal",
|
||||
"mar",
|
||||
"mkd",
|
||||
"mlt",
|
||||
"mon",
|
||||
"mri",
|
||||
"msa",
|
||||
"mya",
|
||||
"nep",
|
||||
"nld",
|
||||
"nor",
|
||||
"oci",
|
||||
"ori",
|
||||
"osd",
|
||||
"pan",
|
||||
"pol",
|
||||
"por",
|
||||
"pus",
|
||||
"que",
|
||||
"ron",
|
||||
"rus",
|
||||
"san",
|
||||
"sin",
|
||||
"slk",
|
||||
"slv",
|
||||
"snd",
|
||||
"snum",
|
||||
"spa",
|
||||
"spa_old",
|
||||
"sqi",
|
||||
"srp",
|
||||
"srp_latn",
|
||||
"sun",
|
||||
"swa",
|
||||
"swe",
|
||||
"syr",
|
||||
"tam",
|
||||
"tat",
|
||||
"tel",
|
||||
"tgk",
|
||||
"tha",
|
||||
"tir",
|
||||
"ton",
|
||||
"tur",
|
||||
"uig",
|
||||
"ukr",
|
||||
"urd",
|
||||
"uzb",
|
||||
"uzb_cyrl",
|
||||
"vie",
|
||||
"yid",
|
||||
"yor",
|
||||
]
|
||||
|
||||
PYTESSERACT_TO_PADDLE_LANG_CODE_MAP = {
|
||||
"afr": "af", # Afrikaans
|
||||
"ara": "ar", # Arabic
|
||||
"aze": "az", # Azerbaijani
|
||||
"bel": "be", # Belarusian
|
||||
"bos": "bs", # Bosnian
|
||||
"bul": "bg", # Bulgarian
|
||||
"ces": "cs", # Czech
|
||||
"chi_sim": "ch", # Simplified Chinese
|
||||
"chi_tra": "chinese_cht", # Traditional Chinese
|
||||
"cym": "cy", # Welsh
|
||||
"dan": "da", # Danish
|
||||
"deu": "german", # German
|
||||
"eng": "en", # English
|
||||
"est": "et", # Estonian
|
||||
"fas": "fa", # Persian
|
||||
"fra": "fr", # French
|
||||
"gle": "ga", # Irish
|
||||
"hin": "hi", # Hindi
|
||||
"hrv": "hr", # Croatian
|
||||
"hun": "hu", # Hungarian
|
||||
"ind": "id", # Indonesian
|
||||
"isl": "is", # Icelandic
|
||||
"ita": "it", # Italian
|
||||
"jpn": "japan", # Japanese
|
||||
"kor": "korean", # Korean
|
||||
"kmr": "ku", # Kurdish
|
||||
"lat": "rs_latin", # Latin
|
||||
"lav": "lv", # Latvian
|
||||
"lit": "lt", # Lithuanian
|
||||
"mar": "mr", # Marathi
|
||||
"mlt": "mt", # Maltese
|
||||
"msa": "ms", # Malay
|
||||
"nep": "ne", # Nepali
|
||||
"nld": "nl", # Dutch
|
||||
"nor": "no", # Norwegian
|
||||
"pol": "pl", # Polish
|
||||
"por": "pt", # Portuguese
|
||||
"ron": "ro", # Romanian
|
||||
"rus": "ru", # Russian
|
||||
"slk": "sk", # Slovak
|
||||
"slv": "sl", # Slovenian
|
||||
"spa": "es", # Spanish
|
||||
"sqi": "sq", # Albanian
|
||||
"srp": "rs_cyrillic", # Serbian
|
||||
"swa": "sw", # Swahili
|
||||
"swe": "sv", # Swedish
|
||||
"tam": "ta", # Tamil
|
||||
"tel": "te", # Telugu
|
||||
"tur": "tr", # Turkish
|
||||
"uig": "ug", # Uyghur
|
||||
"ukr": "uk", # Ukrainian
|
||||
"urd": "ur", # Urdu
|
||||
"uzb": "uz", # Uzbek
|
||||
"vie": "vi", # Vietnamese
|
||||
}
|
||||
|
||||
|
||||
def prepare_languages_for_tesseract(languages: Optional[list[str]] = ["eng"]) -> str:
|
||||
"""
|
||||
Entry point: convert languages (list of strings) into tesseract ocr langcode format (uses +)
|
||||
"""
|
||||
if languages is None:
|
||||
raise ValueError("`languages` can not be `None`")
|
||||
converted_languages = [
|
||||
lang_code
|
||||
for lang_code in (
|
||||
_convert_language_code_to_pytesseract_lang_code(lang) for lang in languages
|
||||
)
|
||||
if lang_code
|
||||
]
|
||||
# Remove duplicates from the list but keep the original order
|
||||
converted_languages = list(dict.fromkeys(converted_languages))
|
||||
if len(converted_languages) == 0:
|
||||
logger.warning(
|
||||
"Failed to find any valid standard language code from "
|
||||
f"languages: {languages}, proceed with `eng` instead.",
|
||||
)
|
||||
return "eng"
|
||||
|
||||
return TESSERACT_LANGUAGES_SPLITTER.join(converted_languages)
|
||||
|
||||
|
||||
def tesseract_to_paddle_language(tesseract_language: str) -> str:
|
||||
"""
|
||||
Convert TesseractOCR language code to PaddleOCR language code.
|
||||
|
||||
:param tesseract_language: str, language code used in TesseractOCR
|
||||
:return: str, corresponding language code for PaddleOCR or None if not found
|
||||
"""
|
||||
|
||||
lang = PYTESSERACT_TO_PADDLE_LANG_CODE_MAP.get(tesseract_language.lower())
|
||||
if not lang:
|
||||
logger.warning(
|
||||
f"{tesseract_language} is not a language code supported by PaddleOCR, "
|
||||
f"proceeding with `en` instead."
|
||||
)
|
||||
return "en"
|
||||
|
||||
return lang
|
||||
|
||||
|
||||
def check_language_args(
|
||||
languages: list[str], ocr_languages: str | list[str] | None
|
||||
) -> list[str] | None:
|
||||
"""Handle users defining both `ocr_languages` and `languages`.
|
||||
|
||||
Give preference to `languages` and convert `ocr_languages` if needed, but default to `None`.
|
||||
|
||||
`ocr_languages` is only a parameter for `auto.partition`, `partition_image`, & `partition_pdf`.
|
||||
`ocr_languages` should not be defined as 'auto' since 'auto' is intended for language detection
|
||||
which is not supported by `partition_image` or `partition_pdf`.
|
||||
"""
|
||||
# --- Clean and update defaults
|
||||
if ocr_languages:
|
||||
ocr_languages = _clean_ocr_languages_arg(ocr_languages)
|
||||
logger.warning(
|
||||
"The ocr_languages kwarg will be deprecated in a future version of unstructured. "
|
||||
"Please use languages instead.",
|
||||
)
|
||||
assert ocr_languages is None or isinstance(ocr_languages, str)
|
||||
|
||||
if ocr_languages and "auto" in ocr_languages:
|
||||
raise ValueError(
|
||||
"`ocr_languages` is deprecated but was used to extract text from pdfs and images."
|
||||
" The 'auto' argument is only for language *detection* when it is assigned"
|
||||
" to `languages` and partitioning documents other than pdfs or images."
|
||||
" Language detection is not currently supported in pdfs or images."
|
||||
)
|
||||
|
||||
if not isinstance(languages, list): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
raise TypeError(
|
||||
"The language parameter must be a list of language codes as strings, ex. ['eng']",
|
||||
)
|
||||
|
||||
# --- If `languages` is a null/default value and `ocr_languages` is defined, use `ocr_languages`
|
||||
if ocr_languages and (languages == ["auto"] or languages == [""] or not languages):
|
||||
languages = ocr_languages.split(TESSERACT_LANGUAGES_SPLITTER)
|
||||
logger.warning(
|
||||
"Only one of languages and ocr_languages should be specified. "
|
||||
"languages is preferred. ocr_languages is marked for deprecation.",
|
||||
)
|
||||
|
||||
# --- Clean `languages`
|
||||
# If "auto" is included in the list of inputs, language detection will be triggered downstream.
|
||||
# The rest of the inputted languages are ignored.
|
||||
if languages:
|
||||
if "auto" not in languages:
|
||||
for i, lang in enumerate(languages):
|
||||
languages[i] = TESSERACT_LANGUAGES_AND_CODES.get(lang.lower(), lang)
|
||||
|
||||
str_languages = _clean_ocr_languages_arg(languages)
|
||||
if not str_languages:
|
||||
return None
|
||||
languages = str_languages.split(TESSERACT_LANGUAGES_SPLITTER)
|
||||
# else, remove the extraneous languages.
|
||||
# NOTE (jennings): "auto" should only be used for partitioners OTHER THAN `_pdf` or `_image`
|
||||
else:
|
||||
# define as 'auto' for language detection when partitioning non-pdfs or -images
|
||||
languages = ["auto"]
|
||||
return languages
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def convert_old_ocr_languages_to_languages(ocr_languages: str) -> list[str]:
|
||||
"""
|
||||
Convert ocr_languages parameter to list of langcode strings.
|
||||
Assumption: ocr_languages is in tesseract plus sign format
|
||||
"""
|
||||
|
||||
return ocr_languages.split(TESSERACT_LANGUAGES_SPLITTER)
|
||||
|
||||
|
||||
def _convert_language_code_to_pytesseract_lang_code(lang: str) -> str:
|
||||
"""
|
||||
Convert a single language code to its tesseract formatted and recognized
|
||||
langcode(s), if supported.
|
||||
"""
|
||||
# if language is already tesseract langcode, return it immediately
|
||||
# this will catch the tesseract special cases equ and osd
|
||||
# NOTE(shreya): this may catch some cases of choosing between tesseract code variants for a lang
|
||||
if lang in PYTESSERACT_LANG_CODES:
|
||||
return lang
|
||||
|
||||
lang_iso639 = _get_iso639_language_object(lang)
|
||||
|
||||
# tesseract uses 3 digit codes (639-3, 639-2b, etc) as prefixes, with suffixes for orthography
|
||||
# use first 3 letters of tesseract codes for matching to standard codes
|
||||
pytesseract_langs_3 = {lang[:3] for lang in PYTESSERACT_LANG_CODES}
|
||||
|
||||
if lang_iso639:
|
||||
# try to match ISO 639-3 code
|
||||
if lang_iso639.part3 in pytesseract_langs_3:
|
||||
matched_langcodes = _get_all_tesseract_langcodes_with_prefix(lang_iso639.part3)
|
||||
return TESSERACT_LANGUAGES_SPLITTER.join(matched_langcodes)
|
||||
|
||||
# try to match ISO 639-2b
|
||||
elif lang_iso639.part2b in pytesseract_langs_3:
|
||||
matched_langcodes = _get_all_tesseract_langcodes_with_prefix(lang_iso639.part2b)
|
||||
return TESSERACT_LANGUAGES_SPLITTER.join(matched_langcodes)
|
||||
|
||||
# try to match ISO 639-2t
|
||||
elif lang_iso639.part2t in pytesseract_langs_3:
|
||||
matched_langcodes = _get_all_tesseract_langcodes_with_prefix(lang_iso639.part2t)
|
||||
return TESSERACT_LANGUAGES_SPLITTER.join(matched_langcodes)
|
||||
|
||||
else:
|
||||
logger.warning(f"{lang} is not a language supported by Tesseract.")
|
||||
return ""
|
||||
logger.warning(f"{lang} is not a language supported by Tesseract.")
|
||||
return ""
|
||||
|
||||
|
||||
def _get_iso639_language_object(lang: str) -> Optional[iso639.Language]:
|
||||
language = _cached_iso639_language_match(lang)
|
||||
if language is not None:
|
||||
return language
|
||||
logger.warning(f"{lang} is not a valid standard language code.")
|
||||
return None
|
||||
|
||||
|
||||
def _get_all_tesseract_langcodes_with_prefix(prefix: str) -> list[str]:
|
||||
"""
|
||||
Get all matching tesseract langcodes with this prefix (may be one or multiple variants).
|
||||
"""
|
||||
return [langcode for langcode in PYTESSERACT_LANG_CODES if langcode.startswith(prefix)]
|
||||
|
||||
|
||||
def detect_languages(
|
||||
text: str,
|
||||
languages: Optional[list[str]] = ["auto"],
|
||||
) -> Optional[list[str]]:
|
||||
"""
|
||||
Detects the list of languages present in the text (in the default "auto" mode),
|
||||
or formats and passes through the user inputted document languages if provided.
|
||||
"""
|
||||
if languages is None:
|
||||
languages = ["auto"]
|
||||
if not isinstance(languages, list):
|
||||
raise TypeError(
|
||||
'The language parameter must be a list of language codes as strings, ex. ["eng"]',
|
||||
)
|
||||
|
||||
# Skip language detection for partitioners that use other partitioners.
|
||||
# For example, partition_msg relies on partition_html and partition_text, but the metadata
|
||||
# gets overwritten after elements have been returned by _html and _text,
|
||||
# so `languages` would be detected twice.
|
||||
# Also return None if there is no text.
|
||||
if languages[0] == "" or text.strip() == "":
|
||||
return None
|
||||
|
||||
# If text contains special characters (like ñ, å, or Korean/Mandarin/etc.) it will NOT default
|
||||
# to English. It will default to English if text is only ascii characters and is short.
|
||||
if _ASCII_RE.match(text) and len(text.split()) < 5:
|
||||
logger.debug(f'short text: "{text}". Defaulting to English.')
|
||||
return ["eng"]
|
||||
|
||||
# set seed for deterministic langdetect outputs
|
||||
DetectorFactory.seed = 0
|
||||
|
||||
doc_languages: list[str] = []
|
||||
|
||||
# user inputted languages:
|
||||
# if "auto" is included in the list of inputs, language detection will be triggered
|
||||
# and the rest of the inputted languages will be ignored
|
||||
if languages and "auto" not in languages:
|
||||
for lang in languages:
|
||||
str_lang = TESSERACT_LANGUAGES_AND_CODES.get(lang.lower(), lang)
|
||||
language = _get_iso639_language_object(str_lang[:3])
|
||||
if language:
|
||||
doc_languages.append(language.part3)
|
||||
|
||||
# language detection:
|
||||
else:
|
||||
# warn if any values other than "auto" were provided
|
||||
if len(languages) > 1:
|
||||
logger.warning(
|
||||
f'Since "auto" is present in the input languages provided ({languages}), '
|
||||
"the language will be auto detected and the rest of the inputted "
|
||||
"languages will be ignored.",
|
||||
)
|
||||
|
||||
try:
|
||||
langdetect_result = detect_langs(text)
|
||||
except lang_detect_exception.LangDetectException as e:
|
||||
logger.warning(e)
|
||||
return None # None as default
|
||||
|
||||
langdetect_langs: list[str] = []
|
||||
|
||||
# NOTE(robinson) - Chinese gets detected with codes zh-cn, zh-tw, zh-hk for various
|
||||
# Chinese variants. We normalizes these because there is a single model for Chinese
|
||||
# machine translation
|
||||
# TODO(shreya): decide how to maintain nonstandard chinese script information
|
||||
for langobj in langdetect_result:
|
||||
lang_val = str(langobj.lang)
|
||||
if lang_val.startswith("zh"): # pyright: ignore
|
||||
langdetect_langs.append("zho")
|
||||
else:
|
||||
language = _get_iso639_language_object(lang_val[:3]) # pyright: ignore
|
||||
if language:
|
||||
langdetect_langs.append(language.part3)
|
||||
|
||||
# remove duplicate chinese (if exists) without modifying order
|
||||
seen = set(doc_languages)
|
||||
for lang in langdetect_langs:
|
||||
if lang not in seen:
|
||||
doc_languages.append(lang)
|
||||
seen.add(lang)
|
||||
|
||||
return doc_languages
|
||||
|
||||
|
||||
def apply_lang_metadata(
|
||||
elements: Iterable[Element],
|
||||
languages: Optional[list[str]],
|
||||
detect_language_per_element: bool = False,
|
||||
) -> Iterator[Element]:
|
||||
"""Detect language and apply it to metadata.languages for each element in `elements`.
|
||||
If languages is None, default to auto detection.
|
||||
If languages is and empty string, skip."""
|
||||
# -- Note this function has a stream interface, but reads the full `elements` stream into memory
|
||||
# -- before emitting the first updated element as output.
|
||||
|
||||
# The auto `partition` function uses `None` as a default because the default for
|
||||
# `partition_pdf` and `partition_img` conflict with the other partitioners that use ["auto"]
|
||||
if languages is None:
|
||||
languages = ["auto"]
|
||||
|
||||
# Skip language detection for partitioners that use other partitioners.
|
||||
# For example, partition_msg relies on partition_html and partition_text, but the metadata
|
||||
# gets overwritten after elements have been returned by _html and _text,
|
||||
# so `languages` would be detected twice.
|
||||
if languages == [""]:
|
||||
yield from elements
|
||||
return
|
||||
|
||||
# Convert elements to a list to get the text, detect the language, and add it to the elements
|
||||
if not isinstance(elements, list):
|
||||
elements = list(elements)
|
||||
|
||||
full_text = " ".join(str(e.text) for e in elements if hasattr(e, "text") and e.text)
|
||||
detected_languages = detect_languages(text=full_text, languages=languages)
|
||||
if (
|
||||
detected_languages is not None
|
||||
and len(detected_languages) == 1
|
||||
and detect_language_per_element is False
|
||||
):
|
||||
# -- apply detected language to each element's metadata --
|
||||
for e in elements:
|
||||
e.metadata.languages = detected_languages
|
||||
yield e
|
||||
else:
|
||||
for e in elements:
|
||||
if hasattr(e, "text"):
|
||||
text_value = str(e.text) if e.text is not None else ""
|
||||
e.metadata.languages = detect_languages(text_value)
|
||||
yield e
|
||||
else:
|
||||
yield e
|
||||
|
||||
|
||||
def _clean_ocr_languages_arg(ocr_languages: list[str] | str) -> str:
|
||||
"""Fix common incorrect definitions for ocr_languages:
|
||||
defining it as a list, adding extra quotation marks, adding brackets.
|
||||
Returns a single string of ocr_languages"""
|
||||
# extract from list
|
||||
if isinstance(ocr_languages, list):
|
||||
ocr_languages = "+".join(ocr_languages)
|
||||
|
||||
# remove extra quotations
|
||||
ocr_languages = re.sub(r"[\"']", "", ocr_languages)
|
||||
# remove brackets
|
||||
ocr_languages = re.sub(r"[\[\]]", "", ocr_languages)
|
||||
|
||||
return ocr_languages
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _cached_iso639_language_match(lang: str) -> Optional[iso639.Language]:
|
||||
try:
|
||||
return iso639.Language.match(lang.lower()) # pyright: ignore[reportUnknownMemberType]
|
||||
except iso639.LanguageNotFoundError:
|
||||
return None
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Helpers used across multiple partitioners to compute metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import datetime as dt
|
||||
import functools
|
||||
import os
|
||||
from typing import Any, Callable, Iterator, Sequence
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from unstructured.documents.elements import Element, ElementMetadata
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.lang import apply_lang_metadata
|
||||
from unstructured.utils import get_call_args_applying_defaults
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
def get_last_modified_date(filename: str) -> str | None:
|
||||
"""Modification time of file at path `filename`, if it exists.
|
||||
|
||||
Returns `None` when `filename` is not a path to a file on the local filesystem.
|
||||
|
||||
Otherwise returns date and time in ISO 8601 string format (YYYY-MM-DDTHH:MM:SS) like
|
||||
"2024-03-05T17:02:53".
|
||||
"""
|
||||
if not os.path.isfile(filename):
|
||||
return None
|
||||
|
||||
modify_date = dt.datetime.fromtimestamp(os.path.getmtime(filename))
|
||||
return modify_date.strftime("%Y-%m-%dT%H:%M:%S%z")
|
||||
|
||||
|
||||
HIERARCHY_RULE_SET = {
|
||||
"Title": [
|
||||
"Text",
|
||||
"UncategorizedText",
|
||||
"NarrativeText",
|
||||
"ListItem",
|
||||
"BulletedText",
|
||||
"Table",
|
||||
"FigureCaption",
|
||||
"CheckBox",
|
||||
"Table",
|
||||
],
|
||||
"Header": [
|
||||
"Title",
|
||||
"Text",
|
||||
"UncategorizedText",
|
||||
"NarrativeText",
|
||||
"ListItem",
|
||||
"BulletedText",
|
||||
"Table",
|
||||
"FigureCaption",
|
||||
"CheckBox",
|
||||
"Table",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def set_element_hierarchy(
|
||||
elements: Sequence[Element], ruleset: dict[str, list[str]] = HIERARCHY_RULE_SET
|
||||
) -> list[Element]:
|
||||
"""Sets `.metadata.parent_id` for each element it applies to.
|
||||
|
||||
`parent_id` assignment is based on the element's category and depth. The importance of an
|
||||
element's category is determined by a rule set. The rule set trumps category_depth. That is,
|
||||
category_depth is only relevant when elements are of the same category.
|
||||
"""
|
||||
stack: list[Element] = []
|
||||
for element in elements:
|
||||
if element.metadata.parent_id is not None:
|
||||
continue
|
||||
parent_id = None
|
||||
element_category = getattr(element, "category", None)
|
||||
element_category_depth = getattr(element.metadata, "category_depth", 0) or 0
|
||||
|
||||
# -- skip elements without a category --
|
||||
if not element_category:
|
||||
continue
|
||||
|
||||
while stack:
|
||||
top_element: Element = stack[-1]
|
||||
top_element_category = getattr(top_element, "category")
|
||||
top_element_category_depth = (
|
||||
getattr(
|
||||
top_element.metadata,
|
||||
"category_depth",
|
||||
0,
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
if (
|
||||
top_element_category == element_category
|
||||
and top_element_category_depth < element_category_depth
|
||||
) or (
|
||||
top_element_category != element_category
|
||||
and element_category in ruleset.get(top_element_category, [])
|
||||
):
|
||||
parent_id = top_element.id
|
||||
break
|
||||
|
||||
stack.pop()
|
||||
|
||||
element.metadata.parent_id = parent_id
|
||||
stack.append(element)
|
||||
|
||||
return list(elements)
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# METADATA POST-PARTITIONING PROCESSING DECORATOR
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
def apply_metadata(
|
||||
file_type: FileType | None = None,
|
||||
) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
|
||||
"""Post-process element-metadata for this document.
|
||||
|
||||
This decorator adds a post-processing step to a partitioner, primarily to apply metadata that
|
||||
is common to all partitioners. It assumes the following responsibilities:
|
||||
|
||||
- Hash element-ids. Computes and applies SHA1 hash element.id when `unique_element_ids`
|
||||
argument is False.
|
||||
|
||||
- Element Hierarchy. Computes and applies `parent_id` metadata based on `category_depth`
|
||||
etc. added by partitioner.
|
||||
|
||||
- Language metadata. Computes and applies `language` metadata based on a language detection
|
||||
model.
|
||||
|
||||
- Apply `filetype` (MIME-type) metadata. There are three cases; first one in this order that
|
||||
applies is used:
|
||||
|
||||
- `metadata_file_type` argument is present in call, use that.
|
||||
- `file_type` decorator argument is populated, use that.
|
||||
- `file_type` decorator argument is omitted or None, don't apply `.metadata.filetype`
|
||||
(assume the partitioner will do that for itself, like `partition_image()`.
|
||||
|
||||
- Replace `filename` with `metadata_filename` when present.
|
||||
|
||||
- Replace `last_modified` with `metadata_last_modified` when present.
|
||||
|
||||
- Apply `url` metadata when present.
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
"""The decorator function itself.
|
||||
|
||||
This function is returned by the `apply_metadata()` function and is the actual decorator.
|
||||
Think of `apply_metadata()` as a factory function that configures this decorator, in
|
||||
particular by setting its `file_type` value.
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
|
||||
elements = func(*args, **kwargs)
|
||||
call_args = get_call_args_applying_defaults(func, *args, **kwargs)
|
||||
|
||||
# ------------------------------------------------------------------------------------
|
||||
# unique-ify elements
|
||||
# ------------------------------------------------------------------------------------
|
||||
# Do this first to ensure all following operations behave as expected. It's easy for a
|
||||
# partitioner to re-use an element or metadata instance when its values are common to
|
||||
# multiple elements. This can lead to very hard-to diagnose bugs downstream when
|
||||
# mutating one element unexpectedly also mutates others (because they are the same
|
||||
# instance).
|
||||
# ------------------------------------------------------------------------------------
|
||||
|
||||
elements = _uniqueify_elements_and_metadata(elements)
|
||||
|
||||
# ------------------------------------------------------------------------------------
|
||||
# apply metadata - do this first because it affects the hash computation.
|
||||
# ------------------------------------------------------------------------------------
|
||||
|
||||
# -- `language` - auto-detect language (e.g. eng, spa) --
|
||||
languages = call_args.get("languages")
|
||||
detect_language_per_element = call_args.get("detect_language_per_element", False)
|
||||
elements = list(
|
||||
apply_lang_metadata(
|
||||
elements=elements,
|
||||
languages=languages,
|
||||
detect_language_per_element=detect_language_per_element,
|
||||
)
|
||||
)
|
||||
|
||||
# == apply filetype, filename, last_modified, and url metadata ===================
|
||||
metadata_kwargs: dict[str, Any] = {}
|
||||
|
||||
# -- `filetype` (MIME-type) metadata --
|
||||
metadata_file_type = call_args.get("metadata_file_type") or file_type
|
||||
if metadata_file_type is not None:
|
||||
metadata_kwargs["filetype"] = metadata_file_type.mime_type
|
||||
|
||||
# -- `filename` metadata - override with metadata_filename when it's present --
|
||||
filename = call_args.get("metadata_filename") or call_args.get("filename")
|
||||
if filename:
|
||||
metadata_kwargs["filename"] = filename
|
||||
|
||||
# -- `last_modified` metadata - override with metadata_last_modified when present --
|
||||
metadata_last_modified = call_args.get("metadata_last_modified")
|
||||
if metadata_last_modified:
|
||||
metadata_kwargs["last_modified"] = metadata_last_modified
|
||||
|
||||
# -- `url` metadata - record url when present --
|
||||
url = call_args.get("url")
|
||||
if url:
|
||||
metadata_kwargs["url"] = url
|
||||
|
||||
# -- update element.metadata in single pass --
|
||||
for element in elements:
|
||||
# NOTE(robinson) - Attached files have already run through this logic in their own
|
||||
# partitioning function
|
||||
if element.metadata.attached_to_filename:
|
||||
continue
|
||||
element.metadata.update(ElementMetadata(**metadata_kwargs))
|
||||
|
||||
# ------------------------------------------------------------------------------------
|
||||
# compute hash ids (when so requestsd)
|
||||
# ------------------------------------------------------------------------------------
|
||||
|
||||
# -- Compute and apply hash-ids if the user does not want UUIDs. Note this mutates the
|
||||
# -- elements themselves, not their metadata.
|
||||
unique_element_ids: bool = call_args.get("unique_element_ids", False)
|
||||
if unique_element_ids is False:
|
||||
elements = _assign_hash_ids(elements)
|
||||
|
||||
# ------------------------------------------------------------------------------------
|
||||
# assign parent-id - do this after hash computation so parent-id is stable.
|
||||
# ------------------------------------------------------------------------------------
|
||||
|
||||
# -- `parent_id` - process category-level etc. to assign parent-id --
|
||||
elements = set_element_hierarchy(elements)
|
||||
|
||||
return elements
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def _assign_hash_ids(elements: list[Element]) -> list[Element]:
|
||||
"""Converts `.id` of each element from UUID to hash.
|
||||
|
||||
The hash is based on the `.text` of the element, but also on its page-number and sequence number
|
||||
on that page. This provides for deterministic results even when the document is split into one
|
||||
or more fragments for parallel processing.
|
||||
"""
|
||||
# -- generate sequence number for each element on a page --
|
||||
page_seq_counts = {}
|
||||
for element in elements:
|
||||
page_number = element.metadata.page_number
|
||||
seq_on_page_counter = page_seq_counts.get(page_number, 0)
|
||||
element.id_to_hash(seq_on_page_counter)
|
||||
page_seq_counts[page_number] = seq_on_page_counter + 1
|
||||
|
||||
return elements
|
||||
|
||||
|
||||
def _uniqueify_elements_and_metadata(elements: list[Element]) -> list[Element]:
|
||||
"""Ensure each of `elements` and their metadata are unique instances.
|
||||
|
||||
This prevents hard-to-diagnose bugs downstream when mutating one element unexpectedly also
|
||||
mutates others because they are the same instance.
|
||||
"""
|
||||
|
||||
def iter_unique_elements(elements: list[Element]) -> Iterator[Element]:
|
||||
"""Substitute deep-copies of any non-unique elements or metadata in `elements`."""
|
||||
seen_elements: set[int] = set()
|
||||
seen_metadata: set[int] = set()
|
||||
|
||||
for element in elements:
|
||||
if id(element) in seen_elements:
|
||||
element = copy.deepcopy(element)
|
||||
if id(element.metadata) in seen_metadata:
|
||||
element.metadata = copy.deepcopy(element.metadata)
|
||||
seen_elements.add(id(element))
|
||||
seen_metadata.add(id(element.metadata))
|
||||
yield element
|
||||
|
||||
return list(iter_unique_elements(elements))
|
||||
@@ -0,0 +1,182 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import csv
|
||||
from typing import IO, Any, Iterator
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.common.html_table import HtmlTable
|
||||
from unstructured.documents.elements import Element, ElementMetadata, Table
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
from unstructured.utils import is_temp_file_path, lazyproperty
|
||||
|
||||
DETECTION_ORIGIN: str = "csv"
|
||||
CSV_FIELD_LIMIT = 10 * 1048576 # 10MiB
|
||||
|
||||
|
||||
@apply_metadata(FileType.CSV)
|
||||
@add_chunking_strategy
|
||||
def partition_csv(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
encoding: str | None = None,
|
||||
include_header: bool = False,
|
||||
infer_table_structure: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions Microsoft Excel Documents in .csv format into its document elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
encoding
|
||||
The encoding method used to decode the text input. If None, utf-8 will be used.
|
||||
include_header
|
||||
Determines whether or not header info info is included in text and medatada.text_as_html.
|
||||
infer_table_structure
|
||||
If True, any Table elements that are extracted will also have a metadata field
|
||||
named "text_as_html" where the table's text content is rendered into an html string.
|
||||
I.e., rows and cells are preserved.
|
||||
Whether True or False, the "text" field is always present in any Table element
|
||||
and is the text content of the table (no structure).
|
||||
"""
|
||||
ctx = _CsvPartitioningContext.load(
|
||||
file_path=filename,
|
||||
file=file,
|
||||
encoding=encoding,
|
||||
include_header=include_header,
|
||||
infer_table_structure=infer_table_structure,
|
||||
)
|
||||
|
||||
csv.field_size_limit(CSV_FIELD_LIMIT)
|
||||
with ctx.open() as file:
|
||||
dataframe = pd.read_csv(file, header=ctx.header, sep=ctx.delimiter, encoding=ctx.encoding)
|
||||
|
||||
html_table = HtmlTable.from_html_text(
|
||||
dataframe.to_html(index=False, header=include_header, na_rep="")
|
||||
)
|
||||
|
||||
metadata = ElementMetadata(
|
||||
filename=filename,
|
||||
last_modified=ctx.last_modified,
|
||||
text_as_html=html_table.html if infer_table_structure else None,
|
||||
)
|
||||
|
||||
# -- a CSV file becomes a single `Table` element --
|
||||
return [Table(text=html_table.text, metadata=metadata, detection_origin=DETECTION_ORIGIN)]
|
||||
|
||||
|
||||
class _CsvPartitioningContext:
|
||||
"""Encapsulates the partitioning-run details.
|
||||
|
||||
Provides access to argument values and especially encapsulates computation of values derived
|
||||
from those values so they don't obscure the core partitioning logic.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str | None = None,
|
||||
file: IO[bytes] | None = None,
|
||||
encoding: str | None = None,
|
||||
include_header: bool = False,
|
||||
infer_table_structure: bool = True,
|
||||
):
|
||||
self._file_path = file_path
|
||||
self._file = file
|
||||
self._encoding = encoding
|
||||
self._include_header = include_header
|
||||
self._infer_table_structure = infer_table_structure
|
||||
|
||||
@classmethod
|
||||
def load(
|
||||
cls,
|
||||
file_path: str | None,
|
||||
file: IO[bytes] | None,
|
||||
encoding: str | None,
|
||||
include_header: bool,
|
||||
infer_table_structure: bool,
|
||||
) -> _CsvPartitioningContext:
|
||||
return cls(
|
||||
file_path=file_path,
|
||||
file=file,
|
||||
encoding=encoding,
|
||||
include_header=include_header,
|
||||
infer_table_structure=infer_table_structure,
|
||||
)._validate()
|
||||
|
||||
@lazyproperty
|
||||
def delimiter(self) -> str | None:
|
||||
"""The CSV delimiter, nominally a comma ",".
|
||||
|
||||
`None` for a single-column CSV file which naturally has no delimiter.
|
||||
"""
|
||||
sniffer = csv.Sniffer()
|
||||
num_bytes = 65536
|
||||
|
||||
with self.open() as file:
|
||||
# -- read whole lines, sniffer can be confused by a trailing partial line --
|
||||
data = "\n".join(
|
||||
ln.decode(self._encoding or "utf-8") for ln in file.readlines(num_bytes)
|
||||
)
|
||||
|
||||
try:
|
||||
return sniffer.sniff(data, delimiters=",;|").delimiter
|
||||
except csv.Error:
|
||||
# -- sniffing will fail on single-column csv as no default can be assumed --
|
||||
return None
|
||||
|
||||
@lazyproperty
|
||||
def header(self) -> int | None:
|
||||
"""Identifies the header row, if any, to Pandas, by idx."""
|
||||
return 0 if self._include_header else None
|
||||
|
||||
@lazyproperty
|
||||
def encoding(self) -> str | None:
|
||||
"""The encoding to use for reading the file."""
|
||||
return self._encoding
|
||||
|
||||
@lazyproperty
|
||||
def last_modified(self) -> str | None:
|
||||
"""The best last-modified date available, None if no sources are available."""
|
||||
return (
|
||||
None
|
||||
if not self._file_path or is_temp_file_path(self._file_path)
|
||||
else get_last_modified_date(self._file_path)
|
||||
)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open(self) -> Iterator[IO[bytes]]:
|
||||
"""Encapsulates complexity of dealing with file-path or file-like-object.
|
||||
|
||||
Provides an `IO[bytes]` object as the "common-denominator" document source.
|
||||
|
||||
Must be used as a context manager using a `with` statement:
|
||||
|
||||
with self._file as file:
|
||||
do things with file
|
||||
|
||||
File is guaranteed to be at read position 0 when called.
|
||||
"""
|
||||
if self._file_path:
|
||||
with open(self._file_path, "rb") as f:
|
||||
yield f
|
||||
else:
|
||||
file = self._file
|
||||
assert file is not None # -- guaranteed by `._validate()` --
|
||||
# -- Be polite on principle. Reset file-pointer both before and after use --
|
||||
file.seek(0)
|
||||
yield file
|
||||
file.seek(0)
|
||||
|
||||
def _validate(self) -> _CsvPartitioningContext:
|
||||
"""Raise on invalid argument values."""
|
||||
if self._file_path is None and self._file is None:
|
||||
raise ValueError("either file-path or file-like object must be provided")
|
||||
return self
|
||||
@@ -0,0 +1,103 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import convert_office_doc, exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.docx import partition_docx
|
||||
|
||||
|
||||
def partition_doc(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
libre_office_filter: Optional[str] = "MS Word 2007 XML",
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions Microsoft Word Documents in .doc format into its document elements.
|
||||
|
||||
All parameters available on `partition_docx()` are also available here.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
libre_office_filter
|
||||
The filter to use when coverting to .doc. The default is the
|
||||
filter that is required when using LibreOffice7. Pass in None
|
||||
if you do not want to apply any filter.
|
||||
languages
|
||||
User defined value for `metadata.languages` if provided. Otherwise language is detected
|
||||
using naive Bayesian filter via `langdetect`. Multiple languages indicates text could be
|
||||
in either language.
|
||||
Additional Parameters:
|
||||
detect_language_per_element
|
||||
Detect language per element instead of at the document level.
|
||||
starting_page_number
|
||||
Indicates what page number should be assigned to the first page in the document.
|
||||
This information will be reflected in elements' metadata and can be be especially
|
||||
useful when partitioning a document that is part of a larger document.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
# -- validate file-path when provided so we can provide a more meaningful error --
|
||||
if filename is not None and not os.path.exists(filename):
|
||||
raise ValueError(f"The file {filename} does not exist.")
|
||||
|
||||
# -- `convert_office_doc` uses a command-line program that ships with LibreOffice to convert
|
||||
# -- from DOC -> DOCX. So both the source and the target need to be file-system files. Put
|
||||
# -- transient files in a temporary directory that is automatically removed so they don't
|
||||
# -- pile up.
|
||||
with tempfile.TemporaryDirectory() as target_dir:
|
||||
source_file_path = f"{target_dir}/document.doc" if file is not None else filename
|
||||
assert source_file_path is not None
|
||||
|
||||
# -- when source is a file-like object, write it to the filesystem so the command-line
|
||||
# -- process can access it (CLI executes in different memory-space).
|
||||
if file is not None:
|
||||
with open(source_file_path, "wb") as f:
|
||||
f.write(file.read())
|
||||
|
||||
# -- convert the .doc file to .docx. The resulting file takes the same base-name as the
|
||||
# -- source file and is written to `target_dir`.
|
||||
convert_office_doc(
|
||||
source_file_path,
|
||||
target_dir,
|
||||
target_format="docx",
|
||||
target_filter=libre_office_filter,
|
||||
)
|
||||
|
||||
# -- compute the path of the resulting .docx document --
|
||||
_, filename_no_path = os.path.split(os.path.abspath(source_file_path))
|
||||
base_filename, _ = os.path.splitext(filename_no_path)
|
||||
target_file_path = os.path.join(target_dir, f"{base_filename}.docx")
|
||||
|
||||
# -- and partition it. Note that `kwargs` is not passed which is a sketchy way to partially
|
||||
# -- disable post-partitioning processing (what the decorators do) so for example the
|
||||
# -- resulting elements are not double-chunked.
|
||||
elements = partition_docx(
|
||||
filename=target_file_path,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.DOC,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# -- Remove temporary document.docx path from metadata when necessary. Note `metadata_filename`
|
||||
# -- defaults to `None` but that's better than a meaningless temporary filename.
|
||||
if file:
|
||||
for element in elements:
|
||||
element.metadata.filename = metadata_filename
|
||||
|
||||
return elements
|
||||
@@ -0,0 +1,987 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import zipfile
|
||||
from typing import IO, Any, Iterator, Protocol, Type
|
||||
|
||||
import docx
|
||||
from docx.document import Document
|
||||
from docx.enum.section import WD_SECTION_START
|
||||
from docx.oxml.table import CT_Tbl
|
||||
from docx.oxml.text.paragraph import CT_P
|
||||
from docx.section import Section, _Footer, _Header
|
||||
from docx.table import Table as DocxTable
|
||||
from docx.table import _Cell, _Row
|
||||
from docx.text.hyperlink import Hyperlink
|
||||
from docx.text.pagebreak import RenderedPageBreak
|
||||
from docx.text.paragraph import Paragraph
|
||||
from docx.text.run import Run
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.cleaners.core import clean_bullets
|
||||
from unstructured.common.html_table import htmlify_matrix_of_cell_texts
|
||||
from unstructured.documents.elements import (
|
||||
Address,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
EmailAddress,
|
||||
Footer,
|
||||
Header,
|
||||
Image,
|
||||
Link,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
PageBreak,
|
||||
Table,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
from unstructured.partition.text_type import (
|
||||
is_bulleted_text,
|
||||
is_email_address,
|
||||
is_possible_narrative_text,
|
||||
is_us_city_state_zip,
|
||||
)
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
from unstructured.utils import is_temp_file_path, lazyproperty
|
||||
|
||||
DETECTION_ORIGIN: str = "docx"
|
||||
# -- CT_* stands for "complex-type", an XML element type in docx parlance --
|
||||
BlockElement: TypeAlias = "CT_P | CT_Tbl"
|
||||
BlockItem: TypeAlias = "Paragraph | DocxTable"
|
||||
|
||||
|
||||
def register_picture_partitioner(picture_partitioner: PicturePartitionerT) -> None:
|
||||
"""Specify a pluggable sub-partitioner to be used for partitioning DOCX images."""
|
||||
DocxPartitionerOptions.register_picture_partitioner(picture_partitioner)
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# DOCX DOMAIN MODEL DEFINITIONS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class PicturePartitionerT(Protocol):
|
||||
"""Defines the interface for a pluggable sub-partitioner for DOCX Picture objects.
|
||||
|
||||
In Microsoft Word parlance, an image is a "picture". We use that term here for an image in a
|
||||
DOCX file both for domain consistency and because it conveniently avoids confusion with an
|
||||
`unstructured` `Image` element.
|
||||
|
||||
A picture can be either *inline* or *floating*. An inline picture is treated like a big
|
||||
character in the text of a paragraph, moving with the text. A floating picture can be moved
|
||||
freely and text flows around it.
|
||||
|
||||
Both inline and floating pictures are defined inside a paragraph in the DOCX file. A paragraph
|
||||
can have zero or more pictures. A DOCX picture partitioner takes a `docx` `Paragraph` object
|
||||
and generates an `Image` element for each picture found in that paragraph.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, paragraph: Paragraph, opts: DocxPartitionerOptions) -> Iterator[Image]:
|
||||
"""Generate an `Image` element for each picture in `paragraph`."""
|
||||
...
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# PARTITIONER
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@apply_metadata(FileType.DOCX)
|
||||
@add_chunking_strategy
|
||||
def partition_docx(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
include_page_breaks: bool = True,
|
||||
infer_table_structure: bool = True,
|
||||
starting_page_number: int = 1,
|
||||
strategy: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions Microsoft Word Documents in .docx format into its document elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
include_page_breaks
|
||||
When True, add a `PageBreak` element to the element-stream when a page-break is detected in
|
||||
the document. Note that not all DOCX files include page-break information.
|
||||
infer_table_structure
|
||||
If True, any Table elements that are extracted will also have a metadata field
|
||||
named "text_as_html" where the table's text content is rendered into an html string.
|
||||
I.e., rows and cells are preserved.
|
||||
Whether True or False, the "text" field is always present in any Table element
|
||||
and is the text content of the table (no structure).
|
||||
metadata_filename
|
||||
The filename to use for the metadata. Relevant because partition_doc converts the document
|
||||
to .docx before partition. We want the original source filename in the metadata.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
starting_page_number
|
||||
Assign this number to the first page of this document and increment the page number from
|
||||
there.
|
||||
"""
|
||||
opts = DocxPartitionerOptions.load(
|
||||
file=file,
|
||||
file_path=filename,
|
||||
include_page_breaks=include_page_breaks,
|
||||
infer_table_structure=infer_table_structure,
|
||||
starting_page_number=starting_page_number,
|
||||
strategy=strategy,
|
||||
)
|
||||
|
||||
elements = _DocxPartitioner.iter_document_elements(opts)
|
||||
|
||||
return list(elements)
|
||||
|
||||
|
||||
class DocxPartitionerOptions:
|
||||
"""Encapsulates partitioning option validation, computation, and application of defaults."""
|
||||
|
||||
_PicturePartitionerCls = None
|
||||
"""Sub-partitioner used to extract pictures from a paragraph as `Image` elements.
|
||||
|
||||
This value has module lifetime and is updated by calling the `register_picture_partitioner()`
|
||||
function defined in this module. The value sent to `register_picture_partitioner()` must be a
|
||||
pluggable sub-partitioner implementing the `PicturePartitionerT` interface. After
|
||||
registration, all paragraphs in subsequently partitioned DOCX documents will be sent to this
|
||||
sub-partitioner to extract images when so configured.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
file: IO[bytes] | None,
|
||||
file_path: str | None,
|
||||
include_page_breaks: bool,
|
||||
infer_table_structure: bool,
|
||||
starting_page_number: int = 1,
|
||||
strategy: str | None = None,
|
||||
):
|
||||
self._file = file
|
||||
self._file_path = file_path
|
||||
self._include_page_breaks = include_page_breaks
|
||||
self._infer_table_structure = infer_table_structure
|
||||
self._strategy = strategy
|
||||
# -- options object maintains page-number state --
|
||||
self._page_counter = starting_page_number
|
||||
|
||||
@classmethod
|
||||
def load(cls, **kwargs: Any) -> DocxPartitionerOptions:
|
||||
"""Construct and validate an instance."""
|
||||
return cls(**kwargs)._validate()
|
||||
|
||||
@classmethod
|
||||
def register_picture_partitioner(cls, picture_partitioner: PicturePartitionerT):
|
||||
"""Specify a pluggable sub-partitioner to extract images from DOCX paragraphs."""
|
||||
cls._PicturePartitionerCls = picture_partitioner
|
||||
|
||||
@lazyproperty
|
||||
def document(self) -> Document:
|
||||
"""The python-docx `Document` object loaded from file or filename."""
|
||||
return docx.Document(self._docx_file)
|
||||
|
||||
@lazyproperty
|
||||
def include_page_breaks(self) -> bool:
|
||||
"""When True, include `PageBreak` elements in element-stream.
|
||||
|
||||
Note that regardless of this setting, page-breaks are detected, and page-number is tracked
|
||||
and included in element metadata. Only the presence of distinct `PageBreak` elements (which
|
||||
contain no text) in the element stream is affected.
|
||||
"""
|
||||
return self._include_page_breaks
|
||||
|
||||
def increment_page_number(self) -> Iterator[PageBreak]:
|
||||
"""Increment page-number by 1 and generate a PageBreak element if enabled."""
|
||||
self._page_counter += 1
|
||||
# -- only emit page-breaks when enabled --
|
||||
if self._include_page_breaks:
|
||||
yield PageBreak("", detection_origin=DETECTION_ORIGIN)
|
||||
|
||||
@lazyproperty
|
||||
def infer_table_structure(self) -> bool:
|
||||
"""True when partitioner should compute and apply `text_as_html` metadata for tables."""
|
||||
return self._infer_table_structure
|
||||
|
||||
@lazyproperty
|
||||
def last_modified(self) -> str | None:
|
||||
"""The best last-modified date available, None if no sources are available."""
|
||||
if not self._file_path:
|
||||
return None
|
||||
|
||||
return (
|
||||
None if is_temp_file_path(self._file_path) else get_last_modified_date(self._file_path)
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def metadata_file_path(self) -> str | None:
|
||||
"""The best available file-path for this document or `None` if unavailable."""
|
||||
return self._file_path
|
||||
|
||||
@property
|
||||
def metadata_page_number(self) -> int | None:
|
||||
"""The current page number to report in metadata, or None if we can't really tell.
|
||||
|
||||
Page numbers are not added to element metadata if we can't find any page-breaks in the
|
||||
document (which may be a common case).
|
||||
|
||||
In the DOCX format, determining page numbers is strictly a best-efforts attempt since
|
||||
actual page-breaks are determined at rendering time (e.g. printing) based on the
|
||||
font-metrics of the target device. Explicit (hard) page-breaks are always recorded in the
|
||||
docx file but the rendered page-breaks are only added optionally.
|
||||
"""
|
||||
return self._page_counter if self._document_contains_pagebreaks else None
|
||||
|
||||
@property
|
||||
def page_number(self) -> int:
|
||||
"""The current page number.
|
||||
|
||||
Note this value may not represent the actual rendered page number when rendered page-break
|
||||
indicators are not present in the document (not uncommon). Use `.metadata_page_number` for
|
||||
metadata purposes, which is `None` when rendered page-breaks are not present in this
|
||||
document.
|
||||
"""
|
||||
return self._page_counter
|
||||
|
||||
@lazyproperty
|
||||
def picture_partitioner(self) -> PicturePartitionerT:
|
||||
"""The sub-partitioner to use for DOCX image extraction."""
|
||||
# -- Note this value has partitioning-run scope. An instance of this options class is
|
||||
# -- instantiated once per partitioning run (each document can have different options).
|
||||
# -- Because this is a lazyproperty, it is computed only on the first reference. All
|
||||
# -- subsequent references during the same partitioning run will get the same value. This
|
||||
# -- ensures image extraction is processed consistently within a single document.
|
||||
return self._PicturePartitionerCls or _NullPicturePartitioner
|
||||
|
||||
@lazyproperty
|
||||
def strategy(self) -> str:
|
||||
"""The partitioning strategy for this document.
|
||||
|
||||
One of "hi_res", "fast", and a few others. These are available as class attributes on
|
||||
`unstructured.partition.utils.constants.PartitionStrategy` but resolve to str values.
|
||||
"""
|
||||
return PartitionStrategy.HI_RES if self._strategy is None else self._strategy
|
||||
|
||||
@lazyproperty
|
||||
def _document_contains_pagebreaks(self) -> bool:
|
||||
"""True when there is at least one page-break detected in the document.
|
||||
|
||||
Only `w:lastRenderedPageBreak` elements reliably indicate a page-break. These are reliably
|
||||
inserted by Microsoft Word, but probably don't appear in documents converted into .docx
|
||||
format from for example .odt format.
|
||||
"""
|
||||
xpath = (
|
||||
# NOTE(scanny) - w:lastRenderedPageBreak (lrpb) is run (w:r) inner content. `w:r` can
|
||||
# appear in a paragraph (w:p). w:r can also appear in a hyperlink (w:hyperlink), which
|
||||
# is w:p inner-content and both of these can occur inside a table-cell as well as the
|
||||
# document body
|
||||
"./w:body/w:p/w:r/w:lastRenderedPageBreak"
|
||||
" | ./w:body/w:p/w:hyperlink/w:r/w:lastRenderedPageBreak"
|
||||
" | ./w:body/w:tbl/w:tr/w:tc/w:p/w:r/w:lastRenderedPageBreak"
|
||||
" | ./w:body/w:tbl/w:tr/w:tc/w:p/w:hyperlink/w:r/w:lastRenderedPageBreak"
|
||||
)
|
||||
|
||||
return bool(self.document.element.xpath(xpath))
|
||||
|
||||
@lazyproperty
|
||||
def _docx_file(self) -> str | IO[bytes]:
|
||||
"""The Word 2007+ document file to be partitioned.
|
||||
|
||||
This is either a `str` path or a file-like object. `python-docx` accepts either for opening
|
||||
a document file.
|
||||
"""
|
||||
if self._file_path:
|
||||
return self._file_path
|
||||
|
||||
# -- In Python <3.11 SpooledTemporaryFile does not implement ".seekable" which triggers an
|
||||
# -- exception when Zipfile tries to open it. The docx format is a zip archive so we need
|
||||
# -- to work around that bug here.
|
||||
if isinstance(self._file, tempfile.SpooledTemporaryFile):
|
||||
self._file.seek(0)
|
||||
return io.BytesIO(self._file.read())
|
||||
|
||||
assert self._file is not None # -- assured by `._validate()` --
|
||||
return self._file
|
||||
|
||||
def _validate(self) -> DocxPartitionerOptions:
|
||||
"""Raise on first invalide option, return self otherwise."""
|
||||
# -- provide distinguished error between "file-not-found" and "not-a-DOCX-file" --
|
||||
if self._file_path:
|
||||
if not os.path.isfile(self._file_path):
|
||||
raise FileNotFoundError(f"no such file or directory: {repr(self._file_path)}")
|
||||
if not zipfile.is_zipfile(self._file_path):
|
||||
raise ValueError(f"not a ZIP archive (so not a DOCX file): {repr(self._file_path)}")
|
||||
elif self._file:
|
||||
if not zipfile.is_zipfile(self._file):
|
||||
raise ValueError(f"not a ZIP archive (so not a DOCX file): {repr(self._file)}")
|
||||
else:
|
||||
raise ValueError(
|
||||
"no DOCX document specified, either `filename` or `file` argument must be provided"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class _DocxPartitioner:
|
||||
"""Provides `.partition()` for MS-Word 2007+ (.docx) files."""
|
||||
|
||||
def __init__(self, opts: DocxPartitionerOptions) -> None:
|
||||
self._opts = opts
|
||||
|
||||
@classmethod
|
||||
def iter_document_elements(cls, opts: DocxPartitionerOptions) -> Iterator[Element]:
|
||||
"""Partition MS Word documents (.docx format) into its document elements."""
|
||||
self = cls(opts)
|
||||
# NOTE(scanny): It's possible for a Word document to have no sections. In particular, a
|
||||
# Microsoft Teams chat transcript exported to DOCX contains no sections. Such a
|
||||
# "section-less" document has to be interated differently and has no headers or footers and
|
||||
# therefore no page-size or margins.
|
||||
return (
|
||||
self._iter_document_elements()
|
||||
if self._document_contains_sections
|
||||
else self._iter_sectionless_document_elements()
|
||||
)
|
||||
|
||||
def _iter_document_elements(self) -> Iterator[Element]:
|
||||
"""Generate each document-element in (docx) `document` in document order."""
|
||||
# -- This implementation composes a collection of iterators into a "combined" iterator
|
||||
# -- return value using `yield from`. You can think of the return value as an Element
|
||||
# -- stream and each `yield from` as "add elements found by this function to the stream".
|
||||
# -- This is functionally analogous to declaring `elements: list[Element] = []` at the top
|
||||
# -- and using `elements.extend()` for the results of each of the function calls, but is
|
||||
# -- more perfomant, uses less memory (avoids producing and then garbage-collecting all
|
||||
# -- those small lists), is more flexible for later iterator operations like filter,
|
||||
# -- chain, map, etc. and is perhaps more elegant and simpler to read once you have the
|
||||
# -- concept of what it's doing. You can see the same pattern repeating in the "sub"
|
||||
# -- functions like `._iter_paragraph_elements()` where the "just return when done"
|
||||
# -- characteristic of a generator avoids repeated code to form interim results into lists.
|
||||
for section_idx, section in enumerate(self._document.sections):
|
||||
yield from self._iter_section_page_breaks(section_idx, section)
|
||||
yield from self._iter_section_headers(section)
|
||||
|
||||
for block_item in section.iter_inner_content():
|
||||
# -- a block-item can be a Paragraph or a Table, maybe others later so elif here.
|
||||
# -- Paragraph is more common so check that first.
|
||||
if isinstance(block_item, Paragraph):
|
||||
yield from self._iter_paragraph_elements(block_item)
|
||||
elif isinstance( # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
block_item, DocxTable
|
||||
):
|
||||
yield from self._iter_table_element(block_item)
|
||||
|
||||
yield from self._iter_section_footers(section)
|
||||
|
||||
def _iter_sectionless_document_elements(self) -> Iterator[Element]:
|
||||
"""Generate each document-element in a docx `document` that has no sections.
|
||||
|
||||
A "section-less" DOCX must be iterated differently. Also it will have no headers or footers
|
||||
(because those live in a section).
|
||||
"""
|
||||
for block_item in self._document.iter_inner_content():
|
||||
if isinstance(block_item, Paragraph):
|
||||
yield from self._iter_paragraph_elements(block_item)
|
||||
# -- can only be a Paragraph or Table so far but more types may come later --
|
||||
elif isinstance(block_item, DocxTable): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
yield from self._iter_table_element(block_item)
|
||||
|
||||
def _classify_paragraph_to_element(self, paragraph: Paragraph) -> Iterator[Element]:
|
||||
"""Generate zero-or-one document element for `paragraph`.
|
||||
|
||||
In Word, an empty paragraph is commonly used for inter-paragraph spacing. An empty paragraph
|
||||
does not contribute to the document-element stream and will not cause an element to be
|
||||
emitted.
|
||||
"""
|
||||
text = "".join(
|
||||
e.text
|
||||
for e in paragraph._p.xpath(
|
||||
"w:r | w:hyperlink | w:r/descendant::wp:inline[ancestor::w:drawing][1]//w:r"
|
||||
)
|
||||
)
|
||||
|
||||
# -- blank paragraphs are commonly used for spacing between paragraphs and do not
|
||||
# -- contribute to the document-element stream
|
||||
if not text.strip():
|
||||
return
|
||||
|
||||
metadata = self._paragraph_metadata(paragraph)
|
||||
|
||||
# -- a list-item gets some special treatment, mutating the text to remove a
|
||||
# -- bullet-character if present
|
||||
if self._is_list_item(paragraph):
|
||||
clean_text = clean_bullets(text).strip()
|
||||
if clean_text:
|
||||
yield ListItem(
|
||||
text=clean_text,
|
||||
metadata=metadata,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
)
|
||||
return
|
||||
|
||||
# -- determine element-type from an explicit Word paragraph-style if possible --
|
||||
TextSubCls = self._style_based_element_type(paragraph)
|
||||
if TextSubCls:
|
||||
yield TextSubCls(text=text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
return
|
||||
|
||||
# -- try to recognize the element type by parsing its text --
|
||||
TextSubCls = self._parse_paragraph_text_for_element_type(paragraph)
|
||||
if TextSubCls:
|
||||
yield TextSubCls(text=text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
return
|
||||
|
||||
# -- if all that fails we give it the default `Text` element-type --
|
||||
yield Text(text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
|
||||
def _convert_table_to_html(self, table: DocxTable) -> str:
|
||||
"""HTML string version of `table`.
|
||||
|
||||
Example:
|
||||
|
||||
<table>
|
||||
<tbody>
|
||||
<tr><th>item </th><th style="text-align: right;"> qty</th></tr>
|
||||
<tr><td>spam </td><td style="text-align: right;"> 42</td></tr>
|
||||
<tr><td>eggs </td><td style="text-align: right;"> 451</td></tr>
|
||||
<tr><td>bacon </td><td style="text-align: right;"> 0</td></tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
`is_nested` is used for recursive calls when a nested table is encountered. Certain
|
||||
behaviors are different in that case, but the caller can safely ignore that parameter and
|
||||
allow it to take its default value.
|
||||
"""
|
||||
|
||||
def iter_cell_block_items(cell: _Cell) -> Iterator[str]:
|
||||
"""Generate the text of each paragraph or table in `cell` as a separate string.
|
||||
|
||||
A table nested in `cell` is converted to the normalized text it contains.
|
||||
"""
|
||||
for block_item in cell.iter_inner_content():
|
||||
if isinstance(paragraph := block_item, Paragraph):
|
||||
# -- all docx content is ultimately in a paragraph; a nested table contributes
|
||||
# -- structure only
|
||||
yield paragraph.text
|
||||
elif isinstance(table := block_item, DocxTable):
|
||||
for row in table.rows:
|
||||
yield from iter_row_cells_as_text(row)
|
||||
|
||||
def iter_row_cells_as_text(row: _Row) -> Iterator[str]:
|
||||
"""Generate the normalized text of each cell in `row` as a separate string.
|
||||
|
||||
The text of each paragraph within a cell is not separated. A table nested in a cell is
|
||||
converted to a normalized string of its contents and combined with the text of the
|
||||
cell that contains the table.
|
||||
"""
|
||||
# -- Each omitted cell at the start of the row (pretty rare) gets the empty string.
|
||||
# -- This preserves column alignment when one or more initial cells are omitted.
|
||||
for _ in range(row.grid_cols_before):
|
||||
yield ""
|
||||
|
||||
try:
|
||||
# -- row.cells may introduce `ValueError: no tc element at grid_offset=X` if the
|
||||
# -- table has merged or malformed cells. always wrap in try/except.
|
||||
for cell in row.cells:
|
||||
cell_text = " ".join(iter_cell_block_items(cell))
|
||||
yield " ".join(cell_text.split())
|
||||
except Exception as e:
|
||||
logging.warning(f"Skipping cell in _iter_row_cells_as_text due to: {e}")
|
||||
yield ""
|
||||
|
||||
# -- Each omitted cell at the end of the row (also rare) gets the empty string. --
|
||||
for _ in range(row.grid_cols_after):
|
||||
yield ""
|
||||
|
||||
return htmlify_matrix_of_cell_texts([list(iter_row_cells_as_text(r)) for r in table.rows])
|
||||
|
||||
@lazyproperty
|
||||
def _document(self) -> Document:
|
||||
"""The python-docx `Document` object loaded from file or filename."""
|
||||
return self._opts.document
|
||||
|
||||
@lazyproperty
|
||||
def _document_contains_sections(self) -> bool:
|
||||
"""True when there is at least one section in the document.
|
||||
|
||||
This is always true for a document produced by Word, but may not always be the case when the
|
||||
document results from conversion or export. In particular, a Microsoft Teams chat-transcript
|
||||
export will have no sections.
|
||||
"""
|
||||
return bool(self._document.sections)
|
||||
|
||||
def _header_footer_text(self, hdrftr: _Header | _Footer) -> str:
|
||||
"""The text enclosed in `hdrftr` as a single string.
|
||||
|
||||
Each paragraph is included along with the text of each table cell. Empty text is omitted.
|
||||
Each paragraph text-item is separated by a newline ("\n") although note that a paragraph
|
||||
that contains a line-break will also include a newline representing that line-break, so
|
||||
newlines do not necessarily distinguish separate paragraphs.
|
||||
|
||||
The entire text of a table is included as a single string with a space separating the text
|
||||
of each cell.
|
||||
|
||||
A header with no text or only whitespace returns the empty string ("").
|
||||
"""
|
||||
|
||||
def iter_hdrftr_texts(hdrftr: _Header | _Footer) -> Iterator[str]:
|
||||
"""Generate each text item in `hdrftr` stripped of leading and trailing whitespace.
|
||||
|
||||
This includes paragraphs as well as table cell contents.
|
||||
"""
|
||||
for block_item in hdrftr.iter_inner_content():
|
||||
if isinstance(block_item, Paragraph):
|
||||
yield block_item.text.strip()
|
||||
# -- can only be a Paragraph or Table so far but more types may come later --
|
||||
elif isinstance( # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
block_item, DocxTable
|
||||
):
|
||||
yield " ".join(self._iter_table_texts(block_item))
|
||||
|
||||
return "\n".join(text for text in iter_hdrftr_texts(hdrftr) if text)
|
||||
|
||||
def _is_list_item(self, paragraph: Paragraph) -> bool:
|
||||
"""True when `paragraph` can be identified as a list-item."""
|
||||
if is_bulleted_text(paragraph.text):
|
||||
return True
|
||||
|
||||
return "<w:numPr>" in paragraph._p.xml
|
||||
|
||||
def _iter_paragraph_elements(self, paragraph: Paragraph) -> Iterator[Element]:
|
||||
"""Generate zero-or-more document elements for `paragraph`.
|
||||
|
||||
The generated elements can be both textual elements and PageBreak elements. An empty
|
||||
paragraph produces no elements.
|
||||
"""
|
||||
|
||||
def iter_paragraph_items(paragraph: Paragraph) -> Iterator[Paragraph | RenderedPageBreak]:
|
||||
"""Generate Paragraph and RenderedPageBreak items from `paragraph`.
|
||||
|
||||
Each generated paragraph is the portion of the paragraph on the same page. When the
|
||||
paragraph contains no page-breaks, it is iterated unchanged and iteration stops. When
|
||||
there is a page-break, in general there one paragraph "fragment" before the page break,
|
||||
the page break, and then the fragment after the page break. However many combinations
|
||||
are possible. The first item can be either a page-break or a paragraph, but the type
|
||||
always alternates throughout the sequence.
|
||||
"""
|
||||
if not paragraph.contains_page_break:
|
||||
yield paragraph
|
||||
return
|
||||
|
||||
page_break = paragraph.rendered_page_breaks[0]
|
||||
|
||||
# -- preceding-fragment is None when first paragraph content is a page-break --
|
||||
preceding_paragraph_fragment = page_break.preceding_paragraph_fragment
|
||||
if preceding_paragraph_fragment:
|
||||
yield preceding_paragraph_fragment
|
||||
|
||||
yield page_break
|
||||
|
||||
# -- following-fragment is None when page-break is last paragraph content. This is
|
||||
# -- probably quite rare (Word moves these to the start of the next paragraph) but
|
||||
# -- easier to check for it than prove it can't happen.
|
||||
following_paragraph_fragment = page_break.following_paragraph_fragment
|
||||
# -- the paragraph fragment following a page-break can itself contain another
|
||||
# -- page-break; this would also be quite rare, but it can happen so we just recurse
|
||||
# -- into the second fragment the same way we handled the original paragraph
|
||||
if following_paragraph_fragment:
|
||||
yield from iter_paragraph_items(following_paragraph_fragment)
|
||||
|
||||
for item in iter_paragraph_items(paragraph):
|
||||
if isinstance(item, Paragraph):
|
||||
yield from self._classify_paragraph_to_element(item)
|
||||
yield from self._iter_paragraph_images(item)
|
||||
else:
|
||||
yield from self._opts.increment_page_number()
|
||||
|
||||
def _iter_paragraph_emphasis(self, paragraph: Paragraph) -> Iterator[dict[str, str]]:
|
||||
"""Generate e.g. {"text": "MUST", "tag": "b"} for each emphasis in `paragraph`."""
|
||||
for run in paragraph.runs:
|
||||
text = run.text.strip() if run.text else ""
|
||||
if not text:
|
||||
continue
|
||||
if run.bold:
|
||||
yield {"text": text, "tag": "b"}
|
||||
if run.italic:
|
||||
yield {"text": text, "tag": "i"}
|
||||
|
||||
def _iter_paragraph_images(self, paragraph: Paragraph) -> Iterator[Image]:
|
||||
"""Generate `Image` element for each picture shape in `paragraph` when so configured."""
|
||||
# -- Delegate this job to the pluggable Picture partitioner. Note the default picture
|
||||
# -- partitioner does not extract images.
|
||||
PicturePartitionerCls = self._opts.picture_partitioner
|
||||
yield from PicturePartitionerCls.iter_elements(paragraph, self._opts)
|
||||
|
||||
def _iter_section_footers(self, section: Section) -> Iterator[Footer]:
|
||||
"""Generate any `Footer` elements defined for this section.
|
||||
|
||||
A Word document has up to three header and footer definition pairs for each document
|
||||
section, a primary, first-page, and even-page header and footer. The first-page pair
|
||||
applies only to the first page of the section (perhaps a title page or chapter start). The
|
||||
even-page pair is used in book-bound documents where there are both recto and verso pages
|
||||
(it is applied to verso (even-numbered) pages). A page where neither more specialized
|
||||
footer applies uses the primary footer.
|
||||
"""
|
||||
|
||||
def iter_footer(footer: _Footer, header_footer_type: str) -> Iterator[Footer]:
|
||||
"""Generate zero-or-one Footer elements for `footer`."""
|
||||
if footer.is_linked_to_previous:
|
||||
return
|
||||
text = self._header_footer_text(footer)
|
||||
if not text:
|
||||
return
|
||||
yield Footer(
|
||||
text=text,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
metadata=ElementMetadata(
|
||||
filename=self._opts.metadata_file_path,
|
||||
header_footer_type=header_footer_type,
|
||||
category_depth=0,
|
||||
),
|
||||
)
|
||||
|
||||
yield from iter_footer(section.footer, "primary")
|
||||
if section.different_first_page_header_footer:
|
||||
yield from iter_footer(section.first_page_footer, "first_page")
|
||||
if self._document.settings.odd_and_even_pages_header_footer:
|
||||
yield from iter_footer(section.even_page_footer, "even_page")
|
||||
|
||||
def _iter_section_headers(self, section: Section) -> Iterator[Header]:
|
||||
"""Generate `Header` elements for this section if it has them.
|
||||
|
||||
See `._iter_section_footers()` docstring for more on docx headers and footers.
|
||||
"""
|
||||
|
||||
def maybe_iter_header(header: _Header, header_footer_type: str) -> Iterator[Header]:
|
||||
"""Generate zero-or-one Header elements for `header`."""
|
||||
if header.is_linked_to_previous:
|
||||
return
|
||||
text = self._header_footer_text(header)
|
||||
if not text:
|
||||
return
|
||||
yield Header(
|
||||
text=text,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
metadata=ElementMetadata(
|
||||
filename=self._opts.metadata_file_path,
|
||||
header_footer_type=header_footer_type,
|
||||
category_depth=0, # -- headers are always at the root level}
|
||||
),
|
||||
)
|
||||
|
||||
yield from maybe_iter_header(section.header, "primary")
|
||||
if section.different_first_page_header_footer:
|
||||
yield from maybe_iter_header(section.first_page_header, "first_page")
|
||||
if self._document.settings.odd_and_even_pages_header_footer:
|
||||
yield from maybe_iter_header(section.even_page_header, "even_page")
|
||||
|
||||
def _iter_section_page_breaks(self, section_idx: int, section: Section) -> Iterator[PageBreak]:
|
||||
"""Generate zero-or-one `PageBreak` document elements for `section`.
|
||||
|
||||
A docx section has a "start" type which can be "continuous" (no page-break), "nextPage",
|
||||
"evenPage", or "oddPage". For the next, even, and odd varieties, a `w:renderedPageBreak`
|
||||
element signals one page break. Here we only need to handle the case where we need to add
|
||||
another, for example to go from one odd page to another odd page and we need a total of
|
||||
two page-breaks.
|
||||
"""
|
||||
|
||||
def page_is_odd() -> bool:
|
||||
return self._opts.page_number % 2 == 1
|
||||
|
||||
start_type = section.start_type
|
||||
|
||||
# -- This method is called upon entering a new section, which happens before any paragraphs
|
||||
# -- in that section are partitioned. A rendered page-break due to a section-start occurs
|
||||
# -- in the first paragraph of the section and so occurs _later_ in the proces. Here we
|
||||
# -- predict when two page breaks will be needed and emit one of them. The second will be
|
||||
# -- emitted by the rendered page-break to follow.
|
||||
|
||||
if start_type == WD_SECTION_START.EVEN_PAGE: # noqa
|
||||
# -- on an even page we need two total, add one to supplement the rendered page break
|
||||
# -- to follow. There is no "first-document-page" special case because 1 is odd.
|
||||
if not page_is_odd():
|
||||
yield from self._opts.increment_page_number()
|
||||
|
||||
elif start_type == WD_SECTION_START.ODD_PAGE:
|
||||
# -- the first page of the document is an implicit "new" odd-page, so no page-break --
|
||||
if section_idx == 0:
|
||||
return
|
||||
if page_is_odd():
|
||||
yield from self._opts.increment_page_number()
|
||||
|
||||
# -- otherwise, start-type is one of "continuous", "new-column", or "next-page", none of
|
||||
# -- which need our help to get the page-breaks right.
|
||||
return
|
||||
|
||||
def _iter_table_element(self, table: DocxTable) -> Iterator[Table]:
|
||||
"""Generate zero-or-one Table element for a DOCX `w:tbl` XML element."""
|
||||
# -- at present, we always generate exactly one Table element, but we might want
|
||||
# -- to skip, for example, an empty table.
|
||||
html_table = (
|
||||
self._convert_table_to_html(table) if self._opts.infer_table_structure else None
|
||||
)
|
||||
text_table = " ".join(self._iter_table_texts(table))
|
||||
emphasized_text_contents, emphasized_text_tags = self._table_emphasis(table)
|
||||
|
||||
yield Table(
|
||||
text_table,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
metadata=ElementMetadata(
|
||||
text_as_html=html_table,
|
||||
filename=self._opts.metadata_file_path,
|
||||
page_number=self._opts.metadata_page_number,
|
||||
last_modified=self._opts.last_modified,
|
||||
emphasized_text_contents=emphasized_text_contents or None,
|
||||
emphasized_text_tags=emphasized_text_tags or None,
|
||||
),
|
||||
)
|
||||
|
||||
def _iter_table_emphasis(self, table: DocxTable) -> Iterator[dict[str, str]]:
|
||||
"""Generate e.g. {"text": "word", "tag": "b"} for each emphasis in `table`."""
|
||||
for row in table.rows:
|
||||
try:
|
||||
# -- row.cells may introduce `ValueError: no tc element at grid_offset=X` if the
|
||||
# -- table has merged or malformed cells. always wrap in try/except.
|
||||
for cell in row.cells:
|
||||
for paragraph in cell.paragraphs:
|
||||
yield from self._iter_paragraph_emphasis(paragraph)
|
||||
except Exception as e:
|
||||
logging.warning(f"Skipping row in _iter_table_emphasis due to: {e}")
|
||||
continue
|
||||
|
||||
def _iter_table_texts(self, table: DocxTable) -> Iterator[str]:
|
||||
"""Generate text of each cell in `table` stripped of leading and trailing whitespace.
|
||||
|
||||
Nested tables are recursed into and their text contributes to the output in depth-first
|
||||
pre-order. Empty strings due to empty or whitespace-only cells are dropped.
|
||||
"""
|
||||
|
||||
def iter_cell_texts(cell: _Cell) -> Iterator[str]:
|
||||
"""Generate each text item in `cell` stripped of leading and trailing whitespace.
|
||||
|
||||
This includes paragraphs as well as table cell contents.
|
||||
"""
|
||||
for block_item in cell.iter_inner_content():
|
||||
if isinstance(block_item, Paragraph):
|
||||
yield block_item.text.strip()
|
||||
# -- can only be a Paragraph or Table so far but more types may come later --
|
||||
elif isinstance( # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
block_item, DocxTable
|
||||
):
|
||||
yield from self._iter_table_texts(block_item)
|
||||
|
||||
for row in table.rows:
|
||||
tr = row._tr
|
||||
for tc in tr.tc_lst:
|
||||
# -- vMerge="continue" indicates a spanned cell in a vertical merge --
|
||||
if tc.vMerge == "continue":
|
||||
continue
|
||||
# -- do not generate empty strings --
|
||||
yield from (text for text in iter_cell_texts(_Cell(tc, table)) if text)
|
||||
|
||||
def _paragraph_emphasis(self, paragraph: Paragraph) -> tuple[list[str], list[str]]:
|
||||
"""[contents, tags] pair describing emphasized text in `paragraph`."""
|
||||
iter_p_emph, iter_p_emph_2 = itertools.tee(self._iter_paragraph_emphasis(paragraph))
|
||||
return ([e["text"] for e in iter_p_emph], [e["tag"] for e in iter_p_emph_2])
|
||||
|
||||
def _paragraph_link_meta(self, paragraph: Paragraph) -> tuple[list[str], list[str], list[Link]]:
|
||||
"""Describes hyperlinks in `paragraph`, if any."""
|
||||
if not paragraph.hyperlinks:
|
||||
return [], [], []
|
||||
|
||||
def iter_paragraph_links() -> Iterator[Link]:
|
||||
"""Generate `Link` typed-dict for each external link in `paragraph`.
|
||||
|
||||
Word uses hyperlinks for internal "jumps" within the document, as well as for web and
|
||||
other external locations. Only generate the external ones.
|
||||
"""
|
||||
offset = 0
|
||||
for item in paragraph.iter_inner_content():
|
||||
if isinstance(item, Run):
|
||||
offset += len(item.text)
|
||||
elif isinstance(item, Hyperlink): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
text = item.text
|
||||
url = item.url
|
||||
start_index = offset
|
||||
offset += len(text)
|
||||
# -- docx hyperlinks include "internal" links, like a table-of-contents
|
||||
# -- (TOC) entry has a jump to the named heading in the document (e.g.
|
||||
# -- '#_Toc147925734'. Such links have a fragment but not an address
|
||||
# -- (URL). Treat those as regular text.
|
||||
if not url:
|
||||
continue
|
||||
# -- all Word hyperlinks should contain text, otherwise they have no
|
||||
# -- visual appearance on the document. Not expected, but technically possible
|
||||
# -- so filter these out too.
|
||||
if not text:
|
||||
continue
|
||||
yield Link(text=text, url=url, start_index=start_index)
|
||||
|
||||
links = list(iter_paragraph_links())
|
||||
# -- link["text"] is allowed to be None by the declared type for `Link`, but never will be
|
||||
# -- here because such a link is filtered out above. Use empty str to satisfy type-checker.
|
||||
link_texts = [link["text"] or "" for link in links]
|
||||
link_urls = [link["url"] for link in links]
|
||||
return link_texts, link_urls, links
|
||||
|
||||
def _paragraph_metadata(self, paragraph: Paragraph) -> ElementMetadata:
|
||||
"""ElementMetadata object describing `paragraph`."""
|
||||
category_depth = self._parse_category_depth_by_style(paragraph)
|
||||
emphasized_text_contents, emphasized_text_tags = self._paragraph_emphasis(paragraph)
|
||||
link_texts, link_urls, links = self._paragraph_link_meta(paragraph)
|
||||
element_metadata = ElementMetadata(
|
||||
category_depth=category_depth,
|
||||
emphasized_text_contents=emphasized_text_contents or None,
|
||||
emphasized_text_tags=emphasized_text_tags or None,
|
||||
filename=self._opts.metadata_file_path,
|
||||
last_modified=self._opts.last_modified,
|
||||
link_texts=link_texts or None,
|
||||
link_urls=link_urls or None,
|
||||
links=links or None,
|
||||
page_number=self._opts.metadata_page_number,
|
||||
)
|
||||
element_metadata.detection_origin = "docx"
|
||||
return element_metadata
|
||||
|
||||
def _parse_category_depth_by_style(self, paragraph: Paragraph) -> int:
|
||||
"""Determine category depth from paragraph metadata"""
|
||||
|
||||
# Determine category depth from paragraph ilvl xpath
|
||||
xpath = paragraph._element.xpath("./w:pPr/w:numPr/w:ilvl/@w:val")
|
||||
if xpath:
|
||||
return round(float(xpath[0]))
|
||||
|
||||
# Determine category depth from style name
|
||||
style_name = (paragraph.style and paragraph.style.name) or "Normal"
|
||||
depth = self._parse_category_depth_by_style_name(style_name)
|
||||
|
||||
if depth > 0:
|
||||
return depth
|
||||
else:
|
||||
# Check if category depth can be determined from style ilvl
|
||||
return self._parse_category_depth_by_style_ilvl()
|
||||
|
||||
def _parse_category_depth_by_style_ilvl(self) -> int:
|
||||
# TODO(newelh) Parsing category depth by style ilvl is not yet implemented
|
||||
return 0
|
||||
|
||||
def _parse_category_depth_by_style_name(self, style_name: str) -> int:
|
||||
"""Parse category-depth from the style-name of `paragraph`.
|
||||
|
||||
Category depth is 0-indexed and relative to the other element types in the document.
|
||||
"""
|
||||
|
||||
def _extract_number(suffix: str) -> int:
|
||||
return int(suffix.split()[-1]) - 1 if suffix.split()[-1].isdigit() else 0
|
||||
|
||||
# Heading styles
|
||||
if style_name.startswith("Heading"):
|
||||
return _extract_number(style_name)
|
||||
|
||||
if style_name == "Subtitle":
|
||||
return 1
|
||||
|
||||
# List styles
|
||||
list_prefixes = ["List", "List Bullet", "List Continue", "List Number"]
|
||||
if any(style_name.startswith(prefix) for prefix in list_prefixes):
|
||||
return _extract_number(style_name)
|
||||
|
||||
# Other styles
|
||||
return 0
|
||||
|
||||
def _parse_paragraph_text_for_element_type(self, paragraph: Paragraph) -> Type[Text] | None:
|
||||
"""Attempt to differentiate the element-type by inspecting the raw text."""
|
||||
text = paragraph.text.strip()
|
||||
|
||||
if len(text) < 2:
|
||||
return None
|
||||
if is_us_city_state_zip(text):
|
||||
return Address
|
||||
if is_email_address(text):
|
||||
return EmailAddress
|
||||
if is_possible_narrative_text(text):
|
||||
return NarrativeText
|
||||
|
||||
return None
|
||||
|
||||
def _style_based_element_type(self, paragraph: Paragraph) -> Type[Text] | None:
|
||||
"""Element-type for `paragraph` based on its paragraph-style.
|
||||
|
||||
Returns `None` when the style doesn't tell us anything useful, including when it
|
||||
is the default "Normal" style.
|
||||
"""
|
||||
# NOTE(robinson) - documentation on built-in styles at the link below:
|
||||
# https://python-docx.readthedocs.io/en/latest/user/styles-understanding.html \
|
||||
# #paragraph-styles-in-default-template
|
||||
STYLE_TO_ELEMENT_MAPPING = {
|
||||
"Caption": Text, # TODO(robinson) - add caption element type
|
||||
"Heading 1": Title,
|
||||
"Heading 2": Title,
|
||||
"Heading 3": Title,
|
||||
"Heading 4": Title,
|
||||
"Heading 5": Title,
|
||||
"Heading 6": Title,
|
||||
"Heading 7": Title,
|
||||
"Heading 8": Title,
|
||||
"Heading 9": Title,
|
||||
"Intense Quote": Text, # TODO(robinson) - add quote element type
|
||||
"List": ListItem,
|
||||
"List 2": ListItem,
|
||||
"List 3": ListItem,
|
||||
"List Bullet": ListItem,
|
||||
"List Bullet 2": ListItem,
|
||||
"List Bullet 3": ListItem,
|
||||
"List Continue": ListItem,
|
||||
"List Continue 2": ListItem,
|
||||
"List Continue 3": ListItem,
|
||||
"List Number": ListItem,
|
||||
"List Number 2": ListItem,
|
||||
"List Number 3": ListItem,
|
||||
"List Paragraph": ListItem,
|
||||
"Macro Text": Text,
|
||||
"No Spacing": Text,
|
||||
"Quote": Text, # TODO(robinson) - add quote element type
|
||||
"Subtitle": Title,
|
||||
"TOCHeading": Title,
|
||||
"Title": Title,
|
||||
}
|
||||
|
||||
# -- paragraph.style can be None in rare cases, so can style.name. That's going
|
||||
# -- to mean default style which is equivalent to "Normal" for our purposes.
|
||||
style_name = (paragraph.style and paragraph.style.name) or "Normal"
|
||||
|
||||
# NOTE(robinson) - The "Normal" style name will return None since it's not
|
||||
# in the mapping. Unknown style names will also return None.
|
||||
return STYLE_TO_ELEMENT_MAPPING.get(style_name)
|
||||
|
||||
def _table_emphasis(self, table: DocxTable) -> tuple[list[str], list[str]]:
|
||||
"""[contents, tags] pair describing emphasized text in `table`."""
|
||||
iter_tbl_emph, iter_tbl_emph_2 = itertools.tee(self._iter_table_emphasis(table))
|
||||
return ([e["text"] for e in iter_tbl_emph], [e["tag"] for e in iter_tbl_emph_2])
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# SUB-PARTITIONERS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class _NullPicturePartitioner:
|
||||
"""Does not parse the provided paragraph for pictures and generates zero `Image` elements."""
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, paragraph: Paragraph, opts: DocxPartitionerOptions) -> Iterator[Image]:
|
||||
"""No-op picture partitioner."""
|
||||
return
|
||||
yield
|
||||
@@ -0,0 +1,435 @@
|
||||
"""Provides `partition_email()` function.
|
||||
|
||||
Suitable for use with `.eml` files, which can be exported from many email clients.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import email
|
||||
import email.policy
|
||||
import email.utils
|
||||
import io
|
||||
import os
|
||||
from email.message import EmailMessage, MIMEPart
|
||||
from typing import IO, Any, Final, Iterator, cast
|
||||
|
||||
from dateutil import parser
|
||||
|
||||
from unstructured.documents.elements import Element, ElementMetadata
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common import UnsupportedFileFormatError
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.partition.text import partition_text
|
||||
from unstructured.utils import lazyproperty
|
||||
|
||||
VALID_CONTENT_SOURCES: Final[tuple[str, ...]] = ("text/html", "text/plain")
|
||||
|
||||
|
||||
def partition_email(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
content_source: str = "text/html",
|
||||
metadata_filename: str | None = None,
|
||||
metadata_last_modified: str | None = None,
|
||||
process_attachments: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an .eml file into document elements.
|
||||
|
||||
Args:
|
||||
filename: str path of the target file.
|
||||
file: A file-like object open for reading bytes (not str) e.g. --> open(filename, "rb").
|
||||
content_source: The preferred message body. Many emails contain both a plain-text and an
|
||||
HTML version of the message body. By default, the HTML version will be used when
|
||||
available. Specifying "text/plain" will cause the plain-text version to be preferred.
|
||||
When the preferred version is not available, the other version will be used.
|
||||
metadata_filename: The file-path to use for metadata purposes. Useful when the target file
|
||||
is specified as a file-like object or when `filename` is a temporary file and the
|
||||
original file-path is known or a more meaningful file-path is desired.
|
||||
metadata_last_modified: The last-modified timestamp to be applied in metadata. Useful when
|
||||
a file-like object (which can have no last-modified date) target is used. The
|
||||
last-modified metadata is otherwise drawn from the filesystem when a path is provided.
|
||||
process_attachments: When True, also partition any attachments in the message after
|
||||
partitioning the message body. All document elements appear in the single returned
|
||||
element list. The filename of the attachment, when available, is used as the
|
||||
`filename` metadata value for elements arising from the attachment.
|
||||
|
||||
Note that all global keyword arguments such as `unique_element_ids`, `language` and
|
||||
`chunking_strategy` can be used and will be passed along to the decorators that implement
|
||||
those functions. Further, any keyword arguments applicable to HTML will be passed along to the
|
||||
HTML partitioner when processing an HTML message body.
|
||||
"""
|
||||
ctx = EmailPartitioningContext.load(
|
||||
file_path=filename,
|
||||
file=file,
|
||||
content_source=content_source,
|
||||
metadata_file_path=metadata_filename,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
process_attachments=process_attachments,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
return list(_EmailPartitioner.iter_elements(ctx=ctx))
|
||||
|
||||
|
||||
class EmailPartitioningContext:
|
||||
"""Encapsulates partitioning option validation, computation, and application of defaults."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str | None = None,
|
||||
file: IO[bytes] | None = None,
|
||||
content_source: str = "text/html",
|
||||
metadata_file_path: str | None = None,
|
||||
metadata_last_modified: str | None = None,
|
||||
process_attachments: bool = False,
|
||||
kwargs: dict[str, Any] = {},
|
||||
):
|
||||
self._file_path = file_path
|
||||
self._file = file
|
||||
self._content_source = content_source
|
||||
self._metadata_file_path = metadata_file_path
|
||||
self._metadata_last_modified = metadata_last_modified
|
||||
self._process_attachments = process_attachments
|
||||
self._kwargs = kwargs
|
||||
|
||||
@classmethod
|
||||
def load(
|
||||
cls,
|
||||
file_path: str | None,
|
||||
file: IO[bytes] | None,
|
||||
content_source: str,
|
||||
metadata_file_path: str | None,
|
||||
metadata_last_modified: str | None,
|
||||
process_attachments: bool,
|
||||
kwargs: dict[str, Any],
|
||||
) -> EmailPartitioningContext:
|
||||
"""Construct and validate an instance."""
|
||||
return cls(
|
||||
file_path=file_path,
|
||||
file=file,
|
||||
content_source=content_source,
|
||||
metadata_file_path=metadata_file_path,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
process_attachments=process_attachments,
|
||||
kwargs=kwargs,
|
||||
)._validate()
|
||||
|
||||
@lazyproperty
|
||||
def bcc_addresses(self) -> list[str] | None:
|
||||
"""The "blind carbon-copy" Bcc: addresses of the message."""
|
||||
bccs = self.msg.get_all("Bcc")
|
||||
if not bccs:
|
||||
return None
|
||||
addrs = email.utils.getaddresses(bccs)
|
||||
return [email.utils.formataddr(addr) for addr in addrs]
|
||||
|
||||
@lazyproperty
|
||||
def body_part(self) -> MIMEPart | None:
|
||||
"""The message part containing the actual textual email message.
|
||||
|
||||
This is as opposed to attachments or "related" parts like an image that appears in the
|
||||
message etc.
|
||||
"""
|
||||
return self.msg.get_body(preferencelist=self.content_type_preference)
|
||||
|
||||
@lazyproperty
|
||||
def cc_addresses(self) -> list[str] | None:
|
||||
"""The "carbon-copy" Cc: addresses of the message."""
|
||||
ccs = self.msg.get_all("Cc")
|
||||
if not ccs:
|
||||
return None
|
||||
addrs = email.utils.getaddresses(ccs)
|
||||
return [email.utils.formataddr(addr) for addr in addrs]
|
||||
|
||||
@lazyproperty
|
||||
def content_type_preference(self) -> tuple[str, ...]:
|
||||
"""Whether to prefer HTML or plain-text body when message-body has both.
|
||||
|
||||
The default order of preference is `("html", "plain")`. The order can be switched by
|
||||
specifying `"text/plain"` as the `content_source` arg value.
|
||||
"""
|
||||
return ("plain", "html") if self._content_source == "text/plain" else ("html", "plain")
|
||||
|
||||
@lazyproperty
|
||||
def email_metadata(self) -> ElementMetadata:
|
||||
"""The email-specific metadata fields for this message.
|
||||
|
||||
Suitable for use with `.metadata.update()` on the base metadata applied to message body
|
||||
elements by delegate partitioners for text and HTML.
|
||||
"""
|
||||
return ElementMetadata(
|
||||
bcc_recipient=self.bcc_addresses,
|
||||
cc_recipient=self.cc_addresses,
|
||||
email_message_id=self.message_id,
|
||||
sent_from=[self.from_address] if self.from_address else None,
|
||||
sent_to=self.to_addresses,
|
||||
subject=self.subject,
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def from_address(self) -> str | None:
|
||||
"""The address of the message sender."""
|
||||
froms = self.msg.get_all("From")
|
||||
if not froms:
|
||||
# -- this should never occur because the From: header is mandatory per RFC 5322 --
|
||||
return None
|
||||
addrs = email.utils.getaddresses(froms)
|
||||
formatted_addrs = [email.utils.formataddr(addr) for addr in addrs]
|
||||
return formatted_addrs[0]
|
||||
|
||||
@lazyproperty
|
||||
def message_id(self) -> str | None:
|
||||
"""The value of the Message-ID: header, when present."""
|
||||
raw_id = self.msg.get("Message-ID")
|
||||
if not raw_id:
|
||||
return None
|
||||
return raw_id.strip().strip("<>")
|
||||
|
||||
@lazyproperty
|
||||
def metadata_file_path(self) -> str | None:
|
||||
"""The best available file-path information for this email message.
|
||||
|
||||
It's value is computed according to these rules, applied in order:
|
||||
|
||||
- The `metadata_filename` arg value when one was provided to `partition_email()`.
|
||||
- The `file_path` value when one was provided.
|
||||
- None otherwise.
|
||||
|
||||
This value is used as the `filename` metadata value for elements produced by partitioning
|
||||
the email message (but not those from its attachments).
|
||||
"""
|
||||
return self._metadata_file_path or self._file_path or None
|
||||
|
||||
@lazyproperty
|
||||
def metadata_last_modified(self) -> str | None:
|
||||
"""The best available last-modified date for this message, as an ISO8601 string.
|
||||
|
||||
It's value is computed according to these rules, applied in order:
|
||||
|
||||
- The `metadata_last_modified` arg value when one was provided to `partition_email()`.
|
||||
- The date-time in the `Sent:` header of the message, when present.
|
||||
- The last-modified date recorded on the filesystem for `file_path` when it was provided.
|
||||
- None otherwise.
|
||||
|
||||
This value is used as the `last_modified` metadata value for all elements produced by
|
||||
partitioning this email message, including any attachments.
|
||||
"""
|
||||
return self._metadata_last_modified or self._sent_date or self._filesystem_last_modified
|
||||
|
||||
@lazyproperty
|
||||
def msg(self) -> EmailMessage:
|
||||
"""The Python stdlib `email.message.EmailMessage` object parsed from the EML file."""
|
||||
if self._file_path is not None:
|
||||
with open(self._file_path, "rb") as f:
|
||||
return cast(
|
||||
EmailMessage, email.message_from_binary_file(f, policy=email.policy.default)
|
||||
)
|
||||
|
||||
assert self._file is not None
|
||||
|
||||
file_bytes = self._file.read()
|
||||
|
||||
return cast(EmailMessage, email.message_from_bytes(file_bytes, policy=email.policy.default))
|
||||
|
||||
@lazyproperty
|
||||
def partitioning_kwargs(self) -> dict[str, Any]:
|
||||
"""The "extra" keyword arguments received by `partition_email()`.
|
||||
|
||||
These are passed along to delegate partitioners which extract keyword args like
|
||||
`chunking_strategy` etc. in their decorators to control metadata behaviors, etc.
|
||||
"""
|
||||
return self._kwargs
|
||||
|
||||
@lazyproperty
|
||||
def process_attachments(self) -> bool:
|
||||
"""When True, partition attachments in addition to the email message body.
|
||||
|
||||
Any attachment having file-format that cannot be partitioned by unstructured is silently
|
||||
skipped.
|
||||
"""
|
||||
return self._process_attachments
|
||||
|
||||
@lazyproperty
|
||||
def subject(self) -> str | None:
|
||||
"""The value of the Subject: header, when present."""
|
||||
subject = self.msg.get("Subject")
|
||||
if not subject:
|
||||
return None
|
||||
return subject
|
||||
|
||||
@lazyproperty
|
||||
def to_addresses(self) -> list[str] | None:
|
||||
"""The To: addresses of the message."""
|
||||
tos = self.msg.get_all("To")
|
||||
if not tos:
|
||||
return None
|
||||
addrs = email.utils.getaddresses(tos)
|
||||
return [email.utils.formataddr(addr) for addr in addrs]
|
||||
|
||||
@lazyproperty
|
||||
def _filesystem_last_modified(self) -> str | None:
|
||||
"""Last-modified retrieved from filesystem when a file-path was provided, None otherwise."""
|
||||
return get_last_modified_date(self._file_path) if self._file_path else None
|
||||
|
||||
@lazyproperty
|
||||
def _sent_date(self) -> str | None:
|
||||
"""ISO-8601 str representation of message sent-date, if available."""
|
||||
date_str = self.msg.get("Date")
|
||||
if not date_str:
|
||||
return None
|
||||
try:
|
||||
sent_date = parser.parse(date_str)
|
||||
except (parser.ParserError, TypeError, ValueError):
|
||||
return None
|
||||
|
||||
return sent_date.astimezone(dt.timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
def _validate(self) -> EmailPartitioningContext:
|
||||
"""Raise on first invalid option, return self otherwise."""
|
||||
if not self._file_path and not self._file:
|
||||
raise ValueError(
|
||||
"no document specified; either a `filename` or `file` argument must be provided."
|
||||
)
|
||||
|
||||
if self._file:
|
||||
if not isinstance( # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
self._file.read(0), bytes
|
||||
):
|
||||
raise ValueError("file object must be opened in binary mode")
|
||||
self._file.seek(0)
|
||||
|
||||
if self._content_source not in VALID_CONTENT_SOURCES:
|
||||
raise ValueError(
|
||||
f"{repr(self._content_source)} is not a valid value for content_source;"
|
||||
f" must be one of: {VALID_CONTENT_SOURCES}",
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class _EmailPartitioner:
|
||||
"""Encapsulates the partitioning logic for email documents."""
|
||||
|
||||
def __init__(self, ctx: EmailPartitioningContext):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, ctx: EmailPartitioningContext) -> Iterator[Element]:
|
||||
"""Generate the document elements for the email described by `ctx`."""
|
||||
return cls(ctx=ctx)._iter_elements()
|
||||
|
||||
def _iter_elements(self) -> Iterator[Element]:
|
||||
"""Generate the document elements for the email described in the partitioning context.
|
||||
|
||||
This optionally includes elements generated by partitioning any partitionable attachments
|
||||
in the message as well.
|
||||
"""
|
||||
for e in self._iter_email_body_elements():
|
||||
e.metadata.update(self._ctx.email_metadata)
|
||||
yield e
|
||||
|
||||
if not self._ctx.process_attachments:
|
||||
return
|
||||
|
||||
for attachment in self._ctx.msg.iter_attachments():
|
||||
yield from _AttachmentPartitioner.iter_elements(attachment, self._ctx)
|
||||
|
||||
def _iter_email_body_elements(self) -> Iterator[Element]:
|
||||
"""Generate document elements from the email body."""
|
||||
body_part = self._ctx.body_part
|
||||
|
||||
# -- it's possible to have no body part; that translates to zero elements --
|
||||
if body_part is None:
|
||||
return
|
||||
|
||||
content_type = body_part.get_content_type()
|
||||
content = body_part.get_content()
|
||||
assert isinstance(content, str)
|
||||
|
||||
if content_type == "text/html":
|
||||
yield from partition_html(
|
||||
text=content,
|
||||
metadata_filename=self._ctx.metadata_file_path,
|
||||
metadata_file_type=FileType.EML,
|
||||
metadata_last_modified=self._ctx.metadata_last_modified,
|
||||
**self._ctx.partitioning_kwargs,
|
||||
)
|
||||
else:
|
||||
yield from partition_text(
|
||||
text=content,
|
||||
metadata_filename=self._ctx.metadata_file_path,
|
||||
metadata_file_type=FileType.EML,
|
||||
metadata_last_modified=self._ctx.metadata_last_modified,
|
||||
**self._ctx.partitioning_kwargs,
|
||||
)
|
||||
|
||||
|
||||
class _AttachmentPartitioner:
|
||||
"""Partitions an attachment to a MSG file."""
|
||||
|
||||
def __init__(self, attachment: EmailMessage, ctx: EmailPartitioningContext):
|
||||
self._attachment = attachment
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def iter_elements(
|
||||
cls, attachment: EmailMessage, ctx: EmailPartitioningContext
|
||||
) -> Iterator[Element]:
|
||||
"""Partition an attachment MIME-part from a MIME email message (.eml file)."""
|
||||
return cls(attachment, ctx)._iter_elements()
|
||||
|
||||
def _iter_elements(self) -> Iterator[Element]:
|
||||
"""Partition the byte-stream in the attachment MIME-part into elements.
|
||||
|
||||
Generates zero elements if the attachment is not partitionable.
|
||||
"""
|
||||
# -- `auto.partition()` imports this module, so we need to defer the import to here to
|
||||
# -- avoid a circular import.
|
||||
from unstructured.partition.auto import partition
|
||||
|
||||
file = io.BytesIO(self._file_bytes)
|
||||
|
||||
# -- partition the attachment --
|
||||
try:
|
||||
elements = partition(
|
||||
file=file,
|
||||
metadata_filename=self._attachment_file_name,
|
||||
metadata_last_modified=self._ctx.metadata_last_modified,
|
||||
**self._ctx.partitioning_kwargs,
|
||||
)
|
||||
except UnsupportedFileFormatError:
|
||||
# -- indicates `auto.partition()` has no partitioner for this file-format;
|
||||
# -- silently skip the attachment
|
||||
return
|
||||
|
||||
for e in elements:
|
||||
e.metadata.attached_to_filename = self._attached_to_filename
|
||||
yield e
|
||||
|
||||
@lazyproperty
|
||||
def _attached_to_filename(self) -> str | None:
|
||||
"""The file-name (no path) of the message. `None` if not available."""
|
||||
file_path = self._ctx.metadata_file_path
|
||||
if file_path is None:
|
||||
return None
|
||||
return os.path.basename(file_path)
|
||||
|
||||
@lazyproperty
|
||||
def _attachment_file_name(self) -> str | None:
|
||||
"""The original name of the attached file, `None` if not present in the MIME part."""
|
||||
return self._attachment.get_filename()
|
||||
|
||||
@lazyproperty
|
||||
def _file_bytes(self) -> bytes:
|
||||
"""The bytes of the attached file."""
|
||||
content = self._attachment.get_content()
|
||||
|
||||
if isinstance(content, str):
|
||||
return content.encode("utf-8")
|
||||
|
||||
assert isinstance(content, bytes)
|
||||
return content
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.file_conversion import convert_file_to_html_text_using_pandoc
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
DETECTION_ORIGIN: str = "epub"
|
||||
|
||||
|
||||
def partition_epub(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
languages: Optional[list[str]] = ["auto"],
|
||||
detect_language_per_element: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an EPUB document. The document is first converted to HTML and then
|
||||
partitioned using partition_html.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
languages
|
||||
User defined value for `metadata.languages` if provided. Otherwise language is detected
|
||||
using naive Bayesian filter via `langdetect`. Multiple languages indicates text could be
|
||||
in either language.
|
||||
Additional Parameters:
|
||||
detect_language_per_element
|
||||
Detect language per element instead of at the document level.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
html_text = convert_file_to_html_text_using_pandoc(
|
||||
source_format="epub",
|
||||
filename=filename,
|
||||
file=file,
|
||||
)
|
||||
|
||||
return partition_html(
|
||||
text=html_text,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.EPUB,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
languages=languages,
|
||||
detect_language_per_element=detect_language_per_element,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
from unstructured.partition.html.partition import partition_html
|
||||
|
||||
__all__ = ["partition_html"]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,320 @@
|
||||
import logging
|
||||
from abc import ABC
|
||||
from collections import defaultdict
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from bs4 import BeautifulSoup, Tag
|
||||
|
||||
from unstructured.documents.elements import Element, ElementType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HTML_PARSER = "html.parser"
|
||||
HTML_TEMPLATE = """
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title></title>
|
||||
</head>
|
||||
<body>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
TABLE_BORDER_STYLE = "border: 1px solid black;"
|
||||
TABLE_BORDER_COLLAPSE_STYLE = "border-collapse: collapse;"
|
||||
|
||||
|
||||
class ElementHtml(ABC):
|
||||
element: Element
|
||||
children: list["ElementHtml"]
|
||||
_html_tag: str = "div"
|
||||
|
||||
def __init__(self, element: Element, children: Optional[list["ElementHtml"]] = None):
|
||||
self.element = element
|
||||
self.children = children or []
|
||||
|
||||
@property
|
||||
def html_tag(self) -> str:
|
||||
return self._html_tag
|
||||
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
return None
|
||||
|
||||
def _inject_html_element_content(self, element_html: Tag, **kwargs: Any) -> None:
|
||||
element_html.string = self.element.text
|
||||
|
||||
def get_text_as_html(self) -> Union[Tag, None]:
|
||||
element_html = BeautifulSoup(self.element.metadata.text_as_html or "", HTML_PARSER).find()
|
||||
if not isinstance(element_html, Tag):
|
||||
return None
|
||||
return element_html
|
||||
|
||||
def _get_children_html(self, soup: BeautifulSoup, element_html: Tag, **kwargs: Any) -> Tag:
|
||||
wrapper = soup.new_tag(name="div")
|
||||
wrapper.append(element_html)
|
||||
for child in self.children:
|
||||
child_html = child.get_html_element(_soup=soup, **kwargs)
|
||||
wrapper.append(child_html)
|
||||
return wrapper
|
||||
|
||||
def get_html_element(self, **kwargs: Any) -> Tag:
|
||||
soup: Optional[BeautifulSoup] = kwargs.pop("_soup", None)
|
||||
if soup is None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
|
||||
element_html = self.get_text_as_html()
|
||||
if element_html is None:
|
||||
element_html = soup.new_tag(name=self.html_tag)
|
||||
self._inject_html_element_content(element_html, **kwargs)
|
||||
element_html["class"] = self.element.category
|
||||
element_html["id"] = self.element.id
|
||||
self._inject_html_element_attrs(element_html)
|
||||
if self.children: # if element has children wrap it with a 'div' tag
|
||||
return self._get_children_html(soup, element_html, **kwargs)
|
||||
return element_html
|
||||
|
||||
def set_children(self, children: list["ElementHtml"]) -> None:
|
||||
self.children = children
|
||||
|
||||
|
||||
class TitleElementHtml(ElementHtml):
|
||||
_html_tag = "h%d"
|
||||
|
||||
@property
|
||||
def html_tag(self) -> str:
|
||||
return self._html_tag % (self.element.metadata.category_depth or 1)
|
||||
|
||||
|
||||
class ImageElementHtml(ElementHtml):
|
||||
_html_tag = "img"
|
||||
|
||||
def _inject_html_element_content(self, element_html: Tag, **kwargs: Any) -> None:
|
||||
exclude_binary_image_data = kwargs.get("exclude_binary_image_data", False)
|
||||
if self.element.metadata.image_base64 and not exclude_binary_image_data:
|
||||
image_mime_type = self.element.metadata.image_mime_type or "image/png"
|
||||
element_html["src"] = (
|
||||
f"data:{image_mime_type};base64,{self.element.metadata.image_base64}"
|
||||
)
|
||||
element_html["alt"] = self.element.text
|
||||
|
||||
|
||||
class TableElementHtml(ElementHtml):
|
||||
_html_tag = "table"
|
||||
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["style"] = f"{TABLE_BORDER_STYLE} {TABLE_BORDER_COLLAPSE_STYLE}"
|
||||
for tag in element_html.find_all(["tr", "th", "td"]):
|
||||
tag["style"] = TABLE_BORDER_STYLE
|
||||
|
||||
|
||||
class LinkElementHtml(ElementHtml):
|
||||
_html_tag = "a"
|
||||
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["href"] = self.element.metadata.url or ""
|
||||
|
||||
|
||||
class TextElementHtml(ElementHtml):
|
||||
_html_tag = "p"
|
||||
|
||||
|
||||
class UnorderedListElementHtml(ElementHtml):
|
||||
_html_tag = "ul"
|
||||
|
||||
def _get_children_html(self, soup: BeautifulSoup, element_html: Tag, **kwargs: Any) -> Tag:
|
||||
for child in self.children:
|
||||
child_html = child.get_html_element(**kwargs)
|
||||
element_html.append(child_html)
|
||||
return element_html
|
||||
|
||||
|
||||
class OrderedListElementHtml(UnorderedListElementHtml):
|
||||
_html_tag = "ol"
|
||||
|
||||
|
||||
class ListItemElementHtml(UnorderedListElementHtml):
|
||||
_html_tag = "li"
|
||||
|
||||
|
||||
class LabelElementHtml(ElementHtml):
|
||||
_html_tag = "label"
|
||||
|
||||
|
||||
class FormElementHtml(ElementHtml):
|
||||
_html_tag = "form"
|
||||
|
||||
|
||||
class InputElementHtml(ElementHtml):
|
||||
_html_tag = "input"
|
||||
|
||||
|
||||
class CheckboxElementHtml(InputElementHtml):
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["type"] = "checkbox"
|
||||
|
||||
|
||||
class CheckboxCheckedElementHtml(InputElementHtml):
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["type"] = "checkbox"
|
||||
element_html["checked"] = "true"
|
||||
|
||||
|
||||
class RadioElementHtml(InputElementHtml):
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["type"] = "radio"
|
||||
|
||||
|
||||
class RadioCheckedElementHtml(InputElementHtml):
|
||||
def _inject_html_element_attrs(self, element_html: Tag) -> None:
|
||||
element_html["type"] = "radio"
|
||||
element_html["checked"] = "true"
|
||||
|
||||
|
||||
LIST_ELEMENTS = [ElementType.LIST_ITEM, ElementType.LIST_ITEM_OTHER]
|
||||
|
||||
TYPE_TO_HTML_MAP = {
|
||||
ElementType.UNCATEGORIZED_TEXT: TextElementHtml,
|
||||
ElementType.TITLE: TitleElementHtml,
|
||||
ElementType.IMAGE: ImageElementHtml,
|
||||
ElementType.TABLE: TableElementHtml,
|
||||
ElementType.LINK: LinkElementHtml,
|
||||
ElementType.TEXT: TextElementHtml,
|
||||
ElementType.PARAGRAPH: TextElementHtml,
|
||||
ElementType.LIST: OrderedListElementHtml,
|
||||
ElementType.LIST_ITEM: ListItemElementHtml,
|
||||
ElementType.LIST_ITEM_OTHER: ListItemElementHtml,
|
||||
ElementType.FIELD_NAME: LabelElementHtml,
|
||||
ElementType.BULLETED_TEXT: ListItemElementHtml,
|
||||
ElementType.FORM: FormElementHtml,
|
||||
ElementType.CAPTION: TextElementHtml,
|
||||
ElementType.CHECKED: CheckboxCheckedElementHtml,
|
||||
ElementType.UNCHECKED: CheckboxElementHtml,
|
||||
ElementType.CHECK_BOX_CHECKED: CheckboxCheckedElementHtml,
|
||||
ElementType.CHECK_BOX_UNCHECKED: CheckboxElementHtml,
|
||||
ElementType.RADIO_BUTTON_CHECKED: RadioCheckedElementHtml,
|
||||
ElementType.RADIO_BUTTON_UNCHECKED: RadioElementHtml,
|
||||
ElementType.NARRATIVE_TEXT: TextElementHtml,
|
||||
ElementType.FIGURE_CAPTION: TextElementHtml,
|
||||
ElementType.VALUE: InputElementHtml,
|
||||
ElementType.ABSTRACT: ElementHtml,
|
||||
ElementType.THREADING: ElementHtml,
|
||||
ElementType.COMPOSITE_ELEMENT: ElementHtml,
|
||||
ElementType.PICTURE: ElementHtml,
|
||||
ElementType.FIGURE: ElementHtml,
|
||||
ElementType.ADDRESS: ElementHtml,
|
||||
ElementType.EMAIL_ADDRESS: ElementHtml,
|
||||
ElementType.PAGE_BREAK: ElementHtml,
|
||||
ElementType.FORMULA: ElementHtml,
|
||||
ElementType.HEADER: ElementHtml,
|
||||
ElementType.HEADLINE: ElementHtml,
|
||||
ElementType.SUB_HEADLINE: ElementHtml,
|
||||
ElementType.PAGE_HEADER: ElementHtml,
|
||||
ElementType.SECTION_HEADER: ElementHtml,
|
||||
ElementType.FOOTER: ElementHtml,
|
||||
ElementType.FOOTNOTE: ElementHtml,
|
||||
ElementType.PAGE_FOOTER: ElementHtml,
|
||||
ElementType.PAGE_NUMBER: ElementHtml,
|
||||
ElementType.CODE_SNIPPET: ElementHtml,
|
||||
ElementType.FORM_KEYS_VALUES: ElementHtml,
|
||||
ElementType.DOCUMENT_DATA: ElementHtml,
|
||||
}
|
||||
|
||||
|
||||
def _group_element_children(children: list[ElementHtml]) -> list[ElementHtml]:
|
||||
grouped_children: list[ElementHtml] = []
|
||||
temp_group: list["ElementHtml"] = []
|
||||
prev_grouping = False
|
||||
for child in children:
|
||||
grouping = child.element.category in LIST_ELEMENTS
|
||||
if grouping:
|
||||
temp_group.append(child)
|
||||
elif prev_grouping:
|
||||
grouped_children.append(OrderedListElementHtml(Element(), temp_group))
|
||||
grouped_children.append(child)
|
||||
temp_group = []
|
||||
else:
|
||||
grouped_children.append(child)
|
||||
prev_grouping = grouping
|
||||
if temp_group:
|
||||
grouped_children.append(OrderedListElementHtml(Element(), temp_group))
|
||||
return grouped_children
|
||||
|
||||
|
||||
def _elements_to_html_tags_by_parent(elements: list[ElementHtml]) -> list[ElementHtml]:
|
||||
parent_to_children_map: dict[str, list[ElementHtml]] = defaultdict(list)
|
||||
for element in elements:
|
||||
if element.element.metadata.parent_id is not None:
|
||||
parent_to_children_map[element.element.metadata.parent_id].append(element)
|
||||
for parent_id, children in parent_to_children_map.items():
|
||||
grouped_children = _group_element_children(children)
|
||||
parent = next((el for el in elements if el.element.id == parent_id), None)
|
||||
if parent is None:
|
||||
logger.warning(f"Parent element with id {parent_id} not found. Skipping.")
|
||||
continue
|
||||
parent.set_children(grouped_children)
|
||||
return [el for el in elements if el.element.metadata.parent_id is None]
|
||||
|
||||
|
||||
def _elements_to_html_tags(
|
||||
elements: list[Element], exclude_binary_image_data: bool = False
|
||||
) -> list[Tag]:
|
||||
elements_html = [
|
||||
TYPE_TO_HTML_MAP.get(element.category, ElementHtml)(element) for element in elements
|
||||
]
|
||||
elements_html = _elements_to_html_tags_by_parent(elements_html)
|
||||
return [
|
||||
element_html.get_html_element(exclude_binary_image_data=exclude_binary_image_data)
|
||||
for element_html in elements_html
|
||||
]
|
||||
|
||||
|
||||
def _elements_to_html_tags_by_page(
|
||||
elements: list[Element], exclude_binary_image_data: bool = False
|
||||
) -> list[Tag]:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
pages_tags: list[Tag] = []
|
||||
grouped_elements = group_elements_by_page(elements)
|
||||
for page, g_elements in enumerate(grouped_elements, start=1):
|
||||
page_html = soup.new_tag(name="div", attrs={"data-page_number": page})
|
||||
elements_html = _elements_to_html_tags(g_elements, exclude_binary_image_data)
|
||||
for element_html in elements_html:
|
||||
page_html.append(element_html)
|
||||
pages_tags.append(page_html)
|
||||
return pages_tags
|
||||
|
||||
|
||||
def group_elements_by_page(
|
||||
unstructured_elements: list[Element],
|
||||
) -> list[list[Element]]:
|
||||
pages_dict: defaultdict[int, list[Element]] = defaultdict(list)
|
||||
|
||||
for element in unstructured_elements:
|
||||
page_number = element.metadata.page_number
|
||||
if page_number is None:
|
||||
logger.warning(f"Page number is not set for an element {element.id}. Skipping.")
|
||||
continue
|
||||
pages_dict[page_number].append(element)
|
||||
|
||||
pages_list = list(pages_dict.values())
|
||||
return pages_list
|
||||
|
||||
|
||||
def elements_to_html(
|
||||
elements: list[Element],
|
||||
exclude_binary_image_data: bool = False,
|
||||
no_group_by_page: bool = False,
|
||||
) -> str:
|
||||
soup = BeautifulSoup(HTML_TEMPLATE, HTML_PARSER)
|
||||
if soup.body is None:
|
||||
raise ValueError("Body tag not found in the HTML template")
|
||||
elements_html = (
|
||||
_elements_to_html_tags(elements, exclude_binary_image_data)
|
||||
if no_group_by_page
|
||||
else _elements_to_html_tags_by_page(elements, exclude_binary_image_data)
|
||||
)
|
||||
for element_html in elements_html:
|
||||
soup.body.append(element_html)
|
||||
return soup.prettify()
|
||||
@@ -0,0 +1,25 @@
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
|
||||
def indent_html(html_string: str, html_parser="html.parser") -> str:
|
||||
"""
|
||||
Formats / indents HTML.
|
||||
|
||||
This function takes an HTML string and formats it using the specified HTML parser.
|
||||
It parses the HTML content and returns a prettified version of it.
|
||||
|
||||
Args:
|
||||
html_string (str): The HTML content to be formatted.
|
||||
html_parser (str, optional): The parser to use for parsing the HTML. Defaults to 'html5lib':
|
||||
- 'html.parser': The built-in HTML parser. Use when you need just parsing
|
||||
- 'html5lib': The slowest. Use when you expect valid HTML parsed
|
||||
the same way a browser does. It adds some extra
|
||||
tags and attributes like <html>, <head>, <body>
|
||||
More in docs https://www.crummy.com/software/BeautifulSoup/bs4/doc/#installing-a-parser
|
||||
|
||||
Returns:
|
||||
str: The formatted and indented HTML content.
|
||||
"""
|
||||
soup = BeautifulSoup(html_string, html_parser)
|
||||
pretty_html = soup.prettify()
|
||||
return pretty_html
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,281 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Provides `partition_html()."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Iterator, List, Literal, Optional, cast
|
||||
|
||||
import requests
|
||||
from lxml import etree
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.documents.elements import Element, ElementType
|
||||
from unstructured.file_utils.encoding import read_txt_file
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
from unstructured.partition.html.parser import Flow, html_parser
|
||||
from unstructured.partition.html.transformations import (
|
||||
ontology_to_unstructured_elements,
|
||||
parse_html_to_ontology,
|
||||
)
|
||||
from unstructured.utils import is_temp_file_path, lazyproperty
|
||||
|
||||
|
||||
@apply_metadata(FileType.HTML)
|
||||
@add_chunking_strategy
|
||||
def partition_html(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
text: Optional[str] = None,
|
||||
encoding: Optional[str] = None,
|
||||
url: Optional[str] = None,
|
||||
headers: dict[str, str] = {},
|
||||
ssl_verify: bool = True,
|
||||
skip_headers_and_footers: bool = False,
|
||||
detection_origin: Optional[str] = None,
|
||||
html_parser_version: Literal["v1", "v2"] = "v1",
|
||||
image_alt_mode: Optional[Literal["to_text"]] = "to_text",
|
||||
extract_image_block_to_payload: bool = False,
|
||||
extract_image_block_types: Optional[list[str]] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an HTML document into its constituent elements.
|
||||
|
||||
HTML source parameters
|
||||
----------------------
|
||||
The HTML to be partitioned can be specified four different ways:
|
||||
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "r" mode --> open(filename, "r").
|
||||
text
|
||||
The string representation of the HTML document.
|
||||
url
|
||||
The URL of a webpage to parse. Only for URLs that return an HTML document.
|
||||
headers
|
||||
The HTTP headers to be used in the HTTP request when `url` is specified.
|
||||
ssl_verify
|
||||
If the URL parameter is set, determines whether or not SSL verification is performed
|
||||
on the HTTP request.
|
||||
encoding
|
||||
The encoding method used to decode the text input. If None, utf-8 will be used.
|
||||
skip_headers_and_footers
|
||||
If True, ignores any content that is within <header> or <footer> tags
|
||||
|
||||
html_parser_version (Literal['v1', 'v2']):
|
||||
The version of the HTML parser to use. The default is 'v1'. For 'v2' the parser will
|
||||
use the ontology schema to parse the HTML document.
|
||||
|
||||
image_alt_mode (Literal['to_text']):
|
||||
When set 'to_text', the v2 parser will include the alternative text of images in the output.
|
||||
"""
|
||||
# -- parser rejects an empty str, nip that edge-case in the bud here --
|
||||
if text is not None and text.strip() == "" and not file and not filename and not url:
|
||||
return []
|
||||
|
||||
opts = HtmlPartitionerOptions(
|
||||
file_path=filename,
|
||||
file=file,
|
||||
text=text,
|
||||
encoding=encoding,
|
||||
url=url,
|
||||
headers=headers,
|
||||
ssl_verify=ssl_verify,
|
||||
skip_headers_and_footers=skip_headers_and_footers,
|
||||
detection_origin=detection_origin,
|
||||
html_parser_version=html_parser_version,
|
||||
image_alt_mode=image_alt_mode,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
)
|
||||
|
||||
return list(_HtmlPartitioner.iter_elements(opts))
|
||||
|
||||
|
||||
class HtmlPartitionerOptions:
|
||||
"""Encapsulates partitioning option validation, computation, and application of defaults."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
file_path: str | None,
|
||||
file: IO[bytes] | None,
|
||||
text: str | None,
|
||||
encoding: str | None,
|
||||
url: str | None,
|
||||
headers: dict[str, str],
|
||||
ssl_verify: bool,
|
||||
skip_headers_and_footers: bool,
|
||||
detection_origin: str | None,
|
||||
html_parser_version: Literal["v1", "v2"] = "v1",
|
||||
image_alt_mode: Optional[Literal["to_text"]] = "to_text",
|
||||
extract_image_block_types: Optional[list[str]] = None,
|
||||
extract_image_block_to_payload: bool = False,
|
||||
):
|
||||
self._file_path = file_path
|
||||
self._file = file
|
||||
self._text = text
|
||||
self._encoding = encoding
|
||||
self._url = url
|
||||
self._headers = headers
|
||||
self._ssl_verify = ssl_verify
|
||||
self._skip_headers_and_footers = skip_headers_and_footers
|
||||
self._detection_origin = detection_origin
|
||||
self._html_parser_version = html_parser_version
|
||||
self._image_alt_mode = image_alt_mode
|
||||
self._extract_image_block_types = extract_image_block_types
|
||||
self._extract_image_block_to_payload = extract_image_block_to_payload
|
||||
|
||||
@lazyproperty
|
||||
def detection_origin(self) -> str | None:
|
||||
"""Trace of initial partitioner to be included in metadata for debugging purposes."""
|
||||
return self._detection_origin
|
||||
|
||||
@lazyproperty
|
||||
def html_text(self) -> str:
|
||||
"""The HTML document as a string, loaded from wherever the caller specified."""
|
||||
if self._file_path:
|
||||
return read_txt_file(filename=self._file_path, encoding=self._encoding)[1]
|
||||
|
||||
if self._file:
|
||||
return read_txt_file(file=self._file, encoding=self._encoding)[1]
|
||||
|
||||
if self._text:
|
||||
return str(self._text)
|
||||
|
||||
if self._url:
|
||||
response = requests.get(self._url, headers=self._headers, verify=self._ssl_verify)
|
||||
if not response.ok:
|
||||
raise ValueError(
|
||||
f"Error status code on GET of provided URL: {response.status_code}"
|
||||
)
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
if not content_type.startswith("text/html"):
|
||||
raise ValueError(f"Expected content type text/html. Got {content_type}.")
|
||||
|
||||
return response.text
|
||||
|
||||
raise ValueError("Exactly one of filename, file, text, or url must be specified.")
|
||||
|
||||
@lazyproperty
|
||||
def last_modified(self) -> str | None:
|
||||
"""The best last-modified date available, None if no sources are available."""
|
||||
return (
|
||||
None
|
||||
if not self._file_path or is_temp_file_path(self._file_path)
|
||||
else get_last_modified_date(self._file_path)
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def skip_headers_and_footers(self) -> bool:
|
||||
"""When True, elements located within a header or footer are pruned."""
|
||||
return self._skip_headers_and_footers
|
||||
|
||||
@lazyproperty
|
||||
def html_parser_version(self) -> Literal["v1", "v2"]:
|
||||
"""When html_parser_version=='v2', HTML elements follow ontology schema."""
|
||||
return self._html_parser_version
|
||||
|
||||
@lazyproperty
|
||||
def add_img_alt_text(self) -> bool:
|
||||
"""When True, the alternative text of images is included in the output."""
|
||||
return self._image_alt_mode == "to_text"
|
||||
|
||||
|
||||
class _HtmlPartitioner:
|
||||
"""Partition HTML document into document-elements."""
|
||||
|
||||
def __init__(self, opts: HtmlPartitionerOptions):
|
||||
self._opts = opts
|
||||
|
||||
def _should_include_image_base64(self, element: Element) -> bool:
|
||||
"""Determines if an image_base64 element should be included in the output."""
|
||||
return (
|
||||
element.category == ElementType.IMAGE
|
||||
and self._opts._extract_image_block_to_payload
|
||||
and self._opts._extract_image_block_types is not None
|
||||
and "Image" in self._opts._extract_image_block_types
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, opts: HtmlPartitionerOptions) -> Iterator[Element]:
|
||||
"""Partition HTML document provided by `opts` into document-elements."""
|
||||
yield from cls(opts)._iter_elements()
|
||||
|
||||
def _iter_elements(self) -> Iterator[Element]:
|
||||
"""Generated document-elements (e.g. Title, NarrativeText, etc.) parsed from document.
|
||||
|
||||
Elements appear in document order.
|
||||
"""
|
||||
# -- handle empty or whitespace-only HTML content --
|
||||
html_text = self._opts.html_text
|
||||
if not html_text or html_text.strip() == "":
|
||||
return
|
||||
|
||||
elements_iter = (
|
||||
self._main.iter_elements()
|
||||
if self._opts.html_parser_version == "v1"
|
||||
else self._from_ontology
|
||||
)
|
||||
|
||||
for e in elements_iter:
|
||||
e.metadata.last_modified = self._opts.last_modified
|
||||
e.metadata.detection_origin = self._opts.detection_origin
|
||||
|
||||
# -- remove <image_base64> if not requested --
|
||||
if not self._should_include_image_base64(e):
|
||||
e.metadata.image_base64 = None
|
||||
e.metadata.image_mime_type = None
|
||||
yield e
|
||||
|
||||
@lazyproperty
|
||||
def _main(self) -> Flow:
|
||||
"""The root HTML element."""
|
||||
# NOTE(scanny) - get `html_text` first so any encoding error raised is not confused with a
|
||||
# recoverable parsing error.
|
||||
html_text = self._opts.html_text
|
||||
|
||||
# NOTE(scanny) - `lxml` will not parse a `str` that includes an XML encoding declaration
|
||||
# and will raise the following error:
|
||||
# ValueError: Unicode strings with encoding declaration are not supported. ...
|
||||
# This is not valid HTML (would be in XHTML), but Chrome accepts it so we work around it
|
||||
# by UTF-8 encoding the str bytes and parsing those.
|
||||
try:
|
||||
root = etree.fromstring(html_text, html_parser)
|
||||
except ValueError:
|
||||
root = etree.fromstring(html_text.encode("utf-8"), html_parser)
|
||||
|
||||
# -- remove a variety of HTML element types like <script> and <style> that we prefer not
|
||||
# -- to encounter while parsing.
|
||||
etree.strip_elements(
|
||||
root, ["del", "link", "meta", "noscript", "script", "style"], with_tail=False
|
||||
)
|
||||
|
||||
# -- remove <header> and <footer> tags if the caller doesn't want their contents --
|
||||
if self._opts.skip_headers_and_footers:
|
||||
etree.strip_elements(root, ["header", "footer"], with_tail=False)
|
||||
|
||||
# -- jump to the core content if the document indicates where it is --
|
||||
if (main := root.find(".//main")) is not None:
|
||||
return cast(Flow, main)
|
||||
if (body := root.find(".//body")) is not None:
|
||||
return cast(Flow, body)
|
||||
return cast(Flow, root)
|
||||
|
||||
@lazyproperty
|
||||
def _from_ontology(self) -> List[Element]:
|
||||
"""Convert an ontology elements represented in HTML to an ontology element."""
|
||||
html_text = self._opts.html_text
|
||||
|
||||
# -- handle empty or whitespace-only HTML content --
|
||||
if not html_text or html_text.strip() == "":
|
||||
return []
|
||||
|
||||
ontology = parse_html_to_ontology(html_text)
|
||||
unstructured_elements = ontology_to_unstructured_elements(
|
||||
ontology, add_img_alt_text=self._opts.add_img_alt_text
|
||||
)
|
||||
return unstructured_elements
|
||||
@@ -0,0 +1,480 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
from collections import OrderedDict
|
||||
from itertools import chain
|
||||
from typing import Sequence, Type
|
||||
|
||||
from bs4 import BeautifulSoup, Tag
|
||||
|
||||
from unstructured.documents import elements, ontology
|
||||
from unstructured.documents.mappings import (
|
||||
CSS_CLASS_TO_ELEMENT_TYPE_MAP,
|
||||
HTML_TAG_AND_CSS_NAME_TO_ELEMENT_TYPE_MAP,
|
||||
HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP,
|
||||
ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE,
|
||||
)
|
||||
|
||||
RECURSION_LIMIT = 50
|
||||
|
||||
|
||||
def ontology_to_unstructured_elements(
|
||||
ontology_element: ontology.OntologyElement,
|
||||
parent_id: str | None = None,
|
||||
page_number: int | None = None,
|
||||
depth: int = 0,
|
||||
filename: str | None = None,
|
||||
add_img_alt_text: bool = True,
|
||||
) -> list[elements.Element]:
|
||||
"""
|
||||
Converts an OntologyElement object to a list of unstructured Element objects.
|
||||
|
||||
To preserve the structure of the ontology, the function is recursive
|
||||
and the tree structure is represented in flatten list by the parent_id
|
||||
attribute in the metadata of each Element object.
|
||||
To preserve all the attributes of the ontology element, the HTML code
|
||||
is injected to unstructured Element in ElementMetadata.text_as_html attribute.
|
||||
|
||||
For Layout elements, the function creates an empty Text Element (with the
|
||||
HTML code injected the same way).
|
||||
|
||||
TODO (Pluto): Better way would be to have special Element type in Unstructured
|
||||
|
||||
Args:
|
||||
ontology_element (OntologyElement): The ontology element to be converted.
|
||||
parent_id (str, optional): The ID of the parent element. Defaults to None.
|
||||
page_number (int, optional): The page number of the element. Defaults to None.
|
||||
depth (int, optional): The depth of the element in the hierarchy. Defaults to 0.
|
||||
filename (str, optional): The name of the file the element comes from. Defaults to None.
|
||||
add_img_alt_text (bool): Whether to include the alternative text of images
|
||||
in the output. Defaults to True.
|
||||
Returns:
|
||||
list[Element]: A list of unstructured Element objects.
|
||||
"""
|
||||
elements_to_return: list[elements.Element] = []
|
||||
if ontology_element.elementType == ontology.ElementTypeEnum.layout and depth <= RECURSION_LIMIT:
|
||||
if page_number is None and isinstance(ontology_element, ontology.Page):
|
||||
page_number = ontology_element.page_number
|
||||
|
||||
if not isinstance(ontology_element, ontology.Document):
|
||||
elements_to_return += [
|
||||
elements.Text(
|
||||
text="",
|
||||
element_id=ontology_element.id,
|
||||
detection_origin="vlm_partitioner",
|
||||
metadata=elements.ElementMetadata(
|
||||
parent_id=parent_id,
|
||||
text_as_html=ontology_element.to_html(add_children=False),
|
||||
page_number=page_number,
|
||||
category_depth=depth,
|
||||
filename=filename,
|
||||
),
|
||||
)
|
||||
]
|
||||
children: list[elements.Element] = []
|
||||
for child in ontology_element.children:
|
||||
child = ontology_to_unstructured_elements(
|
||||
child,
|
||||
parent_id=ontology_element.id,
|
||||
page_number=page_number,
|
||||
depth=0 if isinstance(ontology_element, ontology.Document) else depth + 1,
|
||||
filename=filename,
|
||||
add_img_alt_text=add_img_alt_text,
|
||||
)
|
||||
children += child
|
||||
|
||||
combined_children = combine_inline_elements(children)
|
||||
elements_to_return += combined_children
|
||||
else:
|
||||
element_class: type[elements.Element] = ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE[
|
||||
ontology_element.__class__
|
||||
]
|
||||
html_code_of_ontology_element = ontology_element.to_html()
|
||||
element_text = ontology_element.to_text(add_img_alt_text=add_img_alt_text)
|
||||
|
||||
unstructured_element = element_class(
|
||||
text=element_text, # type: ignore
|
||||
element_id=ontology_element.id,
|
||||
detection_origin="vlm_partitioner",
|
||||
metadata=elements.ElementMetadata(
|
||||
parent_id=parent_id,
|
||||
text_as_html=html_code_of_ontology_element,
|
||||
page_number=page_number,
|
||||
category_depth=depth,
|
||||
filename=filename,
|
||||
),
|
||||
)
|
||||
elements_to_return = [unstructured_element]
|
||||
|
||||
return elements_to_return
|
||||
|
||||
|
||||
def combine_inline_elements(elements: list[elements.Element]) -> list[elements.Element]:
|
||||
"""
|
||||
Combines consecutive inline elements into a single element. Inline elements
|
||||
can be also combined with text elements.
|
||||
|
||||
Combined elements contains multiple HTML tags together eg.
|
||||
{
|
||||
'text': "Text from element 1 Text from element 2",
|
||||
'metadata': {
|
||||
'text_as_html': "<p>Text from element 1</p><a>Text from element 2</a>"
|
||||
}
|
||||
}
|
||||
|
||||
Args:
|
||||
elements (list[Element]): A list of elements to be combined.
|
||||
|
||||
Returns:
|
||||
list[Element]: A list of combined elements.
|
||||
"""
|
||||
result_elements: list[elements.Element] = []
|
||||
|
||||
current_element: elements.Element | None = None
|
||||
for next_element in elements:
|
||||
if current_element is None:
|
||||
current_element = next_element
|
||||
continue
|
||||
|
||||
if can_unstructured_elements_be_merged(current_element, next_element):
|
||||
current_element.text += " " + next_element.text
|
||||
current_element.metadata.text_as_html += next_element.metadata.text_as_html
|
||||
else:
|
||||
result_elements.append(current_element)
|
||||
current_element = next_element
|
||||
|
||||
if current_element is not None:
|
||||
result_elements.append(current_element)
|
||||
|
||||
return result_elements
|
||||
|
||||
|
||||
def can_unstructured_elements_be_merged(
|
||||
current_element: elements.Element, next_element: elements.Element
|
||||
) -> bool:
|
||||
"""
|
||||
Elements can be merged when:
|
||||
- They are on the same level in the HTML tree
|
||||
- Neither of them has children
|
||||
- All elements are inline elements or text element
|
||||
"""
|
||||
if current_element.metadata.category_depth != next_element.metadata.category_depth:
|
||||
return False
|
||||
|
||||
current_html_tags = BeautifulSoup(
|
||||
current_element.metadata.text_as_html, "html.parser"
|
||||
).find_all(recursive=False)
|
||||
next_html_tags = BeautifulSoup(next_element.metadata.text_as_html, "html.parser").find_all(
|
||||
recursive=False
|
||||
)
|
||||
|
||||
ontology_elements = [
|
||||
parse_html_to_ontology_element(html_tag)
|
||||
for html_tag in chain(current_html_tags, next_html_tags)
|
||||
]
|
||||
|
||||
for ontology_element in ontology_elements:
|
||||
if ontology_element.children:
|
||||
return False
|
||||
|
||||
if not (is_inline_element(ontology_element) or is_text_element(ontology_element)):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_text_element(ontology_element: ontology.OntologyElement) -> bool:
|
||||
"""Categories or classes that we want to combine with inline text"""
|
||||
|
||||
text_classes = [
|
||||
ontology.NarrativeText,
|
||||
ontology.Quote,
|
||||
ontology.Paragraph,
|
||||
ontology.Footnote,
|
||||
ontology.FootnoteReference,
|
||||
ontology.Citation,
|
||||
ontology.Bibliography,
|
||||
ontology.Glossary,
|
||||
]
|
||||
text_categories = [ontology.ElementTypeEnum.metadata]
|
||||
|
||||
if any(isinstance(ontology_element, class_) for class_ in text_classes):
|
||||
return True
|
||||
|
||||
return any(ontology_element.elementType == category for category in text_categories)
|
||||
|
||||
|
||||
def is_inline_element(ontology_element: ontology.OntologyElement) -> bool:
|
||||
"""Categories or classes that we want to combine with text elements"""
|
||||
|
||||
inline_classes = [ontology.Hyperlink]
|
||||
inline_categories = [
|
||||
ontology.ElementTypeEnum.specialized_text,
|
||||
ontology.ElementTypeEnum.annotation,
|
||||
]
|
||||
|
||||
if any(isinstance(ontology_element, class_) for class_ in inline_classes):
|
||||
return True
|
||||
|
||||
return any(ontology_element.elementType == category for category in inline_categories)
|
||||
|
||||
|
||||
def unstructured_elements_to_ontology(
|
||||
unstructured_elements: Sequence[elements.Element],
|
||||
) -> ontology.OntologyElement:
|
||||
"""
|
||||
Converts a sequence of unstructured Element objects to an OntologyElement object.
|
||||
|
||||
The function caches the elements in a dictionary and each element is assigned to its parent.
|
||||
At the end the root element is popped from the dictionary and returned.
|
||||
|
||||
Such approach comes with limitations:
|
||||
- The parent element has to be in the list before the child element
|
||||
|
||||
Args:
|
||||
unstructured_elements (Sequence[Element]): The sequence of unstructured Element objects.
|
||||
|
||||
Returns:
|
||||
OntologyElement: The converted OntologyElement object.
|
||||
"""
|
||||
id_to_element_mapping: OrderedDict[str, ontology.OntologyElement] = OrderedDict()
|
||||
|
||||
root_element_id = unstructured_elements[0].metadata.parent_id
|
||||
|
||||
if root_element_id is None:
|
||||
root_element_id = ontology.OntologyElement.generate_unique_id()
|
||||
unstructured_elements[0].metadata.parent_id = root_element_id
|
||||
|
||||
id_to_element_mapping[root_element_id] = ontology.Document(
|
||||
additional_attributes={"id": root_element_id}
|
||||
)
|
||||
|
||||
for element in unstructured_elements:
|
||||
html_as_tags = BeautifulSoup(element.metadata.text_as_html, "html.parser").find_all(
|
||||
recursive=False
|
||||
)
|
||||
element_id = element.id
|
||||
parent_id = element.metadata.parent_id
|
||||
|
||||
if parent_id is None:
|
||||
# Make sure that no element is lost
|
||||
parent_id = root_element_id
|
||||
|
||||
for html_as_tag in html_as_tags:
|
||||
ontology_element = parse_html_to_ontology_element(html_as_tag)
|
||||
id_to_element_mapping[element_id] = ontology_element
|
||||
id_to_element_mapping[parent_id].children.append(ontology_element)
|
||||
|
||||
root_id, root_element = id_to_element_mapping.popitem(last=False)
|
||||
return root_element
|
||||
|
||||
|
||||
def parse_html_to_ontology(html_code: str) -> ontology.OntologyElement:
|
||||
"""
|
||||
Parses the given HTML code and converts it into an Element object.
|
||||
|
||||
Args:
|
||||
html_code (str): The HTML code to be parsed.
|
||||
Parsing HTML will start from <div class="Page">.
|
||||
Returns:
|
||||
OntologyElement: The parsed Element object.
|
||||
|
||||
Raises:
|
||||
ValueError: If no <body class="Document"> element is found in the HTML.
|
||||
"""
|
||||
html_code = remove_empty_divs_from_html_content(html_code)
|
||||
html_code = remove_empty_tags_from_html_content(html_code)
|
||||
soup = BeautifulSoup(html_code, "html.parser")
|
||||
document = soup.find("body", class_="Document")
|
||||
if not document:
|
||||
document = soup.find("div", class_="Page")
|
||||
|
||||
if not document:
|
||||
raise ValueError(
|
||||
"No <body class='Document'> or <div class='Page'> element found in the HTML."
|
||||
)
|
||||
|
||||
document_element = parse_html_to_ontology_element(document)
|
||||
return document_element
|
||||
|
||||
|
||||
def remove_empty_divs_from_html_content(html_content: str) -> str:
|
||||
soup = BeautifulSoup(html_content, "html.parser")
|
||||
divs = soup.find_all("div")
|
||||
for div in reversed(divs):
|
||||
if not div.attrs:
|
||||
div.unwrap()
|
||||
return str(soup)
|
||||
|
||||
|
||||
def remove_empty_tags_from_html_content(html_content: str) -> str:
|
||||
soup = BeautifulSoup(html_content, "html.parser")
|
||||
|
||||
def is_empty(tag):
|
||||
# Remove only specific tags, omit self-closing ones
|
||||
if tag.name not in ["p", "span", "div", "h1", "h2", "h3", "h4", "h5", "h6"]:
|
||||
return False
|
||||
|
||||
if tag.find():
|
||||
return False
|
||||
|
||||
if tag.attrs:
|
||||
return False
|
||||
|
||||
return bool(not tag.get_text(strip=True))
|
||||
|
||||
def remove_empty_tags(soup):
|
||||
for tag in soup.find_all():
|
||||
if is_empty(tag):
|
||||
tag.decompose()
|
||||
|
||||
remove_empty_tags(soup)
|
||||
|
||||
return str(soup)
|
||||
|
||||
|
||||
def parse_html_to_ontology_element(soup: Tag, recursion_depth: int = 1) -> ontology.OntologyElement:
|
||||
"""
|
||||
Converts a BeautifulSoup Tag object into an OntologyElement object. This function is recursive.
|
||||
First tries to recognize a class from Unstructured Ontology, then if class is matched tries
|
||||
to go deeper inside HTML tree. The recursive parsing is ended if the class is not recognized or
|
||||
there are no HTML Tags inside HTML - just text. Then it is parsed to
|
||||
Paragraph or UncategorizedText object.
|
||||
|
||||
Args:
|
||||
soup (Tag): The BeautifulSoup Tag object to be converted.
|
||||
recursion_depth (int): Flag to control limit of recursion depth.
|
||||
Returns:
|
||||
OntologyElement: The converted OntologyElement object.
|
||||
"""
|
||||
ontology_html_tag, ontology_class = extract_tag_and_ontology_class_from_tag(soup)
|
||||
escaped_attrs = get_escaped_attributes(soup)
|
||||
|
||||
if soup.name == "br": # Note(Pluto) should it be <br class="UncategorizedText">?
|
||||
return ontology.Paragraph(
|
||||
text="",
|
||||
css_class_name=None,
|
||||
html_tag_name="br",
|
||||
additional_attributes=escaped_attrs,
|
||||
)
|
||||
|
||||
has_children = (
|
||||
(ontology_class != ontology.UncategorizedText)
|
||||
and any(isinstance(content, Tag) for content in soup.contents)
|
||||
or ontology_class().elementType == ontology.ElementTypeEnum.layout
|
||||
)
|
||||
should_unwrap_html = has_children and recursion_depth <= RECURSION_LIMIT
|
||||
|
||||
if should_unwrap_html:
|
||||
text = ""
|
||||
children = [
|
||||
(
|
||||
parse_html_to_ontology_element(child, recursion_depth=recursion_depth + 1)
|
||||
if isinstance(child, Tag)
|
||||
else ontology.Paragraph(text=str(child).strip())
|
||||
)
|
||||
for child in soup.children
|
||||
if str(child).strip()
|
||||
]
|
||||
else:
|
||||
text = "\n".join([str(content).strip() for content in soup.contents]).strip()
|
||||
children = []
|
||||
|
||||
output_element = ontology_class(
|
||||
text=text,
|
||||
children=children,
|
||||
html_tag_name=ontology_html_tag,
|
||||
additional_attributes=escaped_attrs,
|
||||
)
|
||||
# TODO (Pluto): <input class="FormFieldValue"/> requires being wrapped in <label> tags
|
||||
return output_element
|
||||
|
||||
|
||||
def extract_tag_and_ontology_class_from_tag(
|
||||
soup: Tag,
|
||||
) -> tuple[str, Type[ontology.OntologyElement]]:
|
||||
"""
|
||||
Extracts the HTML tag and corresponding ontology class
|
||||
from a BeautifulSoup Tag object. The CSS class is prioritized over
|
||||
the HTML tag. If not recognized soup.name and UnstructuredText is returned.
|
||||
|
||||
Args:
|
||||
soup (Tag): The BeautifulSoup Tag object to extract information from.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the HTML tag (str) and the ontology class (Type[OntologyElement]).
|
||||
"""
|
||||
html_tag, element_class = None, None
|
||||
|
||||
# Scenario 1: Valid Ontology Element
|
||||
if soup.attrs.get("class"):
|
||||
html_tag, element_class = (
|
||||
soup.name,
|
||||
HTML_TAG_AND_CSS_NAME_TO_ELEMENT_TYPE_MAP.get((soup.name, soup.attrs["class"][0])),
|
||||
)
|
||||
|
||||
# Scenario 2: HTML tag incorrect, CSS class correct
|
||||
# Fallback to css name selector and overwrite html tag
|
||||
if (
|
||||
not element_class
|
||||
and soup.attrs.get("class")
|
||||
and soup.attrs["class"][0] in CSS_CLASS_TO_ELEMENT_TYPE_MAP
|
||||
):
|
||||
element_class = CSS_CLASS_TO_ELEMENT_TYPE_MAP.get(soup.attrs["class"][0])
|
||||
html_tag = element_class().allowed_tags[0]
|
||||
|
||||
# Scenario 3: <input> elements, handled explicitly based on their 'type' attribute
|
||||
if not element_class and soup.name == "input":
|
||||
input_type = (str(soup.get("type")) or "").lower()
|
||||
if input_type == "checkbox":
|
||||
element_class = ontology.Checkbox
|
||||
elif input_type == "radio":
|
||||
element_class = ontology.RadioButton
|
||||
else:
|
||||
# Any other input (including missing type or text/number/etc.) is considered
|
||||
# a generic form field value.
|
||||
element_class = ontology.FormFieldValue
|
||||
html_tag = "input"
|
||||
|
||||
# Scenario 4: CSS class incorrect, but HTML tag correct and exclusive in ontology
|
||||
if not element_class and soup.name in HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP:
|
||||
html_tag, element_class = soup.name, HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP[soup.name]
|
||||
|
||||
# Scenario 5: CSS class incorrect, HTML tag incorrect
|
||||
# Fallback to default UncategorizedText
|
||||
if not element_class:
|
||||
# TODO (Pluto): Sometimes we could infer that from parent type and soup.name
|
||||
# e.g. parent=FormField soup.name=input -> element=FormFieldInput
|
||||
|
||||
html_tag = "span"
|
||||
element_class = ontology.UncategorizedText
|
||||
|
||||
# Scenario 6: UncategorizedText has image and no text
|
||||
# Typically, this happens with a span or div tag with an image inside
|
||||
if element_class == ontology.UncategorizedText and soup.find("img") and not soup.text.strip():
|
||||
element_class = ontology.Image
|
||||
|
||||
return html_tag, element_class
|
||||
|
||||
|
||||
def get_escaped_attributes(soup: Tag) -> dict[str, str | list[str]]:
|
||||
"""
|
||||
Escapes the attributes of a BeautifulSoup Tag object.
|
||||
|
||||
Args:
|
||||
soup (Tag): The BeautifulSoup Tag object whose attributes need to be escaped.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary with escaped attribute names and values.
|
||||
"""
|
||||
escaped_attrs: dict[str, str | list[str]] = {}
|
||||
for key, value in soup.attrs.items():
|
||||
escaped_key = html.escape(key)
|
||||
escaped_value = None
|
||||
if value:
|
||||
if isinstance(value, list):
|
||||
escaped_value = [html.escape(v) for v in value]
|
||||
else:
|
||||
escaped_value = html.escape(value)
|
||||
escaped_attrs[escaped_key] = escaped_value
|
||||
return escaped_attrs
|
||||
@@ -0,0 +1,123 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.documents.elements import Element, process_metadata
|
||||
from unstructured.file_utils.filetype import add_metadata
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.lang import check_language_args
|
||||
from unstructured.partition.pdf import partition_pdf_or_image
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
|
||||
|
||||
@process_metadata() # TODO(shreya): update to use `apply_metadata` decorator
|
||||
@add_metadata
|
||||
@add_chunking_strategy
|
||||
def partition_image(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
include_page_breaks: bool = False,
|
||||
infer_table_structure: bool = False,
|
||||
ocr_languages: Optional[str] = None,
|
||||
languages: Optional[list[str]] = None,
|
||||
detect_language_per_element: bool = False,
|
||||
strategy: str = PartitionStrategy.HI_RES,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
chunking_strategy: Optional[str] = None,
|
||||
hi_res_model_name: Optional[str] = None,
|
||||
extract_images_in_pdf: bool = False,
|
||||
extract_image_block_types: Optional[list[str]] = None,
|
||||
extract_image_block_output_dir: Optional[str] = None,
|
||||
extract_image_block_to_payload: bool = False,
|
||||
starting_page_number: int = 1,
|
||||
extract_forms: bool = False,
|
||||
form_extraction_skip_tables: bool = True,
|
||||
password: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Parses an image into a list of interpreted elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object as bytes --> open(filename, "rb").
|
||||
include_page_breaks
|
||||
If True, includes page breaks at the end of each page in the document.
|
||||
infer_table_structure
|
||||
Only applicable if `strategy=hi_res`.
|
||||
If True, any Table elements that are extracted will also have a metadata field
|
||||
named "text_as_html" where the table's text content is rendered into an html string.
|
||||
I.e., rows and cells are preserved.
|
||||
Whether True or False, the "text" field is always present in any Table element
|
||||
and is the text content of the table (no structure).
|
||||
languages
|
||||
The languages present in the document, for use in partitioning and/or OCR. To use a language
|
||||
with Tesseract, you'll first need to install the appropriate Tesseract language pack.
|
||||
strategy
|
||||
The strategy to use for partitioning the image. Valid strategies are "hi_res" and
|
||||
"ocr_only". When using the "hi_res" strategy, the function uses a layout detection
|
||||
model if to identify document elements. When using the "ocr_only" strategy,
|
||||
partition_image simply extracts the text from the document using OCR and processes it.
|
||||
The default strategy is `hi_res`.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
hi_res_model_name
|
||||
The layout detection model used when partitioning strategy is set to `hi_res`.
|
||||
extract_images_in_pdf
|
||||
Only applicable if `strategy=hi_res`.
|
||||
If True, any detected images will be saved in the path specified by
|
||||
'extract_image_block_output_dir' or stored as base64 encoded data within metadata fields.
|
||||
Deprecation Note: This parameter is marked for deprecation. Future versions will use
|
||||
'extract_image_block_types' for broader extraction capabilities.
|
||||
extract_image_block_types
|
||||
Only applicable if `strategy=hi_res`.
|
||||
Images of the element type(s) specified in this list (e.g., ["Image", "Table"]) will be
|
||||
saved in the path specified by 'extract_image_block_output_dir' or stored as base64 encoded
|
||||
data within metadata fields.
|
||||
extract_image_block_to_payload
|
||||
Only applicable if `strategy=hi_res`.
|
||||
If True, images of the element type(s) defined in 'extract_image_block_types' will be
|
||||
encoded as base64 data and stored in two metadata fields: 'image_base64' and
|
||||
'image_mime_type'.
|
||||
This parameter facilitates the inclusion of element data directly within the payload,
|
||||
especially for web-based applications or APIs.
|
||||
extract_image_block_output_dir
|
||||
Only applicable if `strategy=hi_res` and `extract_image_block_to_payload=False`.
|
||||
The filesystem path for saving images of the element type(s)
|
||||
specified in 'extract_image_block_types'.
|
||||
extract_forms
|
||||
Whether the form extraction logic should be run
|
||||
(results in adding FormKeysValues elements to output).
|
||||
form_extraction_skip_tables
|
||||
Whether the form extraction logic should ignore regions designated as Tables.
|
||||
password
|
||||
The password to decrypt the PDF file.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
languages = check_language_args(languages or [], ocr_languages)
|
||||
|
||||
return partition_pdf_or_image(
|
||||
filename=filename,
|
||||
file=file,
|
||||
is_image=True,
|
||||
include_page_breaks=include_page_breaks,
|
||||
infer_table_structure=infer_table_structure,
|
||||
languages=languages,
|
||||
detect_language_per_element=detect_language_per_element,
|
||||
strategy=strategy,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
hi_res_model_name=hi_res_model_name,
|
||||
extract_images_in_pdf=extract_images_in_pdf,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_output_dir=extract_image_block_output_dir,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
starting_page_number=starting_page_number,
|
||||
extract_forms=extract_forms,
|
||||
form_extraction_skip_tables=form_extraction_skip_tables,
|
||||
password=password,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Provides `partition_json()`.
|
||||
|
||||
Note this does not partition arbitrary JSON. Its only use-case is to "rehydrate" unstructured
|
||||
document elements serialized to JSON, essentially the same function as `elements_from_json()`, but
|
||||
this allows a document of already-partitioned elements to be combined transparently with other
|
||||
documents in a partitioning run. It also allows multiple (low-cost) chunking runs to be performed on
|
||||
a document while only incurring partitioning cost once.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.documents.elements import Element, process_metadata
|
||||
from unstructured.file_utils.filetype import (
|
||||
FileType,
|
||||
add_metadata_with_filetype,
|
||||
is_json_processable,
|
||||
)
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.staging.base import elements_from_dicts
|
||||
|
||||
|
||||
@process_metadata()
|
||||
@add_metadata_with_filetype(FileType.JSON)
|
||||
@add_chunking_strategy
|
||||
def partition_json(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
text: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions serialized Unstructured output into its constituent elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object as bytes --> open(filename, "rb").
|
||||
text
|
||||
The string representation of the .json document.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
if text is not None and text.strip() == "" and not file and not filename:
|
||||
return []
|
||||
|
||||
exactly_one(filename=filename, file=file, text=text)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
file_text = ""
|
||||
if filename is not None:
|
||||
with open(filename, encoding="utf8") as f:
|
||||
file_text = f.read()
|
||||
|
||||
elif file is not None:
|
||||
file_content = file.read()
|
||||
file_text = file_content if isinstance(file_content, str) else file_content.decode()
|
||||
file.seek(0)
|
||||
|
||||
elif text is not None:
|
||||
file_text = str(text)
|
||||
|
||||
if not is_json_processable(file_text=file_text):
|
||||
raise ValueError(
|
||||
"JSON cannot be partitioned. Schema does not match the Unstructured schema.",
|
||||
)
|
||||
|
||||
try:
|
||||
element_dicts = json.loads(file_text)
|
||||
elements = elements_from_dicts(element_dicts)
|
||||
# if we found at least one json element, but no unstructured elements were found, throw 422
|
||||
if len(element_dicts) > 0 and len(elements) == 0:
|
||||
raise ValueError(
|
||||
"JSON cannot be partitioned. Schema does not match the Unstructured schema.",
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError("Not a valid json")
|
||||
|
||||
for element in elements:
|
||||
element.metadata.last_modified = metadata_last_modified or last_modified
|
||||
|
||||
return elements
|
||||
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any
|
||||
|
||||
import markdown
|
||||
import requests
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.encoding import read_txt_file
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
|
||||
def optional_decode(contents: str | bytes) -> str:
|
||||
if isinstance(contents, bytes):
|
||||
return contents.decode("utf-8")
|
||||
return contents
|
||||
|
||||
|
||||
DETECTION_ORIGIN: str = "md"
|
||||
|
||||
|
||||
def partition_md(
|
||||
filename: str | None = None,
|
||||
file: IO[bytes] | None = None,
|
||||
text: str | None = None,
|
||||
url: str | None = None,
|
||||
metadata_filename: str | None = None,
|
||||
metadata_last_modified: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions a markdown file into its constituent elements
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
text
|
||||
The string representation of the markdown document.
|
||||
url
|
||||
The URL of a webpage to parse. Only for URLs that return a markdown document.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
if text is None:
|
||||
text = ""
|
||||
|
||||
# -- verify that only one of the arguments was provided --
|
||||
exactly_one(filename=filename, file=file, text=text, url=url)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
if filename is not None:
|
||||
_, text = read_txt_file(filename=filename)
|
||||
|
||||
elif file is not None:
|
||||
_, text = read_txt_file(file=file)
|
||||
|
||||
elif url is not None:
|
||||
response = requests.get(url)
|
||||
if not response.ok:
|
||||
raise ValueError(f"URL return an error: {response.status_code}")
|
||||
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
if not content_type.startswith("text/markdown"):
|
||||
raise ValueError(
|
||||
f"Expected content type text/markdown. Got {content_type}.",
|
||||
)
|
||||
|
||||
text = response.text
|
||||
|
||||
html = markdown.markdown(text, extensions=["tables", "fenced_code"])
|
||||
|
||||
return partition_html(
|
||||
text=html,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.MD,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
import os
|
||||
|
||||
from unstructured_inference.models.base import get_model
|
||||
|
||||
|
||||
def initialize():
|
||||
"""Download default model or model specified by UNSTRUCTURED_HI_RES_MODEL_NAME environment
|
||||
variable (avoids subprocesses all doing the same)"""
|
||||
|
||||
# If more than one model will be supported and left up to user selection
|
||||
supported_model = os.environ.get("UNSTRUCTURED_HI_RES_SUPPORTED_MODEL", "")
|
||||
if supported_model:
|
||||
for model_name in supported_model.split(","):
|
||||
get_model(model_name=model_name)
|
||||
|
||||
get_model(os.environ.get("UNSTRUCTURED_HI_RES_MODEL_NAME"))
|
||||
@@ -0,0 +1,315 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from typing import IO, Any, Iterator, Optional
|
||||
|
||||
from oxmsg import Message
|
||||
from oxmsg.attachment import Attachment
|
||||
|
||||
from unstructured.documents.elements import Element, ElementMetadata
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.common import UnsupportedFileFormatError
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.partition.text import partition_text
|
||||
from unstructured.utils import is_temp_file_path, lazyproperty
|
||||
|
||||
|
||||
def partition_msg(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
process_attachments: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions a MSFT Outlook .msg file
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_filename
|
||||
The filename to use for the metadata.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
process_attachments
|
||||
If True, partition_email will process email attachments in addition to
|
||||
processing the content of the email itself.
|
||||
"""
|
||||
opts = MsgPartitionerOptions(
|
||||
file=file,
|
||||
file_path=filename,
|
||||
metadata_file_path=metadata_filename,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
partition_attachments=process_attachments,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
return list(_MsgPartitioner.iter_message_elements(opts))
|
||||
|
||||
|
||||
class MsgPartitionerOptions:
|
||||
"""Encapsulates partitioning option validation, computation, and application of defaults."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
file: IO[bytes] | None,
|
||||
file_path: str | None,
|
||||
metadata_file_path: str | None,
|
||||
metadata_last_modified: str | None,
|
||||
partition_attachments: bool,
|
||||
kwargs: dict[str, Any],
|
||||
):
|
||||
self._file = file
|
||||
self._file_path = file_path
|
||||
self._metadata_file_path = metadata_file_path
|
||||
self._metadata_last_modified = metadata_last_modified
|
||||
self._partition_attachments = partition_attachments
|
||||
self._kwargs = kwargs
|
||||
|
||||
@lazyproperty
|
||||
def extra_msg_metadata(self) -> ElementMetadata:
|
||||
"""ElementMetadata suitable for use on an element formed from message content.
|
||||
|
||||
These are only the metadata fields specific to email messages. The remaining metadata
|
||||
fields produced by the delegate partitioner are used as produced.
|
||||
|
||||
None of these metadata fields change based on the element, so we just compute it once.
|
||||
"""
|
||||
msg = self.msg
|
||||
|
||||
sent_from = [s.strip() for s in sender.split(",")] if (sender := msg.sender) else None
|
||||
sent_to = [r.email_address for r in msg.recipients] or None
|
||||
bcc_recipient = (
|
||||
[c.strip() for c in bcc.split(",")] if (bcc := msg.message_headers.get("Bcc")) else None
|
||||
)
|
||||
cc_recipient = (
|
||||
[c.strip() for c in cc.split(",")] if (cc := msg.message_headers.get("Cc")) else None
|
||||
)
|
||||
if email_message_id := msg.message_headers.get("Message-Id"):
|
||||
email_message_id = re.sub(r"^<|>$", "", email_message_id) # Strip angle brackets
|
||||
|
||||
element_metadata = ElementMetadata(
|
||||
bcc_recipient=bcc_recipient,
|
||||
cc_recipient=cc_recipient,
|
||||
email_message_id=email_message_id,
|
||||
sent_from=sent_from,
|
||||
sent_to=sent_to,
|
||||
subject=msg.subject or None,
|
||||
)
|
||||
element_metadata.detection_origin = "msg"
|
||||
|
||||
return element_metadata
|
||||
|
||||
@lazyproperty
|
||||
def is_encrypted(self) -> bool:
|
||||
"""True when message is encrypted."""
|
||||
# NOTE(robinson) - Per RFC 2015, the content type for emails with PGP encrypted content
|
||||
# is multipart/encrypted (ref: https://www.ietf.org/rfc/rfc2015.txt)
|
||||
# NOTE(scanny) - pretty sure we're going to want to dig deeper to discover messages that are
|
||||
# encrypted with something other than PGP.
|
||||
# - might be able to distinguish based on PID_MESSAGE_CLASS = 'IPM.Note.Signed'
|
||||
# - Content-Type header might include "application/pkcs7-mime" for Microsoft S/MIME
|
||||
# encryption.
|
||||
return "encrypted" in self.msg.message_headers.get("Content-Type", "")
|
||||
|
||||
@lazyproperty
|
||||
def metadata_file_path(self) -> str | None:
|
||||
"""Best available path for MSG file.
|
||||
|
||||
The value is the caller supplied `metadata_filename` if present, falling back to the
|
||||
source file-path if that was provided, otherwise `None`.
|
||||
"""
|
||||
return self._metadata_file_path or self._file_path
|
||||
|
||||
@lazyproperty
|
||||
def metadata_last_modified(self) -> str | None:
|
||||
"""Caller override for `.metadata.last_modified` to be applied to all elements."""
|
||||
email_date = sent_date.isoformat() if (sent_date := self.msg.sent_date) else None
|
||||
return self._metadata_last_modified or email_date or self._last_modified
|
||||
|
||||
@lazyproperty
|
||||
def msg(self) -> Message:
|
||||
"""The `oxmsg.Message` object loaded from file or filename."""
|
||||
return Message.load(self._msg_file)
|
||||
|
||||
@lazyproperty
|
||||
def partition_attachments(self) -> bool:
|
||||
"""True when message attachments should also be partitioned."""
|
||||
return self._partition_attachments
|
||||
|
||||
@lazyproperty
|
||||
def partitioning_kwargs(self) -> dict[str, Any]:
|
||||
"""The "extra" keyword arguments received by `partition_msg()`.
|
||||
|
||||
These are passed along to delegate partitioners which extract keyword args like
|
||||
`chunking_strategy` etc. in their decorators to control metadata behaviors, etc.
|
||||
"""
|
||||
return self._kwargs
|
||||
|
||||
@lazyproperty
|
||||
def _last_modified(self) -> str | None:
|
||||
"""The best last-modified date available from source-file, None if not available."""
|
||||
if not self._file_path or is_temp_file_path(self._file_path):
|
||||
return None
|
||||
|
||||
return get_last_modified_date(self._file_path)
|
||||
|
||||
@lazyproperty
|
||||
def _msg_file(self) -> str | IO[bytes]:
|
||||
"""The source for the bytes of the message, either a file-path or a file-like object."""
|
||||
if file_path := self._file_path:
|
||||
return file_path
|
||||
|
||||
if file := self._file:
|
||||
return file
|
||||
|
||||
raise ValueError("one of `file` or `filename` arguments must be provided")
|
||||
|
||||
|
||||
class _MsgPartitioner:
|
||||
"""Partitions Outlook email message (MSG) files."""
|
||||
|
||||
def __init__(self, opts: MsgPartitionerOptions):
|
||||
self._opts = opts
|
||||
|
||||
@classmethod
|
||||
def iter_message_elements(cls, opts: MsgPartitionerOptions) -> Iterator[Element]:
|
||||
"""Partition MS Outlook email messages (.msg files) into elements."""
|
||||
if opts.is_encrypted:
|
||||
logger.warning("Encrypted email detected. Partitioner will return an empty list.")
|
||||
return
|
||||
|
||||
yield from cls(opts)._iter_message_elements()
|
||||
|
||||
def _iter_message_elements(self) -> Iterator[Element]:
|
||||
"""Partition MS Outlook email messages (.msg files) into elements."""
|
||||
yield from self._iter_message_body_elements()
|
||||
|
||||
if not self._opts.partition_attachments:
|
||||
return
|
||||
|
||||
for attachment in self._attachments:
|
||||
yield from _AttachmentPartitioner.iter_elements(attachment, self._opts)
|
||||
|
||||
@lazyproperty
|
||||
def _attachments(self) -> tuple[Attachment, ...]:
|
||||
"""The `oxmsg.attachment.Attachment` objects for this message."""
|
||||
return tuple(self._opts.msg.attachments)
|
||||
|
||||
def _iter_message_body_elements(self) -> Iterator[Element]:
|
||||
"""Partition the message body (but not the attachments)."""
|
||||
msg = self._opts.msg
|
||||
|
||||
if html_body := msg.html_body:
|
||||
elements = partition_html(
|
||||
text=html_body,
|
||||
metadata_filename=self._opts.metadata_file_path,
|
||||
metadata_file_type=FileType.MSG,
|
||||
metadata_last_modified=self._opts.metadata_last_modified,
|
||||
**self._opts.partitioning_kwargs,
|
||||
)
|
||||
elif msg.body:
|
||||
elements = partition_text(
|
||||
text=msg.body,
|
||||
metadata_filename=self._opts.metadata_file_path,
|
||||
metadata_file_type=FileType.MSG,
|
||||
metadata_last_modified=self._opts.metadata_last_modified,
|
||||
**self._opts.partitioning_kwargs,
|
||||
)
|
||||
else:
|
||||
elements: list[Element] = []
|
||||
|
||||
# -- augment the element metadata with email-specific values --
|
||||
email_specific_metadata = self._opts.extra_msg_metadata
|
||||
for e in elements:
|
||||
e.metadata.update(email_specific_metadata)
|
||||
yield e
|
||||
|
||||
|
||||
class _AttachmentPartitioner:
|
||||
"""Partitions an attachment to a MSG file."""
|
||||
|
||||
def __init__(self, attachment: Attachment, opts: MsgPartitionerOptions):
|
||||
self._attachment = attachment
|
||||
self._opts = opts
|
||||
|
||||
@classmethod
|
||||
def iter_elements(
|
||||
cls, attachment: Attachment, opts: MsgPartitionerOptions
|
||||
) -> Iterator[Element]:
|
||||
"""Partition an `oxmsg.attachment.Attachment` from an Outlook email message (.msg file)."""
|
||||
return cls(attachment, opts)._iter_elements()
|
||||
|
||||
def _iter_elements(self) -> Iterator[Element]:
|
||||
"""Partition the file in an `oxmsg.attachment.Attachment` into elements."""
|
||||
from unstructured.partition.auto import partition
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir_path:
|
||||
# -- save attachment as file in this temporary directory --
|
||||
detached_file_path = os.path.join(tmp_dir_path, self._attachment_file_name)
|
||||
with open(detached_file_path, "wb") as f:
|
||||
f.write(self._file_bytes)
|
||||
|
||||
# -- partition the attachment --
|
||||
try:
|
||||
elements = partition(
|
||||
detached_file_path,
|
||||
metadata_filename=self._attachment_file_name,
|
||||
metadata_last_modified=self._attachment_last_modified,
|
||||
**self._opts.partitioning_kwargs,
|
||||
)
|
||||
except UnsupportedFileFormatError:
|
||||
return
|
||||
|
||||
for e in elements:
|
||||
e.metadata.attached_to_filename = self._opts.metadata_file_path
|
||||
yield e
|
||||
|
||||
@lazyproperty
|
||||
def _attachment_file_name(self) -> str:
|
||||
"""The original name of the attached file, no path.
|
||||
|
||||
This value is 'unknown' if it is not present in the MSG file (not expected).
|
||||
The filename is sanitized to prevent path traversal attacks.
|
||||
"""
|
||||
raw_filename = self._attachment.file_name or "unknown"
|
||||
|
||||
# Sanitize the filename to prevent path traversal attacks
|
||||
# Remove any path components for both Unix and Windows paths
|
||||
# Use both separators to handle cross-platform attacks
|
||||
safe_filename = os.path.basename(raw_filename.replace("\\", "/"))
|
||||
|
||||
# Remove null bytes and other control characters
|
||||
safe_filename = safe_filename.replace("\0", "")
|
||||
|
||||
# If the filename becomes empty after sanitization, use a default
|
||||
if not safe_filename or safe_filename in (".", ".."):
|
||||
safe_filename = "unknown"
|
||||
|
||||
return safe_filename
|
||||
|
||||
@lazyproperty
|
||||
def _attachment_last_modified(self) -> str | None:
|
||||
"""ISO8601 string timestamp of attachment last-modified date.
|
||||
|
||||
This value generally available on the attachment and will be the most reliable last-modifed
|
||||
time. There are fallbacks for when it is not present, ultimately `None` if we have no way
|
||||
of telling.
|
||||
"""
|
||||
if last_modified := self._attachment.last_modified:
|
||||
return last_modified.isoformat()
|
||||
return self._opts.metadata_last_modified
|
||||
|
||||
@lazyproperty
|
||||
def _file_bytes(self) -> bytes:
|
||||
"""The bytes of the attached file."""
|
||||
return self._attachment.file_bytes or b""
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Provides `partition_ndjson()`.
|
||||
|
||||
Note this does not partition arbitrary NDJSON. Its only use-case is to "rehydrate" unstructured
|
||||
document elements serialized to JSON, essentially the same function as `elements_from_json()`, but
|
||||
this allows a document of already-partitioned elements to be combined transparently with other
|
||||
documents in a partitioning run. It also allows multiple (low-cost) chunking runs to be performed on
|
||||
a document while only incurring partitioning cost once.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.documents.elements import Element, process_metadata
|
||||
from unstructured.file_utils.filetype import (
|
||||
FileType,
|
||||
add_metadata_with_filetype,
|
||||
is_ndjson_processable,
|
||||
)
|
||||
from unstructured.file_utils.ndjson import loads as ndjson_loads
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.staging.base import elements_from_dicts
|
||||
|
||||
|
||||
@process_metadata()
|
||||
@add_metadata_with_filetype(FileType.NDJSON)
|
||||
@add_chunking_strategy
|
||||
def partition_ndjson(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
text: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions serialized Unstructured output into its constituent elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object as bytes --> open(filename, "rb").
|
||||
text
|
||||
The string representation of the .json document.
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
if text is not None and text.strip() == "" and not file and not filename:
|
||||
return []
|
||||
|
||||
exactly_one(filename=filename, file=file, text=text)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
file_text = ""
|
||||
if filename is not None:
|
||||
with open(filename, encoding="utf8") as f:
|
||||
file_text = f.read()
|
||||
|
||||
elif file is not None:
|
||||
file_content = file.read()
|
||||
file_text = file_content if isinstance(file_content, str) else file_content.decode()
|
||||
file.seek(0)
|
||||
|
||||
elif text is not None:
|
||||
file_text = str(text)
|
||||
|
||||
if not is_ndjson_processable(file_text=file_text):
|
||||
raise ValueError(
|
||||
"NDJSON cannot be partitioned. Schema does not match the Unstructured schema.",
|
||||
)
|
||||
|
||||
try:
|
||||
element_dicts = ndjson_loads(file_text)
|
||||
elements = elements_from_dicts(element_dicts)
|
||||
# if we found at least one json element, but no unstructured elements were found, throw 422
|
||||
if len(element_dicts) > 0 and len(elements) == 0:
|
||||
raise ValueError(
|
||||
"JSON cannot be partitioned. Schema does not match the Unstructured schema.",
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError("Not a valid ndjson")
|
||||
|
||||
for element in elements:
|
||||
element.metadata.last_modified = metadata_last_modified or last_modified
|
||||
|
||||
return elements
|
||||
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from typing import IO, Any, Optional, cast
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.docx import partition_docx
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
|
||||
def partition_odt(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions Open Office Documents in .odt format into its document elements.
|
||||
|
||||
All parameters that are available on `partition_docx()` are also available here.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
infer_table_structure
|
||||
If True, any Table elements that are extracted will also have a metadata field
|
||||
named "text_as_html" where the table's text content is rendered into an html string.
|
||||
I.e., rows and cells are preserved.
|
||||
Whether True or False, the "text" field is always present in any Table element
|
||||
and is the text content of the table (no structure).
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
languages
|
||||
User defined value for `metadata.languages` if provided. Otherwise language is detected
|
||||
using naive Bayesian filter via `langdetect`. Multiple languages indicates text could be
|
||||
in either language.
|
||||
Additional Parameters:
|
||||
detect_language_per_element
|
||||
Detect language per element instead of at the document level.
|
||||
"""
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
with tempfile.TemporaryDirectory() as target_dir:
|
||||
docx_path = _convert_odt_to_docx(target_dir, filename, file)
|
||||
elements = partition_docx(
|
||||
filename=docx_path,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.ODT,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return elements
|
||||
|
||||
|
||||
@requires_dependencies("pypandoc")
|
||||
def _convert_odt_to_docx(
|
||||
target_dir: str, filename: Optional[str], file: Optional[IO[bytes]]
|
||||
) -> str:
|
||||
"""Convert ODT document to DOCX returning the new .docx file's path.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target_dir
|
||||
The str directory-path to use for conversion purposes. The new DOCX file is written to this
|
||||
directory. When passed as a file-like object, a copy of the source file is written here as
|
||||
well. It is the caller's responsibility to remove this directory and its contents when
|
||||
they are no longer needed.
|
||||
filename
|
||||
A str file-path specifying the location of the source ODT file on the local filesystem.
|
||||
file
|
||||
A file-like object open for reading in binary mode ("rb" mode).
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
# -- validate file-path when provided so we can provide a more meaningful error than whatever
|
||||
# -- would come from pandoc.
|
||||
if filename is not None and not os.path.exists(filename):
|
||||
raise ValueError(f"The file {filename} does not exist.")
|
||||
|
||||
# -- Pandoc is a command-line program running in its own memory-space. It can therefore only
|
||||
# -- operate on files on the filesystem. If the source document was passed as `file`, write
|
||||
# -- it to `target_dir/document.odt` and use that path as the source-path.
|
||||
source_file_path = f"{target_dir}/document.odt" if file is not None else cast(str, filename)
|
||||
if file is not None:
|
||||
with open(source_file_path, "wb") as f:
|
||||
f.write(file.read())
|
||||
|
||||
# -- Compute the path of the resulting .docx document. We want its file-name to be preserved
|
||||
# -- if the source-document was provided as `filename`.
|
||||
# -- a/b/foo.odt -> foo.odt --
|
||||
file_name = os.path.basename(source_file_path)
|
||||
# -- foo.odt -> foo --
|
||||
base_name, _ = os.path.splitext(file_name)
|
||||
# -- foo -> foo.docx --
|
||||
target_docx_path = os.path.join(target_dir, f"{base_name}.docx")
|
||||
|
||||
import pypandoc
|
||||
|
||||
pypandoc.convert_file(
|
||||
source_file_path, "docx", format="odt", outputfile=target_docx_path, sandbox=True
|
||||
)
|
||||
|
||||
return target_docx_path
|
||||
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.file_conversion import convert_file_to_html_text_using_pandoc
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
DETECTION_ORIGIN: str = "org"
|
||||
|
||||
|
||||
def partition_org(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
metadata_filename: str | None = None,
|
||||
metadata_last_modified: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an org document. The document is first converted to HTML and then
|
||||
partitioned using partition_html.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
html_text = convert_file_to_html_text_using_pandoc(
|
||||
source_format="org", filename=filename, file=file
|
||||
)
|
||||
|
||||
return partition_html(
|
||||
text=html_text,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.ORG,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
**kwargs,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,689 @@
|
||||
import logging
|
||||
import math
|
||||
import tempfile
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any, Generator, List, Optional, TypeVar, Union
|
||||
|
||||
import numpy as np
|
||||
from matplotlib import colors, font_manager
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from unstructured_inference.constants import ElementType
|
||||
|
||||
from unstructured.partition.pdf_image.analysis.processor import AnalysisProcessor
|
||||
from unstructured.partition.pdf_image.pdf_image_utils import convert_pdf_to_image
|
||||
|
||||
PageImage = TypeVar("PageImage", Image.Image, np.ndarray)
|
||||
|
||||
|
||||
def get_font():
|
||||
preferred_fonts = ["Arial.ttf"]
|
||||
available_fonts = font_manager.findSystemFonts()
|
||||
if not available_fonts:
|
||||
raise ValueError("No fonts available")
|
||||
for font in preferred_fonts:
|
||||
for available_font in available_fonts:
|
||||
if font in available_font:
|
||||
return available_font
|
||||
return available_fonts[0]
|
||||
|
||||
|
||||
COLOR_WHITE = ("white", (255, 255, 255))
|
||||
COLOR_BLACK = ("black", (0, 0, 0))
|
||||
|
||||
|
||||
class TextAlignment(Enum):
|
||||
TOP_LEFT = "top_left"
|
||||
TOP_RIGHT = "top_right"
|
||||
BOTTOM_LEFT = "bottom_left"
|
||||
BOTTOM_RIGHT = "bottom_right"
|
||||
CENTER = "center"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BboxLabels:
|
||||
top_left: Optional[str] = None
|
||||
top_right: Optional[str] = None
|
||||
bottom_left: Optional[str] = None
|
||||
bottom_right: Optional[str] = None
|
||||
center: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BBox:
|
||||
points: tuple[int, int, int, int]
|
||||
labels: Optional[BboxLabels] = None
|
||||
|
||||
|
||||
def get_rgb_color(color: str) -> tuple[int, int, int]:
|
||||
"""Convert a color name to RGB values.
|
||||
|
||||
Args:
|
||||
color: A color name supported by matplotlib.
|
||||
|
||||
Returns:
|
||||
A tuple of three integers representing the RGB values of the color.
|
||||
"""
|
||||
try:
|
||||
rgb_colors = colors.to_rgb(color)
|
||||
except ValueError:
|
||||
print("Error")
|
||||
raise
|
||||
return int(rgb_colors[0] * 255), int(rgb_colors[1] * 255), int(rgb_colors[2] * 255)
|
||||
|
||||
|
||||
def _get_bbox_to_page_ratio(bbox: tuple[int, int, int, int], page_size: tuple[int, int]) -> float:
|
||||
"""Compute the ratio of the bounding box to the page size.
|
||||
|
||||
Args:
|
||||
bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
|
||||
page_size: Tuple containing page size: (width, height).
|
||||
|
||||
Returns:
|
||||
The ratio of the bounding box to the page size.
|
||||
"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
page_width, page_height = page_size
|
||||
page_diagonal = math.sqrt(page_height**2 + page_width**2)
|
||||
bbox_width = x2 - x1
|
||||
bbox_height = y2 - y1
|
||||
bbox_diagonal = math.sqrt(bbox_height**2 + bbox_width**2)
|
||||
return bbox_diagonal / page_diagonal
|
||||
|
||||
|
||||
def _get_optimal_value_for_bbox(
|
||||
bbox: tuple[int, int, int, int],
|
||||
page_size: tuple[int, int],
|
||||
min_value: int,
|
||||
max_value: int,
|
||||
ratio_for_min_value: float = 0.01,
|
||||
ratio_for_max_value: float = 0.5,
|
||||
) -> int:
|
||||
"""Compute the optimal value for a given bounding box using a linear function
|
||||
generated for given min and max values and ratios
|
||||
|
||||
Args:
|
||||
bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
|
||||
page_size: Tuple containing page size: (width, height).
|
||||
min_value: The minimum value returned by the function.
|
||||
max_value: The maximum value returned by the function.
|
||||
ratio_for_min_value: The ratio of the bbox to page size for the min value.
|
||||
ratio_for_max_value: The ratio of the bbox to page size for the max value.
|
||||
|
||||
Returns:
|
||||
The optimal value for the given bounding box and parameters given.
|
||||
"""
|
||||
bbox_to_page_ratio = _get_bbox_to_page_ratio(bbox, page_size)
|
||||
# Direct linear interpolation instead of np.polyfit for better performance
|
||||
slope = (max_value - min_value) / (ratio_for_max_value - ratio_for_min_value)
|
||||
value = int(min_value + slope * (bbox_to_page_ratio - ratio_for_min_value))
|
||||
return max(min_value, min(max_value, value))
|
||||
|
||||
|
||||
def get_bbox_text_size(
|
||||
bbox: tuple[int, int, int, int],
|
||||
page_size: tuple[int, int],
|
||||
min_font_size: int = 16,
|
||||
max_font_size: int = 32,
|
||||
) -> int:
|
||||
"""Compute the optimal font size for a given bounding box.
|
||||
|
||||
Args:
|
||||
bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
|
||||
page_size: Tuple containing page size: (width, height).
|
||||
min_font_size: The minimum font size returned by the function.
|
||||
max_font_size: The maximum font size returned by the function.
|
||||
|
||||
Returns:
|
||||
The optimal font size for the given bounding box.
|
||||
"""
|
||||
return _get_optimal_value_for_bbox(
|
||||
bbox=bbox,
|
||||
page_size=page_size,
|
||||
min_value=min_font_size,
|
||||
max_value=max_font_size,
|
||||
)
|
||||
|
||||
|
||||
def get_bbox_thickness(
|
||||
bbox: tuple[int, int, int, int],
|
||||
page_size: tuple[int, int],
|
||||
min_thickness: int = 1,
|
||||
max_thickness: int = 4,
|
||||
) -> float:
|
||||
"""Compute the optimal thickness for a given bounding box.
|
||||
|
||||
Args:
|
||||
bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
|
||||
page_size: Tuple containing page size: (width, height).
|
||||
min_thickness: The minimum font size returned by the function.
|
||||
max_thickness: The maximum font size returned by the function.
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
return _get_optimal_value_for_bbox(
|
||||
bbox=bbox,
|
||||
page_size=page_size,
|
||||
min_value=min_thickness,
|
||||
max_value=max_thickness,
|
||||
)
|
||||
|
||||
|
||||
def get_text_color(
|
||||
background_color: Union[str, tuple[int, int, int]], brightness_threshold: float = 0.5
|
||||
) -> tuple[str, tuple[int, int, int]]:
|
||||
"""Returns the contrastive text color (black or white) for a given background color.
|
||||
|
||||
Args:
|
||||
background_color: Tuple containing RGB values of the background color.
|
||||
|
||||
Returns:
|
||||
Tuple containing RGB values of the text color.
|
||||
"""
|
||||
if isinstance(background_color, str):
|
||||
background_color = get_rgb_color(background_color)
|
||||
background_brightness = (
|
||||
0.299 * background_color[0] + 0.587 * background_color[1] + 0.114 * background_color[0]
|
||||
) / 255
|
||||
if background_brightness > brightness_threshold:
|
||||
return COLOR_BLACK
|
||||
else:
|
||||
return COLOR_WHITE
|
||||
|
||||
|
||||
def get_label_rect_and_coords(
|
||||
alignment: TextAlignment,
|
||||
bbox_points: tuple[int, int, int, int],
|
||||
text_width: int,
|
||||
text_height: int,
|
||||
):
|
||||
indent = max(int(text_width * 0.2), 10)
|
||||
vertical_correction = max(int(text_height * 0.3), 10)
|
||||
|
||||
# with this the text should be centered in the rectangle
|
||||
rect_width = text_width + indent * 2
|
||||
# we apply a correction to the height to make the text not overlap over the rectangle
|
||||
# (the text height getter looks to return too small value)
|
||||
rect_height = text_height + vertical_correction
|
||||
x1, y1, x2, y2 = bbox_points
|
||||
if alignment is TextAlignment.CENTER:
|
||||
# center:
|
||||
horizontal_half = int(rect_width / 2 * 1.05)
|
||||
vertical_half = int(rect_height / 2 * 1.05)
|
||||
center_point = x1 + (x2 - x1) // 2, y1 + (y2 - y1) // 2
|
||||
# resize rectangle to make it look better in the center
|
||||
label_rectangle = (
|
||||
(
|
||||
center_point[0] - horizontal_half,
|
||||
center_point[1] - vertical_half,
|
||||
),
|
||||
(
|
||||
center_point[0] + horizontal_half,
|
||||
center_point[1] + vertical_half,
|
||||
),
|
||||
)
|
||||
label_coords = (
|
||||
center_point[0] - horizontal_half + int(indent * 1.05),
|
||||
center_point[1] - vertical_half * 1.05,
|
||||
)
|
||||
elif alignment is TextAlignment.TOP_LEFT:
|
||||
label_rectangle = (
|
||||
(x1, y1 - rect_height),
|
||||
(x1 + rect_width, y1),
|
||||
)
|
||||
label_coords = (x1 + indent, y1 - rect_height)
|
||||
elif alignment is TextAlignment.TOP_RIGHT:
|
||||
label_rectangle = (
|
||||
(x2 - rect_width, y1),
|
||||
(x2, y1 + rect_height),
|
||||
)
|
||||
label_coords = (x2 - text_width - indent, y1)
|
||||
elif alignment is TextAlignment.BOTTOM_LEFT:
|
||||
label_rectangle = (
|
||||
(x1, y2 - rect_height),
|
||||
(x1 + rect_width, y2),
|
||||
)
|
||||
label_coords = (x1 + indent, y2 - rect_height)
|
||||
elif alignment is TextAlignment.BOTTOM_RIGHT:
|
||||
label_rectangle = (
|
||||
(x2 - rect_width, y2 - rect_height),
|
||||
(x2, y2),
|
||||
)
|
||||
label_coords = (x2 - text_width - indent, y2 - rect_height)
|
||||
else:
|
||||
raise ValueError(f"Unknown alignment {alignment}")
|
||||
return label_rectangle, label_coords
|
||||
|
||||
|
||||
def draw_bbox_label(
|
||||
image_draw: ImageDraw.ImageDraw,
|
||||
text: str,
|
||||
bbox_points: tuple[int, int, int, int],
|
||||
alignment: TextAlignment,
|
||||
font_size: int,
|
||||
background_color: str,
|
||||
):
|
||||
"""Draw a label stick to a bounding box.
|
||||
The alignment parameter specifies where the label should be placed.
|
||||
|
||||
Args:
|
||||
image_draw: ImageDraw object to draw on the image.
|
||||
text: Text to draw.
|
||||
bbox_points: Bounding box points.
|
||||
alignment: Text alignment.
|
||||
font_size: Font size of the text.
|
||||
background_color: RGB values of the background color.
|
||||
"""
|
||||
font = ImageFont.truetype(get_font(), font_size)
|
||||
text_x1, text_y1, text_x2, text_y2 = image_draw.textbbox(
|
||||
(0, 0), text, font=font, align="center"
|
||||
)
|
||||
text_width = text_x2 - text_x1
|
||||
text_height = text_y2 - text_y1
|
||||
|
||||
label_rectangle, label_coords = get_label_rect_and_coords(
|
||||
alignment, bbox_points, text_width, text_height
|
||||
)
|
||||
|
||||
rgb_background_color = get_rgb_color(background_color)
|
||||
try:
|
||||
image_draw.rectangle(
|
||||
label_rectangle,
|
||||
fill=background_color,
|
||||
outline=background_color,
|
||||
)
|
||||
except TypeError:
|
||||
image_draw.rectangle(
|
||||
label_rectangle,
|
||||
fill=rgb_background_color,
|
||||
outline=rgb_background_color,
|
||||
)
|
||||
text_color, text_color_rgb = get_text_color(background_color)
|
||||
try:
|
||||
image_draw.text(label_coords, text, fill=text_color, font=font, align="center")
|
||||
except TypeError:
|
||||
image_draw.text(label_coords, text, fill=text_color_rgb, font=font, align="center")
|
||||
|
||||
|
||||
def draw_bbox_on_image(
|
||||
image_draw: ImageDraw.ImageDraw,
|
||||
bbox: BBox,
|
||||
color: str,
|
||||
):
|
||||
"""Draw bbox with additional labels on the image..
|
||||
|
||||
Args:
|
||||
image_draw: ImageDraw object to draw on the image.
|
||||
bbox: Bounding box to draw.
|
||||
color: RGB values of the color of the bounding box (edges + label backgrounds).
|
||||
"""
|
||||
x1, y1, x2, y2 = bbox.points
|
||||
if x1 >= x2 or y1 >= y2:
|
||||
print(f"Invalid bbox coordinates: {bbox.points}")
|
||||
return
|
||||
top_left = x1, y1 # the main
|
||||
bottom_right = x2, y2
|
||||
box_thickness = get_bbox_thickness(bbox=bbox.points, page_size=image_draw.im.size)
|
||||
font_size = get_bbox_text_size(bbox=bbox.points, page_size=image_draw.im.size)
|
||||
|
||||
try:
|
||||
image_draw.rectangle((top_left, bottom_right), outline=color, width=box_thickness)
|
||||
except TypeError:
|
||||
rgb_color = get_rgb_color(color)
|
||||
image_draw.rectangle((top_left, bottom_right), outline=rgb_color, width=box_thickness)
|
||||
|
||||
if bbox.labels is not None:
|
||||
if top_left_label := bbox.labels.top_left:
|
||||
draw_bbox_label(
|
||||
image_draw,
|
||||
top_left_label,
|
||||
bbox_points=bbox.points,
|
||||
alignment=TextAlignment.TOP_LEFT,
|
||||
font_size=font_size,
|
||||
background_color=color,
|
||||
)
|
||||
if top_right_label := bbox.labels.top_right:
|
||||
draw_bbox_label(
|
||||
image_draw,
|
||||
top_right_label,
|
||||
bbox_points=bbox.points,
|
||||
alignment=TextAlignment.TOP_RIGHT,
|
||||
font_size=font_size,
|
||||
background_color=color,
|
||||
)
|
||||
if bottom_left_label := bbox.labels.bottom_left:
|
||||
draw_bbox_label(
|
||||
image_draw,
|
||||
bottom_left_label,
|
||||
bbox_points=bbox.points,
|
||||
alignment=TextAlignment.BOTTOM_LEFT,
|
||||
font_size=font_size,
|
||||
background_color=color,
|
||||
)
|
||||
if bottom_right_label := bbox.labels.bottom_right:
|
||||
draw_bbox_label(
|
||||
image_draw,
|
||||
bottom_right_label,
|
||||
bbox_points=bbox.points,
|
||||
alignment=TextAlignment.BOTTOM_RIGHT,
|
||||
font_size=font_size,
|
||||
background_color=color,
|
||||
)
|
||||
if center_label := bbox.labels.center:
|
||||
draw_bbox_label(
|
||||
image_draw,
|
||||
center_label,
|
||||
bbox_points=bbox.points,
|
||||
alignment=TextAlignment.CENTER,
|
||||
font_size=font_size * 2,
|
||||
background_color=color,
|
||||
)
|
||||
|
||||
|
||||
class LayoutDrawer(ABC):
|
||||
layout_source: str = "unknown"
|
||||
laytout_dump: dict
|
||||
|
||||
def __init__(self, layout_dump: dict):
|
||||
self.layout_dump = layout_dump
|
||||
|
||||
def draw_layout_on_page(self, page_image: Image.Image, page_num: int) -> Image.Image:
|
||||
"""Draw the layout bboxes with additional metadata on the image."""
|
||||
layout_pages = self.layout_dump.get("pages")
|
||||
if not layout_pages:
|
||||
print(f"Warning: layout in drawer {self.__class__.__name__} is empty - skipping")
|
||||
return page_image
|
||||
if len(layout_pages) < page_num:
|
||||
print(f"Error! Page {page_num} not found in layout (pages: {len(layout_pages)})")
|
||||
return page_image
|
||||
image_draw = ImageDraw.ImageDraw(page_image)
|
||||
page_layout_dump = layout_pages[page_num - 1]
|
||||
if page_num != page_layout_dump.get("number"):
|
||||
dump_page_num = page_layout_dump.get("number")
|
||||
print(f"Warning: Requested page num {page_num} differs from dump {dump_page_num}")
|
||||
for idx, elements in enumerate(page_layout_dump["elements"], 1):
|
||||
self.render_element_on_page(idx, image_draw, elements)
|
||||
return page_image
|
||||
|
||||
@abstractmethod
|
||||
def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
|
||||
"""Draw a single element on the image."""
|
||||
|
||||
|
||||
class SimpleLayoutDrawer(LayoutDrawer, ABC):
|
||||
color: str
|
||||
show_order: bool = False
|
||||
show_text_length: bool = False
|
||||
|
||||
def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
|
||||
text_len = len(elements["text"]) if elements.get("text") else 0
|
||||
element_prob = elements.get("prob")
|
||||
element_order = f"{idx}" if self.show_order else None
|
||||
text_len = f"len: {text_len}" if self.show_text_length else None
|
||||
bbox = BBox(
|
||||
points=elements["bbox"],
|
||||
labels=BboxLabels(
|
||||
top_right=f"prob: {element_prob:.2f}" if element_prob else None,
|
||||
bottom_left=text_len,
|
||||
center=element_order,
|
||||
),
|
||||
)
|
||||
draw_bbox_on_image(image_draw, bbox, color=self.color)
|
||||
|
||||
|
||||
class PdfminerLayoutDrawer(SimpleLayoutDrawer):
|
||||
layout_source = "pdfminer"
|
||||
|
||||
def __init__(self, layout_dump: dict, color: str = "red"):
|
||||
self.layout_dump = layout_dump
|
||||
self.color = color
|
||||
self.show_order = True
|
||||
super().__init__(layout_dump)
|
||||
|
||||
|
||||
class OCRLayoutDrawer(SimpleLayoutDrawer):
|
||||
layout_source = "ocr"
|
||||
|
||||
def __init__(self, layout_dump: dict, color: str = "red"):
|
||||
self.color = color
|
||||
self.show_order = False
|
||||
self.show_text_length = False
|
||||
super().__init__(layout_dump)
|
||||
|
||||
|
||||
class ODModelLayoutDrawer(LayoutDrawer):
|
||||
layout_source = "od_model"
|
||||
|
||||
color_map = {
|
||||
ElementType.CAPTION: "salmon",
|
||||
ElementType.FOOTNOTE: "orange",
|
||||
ElementType.FORMULA: "mediumpurple",
|
||||
ElementType.LIST_ITEM: "navy",
|
||||
ElementType.PAGE_FOOTER: "deeppink",
|
||||
ElementType.PAGE_HEADER: "green",
|
||||
ElementType.PICTURE: "sienna",
|
||||
ElementType.SECTION_HEADER: "darkorange",
|
||||
ElementType.TABLE: "blue",
|
||||
ElementType.TEXT: "turquoise",
|
||||
ElementType.TITLE: "greenyellow",
|
||||
}
|
||||
|
||||
def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
|
||||
element_type = elements["type"]
|
||||
element_prob = elements.get("prob")
|
||||
bbox_points = elements["bbox"]
|
||||
color = self.get_element_type_color(element_type)
|
||||
bbox = BBox(
|
||||
points=bbox_points,
|
||||
labels=BboxLabels(
|
||||
top_left=f"{element_type}",
|
||||
top_right=f"prob: {element_prob:.2f}" if element_prob else None,
|
||||
),
|
||||
)
|
||||
draw_bbox_on_image(image_draw, bbox, color=color)
|
||||
|
||||
def get_element_type_color(self, element_type: str) -> str:
|
||||
return self.color_map.get(element_type, "cyan")
|
||||
|
||||
|
||||
class FinalLayoutDrawer(LayoutDrawer):
|
||||
layout_source = "final"
|
||||
|
||||
color_map = {
|
||||
"CheckBox": "brown",
|
||||
"ListItem": "red",
|
||||
"Title": "greenyellow",
|
||||
"NarrativeText": "turquoise",
|
||||
"Header": "green",
|
||||
"Footer": "orange",
|
||||
"FigureCaption" "Image": "sienna",
|
||||
"Table": "blue",
|
||||
"Address": "gold",
|
||||
"EmailAddress": "lightskyblue",
|
||||
"Formula": "mediumpurple",
|
||||
"CodeSnippet": "magenta",
|
||||
"PageNumber": "crimson",
|
||||
}
|
||||
|
||||
def __init__(self, layout_dump: dict):
|
||||
self.layout_dump = layout_dump
|
||||
|
||||
def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
|
||||
element_type = elements["type"]
|
||||
element_prob = elements.get("prob")
|
||||
text_len = len(elements["text"]) if elements.get("text") else 0
|
||||
bbox_points = elements["bbox"]
|
||||
color = self.get_element_type_color(element_type)
|
||||
cluster = elements.get("cluster")
|
||||
bbox = BBox(
|
||||
points=bbox_points,
|
||||
labels=BboxLabels(
|
||||
top_left=f"{element_type}",
|
||||
top_right=f"prob: {element_prob:.2f}" if element_prob else None,
|
||||
bottom_right=f"len: {text_len}",
|
||||
bottom_left=f"cl: {cluster}" if cluster else None,
|
||||
center=f"{idx}",
|
||||
),
|
||||
)
|
||||
draw_bbox_on_image(image_draw, bbox, color=color)
|
||||
|
||||
def get_element_type_color(self, element_type: str) -> str:
|
||||
return self.color_map.get(element_type, "cyan")
|
||||
|
||||
|
||||
class AnalysisDrawer(AnalysisProcessor):
|
||||
def __init__(
|
||||
self,
|
||||
filename: Optional[Union[str, Path]],
|
||||
is_image: bool,
|
||||
save_dir: Union[str, Path],
|
||||
file: Optional[BytesIO] = None,
|
||||
draw_caption: bool = True,
|
||||
draw_grid: bool = False,
|
||||
resize: Optional[float] = None,
|
||||
format: str = "png",
|
||||
):
|
||||
self.draw_caption = draw_caption
|
||||
self.draw_grid = draw_grid
|
||||
self.resize = resize
|
||||
self.is_image = is_image
|
||||
self.format = format
|
||||
self.drawers = []
|
||||
self.file = file
|
||||
|
||||
super().__init__(filename, save_dir)
|
||||
|
||||
def add_drawer(self, drawer: LayoutDrawer):
|
||||
self.drawers.append(drawer)
|
||||
|
||||
def process(self):
|
||||
filename_stem = Path(self.filename).stem
|
||||
analysis_save_dir = Path(self.save_dir) / "analysis" / filename_stem / "bboxes"
|
||||
analysis_save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for page_idx, orig_image_page in enumerate(self.load_source_image()):
|
||||
images_for_grid = []
|
||||
page_num = page_idx + 1
|
||||
for drawer in self.drawers:
|
||||
try:
|
||||
image = drawer.draw_layout_on_page(orig_image_page.copy(), page_num=page_num)
|
||||
except: # noqa: E722
|
||||
logging.exception(
|
||||
f"Error while drawing layout for page {page_num} "
|
||||
f"for file {self.filename} with drawer "
|
||||
f"{drawer.__class__.__name__}"
|
||||
)
|
||||
continue
|
||||
if self.draw_caption:
|
||||
image = self.add_caption(
|
||||
image, caption=f"Layout source: {drawer.layout_source}"
|
||||
)
|
||||
if not self.draw_grid:
|
||||
if self.resize is not None:
|
||||
image = image.resize(
|
||||
(int(image.width * self.resize), int(image.height * self.resize)),
|
||||
)
|
||||
image.save(
|
||||
analysis_save_dir / f"page{page_num}"
|
||||
f"_layout_{drawer.layout_source}.{self.format}",
|
||||
optimize=True,
|
||||
quality=85,
|
||||
)
|
||||
image.close()
|
||||
else:
|
||||
images_for_grid.append(image)
|
||||
if images_for_grid:
|
||||
grid_image = self.paste_images_on_grid(images_for_grid)
|
||||
if self.resize is not None:
|
||||
grid_image = grid_image.resize(
|
||||
(int(grid_image.width * self.resize), int(grid_image.height * self.resize))
|
||||
)
|
||||
grid_image.save(
|
||||
analysis_save_dir / f"page{page_num}_layout_all.{self.format}",
|
||||
optimize=True,
|
||||
quality=85,
|
||||
)
|
||||
grid_image.close()
|
||||
|
||||
def add_caption(self, image: Image.Image, caption: str):
|
||||
font = ImageFont.truetype(get_font(), 52)
|
||||
draw = ImageDraw.ImageDraw(image)
|
||||
text_x1, text_y1, text_x2, text_y2 = draw.textbbox(
|
||||
(0, 0), caption, font=font, align="center"
|
||||
)
|
||||
text_width = text_x2 - text_x1
|
||||
text_height = int((text_y2 - text_y1) * 1.5)
|
||||
text_xy = (image.width - text_width) // 2, 10
|
||||
caption_image = Image.new("RGB", (image.width, text_height), color=(255, 255, 255))
|
||||
caption_draw = ImageDraw.ImageDraw(caption_image)
|
||||
caption_draw.text(text_xy, caption, (0, 0, 0), font=font)
|
||||
|
||||
expanded_image = Image.new("RGB", (image.width, image.height + text_height))
|
||||
expanded_image.paste(caption_image, (0, 0))
|
||||
expanded_image.paste(image, (0, text_height))
|
||||
image.close()
|
||||
return expanded_image
|
||||
|
||||
def paste_images_on_grid(self, images: List[Image.Image]) -> Image.Image:
|
||||
"""Creates a single image that presents all the images on a grid 2 x n/2"""
|
||||
|
||||
pairs = []
|
||||
for i in range(0, len(images), 2):
|
||||
left_image = images[i]
|
||||
right_image = images[i + 1] if i < len(images) - 1 else None
|
||||
pairs.append((left_image, right_image))
|
||||
|
||||
max_pair_width = max([pair[0].width + (pair[1].width if pair[1] else 0) for pair in pairs])
|
||||
sum_height = sum([max(pair[0].height, pair[1].height if pair[1] else 0) for pair in pairs])
|
||||
|
||||
new_im = Image.new("RGB", (max_pair_width, sum_height))
|
||||
|
||||
height_shift = 0
|
||||
for image_left, image_right in pairs:
|
||||
new_im.paste(image_left, (0, height_shift))
|
||||
if image_right:
|
||||
new_im.paste(image_right, (image_left.width, height_shift))
|
||||
height_shift += max(image_left.height, image_right.height if image_right else 0)
|
||||
|
||||
for image in images:
|
||||
image.close()
|
||||
return new_im
|
||||
|
||||
def load_source_image(self) -> Generator[Image.Image, None, None]:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
image_paths = []
|
||||
if self.is_image:
|
||||
if self.file:
|
||||
try:
|
||||
image = Image.open(self.file)
|
||||
output_file = Path(temp_dir) / self.filename
|
||||
image.save(output_file, format="PNG")
|
||||
image_paths = [output_file]
|
||||
except Exception as ex: # noqa: E722
|
||||
print(
|
||||
f"Error while converting image to PNG for file {self.filename}, "
|
||||
f"exception: {ex}"
|
||||
)
|
||||
else:
|
||||
image_paths = [self.filename]
|
||||
else:
|
||||
try:
|
||||
image_paths = convert_pdf_to_image(
|
||||
filename=self.filename,
|
||||
file=self.file,
|
||||
output_folder=temp_dir,
|
||||
path_only=True,
|
||||
)
|
||||
except Exception as ex: # noqa: E722
|
||||
print(
|
||||
f"Error while converting pdf to image for file {self.filename}",
|
||||
f"exception: {ex}",
|
||||
)
|
||||
|
||||
for image_path in image_paths:
|
||||
with Image.open(image_path) as image:
|
||||
yield image.convert("RGB")
|
||||
@@ -0,0 +1,203 @@
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from unstructured_inference.inference.elements import ImageTextRegion, TextRegion
|
||||
from unstructured_inference.inference.layout import DocumentLayout
|
||||
from unstructured_inference.models.base import get_model
|
||||
from unstructured_inference.models.detectron2onnx import (
|
||||
DEFAULT_LABEL_MAP as DETECTRON_LABEL_MAP,
|
||||
)
|
||||
from unstructured_inference.models.detectron2onnx import (
|
||||
UnstructuredDetectronONNXModel,
|
||||
)
|
||||
from unstructured_inference.models.yolox import YOLOX_LABEL_MAP, UnstructuredYoloXModel
|
||||
|
||||
from unstructured.documents.elements import Element, Text
|
||||
from unstructured.partition.pdf_image.analysis.processor import AnalysisProcessor
|
||||
from unstructured.partition.utils.sorting import coordinates_to_bbox
|
||||
|
||||
|
||||
class LayoutDumper(ABC):
|
||||
layout_source: str = "unknown"
|
||||
|
||||
@abstractmethod
|
||||
def dump(self) -> dict:
|
||||
"""Transforms the results to a dict convertible structured formats like JSON or YAML"""
|
||||
|
||||
|
||||
def extract_document_layout_info(layout: DocumentLayout) -> dict:
|
||||
pages = []
|
||||
|
||||
for page in layout.pages:
|
||||
size = {
|
||||
"width": page.image_metadata.get("width"),
|
||||
"height": page.image_metadata.get("height"),
|
||||
}
|
||||
elements = []
|
||||
for element in page.elements:
|
||||
bbox = element.bbox
|
||||
elements.append(
|
||||
{
|
||||
"bbox": [bbox.x1, bbox.y1, bbox.x2, bbox.y2],
|
||||
"type": element.type,
|
||||
"prob": element.prob,
|
||||
}
|
||||
)
|
||||
pages.append({"number": page.number, "size": size, "elements": elements})
|
||||
return {"pages": pages}
|
||||
|
||||
|
||||
def object_detection_classes(model_name) -> List[str]:
|
||||
model = get_model(model_name)
|
||||
if isinstance(model, UnstructuredYoloXModel):
|
||||
return list(YOLOX_LABEL_MAP.values())
|
||||
if isinstance(model, UnstructuredDetectronONNXModel):
|
||||
return list(DETECTRON_LABEL_MAP.values())
|
||||
else:
|
||||
raise ValueError(f"Cannot get OD model classes - unknown model type: {model_name}")
|
||||
|
||||
|
||||
class ObjectDetectionLayoutDumper(LayoutDumper):
|
||||
"""Forms the results in COCO format and saves them to a file"""
|
||||
|
||||
layout_source = "object_detection"
|
||||
|
||||
def __init__(self, layout: DocumentLayout, model_name: Optional[str] = None):
|
||||
self.layout: dict = extract_document_layout_info(layout)
|
||||
self.model_name = model_name
|
||||
|
||||
def dump(self) -> dict:
|
||||
"""Transforms the results to COCO format and saves them to a file"""
|
||||
try:
|
||||
classes_dict = {"object_detection_classes": object_detection_classes(self.model_name)}
|
||||
except ValueError:
|
||||
classes_dict = {"object_detection_classes": []}
|
||||
self.layout.update(classes_dict)
|
||||
return self.layout
|
||||
|
||||
|
||||
def _get_info_from_extracted_page(page: List[TextRegion]) -> List[dict]:
|
||||
elements = []
|
||||
for element in page:
|
||||
is_image = isinstance(element, ImageTextRegion)
|
||||
bbox = element.bbox
|
||||
elements.append(
|
||||
{
|
||||
"bbox": [bbox.x1, bbox.y1, bbox.x2, bbox.y2],
|
||||
"text": element.text,
|
||||
"source": str(element.source.value),
|
||||
"is_image": is_image,
|
||||
}
|
||||
)
|
||||
return elements
|
||||
|
||||
|
||||
def extract_text_regions_info(layout: List[List[TextRegion]]) -> dict:
|
||||
pages = []
|
||||
for page_num, page in enumerate(layout, 1):
|
||||
elements = _get_info_from_extracted_page(page)
|
||||
pages.append({"number": page_num, "elements": elements})
|
||||
return {"pages": pages}
|
||||
|
||||
|
||||
class ExtractedLayoutDumper(LayoutDumper):
|
||||
layout_source = "pdfminer"
|
||||
|
||||
def __init__(self, layout: List[List[TextRegion]]):
|
||||
self.layout = extract_text_regions_info(layout)
|
||||
|
||||
def dump(self) -> dict:
|
||||
return self.layout
|
||||
|
||||
|
||||
class OCRLayoutDumper(LayoutDumper):
|
||||
layout_source = "ocr"
|
||||
|
||||
def __init__(self):
|
||||
self.layout = []
|
||||
self.page_number = 1
|
||||
|
||||
def add_ocred_page(self, page: List[TextRegion]):
|
||||
elements = _get_info_from_extracted_page(page)
|
||||
self.layout.append({"number": self.page_number, "elements": elements})
|
||||
self.page_number += 1
|
||||
|
||||
def dump(self) -> dict:
|
||||
return {"pages": self.layout}
|
||||
|
||||
|
||||
def _extract_final_element_info(element: Element) -> dict:
|
||||
element_type = (
|
||||
element.category if isinstance(element, Text) else str(element.__class__.__name__)
|
||||
)
|
||||
element_prob = getattr(element.metadata, "detection_class_prob", None)
|
||||
text = element.text
|
||||
bbox_points = coordinates_to_bbox(element.metadata.coordinates)
|
||||
cluster = getattr(element.metadata, "cluster", None)
|
||||
return {
|
||||
"type": element_type,
|
||||
"prob": element_prob,
|
||||
"text": text,
|
||||
"bbox": bbox_points,
|
||||
"cluster": cluster,
|
||||
}
|
||||
|
||||
|
||||
def _extract_final_element_page_size(element: Element) -> dict:
|
||||
try:
|
||||
return {
|
||||
"width": element.metadata.coordinates.system.width,
|
||||
"height": element.metadata.coordinates.system.height,
|
||||
}
|
||||
except AttributeError:
|
||||
return {
|
||||
"width": None,
|
||||
"height": None,
|
||||
}
|
||||
|
||||
|
||||
class FinalLayoutDumper(LayoutDumper):
|
||||
layout_source = "final"
|
||||
|
||||
def __init__(self, layout: List[Element]):
|
||||
pages = defaultdict(list)
|
||||
for element in layout:
|
||||
element_page_number = element.metadata.page_number
|
||||
pages[element_page_number].append(_extract_final_element_info(element))
|
||||
extracted_pages = [
|
||||
{
|
||||
"number": page_number,
|
||||
"size": (
|
||||
_extract_final_element_page_size(page_elements[0]) if page_elements else None
|
||||
),
|
||||
"elements": page_elements,
|
||||
}
|
||||
for page_number, page_elements in pages.items()
|
||||
]
|
||||
self.layout = {"pages": sorted(extracted_pages, key=lambda x: x["number"])}
|
||||
|
||||
def dump(self) -> dict:
|
||||
return self.layout
|
||||
|
||||
|
||||
class JsonLayoutDumper(AnalysisProcessor):
|
||||
"""Dumps the results of the analysis to a JSON file"""
|
||||
|
||||
def __init__(self, filename: str, save_dir: str):
|
||||
self.dumpers = []
|
||||
super().__init__(filename, save_dir)
|
||||
|
||||
def add_layout_dumper(self, dumper: LayoutDumper):
|
||||
self.dumpers.append(dumper)
|
||||
|
||||
def process(self):
|
||||
filename_stem = Path(self.filename).stem
|
||||
analysis_save_dir = Path(self.save_dir) / "analysis" / filename_stem / "layout_dump"
|
||||
analysis_save_dir.mkdir(parents=True, exist_ok=True)
|
||||
for dumper in self.dumpers:
|
||||
results = dumper.dump()
|
||||
with open(analysis_save_dir / f"{dumper.layout_source}.json", "w") as f:
|
||||
f.write(json.dumps(results, indent=2))
|
||||
@@ -0,0 +1,18 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
|
||||
class AnalysisProcessor(ABC):
|
||||
def __init__(
|
||||
self,
|
||||
filename: Union[str, Path],
|
||||
save_dir: Union[str, Path],
|
||||
):
|
||||
self.filename = filename
|
||||
self.save_dir = save_dir
|
||||
|
||||
@abstractmethod
|
||||
def process(self):
|
||||
"""Performs the analysis and saves the results"""
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,190 @@
|
||||
import json
|
||||
import uuid
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from unstructured.partition.pdf_image.analysis.bbox_visualisation import (
|
||||
AnalysisDrawer,
|
||||
FinalLayoutDrawer,
|
||||
LayoutDrawer,
|
||||
OCRLayoutDrawer,
|
||||
ODModelLayoutDrawer,
|
||||
PdfminerLayoutDrawer,
|
||||
)
|
||||
from unstructured.partition.pdf_image.analysis.layout_dump import (
|
||||
ExtractedLayoutDumper,
|
||||
FinalLayoutDumper,
|
||||
JsonLayoutDumper,
|
||||
LayoutDumper,
|
||||
ObjectDetectionLayoutDumper,
|
||||
OCRLayoutDumper,
|
||||
)
|
||||
|
||||
|
||||
def _get_drawer_for_dumper(dumper: LayoutDumper) -> Optional[LayoutDrawer]:
|
||||
"""For a given layout dumper, return the corresponding layout drawer instance initialized with
|
||||
a dumped layout dict.
|
||||
|
||||
Args:
|
||||
dumper: The layout dumper instance
|
||||
|
||||
Returns:
|
||||
LayoutDrawer: The corresponding layout drawer instance
|
||||
"""
|
||||
if isinstance(dumper, ObjectDetectionLayoutDumper):
|
||||
return ODModelLayoutDrawer(layout_dump=dumper.dump())
|
||||
elif isinstance(dumper, ExtractedLayoutDumper):
|
||||
return PdfminerLayoutDrawer(layout_dump=dumper.dump())
|
||||
elif isinstance(dumper, OCRLayoutDumper):
|
||||
return OCRLayoutDrawer(layout_dump=dumper.dump())
|
||||
elif isinstance(dumper, FinalLayoutDumper):
|
||||
return FinalLayoutDrawer(layout_dump=dumper.dump())
|
||||
else:
|
||||
raise ValueError(f"Unknown dumper type: {dumper}")
|
||||
|
||||
|
||||
def _generate_filename(is_image: bool):
|
||||
"""Generate a filename for the analysis artifacts based on the file type.
|
||||
Adds a random uuid suffix
|
||||
"""
|
||||
suffix = uuid.uuid4().hex[:5]
|
||||
if is_image:
|
||||
return f"image_{suffix}.png"
|
||||
return f"pdf_{suffix}.pdf"
|
||||
|
||||
|
||||
def save_analysis_artifiacts(
|
||||
*layout_dumpers: LayoutDumper,
|
||||
is_image: bool,
|
||||
analyzed_image_output_dir_path: str,
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[BytesIO] = None,
|
||||
skip_bboxes: bool = False,
|
||||
skip_dump_od: bool = False,
|
||||
draw_grid: bool = False,
|
||||
draw_caption: bool = True,
|
||||
resize: Optional[float] = None,
|
||||
format: str = "png",
|
||||
):
|
||||
"""Save the analysis artifacts for a given file. Loads some settings from
|
||||
the environment configuration.
|
||||
|
||||
Args:
|
||||
layout_dumpers: The layout dumpers to save and use for bboxes rendering
|
||||
is_image: Flag for the file type (pdf/image)
|
||||
analyzed_image_output_dir_path: The directory to save the analysis artifacts
|
||||
filename: The filename of the sources analyzed file (pdf/image).
|
||||
Only one of filename or file should be provided.
|
||||
file: The file object for the analyzed file.
|
||||
Only one of filename or file should be provided.
|
||||
draw_grid: Flag for drawing the analysis bboxes on a single image (as grid)
|
||||
draw_caption: Flag for drawing the caption above the analyzed page (for e.g. layout source)
|
||||
resize: Output image resize value. If not provided, the image will not be resized.
|
||||
format: The format for analyzed pages with bboxes drawn on them. Default is 'png'.
|
||||
"""
|
||||
if not filename:
|
||||
filename = _generate_filename(is_image)
|
||||
if skip_bboxes or skip_dump_od:
|
||||
return
|
||||
|
||||
output_path = Path(analyzed_image_output_dir_path)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
if not skip_dump_od:
|
||||
json_layout_dumper = JsonLayoutDumper(
|
||||
filename=filename,
|
||||
save_dir=output_path,
|
||||
)
|
||||
for layout_dumper in layout_dumpers:
|
||||
json_layout_dumper.add_layout_dumper(layout_dumper)
|
||||
json_layout_dumper.process()
|
||||
|
||||
if not skip_bboxes:
|
||||
analysis_drawer = AnalysisDrawer(
|
||||
filename=filename,
|
||||
file=file,
|
||||
is_image=is_image,
|
||||
save_dir=output_path,
|
||||
draw_grid=draw_grid,
|
||||
draw_caption=draw_caption,
|
||||
resize=resize,
|
||||
format=format,
|
||||
)
|
||||
|
||||
for layout_dumper in layout_dumpers:
|
||||
drawer = _get_drawer_for_dumper(layout_dumper)
|
||||
analysis_drawer.add_drawer(drawer)
|
||||
analysis_drawer.process()
|
||||
|
||||
|
||||
def render_bboxes_for_file(
|
||||
filename: str,
|
||||
analyzed_image_output_dir_path: str,
|
||||
renders_output_dir_path: Optional[str] = None,
|
||||
draw_grid: bool = False,
|
||||
draw_caption: bool = True,
|
||||
resize: Optional[float] = None,
|
||||
format: str = "png",
|
||||
):
|
||||
"""Render the bounding boxes for a given layout dimp file.
|
||||
To be used for analysis after the partition is performed for
|
||||
only dumping the layouts - the bboxes can be rendered later.
|
||||
|
||||
Expects that the analyzed_image_output_dir_path keeps the structure
|
||||
that was created by the save_analysis_artifacts function.
|
||||
|
||||
Args:
|
||||
filename: The filename of the sources analyzed file (pdf/image)
|
||||
analyzed_image_output_dir_path: The directory where the analysis artifacts
|
||||
(layout dumps) are saved. It should be the root directory of the structure
|
||||
created by the save_analysis_artifacts function.
|
||||
renders_output_dir_path: Optional directory to save the rendered bboxes -
|
||||
if not provided, it will be saved in the analysis directory.
|
||||
draw_grid: Flag for drawing the analysis bboxes on a single image (as grid)
|
||||
draw_caption: Flag for drawing the caption above the analyzed page (for e.g. layout source)
|
||||
resize: Output image resize value. If not provided, the image will not be resized.
|
||||
format: The format for analyzed pages with bboxes drawn on them. Default is 'png'.
|
||||
"""
|
||||
filename_stem = Path(filename).stem
|
||||
is_image = not Path(filename).suffix.endswith("pdf")
|
||||
analysis_dumps_dir = (
|
||||
Path(analyzed_image_output_dir_path) / "analysis" / filename_stem / "layout_dump"
|
||||
)
|
||||
if not analysis_dumps_dir.exists():
|
||||
return
|
||||
layout_drawers = []
|
||||
for analysis_dump_filename in analysis_dumps_dir.iterdir():
|
||||
if not analysis_dump_filename.is_file():
|
||||
continue
|
||||
with open(analysis_dump_filename) as f:
|
||||
layout_dump = json.load(f)
|
||||
if analysis_dump_filename.stem == "final":
|
||||
layout_drawers.append(FinalLayoutDrawer(layout_dump=layout_dump))
|
||||
if analysis_dump_filename.stem == "object_detection":
|
||||
layout_drawers.append(ODModelLayoutDrawer(layout_dump=layout_dump))
|
||||
if analysis_dump_filename.stem == "ocr":
|
||||
layout_drawers.append(OCRLayoutDrawer(layout_dump=layout_dump))
|
||||
if analysis_dump_filename.stem == "pdfminer":
|
||||
layout_drawers.append(PdfminerLayoutDrawer(layout_dump=layout_dump))
|
||||
|
||||
if layout_drawers:
|
||||
if not renders_output_dir_path:
|
||||
output_path = (
|
||||
Path(analyzed_image_output_dir_path) / "analysis" / filename_stem / "bboxes"
|
||||
)
|
||||
else:
|
||||
output_path = Path(renders_output_dir_path)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
analysis_drawer = AnalysisDrawer(
|
||||
filename=filename,
|
||||
save_dir=output_path,
|
||||
is_image=is_image,
|
||||
draw_grid=draw_grid,
|
||||
draw_caption=draw_caption,
|
||||
resize=resize,
|
||||
format=format,
|
||||
)
|
||||
|
||||
for drawer in layout_drawers:
|
||||
analysis_drawer.add_drawer(drawer)
|
||||
analysis_drawer.process()
|
||||
@@ -0,0 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO
|
||||
|
||||
from unstructured.documents.elements import Element, FormKeysValues
|
||||
|
||||
|
||||
def run_form_extraction(
|
||||
filename: str,
|
||||
file: IO[bytes],
|
||||
model_name: str,
|
||||
elements: list[Element],
|
||||
skip_table_regions: bool,
|
||||
) -> list[FormKeysValues]:
|
||||
raise NotImplementedError("Form extraction not yet available.")
|
||||
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import numpy as np
|
||||
from unstructured_inference.constants import Source
|
||||
from unstructured_inference.inference.elements import TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layoutelement import (
|
||||
LayoutElement,
|
||||
LayoutElements,
|
||||
partition_groups_from_regions,
|
||||
)
|
||||
|
||||
from unstructured.documents.elements import ElementType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import Rectangle
|
||||
|
||||
|
||||
def build_text_region_from_coords(
|
||||
x1: int | float,
|
||||
y1: int | float,
|
||||
x2: int | float,
|
||||
y2: int | float,
|
||||
text: Optional[str] = None,
|
||||
source: Optional[Source] = None,
|
||||
) -> TextRegion:
|
||||
""""""
|
||||
return TextRegion.from_coords(x1, y1, x2, y2, text=text, source=source)
|
||||
|
||||
|
||||
def build_layout_element(
|
||||
bbox: "Rectangle",
|
||||
text: Optional[str] = None,
|
||||
source: Optional[Source] = None,
|
||||
element_type: Optional[str] = None,
|
||||
) -> LayoutElement:
|
||||
""""""
|
||||
|
||||
return LayoutElement(bbox=bbox, text=text, source=source, type=element_type)
|
||||
|
||||
|
||||
def build_layout_elements_from_ocr_regions(
|
||||
ocr_regions: TextRegions,
|
||||
ocr_text: Optional[str] = None,
|
||||
group_by_ocr_text: bool = False,
|
||||
) -> LayoutElements:
|
||||
"""
|
||||
Get layout elements from OCR regions
|
||||
"""
|
||||
|
||||
grouped_regions = []
|
||||
if group_by_ocr_text:
|
||||
text_sections = ocr_text.split("\n\n")
|
||||
mask = np.ones(ocr_regions.texts.shape).astype(bool)
|
||||
indices = np.arange(len(mask))
|
||||
for text_section in text_sections:
|
||||
regions = []
|
||||
words = text_section.replace("\n", " ").split()
|
||||
for i, text in enumerate(ocr_regions.texts[mask]):
|
||||
if not words:
|
||||
break
|
||||
if text in words:
|
||||
regions.append(indices[mask][i])
|
||||
words.remove(text)
|
||||
|
||||
if not regions:
|
||||
continue
|
||||
|
||||
mask[regions] = False
|
||||
grouped_regions.append(ocr_regions.slice(regions))
|
||||
else:
|
||||
grouped_regions = partition_groups_from_regions(ocr_regions)
|
||||
|
||||
merged_regions = TextRegions.from_list([merge_text_regions(group) for group in grouped_regions])
|
||||
return LayoutElements(
|
||||
element_coords=merged_regions.element_coords,
|
||||
texts=merged_regions.texts,
|
||||
sources=merged_regions.sources,
|
||||
element_class_ids=np.zeros(merged_regions.texts.shape),
|
||||
element_class_id_map={0: ElementType.UNCATEGORIZED_TEXT},
|
||||
)
|
||||
|
||||
|
||||
def merge_text_regions(regions: TextRegions) -> TextRegion:
|
||||
"""
|
||||
Merge a list of TextRegion objects into a single TextRegion.
|
||||
|
||||
Parameters:
|
||||
- group (TextRegions): A group of TextRegion objects to be merged.
|
||||
|
||||
Returns:
|
||||
- TextRegion: A single merged TextRegion object.
|
||||
"""
|
||||
|
||||
if not regions:
|
||||
raise ValueError("The text regions to be merged must be provided.")
|
||||
|
||||
min_x1 = regions.x1.min().astype(float)
|
||||
min_y1 = regions.y1.min().astype(float)
|
||||
max_x2 = regions.x2.max().astype(float)
|
||||
max_y2 = regions.y2.max().astype(float)
|
||||
|
||||
merged_text = " ".join([text for text in regions.texts if text])
|
||||
# assumption is the regions has the same source
|
||||
source = regions.sources[0]
|
||||
|
||||
return TextRegion.from_coords(min_x1, min_y1, max_x2, max_y2, merged_text, source)
|
||||
@@ -0,0 +1,492 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from typing import IO, TYPE_CHECKING, Any, List, Optional, cast
|
||||
|
||||
import numpy as np
|
||||
import pdf2image
|
||||
|
||||
# NOTE(yuming): Rename PIL.Image to avoid conflict with
|
||||
# unstructured.documents.elements.Image
|
||||
from PIL import Image as PILImage
|
||||
from PIL import ImageSequence
|
||||
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.metrics.table.table_formats import SimpleTableCell
|
||||
from unstructured.partition.common.lang import tesseract_to_paddle_language
|
||||
from unstructured.partition.pdf_image.analysis.layout_dump import OCRLayoutDumper
|
||||
from unstructured.partition.pdf_image.pdf_image_utils import valid_text
|
||||
from unstructured.partition.pdf_image.pdfminer_processing import (
|
||||
aggregate_embedded_text_by_block,
|
||||
bboxes1_is_almost_subregion_of_bboxes2,
|
||||
)
|
||||
from unstructured.partition.utils.config import env_config
|
||||
from unstructured.partition.utils.constants import OCR_AGENT_PADDLE, OCR_AGENT_TESSERACT, OCRMode
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
|
||||
from unstructured_inference.models.tables import UnstructuredTableTransformerModel
|
||||
|
||||
|
||||
def process_data_with_ocr(
|
||||
data: bytes | IO[bytes],
|
||||
out_layout: "DocumentLayout",
|
||||
extracted_layout: List[List["TextRegion"]],
|
||||
is_image: bool = False,
|
||||
infer_table_structure: bool = False,
|
||||
ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
ocr_languages: str = "eng",
|
||||
ocr_mode: str = OCRMode.FULL_PAGE.value,
|
||||
pdf_image_dpi: int = 200,
|
||||
ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
|
||||
password: Optional[str] = None,
|
||||
table_ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
) -> "DocumentLayout":
|
||||
"""
|
||||
Process OCR data from a given data and supplement the output DocumentLayout
|
||||
from unstructured_inference with ocr.
|
||||
|
||||
Parameters:
|
||||
- data (Union[bytes, BinaryIO]): The input file data,
|
||||
which can be either bytes or a BinaryIO object.
|
||||
|
||||
- out_layout (DocumentLayout): The output layout from unstructured-inference.
|
||||
|
||||
- is_image (bool, optional): Indicates if the input data is an image (True) or not (False).
|
||||
Defaults to False.
|
||||
|
||||
- infer_table_structure (bool, optional): If true, extract the table content.
|
||||
|
||||
- ocr_languages (str, optional): The languages for OCR processing. Defaults to "eng" (English).
|
||||
|
||||
- ocr_mode (str, optional): The OCR processing mode, e.g., "entire_page" or "individual_blocks".
|
||||
Defaults to "entire_page". If choose "entire_page" OCR, OCR processes the entire image
|
||||
page and will be merged with the output layout. If choose "individual_blocks" OCR,
|
||||
OCR is performed on individual elements by cropping the image.
|
||||
|
||||
- pdf_image_dpi (int, optional): DPI (dots per inch) for processing PDF images. Defaults to 200.
|
||||
|
||||
- ocr_layout_dumper (OCRLayoutDumper, optional): The OCR layout dumper to save the OCR layout.
|
||||
|
||||
Returns:
|
||||
DocumentLayout: The merged layout information obtained after OCR processing.
|
||||
"""
|
||||
data_bytes = data if isinstance(data, bytes) else data.read()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir_path:
|
||||
tmp_file_path = os.path.join(tmp_dir_path, "tmp_file")
|
||||
with open(tmp_file_path, "wb") as tmp_file:
|
||||
tmp_file.write(data_bytes)
|
||||
|
||||
merged_layouts = process_file_with_ocr(
|
||||
filename=tmp_file_path,
|
||||
out_layout=out_layout,
|
||||
extracted_layout=extracted_layout,
|
||||
is_image=is_image,
|
||||
infer_table_structure=infer_table_structure,
|
||||
ocr_agent=ocr_agent,
|
||||
ocr_languages=ocr_languages,
|
||||
ocr_mode=ocr_mode,
|
||||
pdf_image_dpi=pdf_image_dpi,
|
||||
ocr_layout_dumper=ocr_layout_dumper,
|
||||
password=password,
|
||||
table_ocr_agent=table_ocr_agent,
|
||||
)
|
||||
|
||||
return merged_layouts
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def process_file_with_ocr(
|
||||
filename: str,
|
||||
out_layout: "DocumentLayout",
|
||||
extracted_layout: List[TextRegions],
|
||||
is_image: bool = False,
|
||||
infer_table_structure: bool = False,
|
||||
ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
ocr_languages: str = "eng",
|
||||
ocr_mode: str = OCRMode.FULL_PAGE.value,
|
||||
pdf_image_dpi: int = 200,
|
||||
ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
|
||||
password: Optional[str] = None,
|
||||
table_ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
) -> "DocumentLayout":
|
||||
"""
|
||||
Process OCR data from a given file and supplement the output DocumentLayout
|
||||
from unstructured-inference with ocr.
|
||||
|
||||
Parameters:
|
||||
- filename (str): The path to the input file, which can be an image or a PDF.
|
||||
|
||||
- out_layout (DocumentLayout): The output layout from unstructured-inference.
|
||||
|
||||
- extracted_layout (List[TextRegions]): a list of text regions extracted by pdfminer, one for
|
||||
each page
|
||||
|
||||
- is_image (bool, optional): Indicates if the input data is an image (True) or not (False).
|
||||
Defaults to False.
|
||||
|
||||
- infer_table_structure (bool, optional): If true, extract the table content.
|
||||
|
||||
- ocr_languages (str, optional): The languages for OCR processing. Defaults to "eng" (English).
|
||||
|
||||
- ocr_mode (str, optional): The OCR processing mode, e.g., "entire_page" or "individual_blocks".
|
||||
Defaults to "entire_page". If choose "entire_page" OCR, OCR processes the entire image
|
||||
page and will be merged with the output layout. If choose "individual_blocks" OCR,
|
||||
OCR is performed on individual elements by cropping the image.
|
||||
|
||||
- pdf_image_dpi (int, optional): DPI (dots per inch) for processing PDF images. Defaults to 200.
|
||||
|
||||
Returns:
|
||||
DocumentLayout: The merged layout information obtained after OCR processing.
|
||||
"""
|
||||
|
||||
from unstructured_inference.inference.layout import DocumentLayout
|
||||
|
||||
merged_page_layouts: list[PageLayout] = []
|
||||
try:
|
||||
if is_image:
|
||||
with PILImage.open(filename) as images:
|
||||
image_format = images.format
|
||||
for i, image in enumerate(ImageSequence.Iterator(images)):
|
||||
image = image.convert("RGB")
|
||||
image.format = image_format
|
||||
extracted_regions = extracted_layout[i] if i < len(extracted_layout) else None
|
||||
merged_page_layout = supplement_page_layout_with_ocr(
|
||||
page_layout=out_layout.pages[i],
|
||||
image=image,
|
||||
infer_table_structure=infer_table_structure,
|
||||
ocr_agent=ocr_agent,
|
||||
ocr_languages=ocr_languages,
|
||||
ocr_mode=ocr_mode,
|
||||
extracted_regions=extracted_regions,
|
||||
ocr_layout_dumper=ocr_layout_dumper,
|
||||
table_ocr_agent=table_ocr_agent,
|
||||
)
|
||||
merged_page_layouts.append(merged_page_layout)
|
||||
return DocumentLayout.from_pages(merged_page_layouts)
|
||||
else:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
_image_paths = pdf2image.convert_from_path(
|
||||
filename,
|
||||
dpi=pdf_image_dpi,
|
||||
output_folder=temp_dir,
|
||||
paths_only=True,
|
||||
userpw=password or "",
|
||||
)
|
||||
image_paths = cast(List[str], _image_paths)
|
||||
for i, image_path in enumerate(image_paths):
|
||||
extracted_regions = extracted_layout[i] if i < len(extracted_layout) else None
|
||||
with PILImage.open(image_path) as image:
|
||||
merged_page_layout = supplement_page_layout_with_ocr(
|
||||
page_layout=out_layout.pages[i],
|
||||
image=image,
|
||||
infer_table_structure=infer_table_structure,
|
||||
ocr_agent=ocr_agent,
|
||||
ocr_languages=ocr_languages,
|
||||
ocr_mode=ocr_mode,
|
||||
extracted_regions=extracted_regions,
|
||||
ocr_layout_dumper=ocr_layout_dumper,
|
||||
table_ocr_agent=table_ocr_agent,
|
||||
)
|
||||
merged_page_layouts.append(merged_page_layout)
|
||||
return DocumentLayout.from_pages(merged_page_layouts)
|
||||
except Exception as e:
|
||||
if os.path.isdir(filename) or os.path.isfile(filename):
|
||||
raise e
|
||||
else:
|
||||
raise FileNotFoundError(f'File "{filename}" not found!') from e
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def supplement_page_layout_with_ocr(
|
||||
page_layout: "PageLayout",
|
||||
image: PILImage.Image,
|
||||
infer_table_structure: bool = False,
|
||||
ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
ocr_languages: str = "eng",
|
||||
ocr_mode: str = OCRMode.FULL_PAGE.value,
|
||||
extracted_regions: Optional[TextRegions] = None,
|
||||
ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
|
||||
table_ocr_agent: str = OCR_AGENT_TESSERACT,
|
||||
) -> "PageLayout":
|
||||
"""
|
||||
Supplement an PageLayout with OCR results depending on OCR mode.
|
||||
If mode is "entire_page", we get the OCR layout for the entire image and
|
||||
merge it with PageLayout.
|
||||
If mode is "individual_blocks", we find the elements from PageLayout
|
||||
with no text and add text from OCR to each element.
|
||||
"""
|
||||
|
||||
language = ocr_languages
|
||||
if ocr_agent == OCR_AGENT_PADDLE:
|
||||
language = tesseract_to_paddle_language(ocr_languages)
|
||||
_ocr_agent = OCRAgent.get_instance(ocr_agent_module=ocr_agent, language=language)
|
||||
if ocr_mode == OCRMode.FULL_PAGE.value:
|
||||
ocr_layout = _ocr_agent.get_layout_from_image(image)
|
||||
if ocr_layout_dumper:
|
||||
ocr_layout_dumper.add_ocred_page(ocr_layout.as_list())
|
||||
page_layout.elements_array = merge_out_layout_with_ocr_layout(
|
||||
out_layout=page_layout.elements_array,
|
||||
ocr_layout=ocr_layout,
|
||||
)
|
||||
elif ocr_mode == OCRMode.INDIVIDUAL_BLOCKS.value:
|
||||
# individual block mode still keeps using the list data structure for elements instead of
|
||||
# the vectorized page_layout.elements_array data structure
|
||||
for i, text in enumerate(page_layout.elements_array.texts):
|
||||
if text:
|
||||
continue
|
||||
padding = env_config.IMAGE_CROP_PAD
|
||||
cropped_image = image.crop(
|
||||
(
|
||||
page_layout.elements_array.x1[i] - padding,
|
||||
page_layout.elements_array.y1[i] - padding,
|
||||
page_layout.elements_array.x2[i] + padding,
|
||||
page_layout.elements_array.y2[i] + padding,
|
||||
),
|
||||
)
|
||||
# Note(yuming): instead of getting OCR layout, we just need
|
||||
# the text extraced from OCR for individual elements
|
||||
text_from_ocr = _ocr_agent.get_text_from_image(cropped_image)
|
||||
page_layout.elements_array.texts[i] = text_from_ocr
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid OCR mode. Parameter `ocr_mode` "
|
||||
"must be set to `entire_page` or `individual_blocks`.",
|
||||
)
|
||||
|
||||
# Note(yuming): use the OCR data from entire page OCR for table extraction
|
||||
if infer_table_structure:
|
||||
language = ocr_languages
|
||||
if table_ocr_agent == OCR_AGENT_PADDLE:
|
||||
language = tesseract_to_paddle_language(ocr_languages)
|
||||
_table_ocr_agent = OCRAgent.get_instance(
|
||||
ocr_agent_module=table_ocr_agent, language=language
|
||||
)
|
||||
from unstructured_inference.models import tables
|
||||
|
||||
tables.load_agent()
|
||||
if tables.tables_agent is None:
|
||||
raise RuntimeError("Unable to load table extraction agent.")
|
||||
|
||||
page_layout.elements_array = supplement_element_with_table_extraction(
|
||||
elements=page_layout.elements_array,
|
||||
image=image,
|
||||
tables_agent=tables.tables_agent,
|
||||
ocr_agent=_table_ocr_agent,
|
||||
extracted_regions=extracted_regions,
|
||||
)
|
||||
|
||||
return page_layout
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def supplement_element_with_table_extraction(
|
||||
elements: LayoutElements,
|
||||
image: PILImage.Image,
|
||||
tables_agent: "UnstructuredTableTransformerModel",
|
||||
ocr_agent,
|
||||
extracted_regions: Optional[TextRegions] = None,
|
||||
) -> List["LayoutElement"]:
|
||||
"""Supplement the existing layout with table extraction. Any Table elements
|
||||
that are extracted will have a metadata fields "text_as_html" where
|
||||
the table's text content is rendered into a html string and "table_as_cells"
|
||||
with the raw table cells output from table agent if env_config.EXTRACT_TABLE_AS_CELLS is True
|
||||
"""
|
||||
from unstructured_inference.models.tables import cells_to_html
|
||||
|
||||
table_id = {v: k for k, v in elements.element_class_id_map.items()}.get(ElementType.TABLE)
|
||||
if table_id is None:
|
||||
# no table found in this page
|
||||
return elements
|
||||
|
||||
table_ele_indices = np.where(elements.element_class_ids == table_id)[0]
|
||||
table_elements = elements.slice(table_ele_indices)
|
||||
padding = env_config.TABLE_IMAGE_CROP_PAD
|
||||
for i, element_coords in enumerate(table_elements.element_coords):
|
||||
cropped_image = image.crop(
|
||||
(
|
||||
element_coords[0] - padding,
|
||||
element_coords[1] - padding,
|
||||
element_coords[2] + padding,
|
||||
element_coords[3] + padding,
|
||||
),
|
||||
)
|
||||
table_tokens = get_table_tokens(
|
||||
table_element_image=cropped_image,
|
||||
ocr_agent=ocr_agent,
|
||||
)
|
||||
tatr_cells = tables_agent.predict(
|
||||
cropped_image, ocr_tokens=table_tokens, result_format="cells"
|
||||
)
|
||||
|
||||
# NOTE(christine): `tatr_cells == ""` means that the table was not recognized
|
||||
text_as_html = "" if tatr_cells == "" else cells_to_html(tatr_cells)
|
||||
elements.text_as_html[table_ele_indices[i]] = text_as_html
|
||||
|
||||
if env_config.EXTRACT_TABLE_AS_CELLS:
|
||||
simple_table_cells = [
|
||||
SimpleTableCell.from_table_transformer_cell(cell).to_dict() for cell in tatr_cells
|
||||
]
|
||||
elements.table_as_cells[table_ele_indices[i]] = simple_table_cells
|
||||
|
||||
return elements
|
||||
|
||||
|
||||
def get_table_tokens(
|
||||
table_element_image: PILImage.Image,
|
||||
ocr_agent: OCRAgent,
|
||||
) -> List[dict[str, Any]]:
|
||||
"""Get OCR tokens from either paddleocr or tesseract"""
|
||||
|
||||
ocr_layout = ocr_agent.get_layout_from_image(image=table_element_image)
|
||||
table_tokens = []
|
||||
for i, text in enumerate(ocr_layout.texts):
|
||||
table_tokens.append(
|
||||
{
|
||||
"bbox": [
|
||||
ocr_layout.x1[i],
|
||||
ocr_layout.y1[i],
|
||||
ocr_layout.x2[i],
|
||||
ocr_layout.y2[i],
|
||||
],
|
||||
"text": text,
|
||||
# 'table_tokens' is a list of tokens
|
||||
# Need to be in a relative reading order
|
||||
"span_num": i,
|
||||
"line_num": 0,
|
||||
"block_num": 0,
|
||||
}
|
||||
)
|
||||
|
||||
return table_tokens
|
||||
|
||||
|
||||
def merge_out_layout_with_ocr_layout(
|
||||
out_layout: LayoutElements,
|
||||
ocr_layout: TextRegions,
|
||||
supplement_with_ocr_elements: bool = True,
|
||||
subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
|
||||
) -> LayoutElements:
|
||||
"""
|
||||
Merge the out layout with the OCR-detected text regions on page level.
|
||||
|
||||
This function iterates over each out layout element and aggregates the associated text from
|
||||
the OCR layout using the specified threshold. The out layout's text attribute is then updated
|
||||
with this aggregated text. If `supplement_with_ocr_elements` is `True`, the out layout will be
|
||||
supplemented with the OCR layout.
|
||||
"""
|
||||
|
||||
if len(out_layout) == 0 or len(ocr_layout) == 0:
|
||||
# what if od model finds nothing but ocr finds something? should we use ocr output at all
|
||||
# currently we require some kind of bounding box, from `out_layout` to aggreaget ocr
|
||||
# results. Can we just use ocr bounding boxes (gonna be many but at least we save
|
||||
# information)
|
||||
return out_layout
|
||||
|
||||
invalid_text_indices = [i for i, text in enumerate(out_layout.texts) if not valid_text(text)]
|
||||
out_layout.texts = out_layout.texts.astype(object)
|
||||
|
||||
for idx in invalid_text_indices:
|
||||
out_layout.texts[idx], _ = aggregate_embedded_text_by_block(
|
||||
target_region=out_layout.slice([idx]),
|
||||
source_regions=ocr_layout,
|
||||
subregion_threshold=subregion_threshold,
|
||||
)
|
||||
|
||||
final_layout = (
|
||||
supplement_layout_with_ocr_elements(out_layout, ocr_layout)
|
||||
if supplement_with_ocr_elements
|
||||
else out_layout
|
||||
)
|
||||
|
||||
return final_layout
|
||||
|
||||
|
||||
def aggregate_ocr_text_by_block(
|
||||
ocr_layout: List["TextRegion"],
|
||||
region: "TextRegion",
|
||||
subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
|
||||
) -> Optional[str]:
|
||||
"""Extracts the text aggregated from the regions of the ocr layout that lie within the given
|
||||
block."""
|
||||
|
||||
extracted_texts = []
|
||||
|
||||
for ocr_region in ocr_layout:
|
||||
ocr_region_is_subregion_of_given_region = ocr_region.bbox.is_almost_subregion_of(
|
||||
region.bbox,
|
||||
subregion_threshold,
|
||||
)
|
||||
if ocr_region_is_subregion_of_given_region and ocr_region.text:
|
||||
extracted_texts.append(ocr_region.text)
|
||||
|
||||
return " ".join(extracted_texts) if extracted_texts else ""
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def supplement_layout_with_ocr_elements(
|
||||
layout: LayoutElements,
|
||||
ocr_layout: TextRegions,
|
||||
subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
|
||||
) -> LayoutElements:
|
||||
"""
|
||||
Supplement the existing layout with additional OCR-derived elements.
|
||||
|
||||
This function takes two lists: one list of pre-existing layout elements (`layout`)
|
||||
and another list of OCR-detected text regions (`ocr_layout`). It identifies OCR regions
|
||||
that are subregions of the elements in the existing layout and removes them from the
|
||||
OCR-derived list. Then, it appends the remaining OCR-derived regions to the existing layout.
|
||||
|
||||
Parameters:
|
||||
- layout (LayoutElements): A collection of existing layout elements in array structures
|
||||
- ocr_layout (TextRegions): A collection of OCR-derived text regions in array structures
|
||||
|
||||
Returns:
|
||||
- List[LayoutElement]: The final combined layout consisting of both the original layout
|
||||
elements and the new OCR-derived elements.
|
||||
|
||||
Note:
|
||||
- The function relies on `is_almost_subregion_of()` method to determine if an OCR region
|
||||
is a subregion of an existing layout element.
|
||||
- It also relies on `build_layout_elements_from_ocr_regions()` to convert OCR regions to
|
||||
layout elements.
|
||||
- The env_config `OCR_LAYOUT_SUBREGION_THRESHOLD` is used to specify the subregion matching
|
||||
threshold.
|
||||
"""
|
||||
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
from unstructured.partition.pdf_image.inference_utils import (
|
||||
build_layout_elements_from_ocr_regions,
|
||||
)
|
||||
|
||||
if len(layout) == 0:
|
||||
if len(ocr_layout) == 0:
|
||||
return layout
|
||||
else:
|
||||
ocr_regions_to_add = ocr_layout
|
||||
else:
|
||||
mask = (
|
||||
~bboxes1_is_almost_subregion_of_bboxes2(
|
||||
ocr_layout.element_coords, layout.element_coords, subregion_threshold
|
||||
)
|
||||
.sum(axis=1)
|
||||
.astype(bool)
|
||||
)
|
||||
|
||||
# add ocr regions that are not covered by layout
|
||||
ocr_regions_to_add = ocr_layout.slice(mask)
|
||||
|
||||
if len(ocr_regions_to_add):
|
||||
ocr_elements_to_add = build_layout_elements_from_ocr_regions(ocr_regions_to_add)
|
||||
final_layout = LayoutElements.concatenate([layout, ocr_elements_to_add])
|
||||
else:
|
||||
final_layout = layout
|
||||
|
||||
return final_layout
|
||||
@@ -0,0 +1,443 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import unicodedata
|
||||
from copy import deepcopy
|
||||
from io import BytesIO
|
||||
from pathlib import Path, PurePath
|
||||
from typing import IO, TYPE_CHECKING, BinaryIO, Iterator, List, Optional, Tuple, Union, cast
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pdf2image
|
||||
from PIL import Image
|
||||
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.common.common import convert_to_bytes, exactly_one
|
||||
from unstructured.partition.utils.config import env_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import TextRegion
|
||||
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
|
||||
|
||||
def write_image(image: Union[Image.Image, np.ndarray], output_image_path: str):
|
||||
"""
|
||||
Write an image to a specified file path, supporting both PIL Image and numpy ndarray formats.
|
||||
|
||||
Parameters:
|
||||
- image (Union[Image.Image, np.ndarray]): The image to be written, which can be in PIL Image
|
||||
format or a numpy ndarray format.
|
||||
- output_image_path (str): The path to which the image will be written.
|
||||
|
||||
Raises:
|
||||
- ValueError: If the provided image type is neither PIL Image nor numpy ndarray.
|
||||
|
||||
Returns:
|
||||
- None: The function writes the image to the specified path but does not return any value.
|
||||
"""
|
||||
|
||||
if isinstance(image, Image.Image):
|
||||
image.save(output_image_path)
|
||||
elif isinstance(image, np.ndarray):
|
||||
cv2.imwrite(output_image_path, image)
|
||||
else:
|
||||
raise ValueError("Unsupported Image Type")
|
||||
|
||||
|
||||
def convert_pdf_to_image(
|
||||
filename: str,
|
||||
file: Optional[Union[bytes, BinaryIO]] = None,
|
||||
dpi: int = 200,
|
||||
output_folder: Optional[Union[str, PurePath]] = None,
|
||||
path_only: bool = False,
|
||||
password: Optional[str] = None,
|
||||
) -> Union[List[Image.Image], List[str]]:
|
||||
"""Get the image renderings of the pdf pages using pdf2image"""
|
||||
|
||||
if path_only and not output_folder:
|
||||
raise ValueError("output_folder must be specified if path_only is true")
|
||||
|
||||
if file is not None:
|
||||
f_bytes = convert_to_bytes(file)
|
||||
images = pdf2image.convert_from_bytes(
|
||||
f_bytes,
|
||||
dpi=dpi,
|
||||
output_folder=output_folder,
|
||||
paths_only=path_only,
|
||||
userpw=password,
|
||||
)
|
||||
else:
|
||||
images = pdf2image.convert_from_path(
|
||||
filename,
|
||||
dpi=dpi,
|
||||
output_folder=output_folder,
|
||||
paths_only=path_only,
|
||||
)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def pad_element_bboxes(
|
||||
element: "LayoutElement",
|
||||
padding: Union[int, float],
|
||||
) -> "LayoutElement":
|
||||
"""Increases (or decreases, if padding is negative) the size of the bounding
|
||||
boxes of the element by extending the boundary outward (resp. inward)"""
|
||||
|
||||
out_element = deepcopy(element)
|
||||
out_element.bbox.x1 -= padding
|
||||
out_element.bbox.x2 += padding
|
||||
out_element.bbox.y1 -= padding
|
||||
out_element.bbox.y2 += padding
|
||||
|
||||
return out_element
|
||||
|
||||
|
||||
def pad_bbox(
|
||||
bbox: Tuple[float, float, float, float],
|
||||
padding: Tuple[Union[int, float], Union[int, float]],
|
||||
) -> Tuple[float, float, float, float]:
|
||||
"""Pads a bounding box (bbox) by a specified horizontal and vertical padding."""
|
||||
|
||||
x1, y1, x2, y2 = bbox
|
||||
h_padding, v_padding = padding
|
||||
x1 -= h_padding
|
||||
x2 += h_padding
|
||||
y1 -= v_padding
|
||||
y2 += v_padding
|
||||
|
||||
return x1, y1, x2, y2
|
||||
|
||||
|
||||
def save_elements(
|
||||
elements: List["Element"],
|
||||
starting_page_number: int,
|
||||
element_category_to_save: str,
|
||||
pdf_image_dpi: int,
|
||||
filename: str = "",
|
||||
file: bytes | IO[bytes] | None = None,
|
||||
is_image: bool = False,
|
||||
extract_image_block_to_payload: bool = False,
|
||||
output_dir_path: str | None = None,
|
||||
password: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Saves specific elements from a PDF as images either to a directory or embeds them in the
|
||||
element's payload.
|
||||
|
||||
This function processes a list of elements partitioned from a PDF file. For each element of
|
||||
a specified category, it extracts and saves the image. The images can either be saved to
|
||||
a specified directory or embedded into the element's payload as a base64-encoded string.
|
||||
"""
|
||||
|
||||
# Determine the output directory path
|
||||
if not extract_image_block_to_payload:
|
||||
output_dir_path = output_dir_path or (
|
||||
str(Path(env_config.GLOBAL_WORKING_PROCESS_DIR) / "figures")
|
||||
if env_config.GLOBAL_WORKING_DIR_ENABLED
|
||||
else str(Path.cwd() / "figures")
|
||||
)
|
||||
|
||||
os.makedirs(output_dir_path, exist_ok=True)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
if is_image:
|
||||
if file is None:
|
||||
image_paths = [filename]
|
||||
else:
|
||||
if isinstance(file, bytes):
|
||||
file_data = file
|
||||
else:
|
||||
file.seek(0)
|
||||
file_data = file.read()
|
||||
|
||||
tmp_file_path = os.path.join(temp_dir, "tmp_file")
|
||||
with open(tmp_file_path, "wb") as tmp_file:
|
||||
tmp_file.write(file_data)
|
||||
image_paths = [tmp_file_path]
|
||||
else:
|
||||
_image_paths = convert_pdf_to_image(
|
||||
filename,
|
||||
file,
|
||||
pdf_image_dpi,
|
||||
output_folder=temp_dir,
|
||||
path_only=True,
|
||||
password=password,
|
||||
)
|
||||
image_paths = cast(List[str], _image_paths)
|
||||
|
||||
figure_number = 0
|
||||
for el in elements:
|
||||
if el.category != element_category_to_save:
|
||||
continue
|
||||
|
||||
coordinates = el.metadata.coordinates
|
||||
if not coordinates or not coordinates.points:
|
||||
continue
|
||||
|
||||
points = coordinates.points
|
||||
x1, y1 = points[0]
|
||||
x2, y2 = points[2]
|
||||
h_padding = env_config.EXTRACT_IMAGE_BLOCK_CROP_HORIZONTAL_PAD
|
||||
v_padding = env_config.EXTRACT_IMAGE_BLOCK_CROP_VERTICAL_PAD
|
||||
padded_bbox = cast(
|
||||
Tuple[int, int, int, int], pad_bbox((x1, y1, x2, y2), (h_padding, v_padding))
|
||||
)
|
||||
|
||||
# The page number in the metadata may have been offset
|
||||
# by starting_page_number. Make sure we use the right
|
||||
# value for indexing!
|
||||
assert el.metadata.page_number
|
||||
metadata_page_number = el.metadata.page_number
|
||||
page_index = metadata_page_number - starting_page_number
|
||||
|
||||
figure_number += 1
|
||||
try:
|
||||
image_path = image_paths[page_index]
|
||||
image = Image.open(image_path)
|
||||
cropped_image = image.crop(padded_bbox)
|
||||
|
||||
# PNG images with transparency need to be converted before saving
|
||||
if cropped_image.mode == "RGBA":
|
||||
cropped_image = cropped_image.convert("RGB")
|
||||
|
||||
if extract_image_block_to_payload:
|
||||
buffered = BytesIO()
|
||||
cropped_image.save(buffered, format="JPEG")
|
||||
img_base64 = base64.b64encode(buffered.getvalue())
|
||||
img_base64_str = img_base64.decode()
|
||||
el.metadata.image_base64 = img_base64_str
|
||||
el.metadata.image_mime_type = "image/jpeg"
|
||||
else:
|
||||
basename = "table" if el.category == ElementType.TABLE else "figure"
|
||||
assert output_dir_path
|
||||
output_f_path = os.path.join(
|
||||
output_dir_path,
|
||||
f"{basename}-{metadata_page_number}-{figure_number}.jpg",
|
||||
)
|
||||
write_image(cropped_image, output_f_path)
|
||||
# add image path to element metadata
|
||||
el.metadata.image_path = output_f_path
|
||||
except (ValueError, IOError):
|
||||
logger.warning("Image Extraction Error: Skipping the failed image", exc_info=True)
|
||||
|
||||
|
||||
def check_element_types_to_extract(
|
||||
extract_image_block_types: Optional[List[str]],
|
||||
) -> List[str]:
|
||||
"""Check and normalize the provided list of element types to extract."""
|
||||
|
||||
if extract_image_block_types is None:
|
||||
return []
|
||||
|
||||
if not isinstance(extract_image_block_types, list):
|
||||
raise TypeError(
|
||||
"The extract_image_block_types parameter must be a list of element types as strings, "
|
||||
"ex. ['Table', 'Image']",
|
||||
)
|
||||
|
||||
available_element_types = {e_type.lower(): e_type for e_type in ElementType.to_dict().values()}
|
||||
normalized_extract_image_block_types = []
|
||||
for el_type in extract_image_block_types:
|
||||
normalized_el_type = available_element_types.get(
|
||||
el_type.lower(), el_type.lower().capitalize()
|
||||
)
|
||||
if normalized_el_type not in available_element_types.values():
|
||||
logger.warning(f"The requested type ({el_type}) doesn't match any available type")
|
||||
normalized_extract_image_block_types.append(normalized_el_type)
|
||||
|
||||
return normalized_extract_image_block_types
|
||||
|
||||
|
||||
def valid_text(text: str) -> bool:
|
||||
"""a helper that determines if the text is valid ascii text"""
|
||||
if not text:
|
||||
return False
|
||||
return "(cid:" not in text
|
||||
|
||||
|
||||
def cid_ratio(text: str) -> float:
|
||||
"""Gets ratio of unknown 'cid' characters extracted from text to all characters."""
|
||||
if not is_cid_present(text):
|
||||
return 0.0
|
||||
cid_pattern = r"\(cid\:(\d+)\)"
|
||||
unmatched, n_cid = re.subn(cid_pattern, "", text)
|
||||
total = n_cid + len(unmatched)
|
||||
return n_cid / total
|
||||
|
||||
|
||||
def is_cid_present(text: str) -> bool:
|
||||
"""Checks if a cid code is present in a text selection."""
|
||||
if len(text) < len("(cid:x)"):
|
||||
return False
|
||||
return text.find("(cid:") != -1
|
||||
|
||||
|
||||
def annotate_layout_elements_with_image(
|
||||
inferred_page_layout: "PageLayout",
|
||||
extracted_page_layout: Optional["PageLayout"],
|
||||
output_dir_path: str,
|
||||
output_f_basename: str,
|
||||
page_number: int,
|
||||
):
|
||||
"""
|
||||
Annotates a page image with both inferred and extracted layout elements.
|
||||
|
||||
This function takes the layout elements of a single page, either extracted from or inferred
|
||||
for the document, and annotates them on the page image. It creates two separate annotated
|
||||
images, one for each set of layout elements: 'inferred' and 'extracted'.
|
||||
These annotated images are saved to a specified directory.
|
||||
"""
|
||||
|
||||
layout_map = {"inferred": {"layout": inferred_page_layout, "color": "blue"}}
|
||||
if extracted_page_layout:
|
||||
layout_map["extracted"] = {"layout": extracted_page_layout, "color": "green"}
|
||||
|
||||
for label, layout_data in layout_map.items():
|
||||
page_layout = layout_data.get("layout")
|
||||
color = layout_data.get("color")
|
||||
|
||||
img = page_layout.annotate(colors=color)
|
||||
output_f_path = os.path.join(
|
||||
output_dir_path, f"{output_f_basename}_{page_number}_{label}.jpg"
|
||||
)
|
||||
write_image(img, output_f_path)
|
||||
print(f"output_image_path: {output_f_path}")
|
||||
|
||||
|
||||
def annotate_layout_elements(
|
||||
inferred_document_layout: "DocumentLayout",
|
||||
extracted_layout: List["TextRegion"],
|
||||
filename: str,
|
||||
output_dir_path: str,
|
||||
pdf_image_dpi: int,
|
||||
is_image: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Annotates layout elements on images extracted from a PDF or an image file.
|
||||
|
||||
This function processes a given document (PDF or image) and annotates layout elements based
|
||||
on the inferred and extracted layout information.
|
||||
It handles both PDF documents and standalone image files. For PDFs, it converts each page
|
||||
into an image, whereas for image files, it processes the single image.
|
||||
"""
|
||||
|
||||
from unstructured_inference.inference.layout import PageLayout
|
||||
|
||||
output_f_basename = os.path.splitext(os.path.basename(filename))[0]
|
||||
images = []
|
||||
try:
|
||||
if is_image:
|
||||
with Image.open(filename) as img:
|
||||
img = img.convert("RGB")
|
||||
images.append(img)
|
||||
|
||||
extracted_page_layout = None
|
||||
if extracted_layout:
|
||||
extracted_page_layout = PageLayout(
|
||||
number=1,
|
||||
image=img,
|
||||
)
|
||||
extracted_page_layout.elements = extracted_layout[0]
|
||||
|
||||
inferred_page_layout = inferred_document_layout.pages[0]
|
||||
inferred_page_layout.image = img
|
||||
|
||||
annotate_layout_elements_with_image(
|
||||
inferred_page_layout=inferred_document_layout.pages[0],
|
||||
extracted_page_layout=extracted_page_layout,
|
||||
output_dir_path=output_dir_path,
|
||||
output_f_basename=output_f_basename,
|
||||
page_number=1,
|
||||
)
|
||||
else:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
_image_paths = pdf2image.convert_from_path(
|
||||
filename,
|
||||
dpi=pdf_image_dpi,
|
||||
output_folder=temp_dir,
|
||||
paths_only=True,
|
||||
)
|
||||
image_paths = cast(List[str], _image_paths)
|
||||
for i, image_path in enumerate(image_paths):
|
||||
with Image.open(image_path) as img:
|
||||
page_number = i + 1
|
||||
|
||||
extracted_page_layout = None
|
||||
if extracted_layout:
|
||||
extracted_page_layout = PageLayout(
|
||||
number=page_number,
|
||||
image=img,
|
||||
)
|
||||
extracted_page_layout.elements = extracted_layout[i]
|
||||
|
||||
inferred_page_layout = inferred_document_layout.pages[i]
|
||||
inferred_page_layout.image = img
|
||||
|
||||
annotate_layout_elements_with_image(
|
||||
inferred_page_layout=inferred_document_layout.pages[i],
|
||||
extracted_page_layout=extracted_page_layout,
|
||||
output_dir_path=output_dir_path,
|
||||
output_f_basename=output_f_basename,
|
||||
page_number=page_number,
|
||||
)
|
||||
except Exception as e:
|
||||
if os.path.isdir(filename) or os.path.isfile(filename):
|
||||
raise e
|
||||
else:
|
||||
raise FileNotFoundError(f'File "{filename}" not found!') from e
|
||||
|
||||
|
||||
def convert_pdf_to_images(
|
||||
filename: str = "",
|
||||
file: Optional[bytes | IO[bytes]] = None,
|
||||
chunk_size: int = 10,
|
||||
password: Optional[str] = None,
|
||||
) -> Iterator[Image.Image]:
|
||||
# Convert a PDF in small chunks of pages at a time (e.g. 1-10, 11-20... and so on)
|
||||
exactly_one(filename=filename, file=file)
|
||||
if file is not None:
|
||||
f_bytes = convert_to_bytes(file)
|
||||
info = pdf2image.pdfinfo_from_bytes(f_bytes, userpw=password)
|
||||
else:
|
||||
f_bytes = None
|
||||
info = pdf2image.pdfinfo_from_path(filename, userpw=password)
|
||||
|
||||
total_pages = info["Pages"]
|
||||
for start_page in range(1, total_pages + 1, chunk_size):
|
||||
end_page = min(start_page + chunk_size - 1, total_pages)
|
||||
if f_bytes is not None:
|
||||
chunk_images = pdf2image.convert_from_bytes(
|
||||
f_bytes,
|
||||
first_page=start_page,
|
||||
last_page=end_page,
|
||||
userpw=password,
|
||||
)
|
||||
else:
|
||||
chunk_images = pdf2image.convert_from_path(
|
||||
filename,
|
||||
first_page=start_page,
|
||||
last_page=end_page,
|
||||
userpw=password,
|
||||
)
|
||||
|
||||
for image in chunk_images:
|
||||
yield image
|
||||
|
||||
|
||||
def remove_control_characters(text: str) -> str:
|
||||
"""Removes control characters from text."""
|
||||
|
||||
# Replace newline character with a space
|
||||
text = text.replace("\t", " ").replace("\n", " ")
|
||||
# Remove other control characters
|
||||
out_text = "".join(c for c in text if unicodedata.category(c)[0] != "C")
|
||||
return out_text
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,127 @@
|
||||
import os
|
||||
import tempfile
|
||||
from typing import BinaryIO, List, Optional, Tuple
|
||||
|
||||
from pdfminer.converter import PDFPageAggregator
|
||||
from pdfminer.layout import LAParams, LTContainer, LTImage, LTItem, LTTextLine
|
||||
from pdfminer.pdfinterp import PDFPageInterpreter, PDFResourceManager
|
||||
from pdfminer.pdfpage import PDFPage
|
||||
from pdfminer.psexceptions import PSSyntaxError
|
||||
from pydantic import BaseModel
|
||||
|
||||
from unstructured.logger import logger
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
|
||||
class PDFMinerConfig(BaseModel):
|
||||
line_overlap: Optional[float] = None
|
||||
word_margin: Optional[float] = None
|
||||
line_margin: Optional[float] = None
|
||||
char_margin: Optional[float] = None
|
||||
|
||||
|
||||
def init_pdfminer(pdfminer_config: Optional[PDFMinerConfig] = None):
|
||||
rsrcmgr = PDFResourceManager()
|
||||
|
||||
laparams_kwargs = pdfminer_config.model_dump(exclude_none=True) if pdfminer_config else {}
|
||||
laparams = LAParams(**laparams_kwargs)
|
||||
|
||||
device = PDFPageAggregator(rsrcmgr, laparams=laparams)
|
||||
interpreter = PDFPageInterpreter(rsrcmgr, device)
|
||||
|
||||
return device, interpreter
|
||||
|
||||
|
||||
def extract_image_objects(parent_object: LTItem) -> List[LTImage]:
|
||||
"""Recursively extracts image objects from a given parent object in a PDF document."""
|
||||
objects = []
|
||||
|
||||
if isinstance(parent_object, LTImage):
|
||||
objects.append(parent_object)
|
||||
elif isinstance(parent_object, LTContainer):
|
||||
for child in parent_object:
|
||||
objects.extend(extract_image_objects(child))
|
||||
|
||||
return objects
|
||||
|
||||
|
||||
def extract_text_objects(parent_object: LTItem) -> List[LTTextLine]:
|
||||
"""Recursively extracts text objects from a given parent object in a PDF document."""
|
||||
objects = []
|
||||
|
||||
if isinstance(parent_object, LTTextLine):
|
||||
objects.append(parent_object)
|
||||
elif isinstance(parent_object, LTContainer):
|
||||
for child in parent_object:
|
||||
objects.extend(extract_text_objects(child))
|
||||
|
||||
return objects
|
||||
|
||||
|
||||
def rect_to_bbox(
|
||||
rect: Tuple[float, float, float, float],
|
||||
height: float,
|
||||
) -> Tuple[float, float, float, float]:
|
||||
"""
|
||||
Converts a PDF rectangle coordinates (x1, y1, x2, y2) to a bounding box in the specified
|
||||
coordinate system where the vertical axis is measured from the top of the page.
|
||||
|
||||
Args:
|
||||
rect (Tuple[float, float, float, float]): A tuple representing a PDF rectangle
|
||||
coordinates (x1, y1, x2, y2).
|
||||
height (float): The height of the page in the specified coordinate system.
|
||||
|
||||
Returns:
|
||||
Tuple[float, float, float, float]: A tuple representing the bounding box coordinates
|
||||
(x1, y1, x2, y2) with the y-coordinates adjusted to be measured from the top of the page.
|
||||
"""
|
||||
x1, y2, x2, y1 = rect
|
||||
y1 = height - y1
|
||||
y2 = height - y2
|
||||
return (x1, y1, x2, y2)
|
||||
|
||||
|
||||
@requires_dependencies(["pikepdf", "pypdf"])
|
||||
def open_pdfminer_pages_generator(
|
||||
fp: BinaryIO, password: Optional[str] = None, pdfminer_config: Optional[PDFMinerConfig] = None
|
||||
):
|
||||
"""Open PDF pages using PDFMiner, handling and repairing invalid dictionary constructs."""
|
||||
|
||||
import pikepdf
|
||||
|
||||
from unstructured.partition.pdf_image.pypdf_utils import get_page_data
|
||||
|
||||
device, interpreter = init_pdfminer(pdfminer_config=pdfminer_config)
|
||||
with tempfile.TemporaryDirectory() as tmp_dir_path:
|
||||
tmp_file_path = os.path.join(tmp_dir_path, "tmp_file")
|
||||
try:
|
||||
pages = PDFPage.get_pages(fp, password=password or "")
|
||||
# Detect invalid dictionary construct for entire PDF
|
||||
for i, page in enumerate(pages):
|
||||
try:
|
||||
# Detect invalid dictionary construct for one page
|
||||
interpreter.process_page(page)
|
||||
page_layout = device.get_result()
|
||||
except PSSyntaxError:
|
||||
logger.info("Detected invalid dictionary construct for PDFminer")
|
||||
logger.info(f"Repairing the PDF page {i + 1} ...")
|
||||
# find the error page from binary data fp
|
||||
error_page_data = get_page_data(fp, page_number=i)
|
||||
# repair the error page with pikepdf
|
||||
with pikepdf.Pdf.open(error_page_data) as pdf:
|
||||
pdf.save(tmp_file_path)
|
||||
page = next(PDFPage.get_pages(open(tmp_file_path, "rb"))) # noqa: SIM115
|
||||
interpreter.process_page(page)
|
||||
page_layout = device.get_result()
|
||||
yield page, page_layout
|
||||
except PSSyntaxError:
|
||||
logger.info("Detected invalid dictionary construct for PDFminer")
|
||||
logger.info("Repairing the PDF document ...")
|
||||
# repair the entire doc with pikepdf
|
||||
with pikepdf.Pdf.open(fp) as pdf:
|
||||
pdf.save(tmp_file_path)
|
||||
pages = PDFPage.get_pages(open(tmp_file_path, "rb")) # noqa: SIM115
|
||||
for page in pages:
|
||||
interpreter.process_page(page)
|
||||
page_layout = device.get_result()
|
||||
yield page, page_layout
|
||||
@@ -0,0 +1,15 @@
|
||||
import io
|
||||
from typing import BinaryIO
|
||||
|
||||
import pypdf
|
||||
|
||||
|
||||
def get_page_data(fp: BinaryIO, page_number: int):
|
||||
"""Find the binary data for a given page number from a PDF binary file."""
|
||||
pdf_reader = pypdf.PdfReader(fp)
|
||||
pdf_writer = pypdf.PdfWriter()
|
||||
page = pdf_reader.pages[page_number]
|
||||
pdf_writer.add_page(page)
|
||||
page_data = io.BytesIO()
|
||||
pdf_writer.write(page_data)
|
||||
return page_data
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import convert_office_doc, exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.pptx import partition_pptx
|
||||
|
||||
|
||||
def partition_ppt(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions Microsoft PowerPoint Documents in .ppt format into their document elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
|
||||
Note that all arguments valid on `partition_pptx()` are also valid here and will be passed
|
||||
along to the `partition_pptx()` function.
|
||||
"""
|
||||
# -- Verify that only one of the arguments was provided
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
if filename:
|
||||
# -- Verify filename.
|
||||
if not os.path.exists(filename):
|
||||
raise ValueError(f"The file {filename} does not exist.")
|
||||
|
||||
else:
|
||||
assert file
|
||||
# -- Create filename.
|
||||
tmp_file_path = os.path.join(tmpdir, "tmp_file")
|
||||
with open(tmp_file_path, "wb") as tmp_file:
|
||||
tmp_file.write(file.read())
|
||||
filename = tmp_file_path
|
||||
|
||||
_, filename_no_path = os.path.split(os.path.abspath(filename))
|
||||
base_filename, _ = os.path.splitext(filename_no_path)
|
||||
|
||||
convert_office_doc(
|
||||
filename,
|
||||
tmpdir,
|
||||
target_format="pptx",
|
||||
target_filter="Impress MS PowerPoint 2007 XML",
|
||||
)
|
||||
pptx_filename = os.path.join(tmpdir, f"{base_filename}.pptx")
|
||||
|
||||
elements = partition_pptx(
|
||||
filename=pptx_filename,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.PPT,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# -- Remove tmp.name from filename if parsing file
|
||||
if file:
|
||||
for element in elements:
|
||||
element.metadata.filename = metadata_filename
|
||||
|
||||
return elements
|
||||
@@ -0,0 +1,511 @@
|
||||
"""Partitioner for PPTX documents.
|
||||
|
||||
PPTX files are PowerPoint 2007+ documents. These are XML-based and "open" (documented ISO standard),
|
||||
unlike the `.ppt` format which was binary and proprietary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from tempfile import SpooledTemporaryFile
|
||||
from typing import IO, Any, Iterator, Protocol, Sequence
|
||||
|
||||
import pptx
|
||||
from pptx.presentation import Presentation
|
||||
from pptx.shapes.autoshape import Shape
|
||||
from pptx.shapes.base import BaseShape
|
||||
from pptx.shapes.graphfrm import GraphicFrame
|
||||
from pptx.shapes.group import GroupShape
|
||||
from pptx.shapes.picture import Picture
|
||||
from pptx.shapes.shapetree import _BaseGroupShapes # pyright: ignore [reportPrivateUsage]
|
||||
from pptx.slide import Slide
|
||||
from pptx.text.text import _Paragraph # pyright: ignore [reportPrivateUsage]
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.common.html_table import HtmlTable, htmlify_matrix_of_cell_texts
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
ElementMetadata,
|
||||
EmailAddress,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
PageBreak,
|
||||
Table,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
from unstructured.partition.text_type import (
|
||||
is_email_address,
|
||||
is_possible_narrative_text,
|
||||
is_possible_title,
|
||||
)
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
from unstructured.utils import is_temp_file_path, lazyproperty
|
||||
|
||||
DETECTION_ORIGIN = "pptx"
|
||||
|
||||
|
||||
def register_picture_partitioner(picture_partitioner: AbstractPicturePartitioner) -> None:
|
||||
"""Specify a pluggable sub-partitioner to be used for partitioning PPTX images."""
|
||||
PptxPartitionerOptions.register_picture_partitioner(picture_partitioner)
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# PPTX DOMAIN MODEL DEFINITIONS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class AbstractPicturePartitioner(Protocol):
|
||||
"""Defines the interface for a pluggable sub-partitioner for PPTX Picture objects.
|
||||
|
||||
A PPTX Picture object generally contains an image (e.g. JPG, PNG) but can also contain other
|
||||
media types like a video or sound file. The interface classmethod generates zero-or-more
|
||||
elements from the specified Picture object. If the media in the picture object is not supported
|
||||
then it will silently return without generating any elements.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, picture: Picture, opts: PptxPartitionerOptions) -> Iterator[Element]:
|
||||
"""Generate document elements derived from `picture`, a PPTX Picture shape."""
|
||||
...
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# PARTITIONER
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@apply_metadata(FileType.PPTX)
|
||||
@add_chunking_strategy
|
||||
def partition_pptx(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
include_page_breaks: bool = True,
|
||||
include_slide_notes: bool | None = None,
|
||||
infer_table_structure: bool = True,
|
||||
starting_page_number: int = 1,
|
||||
strategy: str = PartitionStrategy.FAST,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partition PowerPoint document in .pptx format into its document elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
include_page_breaks
|
||||
If True, includes a PageBreak element between slides
|
||||
include_slide_notes
|
||||
If True, includes the slide notes as element
|
||||
infer_table_structure
|
||||
If True, any Table elements that are extracted will also have a metadata field
|
||||
named "text_as_html" where the table's text content is rendered into an html string.
|
||||
I.e., rows and cells are preserved.
|
||||
Whether True or False, the "text" field is always present in any Table element
|
||||
and is the text content of the table (no structure).
|
||||
starting_page_number
|
||||
Indicates what page number should be assigned to the first slide in the presentation.
|
||||
This information will be reflected in elements' metadata and can be be especially
|
||||
useful when partitioning a document that is part of a larger document.
|
||||
"""
|
||||
opts = PptxPartitionerOptions(
|
||||
file=file,
|
||||
file_path=filename,
|
||||
include_page_breaks=include_page_breaks,
|
||||
include_slide_notes=include_slide_notes,
|
||||
infer_table_structure=infer_table_structure,
|
||||
strategy=strategy,
|
||||
starting_page_number=starting_page_number,
|
||||
)
|
||||
|
||||
return list(_PptxPartitioner.iter_presentation_elements(opts))
|
||||
|
||||
|
||||
class _PptxPartitioner:
|
||||
"""Provides `.partition()` for PowerPoint 2007+ (.pptx) files."""
|
||||
|
||||
def __init__(self, opts: PptxPartitionerOptions):
|
||||
self._opts = opts
|
||||
|
||||
@classmethod
|
||||
def iter_presentation_elements(cls, opts: PptxPartitionerOptions) -> Iterator[Element]:
|
||||
"""Partition MS Word documents (.docx format) into its document elements."""
|
||||
return cls(opts)._iter_presentation_elements()
|
||||
|
||||
def _iter_presentation_elements(self) -> Iterator[Element]:
|
||||
"""Generate each document-element in presentation in document order."""
|
||||
# -- This implementation composes a collection of iterators into a "combined" iterator
|
||||
# -- return value using `yield from`. You can think of the return value as an Element
|
||||
# -- stream and each `yield from` as "add elements found by this function to the stream".
|
||||
# -- This is functionally analogous to declaring `elements: List[Element] = []` at the top
|
||||
# -- and using `elements.extend()` for the results of each of the function calls, but is
|
||||
# -- more perfomant, uses less memory (avoids producing and then garbage-collecting all
|
||||
# -- those small lists), is more flexible for later iterator operations like filter,
|
||||
# -- chain, map, etc. and is perhaps more elegant and simpler to read once you have the
|
||||
# -- concept of what it's doing. You can see the same pattern repeating in the "sub"
|
||||
# -- functions like `._iter_shape_elements()` where the "just return when done"
|
||||
# -- characteristic of a generator avoids repeated code to form interim results into lists.
|
||||
|
||||
for slide in self._presentation.slides:
|
||||
yield from self._opts.increment_page_number()
|
||||
yield from self._iter_maybe_slide_notes(slide)
|
||||
|
||||
title_shape, shapes = self._order_shapes(slide)
|
||||
|
||||
for shape in shapes:
|
||||
if shape.has_table:
|
||||
assert isinstance(shape, GraphicFrame)
|
||||
yield from self._iter_table_element(shape)
|
||||
elif shape.has_text_frame:
|
||||
assert isinstance(shape, Shape)
|
||||
if shape == title_shape:
|
||||
yield from self._iter_title_shape_element(shape)
|
||||
else:
|
||||
yield from self._iter_shape_elements(shape)
|
||||
elif isinstance(shape, Picture):
|
||||
yield from self._iter_picture_elements(shape)
|
||||
|
||||
# -- otherwise ditch it, this would include charts, connectors (lines),
|
||||
# -- and free-form shapes (squiggly lines). Lines don't have text.
|
||||
|
||||
def _is_bulleted_paragraph(self, paragraph: _Paragraph) -> bool:
|
||||
"""True when `paragraph` has a bullet-charcter prefix.
|
||||
|
||||
Bullet characters in the openxml schema are represented by buChar.
|
||||
"""
|
||||
# -- True when XPath returns a non-empty list (nodeset) --
|
||||
return bool(paragraph._p.xpath("./a:pPr/a:buChar"))
|
||||
|
||||
def _iter_maybe_slide_notes(self, slide: Slide) -> Iterator[NarrativeText]:
|
||||
"""Generate zero-or-one NarrativeText element for the slide-notes."""
|
||||
# -- only emit slide-notes elements when enabled --
|
||||
if not self._opts.include_slide_notes:
|
||||
return
|
||||
|
||||
# -- not all slides have a notes slide --
|
||||
if not slide.has_notes_slide:
|
||||
return
|
||||
|
||||
notes_slide = slide.notes_slide
|
||||
notes_text_frame = notes_slide.notes_text_frame
|
||||
|
||||
# -- not all notes slides have a text-frame (it's created on first use) --
|
||||
if not notes_text_frame:
|
||||
return
|
||||
notes_text = notes_text_frame.text.strip()
|
||||
|
||||
# -- not all notes text-frams contain text (if it's all deleted the text-frame remains) --
|
||||
if not notes_text:
|
||||
return
|
||||
|
||||
yield NarrativeText(
|
||||
text=notes_text, metadata=self._opts.text_metadata(), detection_origin=DETECTION_ORIGIN
|
||||
)
|
||||
|
||||
def _iter_picture_elements(self, picture: Picture) -> Iterator[Element]:
|
||||
"""Generate elements derived from the image in `picture`."""
|
||||
# -- delegate this job to the pluggable Picture partitioner --
|
||||
PicturePartitionerCls = self._opts.picture_partitioner
|
||||
yield from PicturePartitionerCls.iter_elements(picture, self._opts)
|
||||
|
||||
def _iter_shape_elements(self, shape: Shape) -> Iterator[Element]:
|
||||
"""Generate Text or subtype element for each paragraph in `shape`."""
|
||||
if self._shape_is_off_slide(shape):
|
||||
return
|
||||
|
||||
for paragraph in shape.text_frame.paragraphs:
|
||||
text = paragraph.text
|
||||
if text.strip() == "":
|
||||
continue
|
||||
|
||||
level = paragraph.level or 0
|
||||
metadata = self._opts.text_metadata(category_depth=level)
|
||||
|
||||
if self._is_bulleted_paragraph(paragraph):
|
||||
yield ListItem(text=text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
elif is_email_address(text):
|
||||
yield EmailAddress(text=text, detection_origin=DETECTION_ORIGIN)
|
||||
elif is_possible_narrative_text(text):
|
||||
yield NarrativeText(
|
||||
text=text,
|
||||
metadata=metadata,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
)
|
||||
elif is_possible_title(text):
|
||||
# If text is a title but not the title shape increment the category depth)
|
||||
metadata = self._opts.text_metadata(category_depth=level + 1)
|
||||
yield Title(text=text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
else:
|
||||
yield Text(text=text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
|
||||
def _iter_table_element(self, graphfrm: GraphicFrame) -> Iterator[Table]:
|
||||
"""Generate zero-or-one Table element for the table in `shape`.
|
||||
|
||||
An empty table does not produce an element.
|
||||
"""
|
||||
if not (rows := list(graphfrm.table.rows)):
|
||||
return
|
||||
|
||||
html_text = htmlify_matrix_of_cell_texts(
|
||||
[[cell.text for cell in row.cells] for row in rows]
|
||||
)
|
||||
html_table = HtmlTable.from_html_text(html_text)
|
||||
|
||||
if not html_table.text:
|
||||
return
|
||||
|
||||
metadata = self._opts.table_metadata(
|
||||
html_table.html if self._opts.infer_table_structure else None
|
||||
)
|
||||
|
||||
yield Table(text=html_table.text, metadata=metadata, detection_origin=DETECTION_ORIGIN)
|
||||
|
||||
def _iter_title_shape_element(self, shape: Shape) -> Iterator[Element]:
|
||||
"""Generate Title element for each paragraph in title `shape`.
|
||||
|
||||
Text is most likely a title, but in the rare case that the title shape was used
|
||||
for the slide body text, also check for bulleted paragraphs."""
|
||||
if self._shape_is_off_slide(shape):
|
||||
return
|
||||
|
||||
depth = 0
|
||||
for paragraph in shape.text_frame.paragraphs:
|
||||
text = paragraph.text
|
||||
if text.strip() == "":
|
||||
continue
|
||||
|
||||
if self._is_bulleted_paragraph(paragraph):
|
||||
bullet_depth = paragraph.level or 0
|
||||
yield ListItem(
|
||||
text=text,
|
||||
metadata=self._opts.text_metadata(category_depth=bullet_depth),
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
)
|
||||
elif is_email_address(text):
|
||||
yield EmailAddress(text=text, detection_origin=DETECTION_ORIGIN)
|
||||
else:
|
||||
# increment the category depth by the paragraph increment in the shape
|
||||
yield Title(
|
||||
text=text,
|
||||
metadata=self._opts.text_metadata(category_depth=depth),
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
)
|
||||
depth += 1 # Cannot enumerate because we want to skip empty paragraphs
|
||||
|
||||
def _order_shapes(self, slide: Slide) -> tuple[Shape | None, Sequence[BaseShape]]:
|
||||
"""Orders the shapes on `slide` from top to bottom and left to right.
|
||||
|
||||
Returns the title shape if it exists and the ordered shapes."""
|
||||
|
||||
def iter_shapes(shapes: _BaseGroupShapes) -> Iterator[BaseShape]:
|
||||
for shape in shapes:
|
||||
if isinstance(shape, GroupShape):
|
||||
yield from iter_shapes(shape.shapes)
|
||||
else:
|
||||
yield shape
|
||||
|
||||
def sort_key(shape: BaseShape) -> tuple[int, int]:
|
||||
return shape.top or 0, shape.left or 0
|
||||
|
||||
return slide.shapes.title, sorted(iter_shapes(slide.shapes), key=sort_key)
|
||||
|
||||
@lazyproperty
|
||||
def _presentation(self) -> Presentation:
|
||||
"""The python-pptx `Presentation` object loaded from the provided source file."""
|
||||
return pptx.Presentation(self._opts.pptx_file)
|
||||
|
||||
def _shape_is_off_slide(self, shape: Shape) -> bool:
|
||||
# NOTE(robinson) - avoid processing shapes that are not on the actual slide
|
||||
# NOTE - skip check if no top or left position (shape displayed top left)
|
||||
return bool((shape.top and shape.left) and (shape.top < 0 or shape.left < 0))
|
||||
|
||||
|
||||
class PptxPartitionerOptions:
|
||||
"""Encapsulates partitioning option validation, computation, and application of defaults."""
|
||||
|
||||
_PicturePartitionerCls = None
|
||||
"""Sub-partitioner used to partition PPTX Picture (Image) shapes.
|
||||
|
||||
This value has module lifetime and is updated by calling the `register_picture_partitioner()`
|
||||
function defined in this module. The value sent to `register_picture_partitioner()` must be a
|
||||
pluggable sub-partitioner implementing the `AbstractPicturePartitioner` interface. After
|
||||
registration, all picture shapes in subsequent PPTX documents will be partitioned by the
|
||||
specified picture sub-partitioner.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
file: IO[bytes] | None,
|
||||
file_path: str | None,
|
||||
include_page_breaks: bool,
|
||||
include_slide_notes: bool | None,
|
||||
infer_table_structure: bool,
|
||||
strategy: str,
|
||||
starting_page_number: int = 1,
|
||||
):
|
||||
self._file = file
|
||||
self._file_path = file_path
|
||||
self._include_page_breaks = include_page_breaks
|
||||
self._include_slide_notes = include_slide_notes
|
||||
self._infer_table_structure = infer_table_structure
|
||||
self._strategy = strategy
|
||||
# -- options object maintains page-number state --
|
||||
self._page_counter = starting_page_number - 1
|
||||
|
||||
@classmethod
|
||||
def register_picture_partitioner(cls, picture_partitioner: AbstractPicturePartitioner):
|
||||
"""Specify a pluggable sub-partitioner to be used for partitioning PPTX images."""
|
||||
cls._PicturePartitionerCls = picture_partitioner
|
||||
|
||||
@lazyproperty
|
||||
def include_page_breaks(self) -> bool:
|
||||
"""When True, include `PageBreak` elements in element-stream.
|
||||
|
||||
Note that regardless of this setting, page-breaks are detected, and page-number is tracked
|
||||
and included in element metadata. Only the presence of distinct `PageBreak` elements (which
|
||||
contain no text) in the element stream is affected.
|
||||
"""
|
||||
return self._include_page_breaks
|
||||
|
||||
@lazyproperty
|
||||
def include_slide_notes(self) -> bool:
|
||||
"""When True, also partition any text found in slide notes as part of each slide."""
|
||||
return False if self._include_slide_notes is None else self._include_slide_notes
|
||||
|
||||
def increment_page_number(self) -> Iterator[PageBreak]:
|
||||
"""Increment page-number by 1 and generate a PageBreak element if enabled."""
|
||||
self._page_counter += 1
|
||||
# -- no page-break before first page --
|
||||
if self._page_counter < 2:
|
||||
return
|
||||
# -- only emit page-breaks when enabled --
|
||||
if self._include_page_breaks:
|
||||
yield PageBreak(
|
||||
"",
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
metadata=ElementMetadata(
|
||||
last_modified=self.last_modified, page_number=self.page_number - 1
|
||||
),
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def infer_table_structure(self) -> bool:
|
||||
"""True when partitioner should compute and apply `text_as_html` metadata for tables."""
|
||||
return self._infer_table_structure
|
||||
|
||||
@lazyproperty
|
||||
def last_modified(self) -> str | None:
|
||||
"""The best last-modified date available, None if no sources are available."""
|
||||
if not self._file_path:
|
||||
return None
|
||||
|
||||
return (
|
||||
None if is_temp_file_path(self._file_path) else get_last_modified_date(self._file_path)
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def metadata_file_path(self) -> str | None:
|
||||
"""The best available file-path for this document or `None` if unavailable."""
|
||||
return self._file_path
|
||||
|
||||
@property
|
||||
def page_number(self) -> int:
|
||||
"""The current page (slide) number."""
|
||||
return self._page_counter
|
||||
|
||||
@lazyproperty
|
||||
def picture_partitioner(self) -> AbstractPicturePartitioner:
|
||||
"""The sub-partitioner to use for PPTX Picture shapes."""
|
||||
# -- Note this value has partitioning-run scope. An instance of this options class is
|
||||
# -- instantiated once per partitioning run (each document can have different options).
|
||||
# -- Because this is a lazyproperty, it is computed only on the first reference. All
|
||||
# -- subsequent references during the same partitioning run will get the same value. This
|
||||
# -- ensures Picture shapes are processed consistently within a single document. The
|
||||
# -- intended use of `register_picture_partitioner()` is that it be called before processing
|
||||
# -- any documents, however there's no reason not to make the mechanism robust against
|
||||
# -- unintended use.
|
||||
return (
|
||||
_NullPicturePartitioner
|
||||
if self._PicturePartitionerCls is None
|
||||
else self._PicturePartitionerCls
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def pptx_file(self) -> str | IO[bytes]:
|
||||
"""The PowerPoint document file to be partitioned.
|
||||
|
||||
This is either a str path or a file-like object. `python-pptx` accepts either for opening a
|
||||
presentation file.
|
||||
"""
|
||||
if self._file_path:
|
||||
return self._file_path
|
||||
|
||||
# -- In Python <3.11 SpooledTemporaryFile does not implement ".seekable" which triggers an
|
||||
# -- exception when Zipfile tries to open it. The pptx format is a zip archive so we need
|
||||
# -- to work around that bug here.
|
||||
if isinstance(self._file, SpooledTemporaryFile):
|
||||
self._file.seek(0)
|
||||
return io.BytesIO(self._file.read())
|
||||
|
||||
if self._file:
|
||||
return self._file
|
||||
|
||||
raise ValueError(
|
||||
"No PPTX document specified, either `filename` or `file` argument must be provided"
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def strategy(self) -> str:
|
||||
"""The requested partitioning strategy.
|
||||
|
||||
This indicates whether the partitioner should undertake expensive operations like inference
|
||||
and OCR to produce a more thorough and/or accurate partitioning of the document.
|
||||
|
||||
Can take several values but for PPTX purposes there is only "hi_res" and not "hi_res".
|
||||
Depending on the picture-partitioner used, images may only be OCR'ed and added to the
|
||||
element-stream when this partitioning strategy is "hi_res".
|
||||
"""
|
||||
return self._strategy
|
||||
|
||||
def table_metadata(self, text_as_html: str | None):
|
||||
"""ElementMetadata instance suitable for use with Table element."""
|
||||
element_metadata = ElementMetadata(
|
||||
filename=self.metadata_file_path,
|
||||
last_modified=self.last_modified,
|
||||
page_number=self.page_number,
|
||||
text_as_html=text_as_html,
|
||||
)
|
||||
element_metadata.detection_origin = DETECTION_ORIGIN
|
||||
return element_metadata
|
||||
|
||||
def text_metadata(self, category_depth: int = 0) -> ElementMetadata:
|
||||
"""ElementMetadata instance suitable for use with Text and subtypes."""
|
||||
element_metadata = ElementMetadata(
|
||||
filename=self.metadata_file_path,
|
||||
last_modified=self.last_modified,
|
||||
page_number=self.page_number,
|
||||
category_depth=category_depth,
|
||||
)
|
||||
element_metadata.detection_origin = DETECTION_ORIGIN
|
||||
return element_metadata
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# SUB-PARTITIONERS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class _NullPicturePartitioner:
|
||||
"""Does not parse the provided Picture element and generates zero elements."""
|
||||
|
||||
@classmethod
|
||||
def iter_elements(cls, picture: Picture, opts: PptxPartitionerOptions) -> Iterator[Element]:
|
||||
"""No-op picture partitioner."""
|
||||
return
|
||||
yield
|
||||
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.file_conversion import convert_file_to_html_text_using_pandoc
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
DETECTION_ORIGIN: str = "rst"
|
||||
|
||||
|
||||
def partition_rst(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an RST document. The document is first converted to HTML and then
|
||||
partitioned using partition_html.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
html_text = convert_file_to_html_text_using_pandoc(
|
||||
source_format="rst", filename=filename, file=file
|
||||
)
|
||||
|
||||
return partition_html(
|
||||
text=html_text,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.RST,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.file_conversion import convert_file_to_html_text_using_pandoc
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import get_last_modified_date
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
DETECTION_ORIGIN: str = "rtf"
|
||||
|
||||
|
||||
def partition_rtf(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
metadata_filename: Optional[str] = None,
|
||||
metadata_last_modified: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions an RTF document. The document is first converted to HTML and then
|
||||
partitioned using partition_html.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
metadata_last_modified
|
||||
The last modified date for the document.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
last_modified = get_last_modified_date(filename) if filename else None
|
||||
|
||||
html_text = convert_file_to_html_text_using_pandoc(
|
||||
source_format="rtf", filename=filename, file=file
|
||||
)
|
||||
|
||||
return partition_html(
|
||||
text=html_text,
|
||||
metadata_filename=metadata_filename or filename,
|
||||
metadata_file_type=FileType.RTF,
|
||||
metadata_last_modified=metadata_last_modified or last_modified,
|
||||
detection_origin=DETECTION_ORIGIN,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,109 @@
|
||||
from typing import List, Optional
|
||||
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
from unstructured.utils import dependency_exists
|
||||
|
||||
|
||||
def validate_strategy(strategy: str, is_image: bool = False):
|
||||
"""Determines if the strategy is valid for the specified filetype."""
|
||||
|
||||
valid_strategies = [
|
||||
PartitionStrategy.AUTO,
|
||||
PartitionStrategy.FAST,
|
||||
PartitionStrategy.OCR_ONLY,
|
||||
PartitionStrategy.HI_RES,
|
||||
]
|
||||
if strategy not in valid_strategies:
|
||||
raise ValueError(f"{strategy} is not a valid strategy.")
|
||||
|
||||
if strategy == PartitionStrategy.FAST and is_image:
|
||||
raise ValueError("The fast strategy is not available for image files.")
|
||||
|
||||
|
||||
def determine_pdf_or_image_strategy(
|
||||
strategy: str,
|
||||
is_image: bool = False,
|
||||
pdf_text_extractable: bool = False,
|
||||
infer_table_structure: bool = False,
|
||||
extract_images_in_pdf: bool = False,
|
||||
extract_image_block_types: Optional[List[str]] = None,
|
||||
):
|
||||
"""Determines what strategy to use for processing PDFs or images, accounting for fallback
|
||||
logic if some dependencies are not available."""
|
||||
pytesseract_installed = dependency_exists("unstructured_pytesseract")
|
||||
unstructured_inference_installed = dependency_exists("unstructured_inference")
|
||||
|
||||
if strategy == PartitionStrategy.AUTO:
|
||||
extract_element = extract_images_in_pdf or bool(extract_image_block_types)
|
||||
if is_image:
|
||||
strategy = _determine_image_auto_strategy()
|
||||
else:
|
||||
strategy = _determine_pdf_auto_strategy(
|
||||
pdf_text_extractable=pdf_text_extractable,
|
||||
infer_table_structure=infer_table_structure,
|
||||
extract_element=extract_element,
|
||||
)
|
||||
|
||||
if all(
|
||||
[not unstructured_inference_installed, not pytesseract_installed, not pdf_text_extractable],
|
||||
):
|
||||
raise ValueError(
|
||||
"unstructured_inference is not installed, pytesseract is not installed "
|
||||
"and the text of the PDF is not extractable. "
|
||||
"To process this file, install unstructured_inference, install pytesseract, "
|
||||
"or remove copy protection from the PDF.",
|
||||
)
|
||||
|
||||
if strategy == PartitionStrategy.HI_RES and not unstructured_inference_installed:
|
||||
logger.warning(
|
||||
"unstructured_inference is not installed. Cannot use the hi_res partitioning "
|
||||
"strategy. Falling back to partitioning with another strategy.",
|
||||
)
|
||||
# NOTE(robinson) - fallback to ocr_only if possible because it is the most
|
||||
# similar to hi_res
|
||||
if pytesseract_installed:
|
||||
logger.warning("Falling back to partitioning with ocr_only.")
|
||||
return PartitionStrategy.OCR_ONLY
|
||||
else:
|
||||
logger.warning("Falling back to partitioning with fast.")
|
||||
return PartitionStrategy.FAST
|
||||
|
||||
elif strategy == PartitionStrategy.OCR_ONLY and not pytesseract_installed:
|
||||
logger.warning(
|
||||
"pytesseract is not installed. Cannot use the ocr_only partitioning "
|
||||
"strategy. Falling back to partitioning with another strategy.",
|
||||
)
|
||||
if pdf_text_extractable:
|
||||
logger.warning("Falling back to partitioning with fast.")
|
||||
return PartitionStrategy.FAST
|
||||
else:
|
||||
logger.warning("Falling back to partitioning with hi_res.")
|
||||
return PartitionStrategy.HI_RES
|
||||
|
||||
return strategy
|
||||
|
||||
|
||||
def _determine_image_auto_strategy():
|
||||
"""If "auto" is passed in as the strategy, determines what strategy to use
|
||||
for images."""
|
||||
# Use hi_res as the only default since images are only about one page
|
||||
return PartitionStrategy.HI_RES
|
||||
|
||||
|
||||
def _determine_pdf_auto_strategy(
|
||||
pdf_text_extractable: bool = False,
|
||||
infer_table_structure: bool = False,
|
||||
extract_element: bool = False,
|
||||
):
|
||||
"""If "auto" is passed in as the strategy, determines what strategy to use
|
||||
for PDFs."""
|
||||
# NOTE(robinson) - Currently "hi_res" is the only strategy where
|
||||
# infer_table_structure and extract_images_in_pdf are used.
|
||||
if infer_table_structure or extract_element:
|
||||
return PartitionStrategy.HI_RES
|
||||
|
||||
if pdf_text_extractable:
|
||||
return PartitionStrategy.FAST
|
||||
else:
|
||||
return PartitionStrategy.OCR_ONLY
|
||||
@@ -0,0 +1,216 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import re
|
||||
from typing import IO, Any, Callable, Literal
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.cleaners.core import (
|
||||
auto_paragraph_grouper,
|
||||
clean_bullets,
|
||||
)
|
||||
from unstructured.documents.coordinates import CoordinateSystem
|
||||
from unstructured.documents.elements import (
|
||||
Address,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
EmailAddress,
|
||||
Footer,
|
||||
Header,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.file_utils.encoding import read_txt_file
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.nlp.patterns import PARAGRAPH_PATTERN, UNICODE_BULLETS_RE
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
from unstructured.partition.text_type import (
|
||||
is_bulleted_text,
|
||||
is_email_address,
|
||||
is_possible_narrative_text,
|
||||
is_possible_numbered_list,
|
||||
is_possible_title,
|
||||
is_us_city_state_zip,
|
||||
)
|
||||
|
||||
|
||||
@apply_metadata(FileType.TXT)
|
||||
@add_chunking_strategy
|
||||
def partition_text(
|
||||
filename: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
encoding: str | None = None,
|
||||
text: str | None = None,
|
||||
paragraph_grouper: Callable[[str], str] | Literal[False] | None = None,
|
||||
detection_origin: str | None = "text",
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partition a .txt documents into its constituent paragraph elements.
|
||||
|
||||
If paragraphs are below "min_partition" or above "max_partition" boundaries,
|
||||
they are combined or split.
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
encoding
|
||||
The encoding method used to decode the input bytes when drawn from `filename` or `file`.
|
||||
Defaults to "utf-8".
|
||||
text
|
||||
The string representation of the .txt document.
|
||||
paragrapher_grouper
|
||||
A str -> str function for fixing paragraphs that are interrupted by line breaks
|
||||
for formatting purposes.
|
||||
"""
|
||||
if text is not None and text.strip() == "" and not file and not filename:
|
||||
return []
|
||||
|
||||
# -- Verify that only one of the arguments was provided --
|
||||
exactly_one(filename=filename, file=file, text=text)
|
||||
|
||||
file_text = ""
|
||||
if filename is not None:
|
||||
encoding, file_text = read_txt_file(filename=filename, encoding=encoding)
|
||||
elif file is not None:
|
||||
encoding, file_text = read_txt_file(file=file, encoding=encoding)
|
||||
elif text is not None:
|
||||
file_text = str(text)
|
||||
|
||||
if paragraph_grouper is False:
|
||||
pass
|
||||
elif paragraph_grouper is not None:
|
||||
file_text = paragraph_grouper(file_text)
|
||||
else:
|
||||
file_text = auto_paragraph_grouper(file_text)
|
||||
|
||||
file_content = _split_by_paragraph(file_text)
|
||||
|
||||
elements: list[Element] = []
|
||||
metadata = ElementMetadata(
|
||||
last_modified=get_last_modified_date(filename) if filename else None,
|
||||
)
|
||||
metadata.detection_origin = detection_origin
|
||||
|
||||
for ctext in file_content:
|
||||
ctext = ctext.strip()
|
||||
|
||||
if ctext and not _is_empty_bullet(ctext):
|
||||
element = element_from_text(ctext)
|
||||
element.metadata = copy.deepcopy(metadata)
|
||||
elements.append(element)
|
||||
|
||||
return elements
|
||||
|
||||
|
||||
def element_from_text(
|
||||
text: str,
|
||||
coordinates: tuple[tuple[float, float], ...] | None = None,
|
||||
coordinate_system: CoordinateSystem | None = None,
|
||||
) -> Element:
|
||||
if _is_in_header_position(coordinates, coordinate_system):
|
||||
return Header(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif _is_in_footer_position(coordinates, coordinate_system):
|
||||
return Footer(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif is_bulleted_text(text):
|
||||
clean_text = clean_bullets(text)
|
||||
return ListItem(
|
||||
text=clean_text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif is_email_address(text):
|
||||
return EmailAddress(text=text)
|
||||
elif is_us_city_state_zip(text):
|
||||
return Address(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif is_possible_numbered_list(text):
|
||||
return ListItem(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif is_possible_narrative_text(text):
|
||||
return NarrativeText(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
elif is_possible_title(text):
|
||||
return Title(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
else:
|
||||
return Text(
|
||||
text=text,
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# HELPER FUNCTIONS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
def _get_height_percentage(
|
||||
coordinates: tuple[tuple[float, float], ...],
|
||||
coordinate_system: CoordinateSystem,
|
||||
) -> float:
|
||||
avg_y = sum(coordinate[1] for coordinate in coordinates) / len(coordinates)
|
||||
return avg_y / coordinate_system.height
|
||||
|
||||
|
||||
def _is_empty_bullet(text: str) -> bool:
|
||||
"""Checks if input text is an empty bullet."""
|
||||
return bool(UNICODE_BULLETS_RE.match(text) and len(text) == 1)
|
||||
|
||||
|
||||
def _is_in_footer_position(
|
||||
coordinates: tuple[tuple[float, float], ...] | None,
|
||||
coordinate_system: CoordinateSystem | None,
|
||||
threshold: float = 0.93,
|
||||
) -> bool:
|
||||
"""Checks to see if the position of the text indicates that the text belongs
|
||||
to a footer."""
|
||||
if coordinates is None or coordinate_system is None:
|
||||
return False
|
||||
|
||||
height_percentage = _get_height_percentage(coordinates, coordinate_system)
|
||||
return height_percentage > threshold
|
||||
|
||||
|
||||
def _is_in_header_position(
|
||||
coordinates: tuple[tuple[float, float], ...] | None,
|
||||
coordinate_system: CoordinateSystem | None,
|
||||
threshold: float = 0.07,
|
||||
) -> bool:
|
||||
"""Checks to see if the position of the text indicates that the text belongs to a header."""
|
||||
if coordinates is None or coordinate_system is None:
|
||||
return False
|
||||
|
||||
height_percentage = _get_height_percentage(coordinates, coordinate_system)
|
||||
return height_percentage < threshold
|
||||
|
||||
|
||||
def _split_by_paragraph(file_text: str) -> list[str]:
|
||||
"""Split text into paragraphs."""
|
||||
return re.split(PARAGRAPH_PATTERN, file_text.strip())
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Provides functions for classifying text for Element selection during partitioning."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Final, List, Optional
|
||||
|
||||
from unstructured.cleaners.core import remove_punctuation
|
||||
from unstructured.logger import trace_logger
|
||||
from unstructured.nlp.english_words import ENGLISH_WORDS
|
||||
from unstructured.nlp.patterns import (
|
||||
EMAIL_ADDRESS_PATTERN_RE,
|
||||
ENDS_IN_PUNCT_RE,
|
||||
NUMBERED_LIST_RE,
|
||||
UNICODE_BULLETS_RE,
|
||||
US_CITY_STATE_ZIP_RE,
|
||||
US_PHONE_NUMBERS_RE,
|
||||
)
|
||||
from unstructured.nlp.tokenize import pos_tag, sent_tokenize, word_tokenize
|
||||
|
||||
POS_VERB_TAGS: Final[List[str]] = ["VB", "VBG", "VBD", "VBN", "VBP", "VBZ"]
|
||||
ENGLISH_WORD_SPLIT_RE = re.compile(r"[\s\-,.!?_\/]+")
|
||||
NON_LOWERCASE_ALPHA_RE = re.compile(r"[^a-z]")
|
||||
|
||||
|
||||
def is_possible_narrative_text(
|
||||
text: str,
|
||||
cap_threshold: float = 0.5,
|
||||
non_alpha_threshold: float = 0.5,
|
||||
languages: List[str] = ["eng"],
|
||||
language_checks: bool = False,
|
||||
) -> bool:
|
||||
"""Checks to see if the text passes all of the checks for a narrative text section.
|
||||
You can change the cap threshold using the cap_threshold kwarg or the
|
||||
NARRATIVE_TEXT_CAP_THRESHOLD environment variable. The environment variable takes
|
||||
precedence over the kwarg.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text
|
||||
The input text to check
|
||||
cap_threshold
|
||||
The percentage of capitalized words necessary to disqualify the segment as narrative
|
||||
non_alpha_threshold
|
||||
The minimum proportion of alpha characters the text needs to be considered
|
||||
narrative text
|
||||
languages
|
||||
The list of languages present in the document. Defaults to ["eng"] for English
|
||||
language_checks
|
||||
If True, conducts checks that are specific to the chosen language. Turn on for more
|
||||
accurate partitioning and off for faster processing.
|
||||
"""
|
||||
_language_checks = os.environ.get("UNSTRUCTURED_LANGUAGE_CHECKS")
|
||||
if _language_checks is not None:
|
||||
language_checks = _language_checks.lower() == "true"
|
||||
|
||||
if len(text) == 0:
|
||||
trace_logger.detail("Not narrative. Text is empty.") # type: ignore
|
||||
return False
|
||||
|
||||
if text.isnumeric():
|
||||
trace_logger.detail(f"Not narrative. Text is all numeric:\n\n{text}") # type: ignore
|
||||
return False
|
||||
|
||||
if "eng" in languages and language_checks and not contains_english_word(text):
|
||||
return False
|
||||
|
||||
# NOTE(robinson): it gets read in from the environment as a string so we need to
|
||||
# cast it to a float
|
||||
cap_threshold = float(
|
||||
os.environ.get("UNSTRUCTURED_NARRATIVE_TEXT_CAP_THRESHOLD", cap_threshold),
|
||||
)
|
||||
if exceeds_cap_ratio(text, threshold=cap_threshold):
|
||||
trace_logger.detail(f"Not narrative. Text exceeds cap ratio {cap_threshold}:\n\n{text}") # type: ignore # noqa: E501
|
||||
return False
|
||||
|
||||
non_alpha_threshold = float(
|
||||
os.environ.get("UNSTRUCTURED_NARRATIVE_TEXT_NON_ALPHA_THRESHOLD", non_alpha_threshold),
|
||||
)
|
||||
if under_non_alpha_ratio(text, threshold=non_alpha_threshold):
|
||||
return False
|
||||
|
||||
if "eng" in languages and (sentence_count(text, 3) < 2) and (not contains_verb(text)):
|
||||
trace_logger.detail(f"Not narrative. Text does not contain a verb:\n\n{text}") # type: ignore # noqa: E501
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_possible_title(
|
||||
text: str,
|
||||
sentence_min_length: int = 5,
|
||||
title_max_word_length: int = 12,
|
||||
non_alpha_threshold: float = 0.5,
|
||||
languages: List[str] = ["eng"],
|
||||
language_checks: bool = False,
|
||||
) -> bool:
|
||||
"""Checks to see if the text passes all of the checks for a valid title.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text
|
||||
The input text to check
|
||||
sentence_min_length
|
||||
The minimum number of words required to consider a section of text a sentence
|
||||
title_max_word_length
|
||||
The maximum number of words a title can contain
|
||||
non_alpha_threshold
|
||||
The minimum number of alpha characters the text needs to be considered a title
|
||||
languages
|
||||
The list of languages present in the document. Defaults to ["eng"] for English
|
||||
language_checks
|
||||
If True, conducts checks that are specific to the chosen language. Turn on for more
|
||||
accurate partitioning and off for faster processing.
|
||||
"""
|
||||
_language_checks = os.environ.get("UNSTRUCTURED_LANGUAGE_CHECKS")
|
||||
if _language_checks is not None:
|
||||
language_checks = _language_checks.lower() == "true"
|
||||
|
||||
if len(text) == 0:
|
||||
trace_logger.detail("Not a title. Text is empty.") # type: ignore
|
||||
return False
|
||||
|
||||
if text.isupper() and ENDS_IN_PUNCT_RE.search(text) is not None:
|
||||
return False
|
||||
|
||||
title_max_word_length = int(
|
||||
os.environ.get("UNSTRUCTURED_TITLE_MAX_WORD_LENGTH", title_max_word_length),
|
||||
)
|
||||
# NOTE(robinson) - splitting on spaces here instead of word tokenizing because it
|
||||
# is less expensive and actual tokenization doesn't add much value for the length check
|
||||
if len(text.split(" ")) > title_max_word_length:
|
||||
return False
|
||||
|
||||
non_alpha_threshold = float(
|
||||
os.environ.get("UNSTRUCTURED_TITLE_NON_ALPHA_THRESHOLD", non_alpha_threshold),
|
||||
)
|
||||
if under_non_alpha_ratio(text, threshold=non_alpha_threshold):
|
||||
return False
|
||||
|
||||
# NOTE(robinson) - Prevent flagging salutations like "To My Dearest Friends," as titles
|
||||
if text.endswith(","):
|
||||
return False
|
||||
|
||||
if "eng" in languages and not contains_english_word(text) and language_checks:
|
||||
return False
|
||||
|
||||
if text.isnumeric():
|
||||
trace_logger.detail(f"Not a title. Text is all numeric:\n\n{text}") # type: ignore
|
||||
return False
|
||||
|
||||
# NOTE(robinson) - The min length is to capture content such as "ITEM 1A. RISK FACTORS"
|
||||
# that sometimes get tokenized as separate sentences due to the period, but are still
|
||||
# valid titles
|
||||
if sentence_count(text, min_length=sentence_min_length) > 1:
|
||||
trace_logger.detail( # type: ignore
|
||||
f"Not a title. Text is longer than {sentence_min_length} sentences:\n\n{text}",
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_bulleted_text(text: str) -> bool:
|
||||
"""Checks to see if the section of text is part of a bulleted list."""
|
||||
return UNICODE_BULLETS_RE.match(text.strip()) is not None
|
||||
|
||||
|
||||
def contains_us_phone_number(text: str) -> bool:
|
||||
"""Checks to see if a section of text contains a US phone number.
|
||||
|
||||
Example
|
||||
-------
|
||||
contains_us_phone_number("867-5309") -> True
|
||||
"""
|
||||
return US_PHONE_NUMBERS_RE.search(text.strip()) is not None
|
||||
|
||||
|
||||
def contains_verb(text: str) -> bool:
|
||||
"""Use a POS tagger to check if a segment contains verbs. If the section does not have verbs,
|
||||
that indicates that it is not narrative text."""
|
||||
if text.isupper():
|
||||
text = text.lower()
|
||||
|
||||
pos_tags = pos_tag(text)
|
||||
return any(tag in POS_VERB_TAGS for _, tag in pos_tags)
|
||||
|
||||
|
||||
def contains_english_word(text: str) -> bool:
|
||||
"""Checks to see if the text contains an English word."""
|
||||
text = text.lower()
|
||||
words = ENGLISH_WORD_SPLIT_RE.split(text)
|
||||
for word in words:
|
||||
# NOTE(Crag): Remove any non-lowercase alphabetical
|
||||
# characters. These removed chars will usually be trailing or
|
||||
# leading characters not already matched in ENGLISH_WORD_SPLIT_RE.
|
||||
# The possessive case is also generally ok:
|
||||
# "beggar's" -> "beggars" (still an english word)
|
||||
# and of course:
|
||||
# "'beggars'"-> "beggars" (also still an english word)
|
||||
word = NON_LOWERCASE_ALPHA_RE.sub("", word)
|
||||
if len(word) > 1 and word in ENGLISH_WORDS:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def sentence_count(text: str, min_length: Optional[int] = None) -> int:
|
||||
"""Checks the sentence count for a section of text. Titles should not be more than one
|
||||
sentence.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text
|
||||
The string of the text to count
|
||||
min_length
|
||||
The min number of words a section needs to be for it to be considered a sentence.
|
||||
"""
|
||||
sentences = sent_tokenize(text)
|
||||
count = 0
|
||||
if min_length:
|
||||
trace_detail = trace_logger.detail # type: ignore
|
||||
for sentence in sentences:
|
||||
stripped = remove_punctuation(sentence)
|
||||
# Fast token count after punctuation is removed: just split on whitespace
|
||||
word_count = sum(1 for token in stripped.split() if token != ".")
|
||||
if word_count < min_length:
|
||||
trace_detail(
|
||||
f"Sentence does not exceed {min_length} word tokens, it will not count toward "
|
||||
"sentence count.\n"
|
||||
f"{stripped}",
|
||||
)
|
||||
continue
|
||||
count += 1
|
||||
else:
|
||||
for sentence in sentences:
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def under_non_alpha_ratio(text: str, threshold: float = 0.5):
|
||||
"""Checks if the proportion of non-alpha characters in the text snippet exceeds a given
|
||||
threshold. This helps prevent text like "-----------BREAK---------" from being tagged
|
||||
as a title or narrative text. The ratio does not count spaces.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text
|
||||
The input string to test
|
||||
threshold
|
||||
If the proportion of non-alpha characters exceeds this threshold, the function
|
||||
returns False
|
||||
"""
|
||||
if not text:
|
||||
return False
|
||||
|
||||
alpha_count = 0
|
||||
total_count = 0
|
||||
for char in text:
|
||||
if not char.isspace():
|
||||
total_count += 1
|
||||
if char.isalpha():
|
||||
alpha_count += 1
|
||||
|
||||
return ((alpha_count / total_count) < threshold) if total_count > 0 else False
|
||||
|
||||
|
||||
def exceeds_cap_ratio(text: str, threshold: float = 0.5) -> bool:
|
||||
"""Checks the title ratio in a section of text. If a sufficient proportion of the words
|
||||
are capitalized, that can be indicated on non-narrative text (i.e. "1A. Risk Factors").
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text
|
||||
The input string to test
|
||||
threshold
|
||||
If the percentage of words beginning with a capital letter exceeds this threshold,
|
||||
the function returns True
|
||||
"""
|
||||
# NOTE(robinson) - Currently limiting this to only sections of text with one sentence.
|
||||
# The assumption is that sections with multiple sentences are not titles.
|
||||
if sentence_count(text, 3) > 1:
|
||||
return False
|
||||
|
||||
if text.isupper():
|
||||
return True
|
||||
|
||||
# NOTE(jay-ylee) - The word_tokenize function also recognizes and separates special characters
|
||||
# into one word, causing problems with ratio measurement.
|
||||
# Therefore, only words consisting of alphabets are used to measure the ratio.
|
||||
# ex. world_tokenize("ITEM 1. Financial Statements (Unaudited)")
|
||||
# = ['ITEM', '1', '.', 'Financial', 'Statements', '(', 'Unaudited', ')'],
|
||||
# however, "ITEM 1. Financial Statements (Unaudited)" is Title, not NarrativeText
|
||||
tokens = [tk for tk in word_tokenize(text) if tk.isalpha()]
|
||||
|
||||
# NOTE(jay-ylee) - If word_tokenize(text) is empty, return must be True to
|
||||
# avoid being misclassified as Narrative Text.
|
||||
if len(tokens) == 0:
|
||||
return True
|
||||
|
||||
capitalized = sum([word.istitle() or word.isupper() for word in tokens])
|
||||
ratio = capitalized / len(tokens)
|
||||
return ratio > threshold
|
||||
|
||||
|
||||
def is_us_city_state_zip(text: str) -> bool:
|
||||
"""Checks if the given text is in the format of US city/state/zip code.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Doylestown, PA 18901
|
||||
Doylestown, Pennsylvania, 18901
|
||||
DOYLESTOWN, PENNSYLVANIA 18901
|
||||
"""
|
||||
return US_CITY_STATE_ZIP_RE.match(text.strip()) is not None
|
||||
|
||||
|
||||
def is_email_address(text: str) -> bool:
|
||||
"""Check if the given text is the email address"""
|
||||
return EMAIL_ADDRESS_PATTERN_RE.match(text.strip()) is not None
|
||||
|
||||
|
||||
def is_possible_numbered_list(text: str) -> bool:
|
||||
"""Checks to see if the text is a potential numbered list."""
|
||||
return NUMBERED_LIST_RE.match(text.strip()) is not None
|
||||
@@ -0,0 +1,64 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import IO, Any, Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy
|
||||
from unstructured.common.html_table import HtmlTable
|
||||
from unstructured.documents.elements import Element, ElementMetadata, Table
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.common import (
|
||||
exactly_one,
|
||||
spooled_to_bytes_io_if_needed,
|
||||
)
|
||||
from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date
|
||||
|
||||
DETECTION_ORIGIN: str = "tsv"
|
||||
|
||||
|
||||
@apply_metadata(FileType.TSV)
|
||||
@add_chunking_strategy
|
||||
def partition_tsv(
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
include_header: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> list[Element]:
|
||||
"""Partitions TSV files into document elements.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename
|
||||
A string defining the target filename path.
|
||||
file
|
||||
A file-like object using "rb" mode --> open(filename, "rb").
|
||||
include_header
|
||||
Determines whether or not header info info is included in text and medatada.text_as_html.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
header = 0 if include_header else None
|
||||
|
||||
if filename:
|
||||
dataframe = pd.read_csv(filename, sep="\t", header=header)
|
||||
else:
|
||||
assert file is not None
|
||||
# -- Note(scanny): `SpooledTemporaryFile` on Python<3.11 does not implement `.readable()`
|
||||
# -- which triggers an exception on `pd.DataFrame.read_csv()` call.
|
||||
f = spooled_to_bytes_io_if_needed(file)
|
||||
dataframe = pd.read_csv(f, sep="\t", header=header)
|
||||
|
||||
html_table = HtmlTable.from_html_text(
|
||||
dataframe.to_html(index=False, header=include_header, na_rep="")
|
||||
)
|
||||
|
||||
metadata = ElementMetadata(
|
||||
filename=filename,
|
||||
last_modified=get_last_modified_date(filename) if filename else None,
|
||||
text_as_html=html_table.html,
|
||||
)
|
||||
metadata.detection_origin = DETECTION_ORIGIN
|
||||
|
||||
return [Table(text=html_table.text, metadata=metadata)]
|
||||
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user