Compare commits
2 Commits
main
...
feat/6-ocr
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
47b65469e1 | ||
|
|
8e4405f0b1 |
@ -1,11 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from dataclasses import dataclass, asdict
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from typing import Optional, Literal, TypedDict
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from src.schemas.enums import StatementType, FslRowKind
|
||||
from .consts import codes_mapping
|
||||
|
||||
@ -52,6 +55,29 @@ class ExtractResult(TypedDict):
|
||||
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:
|
||||
cells: dict[int, dict]
|
||||
layout: ColumnLayout
|
||||
kind: FslRowKind
|
||||
ordinal: int
|
||||
|
||||
|
||||
# column_index в БД: 0 - наименование, 1 - код, значения начинаются с 2
|
||||
_FIRST_VALUE_COLUMN = 2
|
||||
_YEAR_RE = re.compile(r"\d{4}")
|
||||
_BLANK_CHARS_RE = re.compile(r"[\s\-–—]")
|
||||
_HEADER_CODE = 'Код'
|
||||
|
||||
|
||||
_TABLE_TYPE_MAP: dict[str, StatementType] = {
|
||||
"Бухгалтерский баланс": StatementType.BALANCE,
|
||||
"Отчет о финансовых результатах": StatementType.INCOME,
|
||||
@ -106,16 +132,61 @@ def _parse_amount(raw: str) -> tuple[Optional[Decimal], bool]:
|
||||
return (None, is_verified)
|
||||
|
||||
|
||||
def _classify_row(cells: list[dict]) -> FslRowKind | None:
|
||||
if not (str(cells[1]['value']).strip()):
|
||||
def _normalize_header(value) -> str:
|
||||
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
|
||||
try:
|
||||
int(cells[1]["value"])
|
||||
int(code)
|
||||
except (TypeError, ValueError):
|
||||
return FslRowKind.HEADER
|
||||
if cells[0]["value"] == "Баланс":
|
||||
label = _cell_value(cells, layout.label_idx)
|
||||
if label == "Баланс":
|
||||
return FslRowKind.BALANCE
|
||||
elif cells[0]["value"].lower().startswith("итого"):
|
||||
elif label.lower().startswith("итого"):
|
||||
return FslRowKind.TOTAL
|
||||
return FslRowKind.ROW
|
||||
|
||||
@ -129,19 +200,127 @@ def _keep_digits_only(raw: str) -> str:
|
||||
return res.strip()
|
||||
|
||||
|
||||
def _resolve_table_type(rows) -> StatementType:
|
||||
def _is_blank(value: str) -> bool:
|
||||
# прочерк в отчётности означает ноль, считаем его пустым значением
|
||||
return not _BLANK_CHARS_RE.sub('', value)
|
||||
|
||||
|
||||
def _infer_layout(rows_cells: list[dict[int, dict]]) -> Optional[ColumnLayout]:
|
||||
"""Выводит раскладку столбцов по строкам данных, когда заголовка нет.
|
||||
|
||||
Столбец оценивается по всем строкам сразу, поэтому пустые ячейки в отдельных
|
||||
строках (например, в начале таблицы) не влияют на результат.
|
||||
"""
|
||||
columns: dict[int, list[str]] = {}
|
||||
for cells in rows_cells:
|
||||
for idx, cell in cells.items():
|
||||
columns.setdefault(idx, []).append(str(cell['value'] or '').strip())
|
||||
|
||||
# код — столбец с наибольшим числом известных кодов строк
|
||||
code_hits = {idx: sum(1 for v in values if v in codes_mapping) for idx, values in columns.items()}
|
||||
code_idx = min(code_hits, key=lambda idx: (-code_hits[idx], idx), default=None)
|
||||
if code_idx is None or code_hits[code_idx] == 0:
|
||||
return None
|
||||
|
||||
# наименование — столбец левее кода с наибольшим числом текстовых ячеек
|
||||
text_hits = {
|
||||
idx: sum(1 for v in values if any(c.isalpha() for c in v))
|
||||
for idx, values in columns.items() if idx < code_idx
|
||||
}
|
||||
label_idx = min(text_hits, key=lambda idx: (-text_hits[idx], -idx), default=None)
|
||||
if label_idx is not None and text_hits[label_idx] == 0:
|
||||
label_idx = None
|
||||
|
||||
# значения — столбцы правее кода, где суммы преобладают над мусором (полоса скана),
|
||||
# либо полностью пустые (период без данных)
|
||||
value_idxs = []
|
||||
for idx in sorted(columns):
|
||||
if idx <= code_idx:
|
||||
continue
|
||||
filled = [v for v in columns[idx] if not _is_blank(v)]
|
||||
amounts = sum(1 for v in filled if _parse_amount(v)[0] is not None)
|
||||
if amounts >= len(filled) - amounts:
|
||||
value_idxs.append(idx)
|
||||
|
||||
return ColumnLayout(label_idx=label_idx, code_idx=code_idx, value_idxs=tuple(value_idxs))
|
||||
|
||||
|
||||
def _parse_data_row(cells: dict[int, dict], layout: ColumnLayout, ordinal: int) -> Optional[ParsedRow]:
|
||||
kind = _classify_row(cells, layout)
|
||||
if kind is None or kind == FslRowKind.HEADER:
|
||||
return None
|
||||
return ParsedRow(cells, layout, kind, ordinal)
|
||||
|
||||
|
||||
def _synthetic_header_cells(layout: ColumnLayout) -> dict[int, dict]:
|
||||
"""Ячейки восстановленного заголовка: код и наименование, тексты периодов пустые."""
|
||||
cells = {layout.code_idx: {'value': _HEADER_CODE, 'confidence': _FULL_CONFIDENCE}}
|
||||
if layout.label_idx is not None:
|
||||
cells[layout.label_idx] = {'value': codes_mapping[_HEADER_CODE], 'confidence': _FULL_CONFIDENCE}
|
||||
return cells
|
||||
|
||||
|
||||
def _parse_table_rows(table_rows: list[dict]) -> Optional[list[ParsedRow]]:
|
||||
"""Разбирает строки таблицы; первым элементом результата всегда идёт заголовок.
|
||||
|
||||
- Раскладка берётся из строки с ячейкой «Код»; каждый повторный заголовок
|
||||
перезаписывает индексы столбцов.
|
||||
- Строки до первого заголовка разбираются по раскладке, выведенной из них самих
|
||||
(индексы заголовка могут не совпадать с ними). Заголовок пишется в начало таблицы.
|
||||
- Если заголовка нет, раскладка выводится по всем строкам, а заголовок создаётся
|
||||
с кодом и наименованием и пустыми текстами периодов.
|
||||
Если не удалось определить ни заголовок, ни столбец кода, возвращает None.
|
||||
"""
|
||||
rows_cells = [(row['row_index'], _cells_by_index(row)) for row in table_rows]
|
||||
first_header = next(
|
||||
(pos for pos, (_, cells) in enumerate(rows_cells) if _build_layout(cells) is not None),
|
||||
len(rows_cells),
|
||||
)
|
||||
|
||||
before_header = rows_cells[:first_header]
|
||||
inferred = _infer_layout([cells for _, cells in before_header]) if before_header else None
|
||||
leading_rows = []
|
||||
if inferred is not None:
|
||||
for ordinal, cells in before_header:
|
||||
parsed = _parse_data_row(cells, inferred, ordinal)
|
||||
if parsed is not None:
|
||||
leading_rows.append(parsed)
|
||||
|
||||
if first_header == len(rows_cells):
|
||||
if not leading_rows:
|
||||
return None
|
||||
header_ordinal = min(p.ordinal for p in leading_rows) - 1
|
||||
header = ParsedRow(_synthetic_header_cells(inferred), inferred, FslRowKind.HEADER, header_ordinal)
|
||||
return [header, *leading_rows]
|
||||
|
||||
header_row_index, header_cells = rows_cells[first_header]
|
||||
layout = _build_layout(header_cells)
|
||||
header_ordinal = min([header_row_index, *(p.ordinal - 1 for p in leading_rows)])
|
||||
parsed_rows = [ParsedRow(header_cells, layout, FslRowKind.HEADER, header_ordinal), *leading_rows]
|
||||
|
||||
for ordinal, cells in rows_cells[first_header + 1:]:
|
||||
header_layout = _build_layout(cells)
|
||||
if header_layout is not None:
|
||||
layout = header_layout
|
||||
continue
|
||||
parsed = _parse_data_row(cells, layout, ordinal)
|
||||
if parsed is not None:
|
||||
parsed_rows.append(parsed)
|
||||
|
||||
return parsed_rows
|
||||
|
||||
|
||||
def _resolve_table_type(parsed_rows: list[ParsedRow]) -> StatementType:
|
||||
cell_codes = {
|
||||
'1': 0,
|
||||
'2': 0
|
||||
}
|
||||
try:
|
||||
for r in rows:
|
||||
cells = list(sorted(r["cells"], key=lambda x: x['column_index']))
|
||||
kind = _classify_row(cells)
|
||||
if kind is None or kind == FslRowKind.HEADER:
|
||||
for parsed in parsed_rows:
|
||||
if parsed.kind == FslRowKind.HEADER:
|
||||
continue
|
||||
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[first_symbol] += 1
|
||||
if cell_codes['1'] > cell_codes['2']:
|
||||
@ -196,36 +375,34 @@ def _process_file(
|
||||
rows: list[DbRow] = []
|
||||
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(
|
||||
id=uuid.uuid4(),
|
||||
statement_type=_resolve_table_type(table["rows"]),
|
||||
statement_type=_resolve_table_type(parsed_rows),
|
||||
**shared,
|
||||
)
|
||||
metas.append(asdict(meta))
|
||||
header_found = False
|
||||
|
||||
for row in table["rows"]:
|
||||
cells = list(sorted(row["cells"], key=lambda x: x['column_index']))
|
||||
if len(cells) < 3:
|
||||
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]
|
||||
for parsed in parsed_rows:
|
||||
cells, layout, kind = parsed.cells, parsed.layout, parsed.kind
|
||||
code_cell = cells[layout.code_idx]
|
||||
is_line_code_valid = True
|
||||
line_code = _clear_value(code_cell["value"])
|
||||
line_name = codes_mapping.get(line_code, '')
|
||||
if line_name == '':
|
||||
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
|
||||
|
||||
@ -240,20 +417,24 @@ def _process_file(
|
||||
line_name=line_name,
|
||||
line_name_recognized=line_name_recognized,
|
||||
row_kind=kind,
|
||||
ordinal=row["row_index"],
|
||||
ordinal=parsed.ordinal,
|
||||
)
|
||||
|
||||
rows.append(asdict(db_row))
|
||||
|
||||
for cell in cells:
|
||||
if cell["column_index"] < 2:
|
||||
continue
|
||||
raw = cell["value"]
|
||||
for value_number, source_idx in enumerate(layout.value_idxs):
|
||||
cell = cells.get(source_idx)
|
||||
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"]))
|
||||
raw = _clear_value(raw)
|
||||
else:
|
||||
raw = cell["value"]
|
||||
parsed_amount, is_num_recognized = _parse_amount(raw)
|
||||
is_recognized = _is_recognized(cell["confidence"]) and is_num_recognized
|
||||
raw = _clear_value(raw)
|
||||
@ -261,7 +442,7 @@ def _process_file(
|
||||
line_id=db_row.id,
|
||||
amount=parsed_amount,
|
||||
string_value=raw,
|
||||
column_index=cell["column_index"],
|
||||
column_index=_FIRST_VALUE_COLUMN + value_number,
|
||||
is_recognized=is_recognized,
|
||||
)
|
||||
values.append(
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user