diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c5bff9d1c..eaacd630c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,7 +71,7 @@ ## 0.22.16 ### Enhancements - +- **Infer hierarchical heading levels (H1–H6) for PDFs**: `Title` elements now receive accurate `category_depth` values via two strategies: (1) PDF outline/bookmark extraction matches the document's table-of-contents tree to titles by page and text similarity, and (2) font-size analysis clusters distinct character sizes across titles when outlines are unavailable. Works with all partition strategies (`fast`, `hi_res`, `ocr_only`). Resolves #4204. - **Formula markdown export (`element_to_md` / `elements_to_md`)**: New keyword-only `formula_markdown_style` (`"auto"`, `"display_math"`, `"plain"`; default `"auto"`). In `"auto"`, display math (`$$ ... $$`) is used only when the text looks like notation (heuristic score) and contains no `$`/`$$` (avoids breaking Markdown and noisy OCR captions). `"display_math"` wraps whenever safe (still falls back to plain if `$` would corrupt fences). `"plain"` emits text only. Optional `normalize_formula` (default `True`) maps common Unicode operators to LaTeX-like tokens; `normalize_formula` stays before keyword-only options so positional `encoding` / `no_group_by_page` callers are unchanged. Unicode `√` is never mapped to `\\sqrt{}`. Module constants: `FORMULA_MARKDOWN_AUTO`, `FORMULA_MARKDOWN_DISPLAY_MATH`, `FORMULA_MARKDOWN_PLAIN`. ## 0.22.15 diff --git a/test_unstructured/partition/test_pdf_heading_hierarchy.py b/test_unstructured/partition/test_pdf_heading_hierarchy.py new file mode 100644 index 0000000000..f67bc003ca --- /dev/null +++ b/test_unstructured/partition/test_pdf_heading_hierarchy.py @@ -0,0 +1,189 @@ +"""Tests for unstructured.partition.pdf_heading_hierarchy.""" + +from __future__ import annotations + +from typing import Any, cast + +from test_unstructured.unit_utils import example_doc_path +from unstructured.documents.coordinates import PixelSpace +from unstructured.documents.elements import CoordinatesMetadata, ElementMetadata, Title +from unstructured.partition.pdf_heading_hierarchy import ( + _apply_font_size_levels, + _apply_outline_levels, + _best_match, + _open_reader, + infer_heading_levels, +) + + +class _FakeBox: + def __init__(self, width: float = 600, height: float = 800) -> None: + self.width = width + self.height = height + + +class _FakePage: + def __init__(self) -> None: + self.cropbox = _FakeBox() + + +class _FakeDestination: + def __init__( + self, + title: str, + *, + page_index: int | None, + left: float | None = None, + top: float | None = None, + ) -> None: + self.title = title + self.page_index = page_index + self.left = left + self.top = top + self.right = None + self.bottom = None + self.page = page_index + + +class _FakeReader: + def __init__(self, outline: list[object]) -> None: + self.outline = outline + self.pages = [_FakePage()] + + def get_destination_page_number(self, destination: _FakeDestination) -> int: + return destination.page_index if destination.page_index is not None else -1 + + +def _title_with_bbox(text: str, y1: float, y2: float) -> Title: + title = Title(text) + title.metadata = ElementMetadata( + page_number=1, + coordinates=CoordinatesMetadata( + points=((100, y1), (100, y2), (250, y2), (250, y1)), + system=PixelSpace(width=600, height=800), + ), + ) + return title + + +class Describe_best_match: + def it_picks_strongest_match_not_first(self): + """Prevents PR #4222 regression: greedy early-break selecting wrong match.""" + t1 = Title("Part I Overview") + t2 = Title("Part II Details") + best, ratio = _best_match("Part II Details", [t1, t2]) + assert best is t2 + assert ratio == 1.0 + + def it_skips_already_assigned_titles(self): + t1 = Title("Introduction") + t1.metadata.category_depth = 0 + t2 = Title("Introduction") + best, _ = _best_match("Introduction", [t1, t2]) + assert best is t2 + + +class Describe_apply_outline_levels: + def it_matches_on_correct_page(self): + """Prevents PR #4222 regression: 0-based vs 1-based page mismatch.""" + reader = _open_reader(filename=example_doc_path("pdf/DA-1p.pdf")) + assert reader is not None + # Title on page 1 should match; title on wrong page should not + t_right_page = Title("Codex DAO") + t_right_page.metadata = ElementMetadata(page_number=1) + t_wrong_page = Title("Codex DAO") + t_wrong_page.metadata = ElementMetadata(page_number=999) + matched = _apply_outline_levels([t_right_page, t_wrong_page], reader) + assert matched >= 1 + assert t_right_page.metadata.category_depth == 0 + + def it_uses_outline_destination_when_same_title_repeats_on_same_page(self): + upper_title = _title_with_bbox("Introduction", y1=90, y2=120) + lower_title = _title_with_bbox("Introduction", y1=290, y2=320) + reader = cast( + Any, + _FakeReader( + outline=[ + _FakeDestination("Introduction", page_index=0, left=100, top=700), + [_FakeDestination("Introduction", page_index=0, left=100, top=500)], + ], + ), + ) + + matched = _apply_outline_levels([lower_title, upper_title], reader) + + assert matched == 2 + assert upper_title.metadata.category_depth == 0 + assert lower_title.metadata.category_depth == 1 + + def it_skips_external_outline_entries(self): + title = _title_with_bbox("External Resource", y1=90, y2=120) + reader = cast( + Any, + _FakeReader( + outline=[ + _FakeDestination("External Resource", page_index=None, left=100, top=700), + ], + ), + ) + + matched = _apply_outline_levels([title], reader) + + assert matched == 0 + assert title.metadata.category_depth is None + + +class Describe_apply_font_size_levels: + def it_does_not_overwrite_existing_depths(self): + t1 = Title("Introduction") + t1.metadata = ElementMetadata(page_number=1, category_depth=5) + _apply_font_size_levels([t1], filename=example_doc_path("pdf/layout-parser-paper.pdf")) + assert t1.metadata.category_depth == 5 + + +class Describe_infer_heading_levels: + def it_handles_missing_file_gracefully(self): + t1 = Title("Heading") + t1.metadata = ElementMetadata(page_number=1) + result = infer_heading_levels([t1]) + assert result is not None + assert t1.metadata.category_depth is None + + def it_works_with_file_object(self): + t1 = Title("Codex DAO") + t1.metadata = ElementMetadata(page_number=1) + with open(example_doc_path("pdf/DA-1p.pdf"), "rb") as f: + infer_heading_levels([t1], file=f) + assert t1.metadata.category_depth == 0 + + +class Describe_partition_pdf_end_to_end: + """End-to-end tests calling partition_pdf() and verifying category_depth.""" + + def it_assigns_outline_depths_via_fast_strategy(self): + from unstructured.partition.pdf import partition_pdf + + elements = partition_pdf(example_doc_path("pdf/DA-1p.pdf"), strategy="fast") + titles = { + el.text.strip(): el.metadata.category_depth for el in elements if el.category == "Title" + } + assert titles["MAIN GAME"] == 1 + assert titles["CREATURES"] == 2 + assert titles["Abomination"] == 3 + + def it_assigns_font_size_depths_for_pdf_without_outline(self): + from unstructured.partition.pdf import partition_pdf + + elements = partition_pdf(example_doc_path("pdf/loremipsum-flat.pdf"), strategy="fast") + titles = [el for el in elements if el.category == "Title"] + for t in titles: + if t.metadata.category_depth is not None: + assert 0 <= t.metadata.category_depth <= 5 + + def it_does_not_set_depths_on_non_title_elements(self): + from unstructured.partition.pdf import partition_pdf + + elements = partition_pdf(example_doc_path("pdf/DA-1p.pdf"), strategy="fast") + non_titles = [el for el in elements if el.category != "Title"] + for el in non_titles: + assert el.metadata.category_depth is None diff --git a/unstructured/partition/pdf.py b/unstructured/partition/pdf.py index b5a5af2e53..ef362f139e 100644 --- a/unstructured/partition/pdf.py +++ b/unstructured/partition/pdf.py @@ -52,6 +52,7 @@ ) from unstructured.partition.common.lang import check_language_args, prepare_languages_for_tesseract from unstructured.partition.common.metadata import apply_metadata, get_last_modified_date +from unstructured.partition.pdf_heading_hierarchy import infer_heading_levels from unstructured.partition.pdf_image.pdfminer_processing import ( check_annotations_within_element, get_uris, @@ -347,7 +348,7 @@ def partition_pdf_or_image( # NOTE(robinson): Catches a UserWarning that occurs when detection is called with warnings.catch_warnings(): warnings.simplefilter("ignore") - return _partition_pdf_or_image_local( + elements = _partition_pdf_or_image_local( filename=filename, file=spooled_to_bytes_io_if_needed(file), is_image=is_image, @@ -376,7 +377,7 @@ def partition_pdf_or_image( # are likely not Titles and should not be identified as such. elif strategy == PartitionStrategy.FAST: - return _partition_pdf_with_pdfparser( + elements = _partition_pdf_with_pdfparser( extracted_elements=extracted_elements, include_page_breaks=include_page_breaks, **kwargs, @@ -385,7 +386,7 @@ def partition_pdf_or_image( elif strategy == PartitionStrategy.OCR_ONLY: # NOTE(robinson): Catches file conversion warnings when running with PDFs with warnings.catch_warnings(): - elements = _partition_pdf_or_image_with_ocr( + ocr_elements = _partition_pdf_or_image_with_ocr( filename=filename, file=file, include_page_breaks=include_page_breaks, @@ -397,9 +398,23 @@ def partition_pdf_or_image( password=password, **kwargs, ) - return _process_uncategorized_text_elements(elements) + elements = _process_uncategorized_text_elements(ocr_elements) - raise ValueError(f"Unsupported partitioning strategy: {strategy}") + else: + raise ValueError(f"Unsupported partitioning strategy: {strategy}") + + # Infer hierarchical heading levels for PDF documents (not images). + if not is_image: + if file is not None: + file.seek(0) + infer_heading_levels( + elements, + filename=filename or "", + file=file, + password=password, + ) + + return elements def extractable_elements( diff --git a/unstructured/partition/pdf_heading_hierarchy.py b/unstructured/partition/pdf_heading_hierarchy.py new file mode 100644 index 0000000000..88c40892ee --- /dev/null +++ b/unstructured/partition/pdf_heading_hierarchy.py @@ -0,0 +1,422 @@ +"""Infer hierarchical heading levels (H1-H6) for PDF elements. + +Two strategies in order of preference: +1. Outline extraction - match PDF bookmarks to Title elements by destination + text similarity. +2. Font-size analysis - cluster distinct font sizes across Titles, largest font -> depth 0. + +Sets element.metadata.category_depth using the 0-indexed convention (0=H1, 1=H2, ... 5=H6). +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from difflib import SequenceMatcher +from typing import IO, Any, BinaryIO, Optional, cast + +from pypdf import PdfReader + +from unstructured.documents.coordinates import CoordinateSystem, PointSpace +from unstructured.documents.elements import Element, Title +from unstructured.logger import logger + +_SIMILARITY_THRESHOLD = 0.85 +_MAX_HEADING_DEPTH = 6 + + +@dataclass(frozen=True) +class _OutlineEntry: + text: str + depth: int + page_number: int + left: Optional[float] = None + top: Optional[float] = None + right: Optional[float] = None + bottom: Optional[float] = None + + @property + def has_destination_coordinates(self) -> bool: + return any( + coordinate is not None for coordinate in (self.left, self.top, self.right, self.bottom) + ) + + +def infer_heading_levels( + elements: list[Element], + *, + filename: str = "", + file: Optional[IO[bytes]] = None, + password: Optional[str] = None, +) -> list[Element]: + """Assign ``category_depth`` to ``Title`` elements based on PDF structure. + Mutates *elements* in-place and returns for convenience. + """ + titles = [el for el in elements if isinstance(el, Title)] + if not titles: + return elements + reader = _open_reader(filename=filename, file=file, password=password) + if reader is None: + return elements + _apply_outline_levels(titles, reader) + if any(t.metadata.category_depth is None for t in titles): + _apply_font_size_levels(titles, filename=filename, file=file, password=password) + return elements + + +# -- Strategy 1: PDF outline / bookmarks -- + + +def _apply_outline_levels(titles: list[Title], reader: PdfReader) -> int: + """Walk the PDF outline tree and set ``category_depth`` on matching titles. + Returns the number of matched titles. + """ + try: + outline = reader.outline + except Exception: + logger.debug("Failed to read PDF outline, skipping outline strategy.") + return 0 + if not outline: + return 0 + + outline_entries: list[_OutlineEntry] = [] + _walk_outline(outline, reader, depth=0, entries=outline_entries) + if not outline_entries: + return 0 + + titles_by_page: dict[int, list[Title]] = defaultdict(list) + for title in titles: + if title.metadata.page_number is not None: + titles_by_page[title.metadata.page_number].append(title) + + matched = 0 + for entry in outline_entries: + capped_depth = min(entry.depth, _MAX_HEADING_DEPTH - 1) + candidates = titles_by_page.get(entry.page_number, []) + best_title, best_ratio = _best_destination_match(entry, candidates, reader) + if best_title is None: + best_title, best_ratio = _best_match(entry.text, candidates) + if best_title is not None and best_ratio >= _SIMILARITY_THRESHOLD: + best_title.metadata.category_depth = capped_depth + matched += 1 + return matched + + +def _walk_outline( + items: list[Any], + reader: PdfReader, + depth: int, + entries: list[_OutlineEntry], +) -> None: + """Recursively walk an outline tree produced by ``PdfReader.outline``.""" + for item in items: + if isinstance(item, list): + _walk_outline(item, reader, depth + 1, entries) + else: + title_text = (getattr(item, "title", None) or "").strip() + if not title_text: + continue + page_number = _resolve_page_number(item, reader) + if page_number is None: + continue + entries.append( + _OutlineEntry( + text=title_text, + depth=depth, + page_number=page_number, + left=_as_float(getattr(item, "left", None)), + top=_as_float(getattr(item, "top", None)), + right=_as_float(getattr(item, "right", None)), + bottom=_as_float(getattr(item, "bottom", None)), + ), + ) + + +def _resolve_page_number(destination: object, reader: PdfReader) -> Optional[int]: + """Return the 1-based page number for an outline destination.""" + try: + page_idx = reader.get_destination_page_number(cast(Any, destination)) + if page_idx is not None and page_idx >= 0: + return page_idx + 1 + except Exception: + pass + + try: + page_obj = getattr(destination, "page", None) + if page_obj is None: + return None + if isinstance(page_obj, int): + return page_obj + 1 if page_obj >= 0 else None + page_obj = page_obj.get_object() if hasattr(page_obj, "get_object") else page_obj + for idx, page in enumerate(reader.pages): + if page.get_object() == page_obj: + return idx + 1 + except Exception: + pass + return None + + +def _as_float(value: Any) -> Optional[float]: + """Return *value* as a float when it is a numeric PDF destination coordinate.""" + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _best_destination_match( + entry: _OutlineEntry, + candidates: list[Title], + reader: PdfReader, +) -> tuple[Optional[Title], float]: + """Find the unassigned candidate nearest the outline destination on the target page.""" + if not entry.has_destination_coordinates: + return None, 0.0 + + best_title: Optional[Title] = None + best_ratio = 0.0 + best_distance = float("inf") + query_lower = entry.text.lower().strip() + for title in candidates: + if title.metadata.category_depth is not None: + continue + title_text = (title.text or "").strip().lower() + if not title_text: + continue + ratio = SequenceMatcher(None, query_lower, title_text).ratio() + if ratio < _SIMILARITY_THRESHOLD: + continue + distance = _distance_to_destination(entry, title, reader) + if distance is None: + continue + if distance < best_distance or (distance == best_distance and ratio > best_ratio): + best_distance = distance + best_ratio = ratio + best_title = title + return best_title, best_ratio + + +def _distance_to_destination( + entry: _OutlineEntry, + title: Title, + reader: PdfReader, +) -> Optional[float]: + """Return distance from the outline destination point to a Title bounding box.""" + coordinates = title.metadata.coordinates + if coordinates is None or coordinates.points is None or coordinates.system is None: + return None + + destination_point = _destination_point_in_coordinate_system(entry, coordinates.system, reader) + if destination_point is None: + return None + destination_x, destination_y = destination_point + if destination_x is None and destination_y is None: + return None + + xs = [point[0] for point in coordinates.points] + ys = [point[1] for point in coordinates.points] + x_min, x_max = min(xs), max(xs) + y_min, y_max = min(ys), max(ys) + + dx = _axis_distance(destination_x, x_min, x_max) + dy = _axis_distance(destination_y, y_min, y_max) + return (dx**2 + dy**2) ** 0.5 + + +def _destination_point_in_coordinate_system( + entry: _OutlineEntry, + coordinate_system: CoordinateSystem, + reader: PdfReader, +) -> Optional[tuple[Optional[float], Optional[float]]]: + """Convert the PDF outline destination point into an element coordinate system.""" + pdf_x = entry.left if entry.left is not None else entry.right + pdf_y = entry.top if entry.top is not None else entry.bottom + if pdf_x is None and pdf_y is None: + return None + + page_space = _page_point_space(reader, entry.page_number) + if page_space is None: + return None + + converted_x, converted_y = page_space.convert_coordinates_to_new_system( + new_system=coordinate_system, + x=pdf_x if pdf_x is not None else 0, + y=pdf_y if pdf_y is not None else 0, + ) + return ( + converted_x if pdf_x is not None else None, + converted_y if pdf_y is not None else None, + ) + + +def _page_point_space(reader: PdfReader, page_number: int) -> Optional[PointSpace]: + """Return the PDF point-space for *page_number*.""" + try: + page = reader.pages[page_number - 1] + box = getattr(page, "cropbox", None) or page.mediabox + return PointSpace(width=float(box.width), height=float(box.height)) + except Exception: + return None + + +def _axis_distance(point: Optional[float], lower: float, upper: float) -> float: + """Return zero when *point* falls inside the axis range, else edge distance.""" + if point is None: + return 0.0 + if point < lower: + return lower - point + if point > upper: + return point - upper + return 0.0 + + +def _best_match(query: str, candidates: list[Title]) -> tuple[Optional[Title], float]: + """Find the unassigned candidate Title most similar to *query*. + Always evaluates all candidates (never breaks early on a weaker match). + """ + best_title: Optional[Title] = None + best_ratio = 0.0 + query_lower = query.lower().strip() + for title in candidates: + if title.metadata.category_depth is not None: + continue + title_text = (title.text or "").strip().lower() + if not title_text: + continue + ratio = SequenceMatcher(None, query_lower, title_text).ratio() + if ratio > best_ratio: + best_ratio = ratio + best_title = title + return best_title, best_ratio + + +# -- Strategy 2: Font-size clustering -- + + +def _apply_font_size_levels( + titles: list[Title], + *, + filename: str = "", + file: Optional[IO[bytes]] = None, + password: Optional[str] = None, +) -> None: + """Assign ``category_depth`` based on font-size clusters. + Ranks distinct sizes largest-first -> depth 0, 1, ... + """ + try: + page_font_data = _extract_page_font_data(filename=filename, file=file, password=password) + except Exception: + logger.debug("Font-size analysis failed, skipping fallback.", exc_info=True) + return + if not page_font_data: + return + + # Map ALL titles to font sizes (for full size->depth mapping). + all_title_sizes: dict[int, float] = {} + for idx, title in enumerate(titles): + page_num = title.metadata.page_number + if page_num is None or page_num not in page_font_data: + continue + title_text = (title.text or "").strip() + if not title_text: + continue + best_size: Optional[float] = None + best_ratio = 0.0 + for box_text, dominant_size in page_font_data[page_num]: + ratio = SequenceMatcher(None, title_text.lower(), box_text.lower()).ratio() + if ratio > best_ratio: + best_ratio = ratio + best_size = dominant_size + if best_size is not None and best_ratio >= 0.5: + all_title_sizes[idx] = best_size + if not all_title_sizes: + return + + distinct_sizes = sorted(set(all_title_sizes.values()), reverse=True) + size_to_depth = {size: min(i, _MAX_HEADING_DEPTH - 1) for i, size in enumerate(distinct_sizes)} + for idx, size in all_title_sizes.items(): + if titles[idx].metadata.category_depth is None: + titles[idx].metadata.category_depth = size_to_depth[size] + + +def _extract_page_font_data( + *, + filename: str = "", + file: Optional[IO[bytes]] = None, + password: Optional[str] = None, +) -> dict[int, list[tuple[str, float]]]: + """Extract per-page (text, dominant_font_size) pairs using pdfminer.""" + page_font_data: dict[int, list[tuple[str, float]]] = {} + if filename: + with open(filename, "rb") as fp: + _collect_font_data(fp, password, page_font_data) + elif file: + if hasattr(file, "seek"): + file.seek(0) + _collect_font_data(file, password, page_font_data) + return page_font_data + + +def _collect_font_data( + fp: IO[bytes], + password: Optional[str], + page_font_data: dict[int, list[tuple[str, float]]], +) -> None: + """Iterate pdfminer pages and collect (text, dominant_font_size) pairs.""" + from pdfminer.layout import LTChar, LTTextBox, LTTextLine + + from unstructured.partition.pdf_image.pdfminer_utils import open_pdfminer_pages_generator + + for page_number, (_page, page_layout) in enumerate( + open_pdfminer_pages_generator(cast(BinaryIO, fp), password=password), + start=1, + ): + page_entries: list[tuple[str, float]] = [] + for element in page_layout: + if not isinstance(element, LTTextBox): + continue + text = element.get_text().strip() + if not text: + continue + char_sizes: list[float] = [] + for line in element: + if isinstance(line, LTTextLine): + for char in line: + if isinstance(char, LTChar) and char.size > 0: + char_sizes.append(char.size) + if char_sizes: + page_entries.append((text, _mode_font_size(char_sizes))) + page_font_data[page_number] = page_entries + + +def _mode_font_size(sizes: list[float]) -> float: + """Return the most common font size (rounded to 1 decimal).""" + counts: dict[float, int] = {} + for s in sizes: + r = round(s, 1) + counts[r] = counts.get(r, 0) + 1 + return max(counts, key=counts.get) # type: ignore[arg-type] + + +# -- Helpers -- + + +def _open_reader( + *, + filename: str = "", + file: Optional[IO[bytes]] = None, + password: Optional[str] = None, +) -> Optional[PdfReader]: + """Open a PdfReader, returning None on failure.""" + try: + if filename: + return PdfReader(filename, password=password) if password else PdfReader(filename) + elif file: + if hasattr(file, "seek"): + file.seek(0) + return PdfReader(file, password=password) if password else PdfReader(file) + return None + except Exception: + logger.debug("Could not open PdfReader for heading inference.", exc_info=True) + return None