documentExtractionag051/ExtractDocument

1import logging2from .invoice_post_processing import extract_invoice_tables3 4logging.basicConfig(level=logging.INFO)5logger = logging.getLogger(__name__)6 7"""8Custom Post-Processing Module9==============================10 11This file lets you add custom post-processing steps for different document layouts.12 13HOW TO USE:141. Write your own Python functions in the `custom/` folder.15   - Each function receives:16       - ocr_elems: list of TextRect objects from OCR17       - validation_data: current validation data dict18       - page_dimensions: dict of page dimensions {page_num: (width, height)}19   - Your function must return the updated validation_data20 212. Register your functions in the LAYOUT_PROCESSORS dictionary.22   - Use the layout name as the key23   - The value is a LIST of functions to run in order24   25   Example:26       from custom.my_layout import fix_totals, fix_header27       28       LAYOUT_PROCESSORS = {29           "MY_LAYOUT": [fix_header, fix_totals]30       }31"""32 33#34# Layout -> List of handler functions35#36LAYOUT_PROCESSORS = {37    "Invoice": [extract_invoice_tables]38}39 40 41# -----------------------------------------------------------------------------42# Main routing logic43# -----------------------------------------------------------------------------44def do_custom_post_processing(45    input_data,46    validation_data,47    ocr_elems,48    column_headers=None,49    anchor_column=None,50    anchor_pattern=None,51    advanced_options=None52):53    """54    Main entry point for layout-based post-processing.55    56    Args:57        input_data (dict): Original input JSON data58        validation_data (dict): Validation data structure59        ocr_elems (list): List of TextRect objects60        column_headers (dict): Optional custom column headers61        anchor_column (str): Optional anchor column name62        anchor_pattern (str): Optional anchor pattern regex63        advanced_options (dict): Optional advanced options64    65    Returns:66        dict: Updated validation data67    """68    # Extract page dimensions from input data69    page_dimensions = _extract_page_dimensions(input_data)70    71    # For now, always run invoice table extraction72    # You can add layout detection logic here if needed73    handler_list = [extract_invoice_tables]74    75    logger.info(f"Running {len(handler_list)} handler(s)")76    77    for handler in handler_list:78        logger.info(f"Running handler: {handler.__name__}")79        80        try:81            validation_data = handler(82                ocr_elems=ocr_elems,83                validation_data=validation_data,84                page_dimensions=page_dimensions,85                column_headers=column_headers,86                anchor_column=anchor_column,87                anchor_pattern=anchor_pattern,88                advanced_options=advanced_options89            )90        except Exception as e:91            logger.error(f"Handler {handler.__name__} failed: {e}")92            import traceback93            traceback.print_exc()94    95    return validation_data96 97 98def _extract_page_dimensions(data):99    """100    Extract page dimensions from input data.101    102    Supports both simplified schema and original engine output format.103    """104    # Simplified schema105    if "pages" in data and isinstance(data.get("pages"), list):106        pages = data["pages"]107        if pages and "pageNum" in pages[0]:108            return {109                p.get("pageNum", idx + 1): (p.get("width", 0), p.get("height", 0))110                for idx, p in enumerate(pages)111            }112    113    # Original engine output114    if "engineData" in data:115        pages = (116            data.get("engineData", {})117            .get("imagePreprocessingResult", {})118            .get("pages", [])119        )120        return {121            idx + 1: (p.get("width", 0), p.get("height", 0))122            for idx, p in enumerate(pages)123        }124    125    logger.warning("Could not extract page dimensions")126    return {}127