documentExtractionag051/ExtractDocument

1# pp_handler.py > post_process_data2import argparse3import json4import logging5import os6 7from utils.pp_utils import open_json, create_text_rects8from custom_processor import do_custom_post_processing9 10# Set up logging11logging.basicConfig(level=logging.INFO, format='%(levelname)s:%(name)s:%(message)s')12logger = logging.getLogger(__name__)13 14"""15Entry script for table extraction post-processing.16 17Supports two input formats:181. Simplified schema (pages + ocrBlocks)192. Original DA engine output (engineData structure)20"""21 22 23def create_validation_data_from_input(input_data):24    """25    Create validation data structure from input JSON.26    27    Supports both simplified schema and original engine output format.28    29    Args:30        input_data (dict): Input JSON data31    32    Returns:33        tuple: (validation_data, ocr_elems, page_dimensions)34    """35    # Check if this is simplified schema or original engine output36    if "ocrBlocks" in input_data:37        # Simplified schema38        return _parse_simplified_schema(input_data)39    elif "engineData" in input_data:40        # Original engine output41        return _parse_engine_output(input_data)42    else:43        raise ValueError("Invalid input format. Expected 'ocrBlocks' or 'engineData' in input JSON")44 45 46def _parse_simplified_schema(input_data):47    """48    Parse simplified input schema.49    50    Expected format:51    {52        "pages": [{"pageNum": 1, "width": 2479, "height": 3508, ...}],53        "ocrBlocks": [{"pageNum": 1, "blockType": "LINE", "text": "...", "geometry": {...}}, ...]54    }55    """56    # Extract page dimensions57    page_dimensions = {}58    for page in input_data.get("pages", []):59        page_num = page.get("pageNum", len(page_dimensions) + 1)60        page_dimensions[page_num] = (page.get("width", 0), page.get("height", 0))61    62    # Create validation data structure63    validation_data = {64        "pages": [65            {"width": page.get("width", 0), "height": page.get("height", 0)}66            for page in input_data.get("pages", [])67        ],68        "fields": {},69        "tables": {"table": []}  # Empty, will be filled by post-processor70    }71    72    # Create OCR TextRect objects73    ocr_elems = create_text_rects(input_data.get("ocrBlocks", []))74    75    logger.info(f"Parsed simplified schema: {len(page_dimensions)} pages, {len(ocr_elems)} OCR elements")76    77    return validation_data, ocr_elems, page_dimensions78 79 80def _parse_engine_output(input_data):81    """82    Parse original DA engine output format.83    """84    # Extract page dimensions85    pages = (86        input_data.get("engineData", {})87        .get("imagePreprocessingResult", {})88        .get("pages", [])89    )90    91    page_dimensions = {92        idx + 1: (p.get("width", 0), p.get("height", 0))93        for idx, p in enumerate(pages)94    }95    96    # Create validation data structure97    validation_data = {98        "pages": [99            {"width": page.get("width", 0), "height": page.get("height", 0)}100            for page in pages101        ],102        "fields": {},103        "tables": {"table": []}  # Empty, will be filled by post-processor104    }105    106    # Create OCR TextRect objects107    ocr_blocks = input_data.get("engineData", {}).get("ocrResult", {}).get("blocks", [])108    ocr_elems = create_text_rects(ocr_blocks)109    110    logger.info(f"Parsed engine output: {len(page_dimensions)} pages, {len(ocr_elems)} OCR elements")111    112    return validation_data, ocr_elems, page_dimensions113 114 115def post_process_data(input_dict, standalone=False):116    """117    Execute post-processing using the input JSON.118    119    Args:120        input_dict (dict):121            - 'input_dir' (str): Directory containing input JSON.122            - 'column_headers' (optional): Dict of column field names to header text.123            - 'anchor_column' (optional): Column name for row detection.124            - 'anchor_pattern' (optional): Regex pattern for anchor validation.125            - 'advanced_options' (optional): Dict with special logic options.126        standalone (bool): If True, do not delete input files.127    128    Returns:129        str: Final validation_data JSON (stringified).130    """131    input_dir = input_dict["input_dir"]132    133    # Get optional dynamic parameters134    column_headers = input_dict.get("column_headers")135    anchor_column = input_dict.get("anchor_column")136    anchor_pattern = input_dict.get("anchor_pattern")137    advanced_options = input_dict.get("advanced_options")138    139    files = os.listdir(input_dir)140    if not files:141        logger.info(f"No files found in directory: {input_dir}")142        return None143    144    for file in files:145        try:146            if not file.endswith(".json"):147                continue148            149            path = os.path.join(input_dir, file)150            input_data = open_json(path)151            152            # Parse input and create validation data153            validation_data, ocr_elems, page_dimensions = create_validation_data_from_input(input_data)154            155            logger.info(f"Created {len(ocr_elems)} OCR TextRect objects")156            157            # Apply custom post-processing logic158            validation_data = do_custom_post_processing(159                input_data,160                validation_data,161                ocr_elems,162                column_headers=column_headers,163                anchor_column=anchor_column,164                anchor_pattern=anchor_pattern,165                advanced_options=advanced_options166            )167            168            print(f"validation_data: {validation_data}")169            170            # Cleanup171            if not standalone:172                os.remove(path)173                logger.info(f"Removed processed file: {path}")174            175            return json.dumps(validation_data)176        177        except Exception as e:178            logger.error(f"Error processing file {file}: {e}")179            import traceback180            traceback.print_exc()181            182            if not standalone:183                os.remove(path)184            185            return None186    187    return None188 189 190def run_post_processing(input_dir: str, standalone: bool = True,191                       column_headers: dict = None, anchor_column: str = None,192                       anchor_pattern: str = None, advanced_options: dict = None):193    """194    Wrapper for running post-processing directly from a directory.195    196    Args:197        input_dir (str): Folder path containing input JSON files.198        standalone (bool): If True, keeps input files.199        column_headers (dict): Optional column headers mapping.200        anchor_column (str): Optional anchor column name.201        anchor_pattern (str): Optional anchor pattern regex.202        advanced_options (dict): Optional advanced options for special logic.203    204    Returns:205        str: JSON string of validation_data or None.206    """207    return post_process_data({208        "input_dir": input_dir,209        "column_headers": column_headers,210        "anchor_column": anchor_column,211        "anchor_pattern": anchor_pattern,212        "advanced_options": advanced_options213    }, standalone=standalone)214 215 216def parse_args():217    """Parse command line arguments."""218    parser = argparse.ArgumentParser(219        description="Run the table extraction post-processor."220    )221    parser.add_argument(222        "--input_dir",223        type=str,224        required=True,225        help="Directory containing input JSON files."226    )227    parser.add_argument(228        "--standalone",229        action="store_true",230        help="Enable standalone mode (preserve input files)."231    )232    233    return parser.parse_args()234 235 236if __name__ == "__main__":237    # Example usage - update path as needed238    input_dir = "C:\\Users\\...\\Documents\\...\\Customer_Issue1\\CS\\CS-35329\\ceva_post_processor\\ceva_post_processor"239    run_post_processing(input_dir, standalone=True)240