# pyright: ignore """ Docling Document Header Consolidator with Image Support Processes PDF documents using Docling with the same options as CLI: --image-export-mode referenced --table-mode accurate Then consolidates consecutive headers in the document object before export. Usage: python docling_consolidator.py input.pdf [output_dir] """ import argparse from pathlib import Path from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode from docling.document_converter import DocumentConverter, PdfFormatOption from docling_core.types.doc.base import ImageRefMode def setup_docling_converter(): """ Setup DocumentConverter with CLI-equivalent options: - image-export-mode referenced - table-mode accurate """ # Configure pipeline options to match CLI behavior pipeline_options = PdfPipelineOptions() # Enable table structure with accurate mode (equivalent to --table-mode accurate) pipeline_options.do_table_structure = True pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE # Enable image generation for referenced export mode # (equivalent to --image-export-mode referenced) pipeline_options.generate_page_images = True pipeline_options.generate_picture_images = True pipeline_options.images_scale = 2.0 # Higher resolution images # Create converter with PDF-specific options converter = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options) } ) return converter def identify_header_elements(document): """ Identify elements that are headers in the document structure. Returns: List of (element, index) tuples for header elements """ header_elements = [] # Iterate through document body children to find headers for i, child_ref in enumerate(document.body.children): # Resolve the reference to get the actual element element = child_ref.resolve(document) if element and hasattr(element, "text") and element.text: text = element.text.strip() # Heuristic: identify headers by common patterns # This is a simplified approach - you might need to refine based on your documents is_header = ( # Short lines that look like titles len(text.split(" ")) <= 2 and len(text) < 200 and ( # All caps patterns text.isupper() or # Title case patterns text.istitle() or # Contains common header keywords any( keyword in text.upper() for keyword in [ "REVENUE", "INCOME", "QUARTER", "RESULTS", "SUMMARY", "OPERATIONS", "FINANCIAL", "SEGMENT", "OUTLOOK", "EARNINGS", "CASH FLOW", "BALANCE SHEET", ] ) or # Ends with colon (often section headers) text.endswith(":") ) ) if is_header: header_elements.append((element, i)) return header_elements def consolidate_consecutive_headers(document): """ Consolidate consecutive header elements using ^ separator. Modifies the document in place. """ print("=== Header Consolidation Process ===") header_elements = identify_header_elements(document) if not header_elements: print("No header elements identified for consolidation") return 0 print(f"Found {len(header_elements)} potential header elements") # Group consecutive headers consolidated_count = 0 i = 0 elements_to_remove = [] while i < len(header_elements): current_element, current_index = header_elements[i] consecutive_headers = [current_element.text.strip()] consecutive_indices = [current_index] # Look ahead for consecutive headers j = i + 1 while j < len(header_elements): next_element, next_index = header_elements[j] # Check if indices are consecutive (or nearly consecutive) # Allow for small gaps that might be empty elements if next_index - consecutive_indices[-1] <= 3: consecutive_headers.append(next_element.text.strip()) consecutive_indices.append(next_index) j += 1 else: break # If we found multiple consecutive headers, consolidate them if len(consecutive_headers) > 1: consolidated_text = " ^ ".join(consecutive_headers) print(f"Consolidating {len(consecutive_headers)} headers:") for idx, header in enumerate(consecutive_headers): print(f" {idx + 1}: {header[:50]}{'...' if len(header) > 50 else ''}") print( f" -> {consolidated_text[:100]}{'...' if len(consolidated_text) > 100 else ''}" ) # Update the first element with consolidated text current_element.text = consolidated_text # Mark subsequent elements for removal for remove_idx in consecutive_indices[1:]: elements_to_remove.append(remove_idx) consolidated_count += len(consecutive_headers) - 1 i = j if j > i + 1 else i + 1 # Remove the marked elements (in reverse order to maintain indices) elements_to_remove.sort(reverse=True) for remove_idx in elements_to_remove: if remove_idx < len(document.body.children): del document.body.children[remove_idx] print(f"Consolidated {consolidated_count} header elements") return consolidated_count def save_document_with_images(document, output_dir: Path, doc_filename: str): """ Save document with images using referenced mode (equivalent to CLI --image-export-mode referenced) """ output_dir.mkdir(parents=True, exist_ok=True) # Save page images (equivalent to --image-export-mode referenced) image_count = 0 for page_no, page in document.pages.items(): # noqa: B007 if hasattr(page, "image") and page.image and hasattr(page.image, "pil_image"): page_image_filename = output_dir / f"{doc_filename}-page-{page.page_no}.png" with page_image_filename.open("wb") as fp: page.image.pil_image.save(fp, format="PNG") image_count += 1 # Save figure and table images table_counter = 0 picture_counter = 0 for element, _level in document.iterate_items(): # Handle table images if hasattr(element, "__class__") and "Table" in element.__class__.__name__: if ( hasattr(element, "image") and element.image and hasattr(element.image, "pil_image") ): table_counter += 1 table_image_filename = ( output_dir / f"{doc_filename}-table-{table_counter}.png" ) with table_image_filename.open("wb") as fp: element.image.pil_image.save(fp, format="PNG") image_count += 1 # Handle picture images elif hasattr(element, "__class__") and "Picture" in element.__class__.__name__: if ( hasattr(element, "image") and element.image and hasattr(element.image, "pil_image") ): picture_counter += 1 picture_image_filename = ( output_dir / f"{doc_filename}-picture-{picture_counter}.png" ) with picture_image_filename.open("wb") as fp: element.image.pil_image.save(fp, format="PNG") image_count += 1 # Save markdown with referenced images markdown_filename = output_dir / f"{doc_filename}.md" with open(markdown_filename, "w", encoding="utf-8") as f: f.write(document.export_to_markdown()) # Save HTML with referenced images html_filename = output_dir / f"{doc_filename}.html" document.save_as_html(html_filename, image_mode=ImageRefMode.REFERENCED) print(f"Saved {image_count} images to {output_dir}") print(f"Saved markdown: {markdown_filename}") print(f"Saved HTML: {html_filename}") def process_document(pdf_path: Path, output_dir: Path = None): """ Process a PDF document with header consolidation and image export. """ if output_dir is None: output_dir = Path("output") print(f"Processing: {pdf_path}") print(f"Output directory: {output_dir}") # Setup converter with CLI-equivalent options converter = setup_docling_converter() # Convert document print("Converting document...") result = converter.convert(pdf_path.as_posix()) document = result.document print( f"Document loaded: {len(document.pages)} pages, {len(document.tables)} tables" ) # Consolidate headers consolidated_count = consolidate_consecutive_headers(document) # Save with images doc_filename = pdf_path.stem save_document_with_images(document, output_dir, doc_filename) print("Processing complete!") print(f"- Headers consolidated: {consolidated_count}") print(f"- Output saved to: {output_dir}") return document def main(): """Command line interface""" parser = argparse.ArgumentParser( description="Process PDF with Docling, consolidate headers, and export with images" ) parser.add_argument("pdf_path", help="Path to PDF file") parser.add_argument("output_dir", nargs="?", help="Output directory (optional)") args = parser.parse_args() pdf_path = Path(args.pdf_path) output_dir = Path(args.output_dir) if args.output_dir else None if not pdf_path.exists(): print(f"Error: File {pdf_path} does not exist") return 1 try: process_document(pdf_path, output_dir) return 0 except Exception as e: print(f"Error processing document: {e}") return 1 if __name__ == "__main__": # Example usage when run directly if len(__import__("sys").argv) == 1: print("Example usage:") print("python docling_consolidator.py documents/TMUS_992_Q423.pdf output/") print(" This script replicates:") print( "docling --to md --image-export-mode referenced --table-mode accurate documents/TMUS_992_Q423.pdf" ) print("Plus adds header consolidation with ^ separator") else: exit(main())