110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
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
|