Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from collections.abc import Iterator | |
| from pathlib import Path | |
| from typing import ClassVar | |
| import numpy as np | |
| import pdfplumber | |
| import pypdfium2 as pdfium | |
| from rapidocr_onnxruntime import RapidOCR | |
| from universal_parser.core.router import register | |
| from universal_parser.core.schema import BBox, Element | |
| from universal_parser.core.sniffer import FileType | |
| from universal_parser.extractors.base import BaseExtractor | |
| from universal_parser.extractors.pdf.tables import PDFTableExtractor | |
| class NativePDFExtractor(BaseExtractor): | |
| """Extractor for PDF files (both native text and scanned pages). | |
| Handles: | |
| - Multi-column reading order via column-band clustering | |
| - Dynamic header hierarchy using font-size percentiles | |
| - Bordered & borderless table extraction integrated into the reading flow | |
| - Automatic OCR fallback (via RapidOCR ONNX) for scanned pages with no text | |
| """ | |
| supported_types: ClassVar[list[FileType]] = [FileType.PDF] | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self._ocr = RapidOCR() | |
| def stream(self, path: str | Path) -> Iterator[Element]: | |
| """Stream elements from a PDF page by page.""" | |
| path_str = str(path) | |
| try: | |
| plumber_doc = pdfplumber.open(path_str) | |
| pdfium_doc = pdfium.PdfDocument(path_str) | |
| except Exception: # noqa: BLE001 | |
| return | |
| try: | |
| # Collect font sizes for header hierarchy | |
| font_sizes = self._collect_font_sizes(plumber_doc, max_pages=10) | |
| thresholds = self._compute_header_thresholds(font_sizes) | |
| for page_num in range(len(plumber_doc.pages)): | |
| plumber_page = plumber_doc.pages[page_num] | |
| page_width = float(plumber_page.width) | |
| # A. Extract tables and their bounding boxes using pdfplumber | |
| try: | |
| table_extractor = PDFTableExtractor(plumber_page) | |
| page_tables = table_extractor.extract_tables() | |
| except Exception: # noqa: BLE001 | |
| page_tables = [] | |
| table_bboxes = [table["bbox"] for table in page_tables] | |
| # B. Extract native text words/spans | |
| words = plumber_page.extract_words(extra_attrs=["size"], keep_blank_chars=False) | |
| spans = [] | |
| # Group adjacent words on same line into phrase spans | |
| current_line: list[dict] = [] | |
| for word in words: | |
| bbox = ( | |
| float(word["x0"]), | |
| float(word["top"]), | |
| float(word["x1"]), | |
| float(word["bottom"]), | |
| ) | |
| if self._is_inside_any_bbox(bbox, table_bboxes): | |
| continue | |
| if not current_line: | |
| current_line.append(word) | |
| else: | |
| prev = current_line[-1] | |
| # Check if on same line (vertical overlap) and close horizontal gap | |
| same_line = abs(float(word["top"]) - float(prev["top"])) < 4.0 | |
| gap = float(word["x0"]) - float(prev["x1"]) | |
| if same_line and gap < 12.0: | |
| current_line.append(word) | |
| else: | |
| span_text = " ".join(w["text"] for w in current_line).strip() | |
| if span_text: | |
| spans.append( | |
| { | |
| "is_table": False, | |
| "text": span_text, | |
| "size": round(float(current_line[0].get("size", 10.0)), 1), | |
| "bbox": ( | |
| float(current_line[0]["x0"]), | |
| float(min(w["top"] for w in current_line)), | |
| float(current_line[-1]["x1"]), | |
| float(max(w["bottom"] for w in current_line)), | |
| ), | |
| } | |
| ) | |
| current_line = [word] | |
| if current_line: | |
| span_text = " ".join(w["text"] for w in current_line).strip() | |
| if span_text: | |
| spans.append( | |
| { | |
| "is_table": False, | |
| "text": span_text, | |
| "size": round(float(current_line[0].get("size", 10.0)), 1), | |
| "bbox": ( | |
| float(current_line[0]["x0"]), | |
| float(min(w["top"] for w in current_line)), | |
| float(current_line[-1]["x1"]), | |
| float(max(w["bottom"] for w in current_line)), | |
| ), | |
| } | |
| ) | |
| # C. Wrap extracted tables | |
| for table in page_tables: | |
| spans.append( | |
| { | |
| "is_table": True, | |
| "table_data": table["data"], | |
| "confidence": table["confidence"], | |
| "bbox": table["bbox"], | |
| } | |
| ) | |
| # D. SCANNED PAGE FALLBACK: If page has NO native text or tables, run OCR | |
| if not spans: | |
| if page_num < len(pdfium_doc): | |
| yield from self._ocr_scanned_page( | |
| pdfium_doc[page_num], | |
| page_num + 1, | |
| page_width, | |
| float(plumber_page.height), | |
| ) | |
| continue | |
| # Step 3: Sort in natural reading order | |
| ordered_elements = self._sort_reading_order(spans, page_width) | |
| # Step 4: Yield sorted elements | |
| for el in ordered_elements: | |
| bbox_coords = el["bbox"] | |
| if el["is_table"]: | |
| headers = el["table_data"].headers | |
| rows = el["table_data"].rows | |
| md_header = "| " + " | ".join(headers) + " |" | |
| md_separator = "| " + " | ".join(["---"] * len(headers)) + " |" | |
| md_rows = ["| " + " | ".join(row) + " |" for row in rows] | |
| markdown_repr = "\n".join([md_header, md_separator] + md_rows) | |
| yield Element( | |
| type="table", | |
| page=page_num + 1, | |
| bbox=BBox( | |
| x0=bbox_coords[0], | |
| y0=bbox_coords[1], | |
| x1=bbox_coords[2], | |
| y1=bbox_coords[3], | |
| ), | |
| data=el["table_data"], | |
| markdown_repr=markdown_repr, | |
| confidence=el["confidence"], | |
| ) | |
| else: | |
| size = el["size"] | |
| text = el["text"] | |
| el_type = "paragraph" | |
| level = None | |
| if size >= thresholds["h1"]: | |
| el_type = "heading" | |
| level = 1 | |
| elif size >= thresholds["h2"]: | |
| el_type = "heading" | |
| level = 2 | |
| elif size >= thresholds["h3"]: | |
| el_type = "heading" | |
| level = 3 | |
| yield Element( | |
| type=el_type, | |
| level=level, | |
| text=text, | |
| page=page_num + 1, | |
| bbox=BBox( | |
| x0=bbox_coords[0], | |
| y0=bbox_coords[1], | |
| x1=bbox_coords[2], | |
| y1=bbox_coords[3], | |
| ), | |
| ) | |
| finally: | |
| pdfium_doc.close() | |
| plumber_doc.close() | |
| def _ocr_scanned_page( | |
| self, page: pdfium.PdfPage, page_num: int, page_w: float, page_h: float | |
| ) -> Iterator[Element]: | |
| """Render a scanned PDF page to an image and run RapidOCR.""" | |
| try: | |
| scale = 200.0 / 72.0 | |
| bitmap = page.render(scale=scale) | |
| pil_img = bitmap.to_pil().convert("RGB") | |
| img_np = np.array(pil_img) | |
| bitmap.close() | |
| ocr_results, _ = self._ocr(img_np) | |
| if not ocr_results: | |
| return | |
| scale_x = page_w / pil_img.width | |
| scale_y = page_h / pil_img.height | |
| for item in ocr_results: | |
| dt_boxes, text, score = item | |
| clean_text = text.strip() | |
| if not clean_text: | |
| continue | |
| pts = np.array(dt_boxes, dtype=np.float32) | |
| x0 = float(np.min(pts[:, 0])) * scale_x | |
| y0 = float(np.min(pts[:, 1])) * scale_y | |
| x1 = float(np.max(pts[:, 0])) * scale_x | |
| y1 = float(np.max(pts[:, 1])) * scale_y | |
| yield Element( | |
| type="paragraph", | |
| text=clean_text, | |
| page=page_num, | |
| bbox=BBox(x0=x0, y0=y0, x1=x1, y1=y1), | |
| markdown_repr=clean_text, | |
| confidence=round(float(score), 3), | |
| ) | |
| except Exception: # noqa: BLE001 | |
| return | |
| def _is_inside_any_bbox( | |
| self, | |
| span_bbox: tuple[float, float, float, float], | |
| table_bboxes: list[tuple[float, float, float, float]], | |
| ) -> bool: | |
| """Check if a text span falls inside any table bounding box.""" | |
| sx0, sy0, sx1, sy1 = span_bbox | |
| for tx0, ty0, tx1, ty1 in table_bboxes: | |
| if sx0 >= (tx0 - 2) and sy0 >= (ty0 - 2) and sx1 <= (tx1 + 2) and sy1 <= (ty1 + 2): | |
| return True | |
| return False | |
| def _collect_font_sizes(self, doc: pdfplumber.PDF, max_pages: int) -> list[float]: | |
| """Collect all font sizes from the start of the document.""" | |
| sizes = [] | |
| pages_to_scan = min(len(doc.pages), max_pages) | |
| for i in range(pages_to_scan): | |
| page = doc.pages[i] | |
| words = page.extract_words(extra_attrs=["size"]) | |
| for w in words: | |
| sizes.append(float(w.get("size", 10.0))) | |
| return sizes | |
| def _compute_header_thresholds(self, font_sizes: list[float]) -> dict[str, float]: | |
| """Compute font size cutoffs for H1, H2, H3 using percentiles.""" | |
| if not font_sizes: | |
| return {"h1": 16.0, "h2": 14.0, "h3": 12.0} | |
| arr = np.array(font_sizes) | |
| h1_val = float(np.percentile(arr, 95)) | |
| h2_val = float(np.percentile(arr, 85)) | |
| h3_val = float(np.percentile(arr, 75)) | |
| median = float(np.median(arr)) | |
| h1_val = max(h1_val, median + 3.0) | |
| h2_val = max(h2_val, median + 1.5) | |
| h3_val = max(h3_val, median + 0.5) | |
| return {"h1": h1_val, "h2": h2_val, "h3": h3_val} | |
| def _sort_reading_order(self, elements: list[dict], page_width: float) -> list[dict]: | |
| """Sort elements to preserve natural multi-column reading order.""" | |
| if len(elements) < 5: | |
| return sorted(elements, key=lambda e: (e["bbox"][1], e["bbox"][0])) | |
| midpoint = page_width / 2.0 | |
| left_col = [] | |
| right_col = [] | |
| full_width = [] | |
| for el in elements: | |
| x0, _, x1, _ = el["bbox"] | |
| if x1 <= (midpoint + 15): | |
| left_col.append(el) | |
| elif x0 >= (midpoint - 15): | |
| right_col.append(el) | |
| else: | |
| full_width.append(el) | |
| def key_y(e): | |
| return (e["bbox"][1], e["bbox"][0]) | |
| left_sorted = sorted(left_col, key=key_y) | |
| right_sorted = sorted(right_col, key=key_y) | |
| if len(left_sorted) >= 2 and len(right_sorted) >= 2: | |
| col_top = min(left_sorted[0]["bbox"][1], right_sorted[0]["bbox"][1]) | |
| col_bottom = max(left_sorted[-1]["bbox"][3], right_sorted[-1]["bbox"][3]) | |
| top_full = [e for e in full_width if e["bbox"][3] <= col_top + 10] | |
| bottom_full = [e for e in full_width if e["bbox"][1] >= col_bottom - 10] | |
| mid_full = [e for e in full_width if e not in top_full and e not in bottom_full] | |
| return ( | |
| sorted(top_full, key=key_y) | |
| + sorted(mid_full, key=key_y) | |
| + left_sorted | |
| + right_sorted | |
| + sorted(bottom_full, key=key_y) | |
| ) | |
| return sorted(elements, key=key_y) | |