296 lines
14 KiB
Python
296 lines
14 KiB
Python
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
|