feat(pdf-worker): determine OCR table columns from header row
Column indices for label, code and values are taken from the header row (cell «Код») instead of fixed positions 0/1/2+. Every repeated header overwrites the layout. Value column_index is normalized to 2..n, missing cells become empty unrecognized values, tables without a header are skipped with a warning. Refs #6 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
29512f98f5
commit
8e4405f0b1
@ -1,11 +1,14 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
import uuid
|
import uuid
|
||||||
from dataclasses import dataclass, asdict
|
from dataclasses import dataclass, asdict
|
||||||
from datetime import date, datetime
|
from datetime import date, datetime
|
||||||
from decimal import Decimal, InvalidOperation
|
from decimal import Decimal, InvalidOperation
|
||||||
from typing import Optional, Literal, TypedDict
|
from typing import Optional, Literal, TypedDict
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
from src.schemas.enums import StatementType, FslRowKind
|
from src.schemas.enums import StatementType, FslRowKind
|
||||||
from .consts import codes_mapping
|
from .consts import codes_mapping
|
||||||
|
|
||||||
@ -52,6 +55,27 @@ class ExtractResult(TypedDict):
|
|||||||
skipped_files: list[tuple[str, dict]]
|
skipped_files: list[tuple[str, dict]]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class ColumnLayout:
|
||||||
|
"""Индексы значимых столбцов таблицы OCR, определённые по строке заголовка."""
|
||||||
|
label_idx: Optional[int]
|
||||||
|
code_idx: int
|
||||||
|
value_idxs: tuple[int, ...]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ParsedRow:
|
||||||
|
row: dict
|
||||||
|
cells: dict[int, dict]
|
||||||
|
layout: ColumnLayout
|
||||||
|
kind: FslRowKind
|
||||||
|
|
||||||
|
|
||||||
|
# column_index в БД: 0 - наименование, 1 - код, значения начинаются с 2
|
||||||
|
_FIRST_VALUE_COLUMN = 2
|
||||||
|
_YEAR_RE = re.compile(r"\d{4}")
|
||||||
|
|
||||||
|
|
||||||
_TABLE_TYPE_MAP: dict[str, StatementType] = {
|
_TABLE_TYPE_MAP: dict[str, StatementType] = {
|
||||||
"Бухгалтерский баланс": StatementType.BALANCE,
|
"Бухгалтерский баланс": StatementType.BALANCE,
|
||||||
"Отчет о финансовых результатах": StatementType.INCOME,
|
"Отчет о финансовых результатах": StatementType.INCOME,
|
||||||
@ -106,16 +130,61 @@ def _parse_amount(raw: str) -> tuple[Optional[Decimal], bool]:
|
|||||||
return (None, is_verified)
|
return (None, is_verified)
|
||||||
|
|
||||||
|
|
||||||
def _classify_row(cells: list[dict]) -> FslRowKind | None:
|
def _normalize_header(value) -> str:
|
||||||
if not (str(cells[1]['value']).strip()):
|
text = str(value or '').lower()
|
||||||
|
text = ''.join(c if c.isalnum() else ' ' for c in text)
|
||||||
|
return ' '.join(text.split())
|
||||||
|
|
||||||
|
|
||||||
|
def _is_code_header(text: str) -> bool:
|
||||||
|
return text == 'код' or text.startswith('код ')
|
||||||
|
|
||||||
|
|
||||||
|
def _is_value_header(text: str) -> bool:
|
||||||
|
return bool(_YEAR_RE.search(text)) or text.startswith(('на ', 'за '))
|
||||||
|
|
||||||
|
|
||||||
|
def _cells_by_index(row: dict) -> dict[int, dict]:
|
||||||
|
return {cell['column_index']: cell for cell in row['cells']}
|
||||||
|
|
||||||
|
|
||||||
|
def _cell_value(cells: dict[int, dict], idx: Optional[int]) -> str:
|
||||||
|
if idx is None or idx not in cells:
|
||||||
|
return ''
|
||||||
|
return str(cells[idx]['value'] or '')
|
||||||
|
|
||||||
|
|
||||||
|
def _build_layout(cells: dict[int, dict]) -> Optional[ColumnLayout]:
|
||||||
|
"""Возвращает раскладку столбцов, если строка является заголовком (есть ячейка «Код»)."""
|
||||||
|
headers = {idx: _normalize_header(cells[idx]['value']) for idx in sorted(cells)}
|
||||||
|
code_idx = next((idx for idx, text in headers.items() if _is_code_header(text)), None)
|
||||||
|
if code_idx is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
label_idx = next((idx for idx, text in headers.items() if 'наименование' in text), None)
|
||||||
|
if label_idx is None:
|
||||||
|
left_filled = [idx for idx, text in headers.items() if idx < code_idx and text]
|
||||||
|
label_idx = left_filled[-1] if left_filled else None
|
||||||
|
|
||||||
|
value_idxs = tuple(
|
||||||
|
idx for idx, text in headers.items()
|
||||||
|
if idx > code_idx and _is_value_header(text)
|
||||||
|
)
|
||||||
|
return ColumnLayout(label_idx=label_idx, code_idx=code_idx, value_idxs=value_idxs)
|
||||||
|
|
||||||
|
|
||||||
|
def _classify_row(cells: dict[int, dict], layout: ColumnLayout) -> FslRowKind | None:
|
||||||
|
code = _cell_value(cells, layout.code_idx).strip()
|
||||||
|
if not code:
|
||||||
return None
|
return None
|
||||||
try:
|
try:
|
||||||
int(cells[1]["value"])
|
int(code)
|
||||||
except (TypeError, ValueError):
|
except (TypeError, ValueError):
|
||||||
return FslRowKind.HEADER
|
return FslRowKind.HEADER
|
||||||
if cells[0]["value"] == "Баланс":
|
label = _cell_value(cells, layout.label_idx)
|
||||||
|
if label == "Баланс":
|
||||||
return FslRowKind.BALANCE
|
return FslRowKind.BALANCE
|
||||||
elif cells[0]["value"].lower().startswith("итого"):
|
elif label.lower().startswith("итого"):
|
||||||
return FslRowKind.TOTAL
|
return FslRowKind.TOTAL
|
||||||
return FslRowKind.ROW
|
return FslRowKind.ROW
|
||||||
|
|
||||||
@ -129,19 +198,46 @@ def _keep_digits_only(raw: str) -> str:
|
|||||||
return res.strip()
|
return res.strip()
|
||||||
|
|
||||||
|
|
||||||
def _resolve_table_type(rows) -> StatementType:
|
def _parse_table_rows(table_rows: list[dict]) -> Optional[list[ParsedRow]]:
|
||||||
|
"""Разбирает строки таблицы по раскладке из заголовка.
|
||||||
|
|
||||||
|
Раскладка берётся из строки с ячейкой «Код»; каждый повторный заголовок
|
||||||
|
перезаписывает индексы столбцов. В результат попадает только первый заголовок,
|
||||||
|
строки до первого заголовка игнорируются. Если заголовка нет, возвращает None.
|
||||||
|
"""
|
||||||
|
layout: Optional[ColumnLayout] = None
|
||||||
|
parsed_rows: list[ParsedRow] = []
|
||||||
|
|
||||||
|
for row in table_rows:
|
||||||
|
cells = _cells_by_index(row)
|
||||||
|
header_layout = _build_layout(cells)
|
||||||
|
if header_layout is not None:
|
||||||
|
if layout is None:
|
||||||
|
parsed_rows.append(ParsedRow(row, cells, header_layout, FslRowKind.HEADER))
|
||||||
|
layout = header_layout
|
||||||
|
continue
|
||||||
|
if layout is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
kind = _classify_row(cells, layout)
|
||||||
|
if kind is None or kind == FslRowKind.HEADER:
|
||||||
|
continue
|
||||||
|
parsed_rows.append(ParsedRow(row, cells, layout, kind))
|
||||||
|
|
||||||
|
return parsed_rows if layout is not None else None
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_table_type(parsed_rows: list[ParsedRow]) -> StatementType:
|
||||||
cell_codes = {
|
cell_codes = {
|
||||||
'1': 0,
|
'1': 0,
|
||||||
'2': 0
|
'2': 0
|
||||||
}
|
}
|
||||||
try:
|
try:
|
||||||
for r in rows:
|
for parsed in parsed_rows:
|
||||||
cells = list(sorted(r["cells"], key=lambda x: x['column_index']))
|
if parsed.kind == FslRowKind.HEADER:
|
||||||
kind = _classify_row(cells)
|
|
||||||
if kind is None or kind == FslRowKind.HEADER:
|
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
first_symbol = str(cells[1]['value'])[0]
|
first_symbol = _cell_value(parsed.cells, parsed.layout.code_idx).strip()[0]
|
||||||
cell_codes.setdefault(first_symbol, 0)
|
cell_codes.setdefault(first_symbol, 0)
|
||||||
cell_codes[first_symbol] += 1
|
cell_codes[first_symbol] += 1
|
||||||
if cell_codes['1'] > cell_codes['2']:
|
if cell_codes['1'] > cell_codes['2']:
|
||||||
@ -196,36 +292,34 @@ def _process_file(
|
|||||||
rows: list[DbRow] = []
|
rows: list[DbRow] = []
|
||||||
values: list[DbValue] = []
|
values: list[DbValue] = []
|
||||||
|
|
||||||
for table in data["table_structure"]:
|
for table_number, table in enumerate(data["table_structure"]):
|
||||||
|
parsed_rows = _parse_table_rows(table["rows"])
|
||||||
|
if parsed_rows is None:
|
||||||
|
logger.warning(
|
||||||
|
"В таблице не найден заголовок со столбцом «Код», таблица пропущена. "
|
||||||
|
"file_uid={file_uid}, table_number={table_number}, table_label={table_label}",
|
||||||
|
file_uid=file.get("file_uid"),
|
||||||
|
table_number=table_number,
|
||||||
|
table_label=(table.get("label") or {}).get("form_name"),
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
meta = Meta(
|
meta = Meta(
|
||||||
id=uuid.uuid4(),
|
id=uuid.uuid4(),
|
||||||
statement_type=_resolve_table_type(table["rows"]),
|
statement_type=_resolve_table_type(parsed_rows),
|
||||||
**shared,
|
**shared,
|
||||||
)
|
)
|
||||||
metas.append(asdict(meta))
|
metas.append(asdict(meta))
|
||||||
header_found = False
|
|
||||||
|
|
||||||
for row in table["rows"]:
|
for parsed in parsed_rows:
|
||||||
cells = list(sorted(row["cells"], key=lambda x: x['column_index']))
|
row, cells, layout, kind = parsed.row, parsed.cells, parsed.layout, parsed.kind
|
||||||
if len(cells) < 3:
|
code_cell = cells[layout.code_idx]
|
||||||
raise ValueError(
|
|
||||||
f"Row {row['row_index']} in file {file['file_uid']} "
|
|
||||||
f"has {len(cells)} cells (< 3)."
|
|
||||||
)
|
|
||||||
|
|
||||||
kind = _classify_row(cells)
|
|
||||||
if kind is None or (kind == FslRowKind.HEADER and header_found):
|
|
||||||
continue
|
|
||||||
elif (kind == FslRowKind.HEADER and not(header_found)):
|
|
||||||
header_found = True
|
|
||||||
label_cell = cells[0]
|
|
||||||
code_cell = cells[1]
|
|
||||||
is_line_code_valid = True
|
is_line_code_valid = True
|
||||||
line_code = _clear_value(code_cell["value"])
|
line_code = _clear_value(code_cell["value"])
|
||||||
line_name = codes_mapping.get(line_code, '')
|
line_name = codes_mapping.get(line_code, '')
|
||||||
if line_name == '':
|
if line_name == '':
|
||||||
is_line_code_valid = False
|
is_line_code_valid = False
|
||||||
line_name = label_cell['value']
|
line_name = _cell_value(cells, layout.label_idx)
|
||||||
|
|
||||||
line_code_recognized = _is_recognized(code_cell["confidence"]) and _clear_value(code_cell['value']) == line_code and is_line_code_valid
|
line_code_recognized = _is_recognized(code_cell["confidence"]) and _clear_value(code_cell['value']) == line_code and is_line_code_valid
|
||||||
|
|
||||||
@ -245,15 +339,19 @@ def _process_file(
|
|||||||
|
|
||||||
rows.append(asdict(db_row))
|
rows.append(asdict(db_row))
|
||||||
|
|
||||||
for cell in cells:
|
for value_number, source_idx in enumerate(layout.value_idxs):
|
||||||
if cell["column_index"] < 2:
|
cell = cells.get(source_idx)
|
||||||
continue
|
|
||||||
raw = cell["value"]
|
|
||||||
parsed_amount = None
|
parsed_amount = None
|
||||||
if kind == FslRowKind.HEADER:
|
if cell is None:
|
||||||
|
# ячейки под столбцом заголовка нет — пишем пустое значение, чтобы не сдвигать столбцы
|
||||||
|
raw = ''
|
||||||
|
is_recognized = False
|
||||||
|
elif kind == FslRowKind.HEADER:
|
||||||
|
raw = cell["value"]
|
||||||
is_recognized = (_clear_value(raw) == _remove_spec_space(raw) and _is_recognized(cell["confidence"]))
|
is_recognized = (_clear_value(raw) == _remove_spec_space(raw) and _is_recognized(cell["confidence"]))
|
||||||
raw = _clear_value(raw)
|
raw = _clear_value(raw)
|
||||||
else:
|
else:
|
||||||
|
raw = cell["value"]
|
||||||
parsed_amount, is_num_recognized = _parse_amount(raw)
|
parsed_amount, is_num_recognized = _parse_amount(raw)
|
||||||
is_recognized = _is_recognized(cell["confidence"]) and is_num_recognized
|
is_recognized = _is_recognized(cell["confidence"]) and is_num_recognized
|
||||||
raw = _clear_value(raw)
|
raw = _clear_value(raw)
|
||||||
@ -261,7 +359,7 @@ def _process_file(
|
|||||||
line_id=db_row.id,
|
line_id=db_row.id,
|
||||||
amount=parsed_amount,
|
amount=parsed_amount,
|
||||||
string_value=raw,
|
string_value=raw,
|
||||||
column_index=cell["column_index"],
|
column_index=_FIRST_VALUE_COLUMN + value_number,
|
||||||
is_recognized=is_recognized,
|
is_recognized=is_recognized,
|
||||||
)
|
)
|
||||||
values.append(
|
values.append(
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user