from __future__ import annotations import uuid from typing import Any from sqlalchemy.orm import Session from app.core.logging_config import get_logger from app.repositories.document_repository import DocumentRepository from app.repositories.template_repository import TemplateRepository from app.services.matching_service import MatchingService logger = get_logger(__name__) class ExtractionService: """Extract document values using matched template coordinates and structures.""" def __init__(self, db: Session) -> None: self.db = db self.matching_service = MatchingService(db) self.doc_repo = DocumentRepository(db) self.template_repo = TemplateRepository(db) def extract_document_data(self, document_id: uuid.UUID) -> dict[str, Any]: """Perform template matching on the document and extract structured field values.""" document = self.doc_repo.get_with_pages(document_id) if not document: raise ValueError(f"Document '{document_id}' not found") # 1. Match document against existing templates matches = self.matching_service.match_document(document_id, min_confidence=0.75) if not matches: logger.info("extraction_failed_no_match", document_id=str(document_id)) return { "template_matched": False, "template_id": None, "template_name": None, "confidence_score": 0.0, "extracted_data": {}, } best_match = matches[0] template = self.template_repo.get_by_id(best_match.format_id) if not template: raise ValueError(f"Matched template '{best_match.format_id}' not found") logger.info( "extraction_matched_template", document_id=str(document_id), template_id=str(template.id), template_name=template.name, score=best_match.confidence_score, ) extracted_data: dict[str, Any] = {} # Load all mapped regions for the template regions = [r for r in template.regions if r.region_type == "field_mapping"] regions_by_field: dict[str, list[Any]] = {} for r in regions: field_name = r.content.get("field_name") if r.content else None if field_name: regions_by_field.setdefault(field_name, []).append(r) # 2. Divide fields into Scalar vs Table Column types scalar_cells = [cell for cell in template.cells if cell.data_type != "TABLE_COLUMN"] table_column_cells = [cell for cell in template.cells if cell.data_type == "TABLE_COLUMN"] # 3. Extract Scalar Fields for cell in scalar_cells: field_name = cell.field_name if not field_name: continue field_regions = regions_by_field.get(field_name, []) if not field_regions: extracted_data[field_name] = "" continue extracted_values = [] extracted_block_ids = set() sorted_regions = sorted(field_regions, key=lambda r: (r.page_number, r.sequence, r.y, r.x)) for region in sorted_regions: page = next((p for p in document.pages if p.page_number == region.page_number), None) if not page: continue best_block = self._find_best_overlapping_block(page, region) if best_block and best_block.id not in extracted_block_ids: extracted_block_ids.add(best_block.id) val = best_block.text.strip() if val: extracted_values.append(val) extracted_data[field_name] = " ".join(extracted_values) # 4. Extract Table Column Fields if table_column_cells: table_column_names = {cell.field_name for cell in table_column_cells if cell.field_name} table_regions = [ r for r in regions if r.content and r.content.get("field_name") in table_column_names ] if table_regions: # Determine vertical boundaries of the table area table_start_y = min((r.y for r in table_regions), default=0.0) table_start_y = max(0.0, table_start_y - 10.0) # subtract buffer # Process page where the table coordinates are mapped page_number = min((r.page_number for r in table_regions), default=1) page = next((p for p in document.pages if p.page_number == page_number), None) # Determine table vertical end Y by finding any summary scalar fields below the table summary_keywords = {"tax", "total", "shipping", "discount", "vat", "handling", "duty", "subtotal", "grand"} summary_regions = [] for col_name, regs in regions_by_field.items(): if col_name in table_column_names: continue for r in regs: if r.y > table_start_y and any(kw in col_name.lower() for kw in summary_keywords): # Skip left-aligned metadata fields (like Shipping Method) if page and r.x < page.width * 0.4: continue summary_regions.append(r) table_end_y = min((r.y for r in summary_regions), default=99999.0) # Check if template has a footer to define the end vertical boundary footer_regions = [r for r in template.regions if r.region_type == "footer"] if footer_regions: table_end_y = min(table_end_y, min(r.y for r in footer_regions)) # Determine horizontal ranges (X span) for each column in the template col_x_spans: dict[str, tuple[float, float]] = {} for col_name in table_column_names: col_regs = [r for r in table_regions if r.content and r.content.get("field_name") == col_name] if col_regs: x_min = min(r.x for r in col_regs) x_max = max(r.x + r.width for r in col_regs) # Add a 15px margin to accommodate layout differences col_x_spans[col_name] = (max(0.0, x_min - 15.0), x_max + 15.0) else: col_x_spans[col_name] = (0.0, 0.0) if page: # Check for summary keyword text blocks (e.g. Total, Tax) as vertical boundary fallback summary_labels = {"subtotal", "total", "grand total", "tax", "shipping & handling", "discount"} for block in page.text_blocks: if block.y > table_start_y: # Skip left-aligned text blocks if block.x < page.width * 0.4: continue block_text_clean = block.text.strip().lower() if any(label in block_text_clean for label in summary_labels): table_end_y = min(table_end_y, block.y) # Collect all document blocks inside Y range candidate_blocks = [ b for b in page.text_blocks if b.y >= table_start_y and b.y < table_end_y ] # Sort blocks by Y coordinate sorted_blocks = sorted(candidate_blocks, key=lambda b: b.y) # Group blocks into rows by vertical center alignment (15px threshold) rows: list[list[Any]] = [] current_row: list[Any] = [] current_y_center = None for block in sorted_blocks: if not block.text.strip(): continue block_y_center = block.y + block.height / 2 if current_y_center is None: current_row.append(block) current_y_center = block_y_center elif abs(block_y_center - current_y_center) < 15.0: current_row.append(block) else: rows.append(current_row) current_row = [block] current_y_center = block_y_center if current_row: rows.append(current_row) # Map candidate blocks in each row to columns rows_data: list[dict[str, str]] = [] for row_blocks in rows: row_dict: dict[str, list[Any]] = {name: [] for name in table_column_names} for block in row_blocks: best_col = None best_overlap = 0.0 for col_name, (x_min, x_max) in col_x_spans.items(): overlap = max( 0.0, min(x_max, block.x + block.width) - max(x_min, block.x) ) overlap_ratio = overlap / block.width if block.width > 0 else 0.0 if overlap_ratio > 0.2 and overlap_ratio > best_overlap: best_overlap = overlap_ratio best_col = col_name if best_col: row_dict[best_col].append(block) # Construct row values row_values: dict[str, str] = {} for col_name, blocks in row_dict.items(): sorted_blocks_in_col = sorted(blocks, key=lambda b: b.x) row_values[col_name] = " ".join( b.text.strip() for b in sorted_blocks_in_col ) # Filter out table header rows is_header = False for col_name, val in row_values.items(): val_lower = val.lower() if ( val_lower == col_name.lower() or val_lower in [ "item", "items", "qty", "quantity", "price", "amount", "total", "subtotal", "description" ] ): is_header = True break # Append if not header and at least one cell has a value if not is_header and any(row_values.values()): rows_data.append(row_values) # Merge continuation lines (descriptions spanning multiple rows with empty sibling columns) merged_rows: list[dict[str, str]] = [] for row in rows_data: non_empty_cols = [k for k, v in row.items() if v.strip()] if len(merged_rows) > 0 and len(non_empty_cols) == 1: col_name = non_empty_cols[0] last_row = merged_rows[-1] if last_row.get(col_name): last_row[col_name] = last_row[col_name] + " " + row[col_name] else: last_row[col_name] = row[col_name] else: merged_rows.append(row.copy()) rows_data = merged_rows # Pivot list of rows into parallel arrays under column names for col_name in table_column_names: extracted_data[col_name] = [] for row in rows_data: for col_name in table_column_names: extracted_data[col_name].append(row.get(col_name, "")) return { "template_matched": True, "template_id": str(template.id), "template_name": template.name, "confidence_score": best_match.confidence_score, "extracted_data": extracted_data, } def _find_best_overlapping_block(self, page: Any, region: Any) -> Any | None: """Find the document text block overlapping most with the template region.""" best_block = None best_overlap = 0.0 for block in page.text_blocks: x_overlap = max( 0.0, min(region.x + region.width, block.x + block.width) - max(region.x, block.x) ) y_overlap = max( 0.0, min(region.y + region.height, block.y + block.height) - max(region.y, block.y) ) overlap = x_overlap * y_overlap if overlap > best_overlap: best_overlap = overlap best_block = block # Fallback: if no overlapping block, find the closest center-to-center if not best_block: min_dist = 100.0 # max 100px center-to-center distance for block in page.text_blocks: c_rx = region.x + region.width / 2 c_ry = region.y + region.height / 2 c_bx = block.x + block.width / 2 c_by = block.y + block.height / 2 dist = ((c_rx - c_bx) ** 2 + (c_ry - c_by) ** 2) ** 0.5 if dist < min_dist: min_dist = dist best_block = block return best_block