370 lines
11 KiB
Python
370 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import zlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from .db import PREVIEW_DIR
|
|
|
|
|
|
PDF_RENDER_SCALE = 3.0
|
|
MAX_BLOCKS_PER_PAGE = 12
|
|
DATE_RE = re.compile(r"\b([0-3]?\d)[/-]([01]?\d)[/-]((?:20)?\d{2})\b")
|
|
MONEY_RE = re.compile(r"(?<!\d)(?:R\$\s*)?(\d{1,3}(?:\.\d{3})*,\d{2}|\d+\.\d{2})(?!\d)")
|
|
CNPJ_RE = re.compile(r"\b\d{2}\.?\d{3}\.?\d{3}/?\d{4}-?\d{2}\b")
|
|
ANCHOR_RE = re.compile(
|
|
r"\b(CNPJ|CUPOM|NFC-?E|SAT|EXTRATO|DANFE|VALOR\s+TOTAL|TOTAL\s+A\s+PAGAR)\b",
|
|
re.IGNORECASE,
|
|
)
|
|
START_ANCHOR_RE = re.compile(r"\b(CNPJ|CUPOM|NFC-?E|SAT|EXTRATO|DANFE)\b", re.IGNORECASE)
|
|
LABEL_RE = re.compile(
|
|
r"(valor\s+total|total\s+a\s+pagar|valor\s+pago|total\s+pago|total\s+r\$|total)",
|
|
re.IGNORECASE,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class PageSource:
|
|
page_number: int
|
|
text: str
|
|
image_path: Path | None
|
|
source_kind: str
|
|
|
|
|
|
@dataclass
|
|
class DetectedDocumentCandidate:
|
|
source_file_name: str
|
|
source_page: int | None
|
|
source_location: str
|
|
raw_text: str
|
|
purchase_date: str | None
|
|
supplier_name: str | None
|
|
total_paid: float | None
|
|
confidence: str
|
|
field_confidence: dict[str, str]
|
|
|
|
|
|
def normalize_file(path: Path, original_name: str) -> list[PageSource]:
|
|
suffix = path.suffix.lower()
|
|
if suffix == ".pdf":
|
|
return normalize_pdf(path)
|
|
if suffix in {".jpg", ".jpeg", ".png"}:
|
|
return [normalize_image(path, page_number=1)]
|
|
return [PageSource(1, "", None, "unsupported")]
|
|
|
|
|
|
def normalize_pdf(path: Path) -> list[PageSource]:
|
|
pages = extract_pdf_text_with_pypdf(path)
|
|
if not pages:
|
|
pages = extract_pdf_text_naive(path)
|
|
|
|
preview_images = render_pdf_pages_to_images(path)
|
|
page_count = max(len(pages), len(preview_images), count_pdf_pages(path), 1)
|
|
normalized: list[PageSource] = []
|
|
for index in range(page_count):
|
|
text = pages[index] if index < len(pages) else ""
|
|
if not has_usable_text(text) and index < len(preview_images):
|
|
ocr_text = extract_image_text(preview_images[index])
|
|
if has_usable_text(ocr_text):
|
|
text = ocr_text
|
|
normalized.append(
|
|
PageSource(
|
|
page_number=index + 1,
|
|
text=text if has_usable_text(text) else "",
|
|
image_path=preview_images[index] if index < len(preview_images) else None,
|
|
source_kind="pdf",
|
|
)
|
|
)
|
|
return normalized
|
|
|
|
|
|
def normalize_image(path: Path, page_number: int) -> PageSource:
|
|
text = extract_image_text(path)
|
|
return PageSource(page_number, text, path, "image")
|
|
|
|
|
|
def extract_pdf_text_with_pypdf(path: Path) -> list[str]:
|
|
try:
|
|
from pypdf import PdfReader # type: ignore
|
|
except Exception:
|
|
return []
|
|
|
|
try:
|
|
reader = PdfReader(str(path))
|
|
return [
|
|
text if has_usable_text(text) else ""
|
|
for text in ((page.extract_text() or "") for page in reader.pages)
|
|
]
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
|
|
def extract_pdf_text_naive(path: Path) -> list[str]:
|
|
data = path.read_bytes()
|
|
chunks: list[bytes] = []
|
|
for match in re.finditer(rb"stream\r?\n(.*?)\r?\nendstream", data, re.S):
|
|
stream = match.group(1)
|
|
try:
|
|
chunks.append(zlib.decompress(stream))
|
|
except Exception:
|
|
chunks.append(stream)
|
|
|
|
decoded = "\n".join(chunk.decode("latin-1", errors="ignore") for chunk in chunks)
|
|
if not decoded:
|
|
decoded = data.decode("latin-1", errors="ignore")
|
|
|
|
text_tokens = re.findall(r"\(([^()]*)\)", decoded)
|
|
text = "\n".join(token.replace(r"\)", ")").replace(r"\(", "(") for token in text_tokens)
|
|
|
|
page_count = count_pdf_pages(path)
|
|
if not has_usable_text(text):
|
|
return ["" for _ in range(max(page_count, 1))]
|
|
if "\f" in text and page_count > 1:
|
|
parts = [clean_text(part) for part in text.split("\f")]
|
|
if len(parts) <= page_count * 2 and any(has_usable_text(part) for part in parts):
|
|
return [part if has_usable_text(part) else "" for part in parts]
|
|
if page_count <= 1:
|
|
return [clean_text(text)]
|
|
return [clean_text(text)] + ["" for _ in range(max(page_count - 1, 0))]
|
|
|
|
def render_pdf_pages_to_images(path: Path) -> list[Path]:
|
|
try:
|
|
import fitz # type: ignore
|
|
except Exception:
|
|
return []
|
|
|
|
rendered: list[Path] = []
|
|
try:
|
|
PREVIEW_DIR.mkdir(exist_ok=True)
|
|
document = fitz.open(path)
|
|
for index, page in enumerate(document):
|
|
pix = page.get_pixmap(matrix=fitz.Matrix(PDF_RENDER_SCALE, PDF_RENDER_SCALE), alpha=False)
|
|
output = PREVIEW_DIR / f"{path.stem}-page-{index + 1}.png"
|
|
pix.save(output)
|
|
rendered.append(output)
|
|
except Exception:
|
|
return rendered
|
|
return rendered
|
|
|
|
|
|
def extract_image_text(path: Path) -> str:
|
|
try:
|
|
from PIL import Image # type: ignore
|
|
import pytesseract # type: ignore
|
|
except Exception:
|
|
return ""
|
|
|
|
try:
|
|
image = Image.open(path)
|
|
for language in ("por+eng", "eng"):
|
|
try:
|
|
text = pytesseract.image_to_string(image, lang=language)
|
|
except Exception:
|
|
continue
|
|
if has_usable_text(text):
|
|
return text
|
|
return ""
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
def count_pdf_pages(path: Path) -> int:
|
|
try:
|
|
data = path.read_bytes()
|
|
except OSError:
|
|
return 0
|
|
matches = re.findall(rb"/Type\s*/Page\b", data)
|
|
return len(matches)
|
|
|
|
|
|
def detect_documents(pages: list[PageSource], original_name: str) -> list[DetectedDocumentCandidate]:
|
|
candidates: list[DetectedDocumentCandidate] = []
|
|
for page in pages:
|
|
blocks = segment_blocks(page.text)
|
|
if not blocks:
|
|
if page.source_kind == "pdf" and not page.text.strip():
|
|
continue
|
|
location = f"pagina {page.page_number}"
|
|
if page.image_path:
|
|
location += f" ({page.image_path.name})"
|
|
candidates.append(
|
|
build_candidate(
|
|
original_name,
|
|
page.page_number,
|
|
location,
|
|
"",
|
|
forced_confidence="low",
|
|
)
|
|
)
|
|
continue
|
|
|
|
for block_index, block in enumerate(blocks, start=1):
|
|
location = f"pagina {page.page_number}, bloco {block_index}"
|
|
candidates.append(build_candidate(original_name, page.page_number, location, block))
|
|
return candidates
|
|
|
|
|
|
def segment_blocks(text: str) -> list[str]:
|
|
cleaned = clean_text(text)
|
|
if not has_usable_text(cleaned):
|
|
return []
|
|
|
|
lines = [line.strip() for line in cleaned.splitlines() if line.strip()]
|
|
anchor_indexes = [
|
|
index for index, line in enumerate(lines)
|
|
if START_ANCHOR_RE.search(line) or CNPJ_RE.search(line)
|
|
]
|
|
|
|
starts: list[int] = []
|
|
for index in anchor_indexes:
|
|
start = max(index - 1, 0) if CNPJ_RE.search(lines[index]) else index
|
|
if not starts or start - starts[-1] > 3:
|
|
starts.append(start)
|
|
|
|
if len(starts) <= 1:
|
|
return [cleaned]
|
|
|
|
blocks: list[str] = []
|
|
for position, start in enumerate(starts):
|
|
end = starts[position + 1] if position + 1 < len(starts) else len(lines)
|
|
block = "\n".join(lines[start:end]).strip()
|
|
if block:
|
|
blocks.append(block)
|
|
return blocks[:MAX_BLOCKS_PER_PAGE]
|
|
|
|
|
|
def build_candidate(
|
|
original_name: str,
|
|
page_number: int | None,
|
|
location: str,
|
|
raw_text: str,
|
|
forced_confidence: str | None = None,
|
|
) -> DetectedDocumentCandidate:
|
|
purchase_date = extract_date(raw_text)
|
|
supplier = extract_supplier(raw_text)
|
|
total = extract_total(raw_text)
|
|
|
|
fields = {
|
|
"purchase_date": "high" if purchase_date else "low",
|
|
"supplier_name": "medium" if supplier else "low",
|
|
"total_paid": "high" if total is not None and has_total_label(raw_text) else ("medium" if total is not None else "low"),
|
|
}
|
|
confidence = forced_confidence or overall_confidence(fields)
|
|
return DetectedDocumentCandidate(
|
|
source_file_name=original_name,
|
|
source_page=page_number,
|
|
source_location=location,
|
|
raw_text=raw_text,
|
|
purchase_date=purchase_date,
|
|
supplier_name=supplier,
|
|
total_paid=total,
|
|
confidence=confidence,
|
|
field_confidence=fields,
|
|
)
|
|
|
|
|
|
def extract_date(text: str) -> str | None:
|
|
match = DATE_RE.search(text)
|
|
if not match:
|
|
return None
|
|
day, month, year = match.groups()
|
|
if len(year) == 2:
|
|
year = f"20{year}"
|
|
return f"{int(year):04d}-{int(month):02d}-{int(day):02d}"
|
|
|
|
|
|
def extract_supplier(text: str) -> str | None:
|
|
for line in clean_text(text).splitlines()[:8]:
|
|
value = line.strip(" :-")
|
|
if not value:
|
|
continue
|
|
if should_skip_supplier_line(value):
|
|
continue
|
|
if any(ch.isalpha() for ch in value):
|
|
return value[:120]
|
|
return None
|
|
|
|
|
|
def should_skip_supplier_line(line: str) -> bool:
|
|
upper = line.upper()
|
|
if CNPJ_RE.search(line) or DATE_RE.search(line) or MONEY_RE.search(line):
|
|
return True
|
|
blocked = ["CUPOM", "NFC", "SAT", "DANFE", "EXTRATO", "VALOR", "TOTAL", "CHAVE", "ENDERECO"]
|
|
return any(token in upper for token in blocked)
|
|
|
|
|
|
def extract_total(text: str) -> float | None:
|
|
lines = [line.strip() for line in clean_text(text).splitlines() if line.strip()]
|
|
labelled: list[float] = []
|
|
all_values: list[float] = []
|
|
for line in lines:
|
|
values = [parse_money(match.group(1)) for match in MONEY_RE.finditer(line)]
|
|
values = [value for value in values if value is not None]
|
|
all_values.extend(values)
|
|
if LABEL_RE.search(line):
|
|
labelled.extend(values)
|
|
|
|
if labelled:
|
|
return labelled[-1]
|
|
if all_values:
|
|
return all_values[-1]
|
|
return None
|
|
|
|
|
|
def parse_money(value: str) -> float | None:
|
|
normalized = value.strip()
|
|
if "," in normalized:
|
|
normalized = normalized.replace(".", "").replace(",", ".")
|
|
try:
|
|
return round(float(normalized), 2)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def has_total_label(text: str) -> bool:
|
|
return LABEL_RE.search(text) is not None
|
|
|
|
|
|
def overall_confidence(fields: dict[str, str]) -> str:
|
|
values = list(fields.values())
|
|
if all(value == "high" for value in values):
|
|
return "high"
|
|
if values.count("low") >= 2:
|
|
return "low"
|
|
return "medium"
|
|
|
|
|
|
def clean_text(text: str) -> str:
|
|
text = text.replace("\x00", " ")
|
|
lines = [re.sub(r"\s+", " ", line).strip() for line in text.splitlines()]
|
|
return "\n".join(line for line in lines if line)
|
|
|
|
|
|
def has_usable_text(text: str) -> bool:
|
|
cleaned = clean_text(text)
|
|
if len(cleaned) < 12:
|
|
return False
|
|
|
|
allowed_extra = set("???????????????????????????????????????????????????$??")
|
|
ordinary = sum(
|
|
1
|
|
for ch in cleaned
|
|
if ch.isascii() or ch in allowed_extra
|
|
)
|
|
printable = sum(1 for ch in cleaned if ch.isprintable())
|
|
alnum = sum(1 for ch in cleaned if ch.isalnum())
|
|
letters = sum(1 for ch in cleaned if ch.isalpha())
|
|
length = max(len(cleaned), 1)
|
|
if ordinary / length < 0.85:
|
|
return False
|
|
if printable / length < 0.9:
|
|
return False
|
|
if alnum / length < 0.35:
|
|
return False
|
|
if letters < 3 and not (DATE_RE.search(cleaned) or MONEY_RE.search(cleaned) or CNPJ_RE.search(cleaned)):
|
|
return False
|
|
|
|
return True
|