454 lines
16 KiB
Python
454 lines
16 KiB
Python
"""
|
||
OCR 刷卡信息提取
|
||
|
||
从支付截图中识别刷卡记录（姓名、日期、金额），回填到发票数据中。
|
||
|
||
匹配策略:
|
||
1. 先按文件名匹配（PDF 和图片同名）
|
||
2. 未匹配的通过金额近邻匹配
|
||
|
||
对外接口:
|
||
enrich_with_ocr(rows, directory) -> list[dict] 用 OCR 识别结果丰富发票数据
|
||
"""
|
||
|
||
import csv
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
|
||
from . import get_logger
|
||
|
||
log = get_logger("ocr")
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 懒加载 OCR
|
||
# ------------------------------------------------------------------
|
||
|
||
_ocr_instance = None
|
||
|
||
|
||
def _get_ocr():
|
||
"""懒加载 PaddleOCR 实例（兼容 2.x / 3.x）"""
|
||
global _ocr_instance
|
||
if _ocr_instance is not None:
|
||
return _ocr_instance
|
||
|
||
os.environ.setdefault("FLAGS_use_mkldnn", "0")
|
||
os.environ.setdefault("FLAGS_mkldnn_cache_enabled", "0")
|
||
|
||
from paddleocr import PaddleOCR
|
||
|
||
try:
|
||
_ocr_instance = PaddleOCR(use_textline_orientation=True, lang="ch")
|
||
except TypeError:
|
||
try:
|
||
_ocr_instance = PaddleOCR(lang="ch")
|
||
except TypeError:
|
||
_ocr_instance = PaddleOCR()
|
||
|
||
return _ocr_instance
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# OCR 识别
|
||
# ------------------------------------------------------------------
|
||
|
||
def ocr_image(image_path: Path) -> list[dict]:
|
||
"""对单张图片执行 OCR，返回 [{"text": str, "confidence": float}, ...]"""
|
||
ocr = _get_ocr()
|
||
texts = []
|
||
|
||
try:
|
||
results = ocr.ocr(str(image_path), cls=True)
|
||
if results and isinstance(results, list):
|
||
for page_result in results:
|
||
if not page_result:
|
||
continue
|
||
for line in page_result:
|
||
if isinstance(line, (list, tuple)) and len(line) >= 2:
|
||
_, text_info = line[0], line[1]
|
||
if isinstance(text_info, (list, tuple)) and len(text_info) >= 2:
|
||
texts.append({
|
||
"text": str(text_info[0]),
|
||
"confidence": float(text_info[1]),
|
||
})
|
||
except Exception:
|
||
try:
|
||
if hasattr(ocr, "predict"):
|
||
results = ocr.predict(str(image_path))
|
||
if results:
|
||
for result in results:
|
||
if hasattr(result, "rec_result_list"):
|
||
for line in result.rec_result_list:
|
||
t = getattr(line, "text", "") or ""
|
||
s = getattr(line, "score", 0.0) or 0.0
|
||
texts.append({"text": str(t), "confidence": float(s)})
|
||
elif isinstance(result, list):
|
||
for line in result:
|
||
if isinstance(line, (list, tuple)) and len(line) >= 2:
|
||
t = line[1][0] if isinstance(line[1], (list, tuple)) else str(line[1])
|
||
s = line[1][1] if isinstance(line[1], (list, tuple)) and len(line[1]) > 1 else 0.0
|
||
texts.append({"text": str(t), "confidence": float(s)})
|
||
except Exception as e:
|
||
log.error(f"OCR 识别失败: {e}")
|
||
|
||
return texts
|
||
|
||
|
||
def extract_card_info(texts: list[dict]) -> dict:
|
||
"""从 OCR 文本中提取刷卡信息（日期 / 金额 / 姓名）"""
|
||
info = {"刷卡日期": "", "刷卡金额": "", "人员姓名": ""}
|
||
|
||
valid = [t for t in texts if t["confidence"] > 0.5]
|
||
full_text = " ".join(t["text"] for t in valid)
|
||
if not full_text:
|
||
return info
|
||
|
||
# 日期（优先级匹配，避免误抓发票开票日期）
|
||
date_candidates = []
|
||
for pattern, priority in [
|
||
(r"记账时间[：:\s]*(\d{4}[-/]\d{1,2}[-/]\d{1,2})", 10),
|
||
(r"交易时间[：:\s]*(\d{4}[-/]\d{1,2}[-/]\d{1,2})", 9),
|
||
(r"刷卡日期[：:\s]*(\d{4}[-/]\d{1,2}[-/]\d{1,2})", 9),
|
||
(r"日期[：:\s]*(\d{4}[-/]\d{1,2}[-/]\d{1,2})", 5),
|
||
]:
|
||
for m in re.finditer(pattern, full_text):
|
||
date_candidates.append((priority, m.start(), m.group(1).replace("-", "/")))
|
||
if date_candidates:
|
||
date_candidates.sort(key=lambda x: (-x[0], x[1]))
|
||
info["刷卡日期"] = date_candidates[0][2]
|
||
|
||
# 金额
|
||
amount_candidates = []
|
||
for pattern, priority in [
|
||
(r"交易金额[：:\s]*([+-]?[\d,]+\.?\d*)", 10),
|
||
(r"刷卡金额[：:\s]*([+-]?[\d,]+\.?\d*)", 10),
|
||
(r"金额[：:\s]*([+-]?[\d,]+\.?\d*)", 5),
|
||
]:
|
||
for m in re.finditer(pattern, full_text):
|
||
amt_str = m.group(1).replace(",", "").replace("+", "")
|
||
try:
|
||
val = float(amt_str)
|
||
if 0 < val < 999999:
|
||
amount_candidates.append((priority, m.start(), amt_str))
|
||
except ValueError:
|
||
continue
|
||
if amount_candidates:
|
||
amount_candidates.sort(key=lambda x: (-x[0], x[1]))
|
||
info["刷卡金额"] = amount_candidates[0][2]
|
||
|
||
# 姓名（排除公司/机构后缀）
|
||
EXCLUDE_SUFFIXES = ("公司", "银行", "中心", "支行", "商户", "网点", "有限", "责任")
|
||
name_candidates = []
|
||
for pattern, priority in [
|
||
(r"交易户名[：:\s]*([\u4e00-\u9fff]{2,6})", 10),
|
||
(r"户名[：:\s]*([\u4e00-\u9fff]{2,6})", 8),
|
||
(r"持卡人[：:\s]*([\u4e00-\u9fff]{2,6})", 8),
|
||
(r"姓名[：:\s]*([\u4e00-\u9fff]{2,6})", 8),
|
||
]:
|
||
for m in re.finditer(pattern, full_text):
|
||
name = m.group(1)
|
||
if not any(name.endswith(s) for s in EXCLUDE_SUFFIXES):
|
||
name_candidates.append((priority, m.start(), name))
|
||
if name_candidates:
|
||
name_candidates.sort(key=lambda x: (-x[0], x[1]))
|
||
info["人员姓名"] = name_candidates[0][2]
|
||
|
||
return info
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# PDF 发票号提取
|
||
# ------------------------------------------------------------------
|
||
|
||
def extract_invoice_number(pdf_path: Path) -> str:
|
||
"""从 PDF 中提取发票号码"""
|
||
try:
|
||
import pdfplumber
|
||
except ImportError:
|
||
log.warning("缺少 pdfplumber，跳过发票号提取")
|
||
return ""
|
||
|
||
try:
|
||
with pdfplumber.open(str(pdf_path)) as pdf_file:
|
||
page_text = ""
|
||
for page in pdf_file.pages:
|
||
page_text += page.extract_text() or ""
|
||
|
||
for pattern in [
|
||
r"发票号码[：:\s]*([A-Za-z0-9]{8,20})",
|
||
r"发票代码[：:\s]*([A-Za-z0-9]{10,12})",
|
||
r"号码[：:\s]*([A-Za-z0-9]{8,20})",
|
||
]:
|
||
m = re.search(pattern, page_text)
|
||
if m:
|
||
return m.group(1)
|
||
except Exception as e:
|
||
log.warning(f"PDF 读取失败 ({pdf_path.name}): {e}")
|
||
|
||
return ""
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 图片配对
|
||
# ------------------------------------------------------------------
|
||
|
||
def _extract_amount_from_pdf(pdf_path: Path) -> float | None:
|
||
"""从 PDF 中提取价税合计金额"""
|
||
try:
|
||
import pdfplumber
|
||
with pdfplumber.open(str(pdf_path)) as pdf:
|
||
text = ""
|
||
for page in pdf.pages:
|
||
t = page.extract_text()
|
||
if t:
|
||
text += t + "\n"
|
||
m = re.search(r"价税合计.*?（小写）[¥￥]?\s*(\d+\.?\d*)", text)
|
||
if m:
|
||
return float(m.group(1))
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
def _extract_amount_from_image(img_path: Path) -> float | None:
|
||
"""从图片 OCR 中提取刷卡金额"""
|
||
texts = ocr_image(img_path)
|
||
if not texts:
|
||
return None
|
||
info = extract_card_info(texts)
|
||
amt_str = info.get("刷卡金额", "")
|
||
if amt_str:
|
||
try:
|
||
return float(amt_str)
|
||
except ValueError:
|
||
pass
|
||
return None
|
||
|
||
|
||
def find_image_pairs(directory: str = ".") -> list[tuple[Path, Path]]:
|
||
"""查找 PDF 和对应图片的配对
|
||
|
||
1. 先按文件名匹配（PDF 和图片同名）
|
||
2. 未匹配的通过金额近邻匹配
|
||
"""
|
||
base = Path(directory)
|
||
pdfs = sorted(base.glob("*.pdf"))
|
||
image_exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||
|
||
all_images = sorted(
|
||
f for ext in image_exts for f in base.glob(f"*{ext}")
|
||
)
|
||
|
||
# ---- Phase 1: 文件名匹配 ----
|
||
pairs: list[tuple[Path, Path]] = []
|
||
matched_pdfs: set[Path] = set()
|
||
matched_imgs: set[Path] = set()
|
||
|
||
for pdf in pdfs:
|
||
for ext in image_exts:
|
||
img = base / f"{pdf.stem}{ext}"
|
||
if img.exists():
|
||
pairs.append((pdf, img))
|
||
matched_pdfs.add(pdf)
|
||
matched_imgs.add(img)
|
||
break
|
||
|
||
unmatched_pdfs = [p for p in pdfs if p not in matched_pdfs]
|
||
unmatched_imgs = [i for i in all_images if i not in matched_imgs]
|
||
|
||
if not unmatched_pdfs or not unmatched_imgs:
|
||
return pairs
|
||
|
||
# ---- Phase 2: 金额近邻匹配 ----
|
||
if len(unmatched_pdfs) > 0 and len(unmatched_imgs) > 0:
|
||
log.info(f"文件名匹配 {len(pairs)} 组，剩余 {len(unmatched_pdfs)} 个 PDF、{len(unmatched_imgs)} 张图片，尝试金额匹配...")
|
||
|
||
pdf_amounts: dict[Path, float] = {}
|
||
for pdf in unmatched_pdfs:
|
||
amt = _extract_amount_from_pdf(pdf)
|
||
if amt is not None:
|
||
pdf_amounts[pdf] = amt
|
||
|
||
img_amounts: dict[Path, float] = {}
|
||
for img in unmatched_imgs:
|
||
amt = _extract_amount_from_image(img)
|
||
if amt is not None:
|
||
img_amounts[img] = amt
|
||
|
||
# 贪婪匹配：每张图片找金额差最小的 PDF
|
||
used_pdfs: set[Path] = set()
|
||
for img, img_amt in sorted(img_amounts.items(), key=lambda x: x[0].name):
|
||
best_pdf: Path | None = None
|
||
best_diff: float = float("inf")
|
||
|
||
for pdf, pdf_amt in pdf_amounts.items():
|
||
if pdf in used_pdfs:
|
||
continue
|
||
diff = abs(pdf_amt - img_amt)
|
||
if diff < best_diff:
|
||
best_diff = diff
|
||
best_pdf = pdf
|
||
|
||
if best_pdf is not None:
|
||
pairs.append((best_pdf, img))
|
||
used_pdfs.add(best_pdf)
|
||
|
||
log.info(f"金额匹配完成，共 {len(pairs)} 组配对")
|
||
|
||
return pairs
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# CSV 读写
|
||
# ------------------------------------------------------------------
|
||
|
||
from .extractor import CSV_COLUMNS
|
||
|
||
|
||
def _load_csv(csv_path: Path) -> list[dict] | None:
|
||
"""读取现有 CSV 为 dict 列表，失败返回 None"""
|
||
try:
|
||
with open(csv_path, encoding="utf-8", newline="") as f:
|
||
reader = csv.DictReader(f)
|
||
fieldnames = reader.fieldnames or []
|
||
missing = [c for c in CSV_COLUMNS if c not in fieldnames]
|
||
if missing:
|
||
log.error(f"CSV 缺少必要列: {missing}")
|
||
return None
|
||
return [row for row in reader]
|
||
except FileNotFoundError:
|
||
log.error(f"CSV 文件不存在: {csv_path.name}")
|
||
return None
|
||
except Exception as e:
|
||
log.error(f"CSV 读取失败: {e}")
|
||
return None
|
||
|
||
|
||
def _save_csv(csv_path: Path, rows: list[dict]):
|
||
"""保存 CSV"""
|
||
with open(csv_path, "w", encoding="utf-8", newline="") as f:
|
||
writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS)
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# Markdown 同步
|
||
# ------------------------------------------------------------------
|
||
|
||
def save_markdown_from_csv(csv_path: Path, rows: list[dict]):
|
||
"""根据最新 CSV 数据生成 Markdown 汇总表"""
|
||
md_path = csv_path.with_suffix(".md")
|
||
columns = [
|
||
("序号", "序号"), ("发票号码", "发票号码"), ("开票日期", "开票日期"),
|
||
("项目名称", "项目名称"), ("规格型号", "规格型号"), ("价税合计", "价税合计"),
|
||
("销售方名称", "销售方名称"), ("人员姓名", "人员姓名"),
|
||
("刷卡日期", "刷卡日期"), ("公务卡号", "公务卡号"),
|
||
("刷卡金额", "刷卡金额"), ("备注", "备注"), ("工号", "工号"),
|
||
]
|
||
|
||
lines = ["# 发票信息汇总表", ""]
|
||
header = " | ".join(col[1] for col in columns)
|
||
separator = "|".join(["------" for _ in columns])
|
||
lines.append(f"| {header} |")
|
||
lines.append(f"|{separator}|")
|
||
|
||
total_price = 0.0
|
||
total_card = 0.0
|
||
|
||
for row in rows:
|
||
cells = []
|
||
for key, _ in columns:
|
||
value = row.get(key, "").strip()
|
||
|
||
if key == "价税合计" and value:
|
||
try:
|
||
total_price += float(value.replace(",", ""))
|
||
cells.append(f"¥{float(value.replace(',', '')):,.2f}")
|
||
except (ValueError, TypeError):
|
||
cells.append(value)
|
||
elif key == "刷卡金额" and value:
|
||
try:
|
||
total_card += float(value.replace(",", ""))
|
||
cells.append(f"¥{float(value.replace(',', '')):,.2f}")
|
||
except (ValueError, TypeError):
|
||
cells.append(value)
|
||
else:
|
||
cells.append(value if value else "")
|
||
|
||
lines.append("| " + " | ".join(cells) + " |")
|
||
|
||
lines.append("")
|
||
lines.append(f"**价税合计总计: ¥{total_price:,.2f}**")
|
||
lines.append(f"**刷卡金额总计: ¥{total_card:,.2f}**")
|
||
lines.append("")
|
||
|
||
with open(md_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(lines))
|
||
|
||
log.info(f"Markdown 已同步: {md_path.name}")
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 主入口
|
||
# ------------------------------------------------------------------
|
||
|
||
def enrich_with_ocr(rows: list[dict], directory: str = ".") -> list[dict]:
|
||
"""用 OCR 识别结果丰富发票数据，返回更新后的行列表
|
||
|
||
rows 应包含「发票号码」列，已存在的字段不会覆盖。
|
||
"""
|
||
pairs = find_image_pairs(directory)
|
||
if not pairs:
|
||
log.warning("未找到 PDF-图片配对文件，跳过 OCR")
|
||
return rows
|
||
|
||
log.info(f"找到 {len(pairs)} 组 PDF-图片配对")
|
||
|
||
ocr_by_invoice: dict[str, dict] = {}
|
||
|
||
for idx, (pdf, img) in enumerate(pairs, 1):
|
||
inv_num = extract_invoice_number(pdf)
|
||
if not inv_num:
|
||
inv_num = pdf.stem
|
||
|
||
texts = ocr_image(img)
|
||
if not texts:
|
||
log.warning(f"OCR 未识别到文本: {img.name}")
|
||
continue
|
||
|
||
info = extract_card_info(texts)
|
||
ocr_by_invoice[inv_num] = info
|
||
|
||
# 更新行数据
|
||
updated = 0
|
||
matched = 0
|
||
|
||
for i, row in enumerate(rows):
|
||
inv_num = row.get("发票号码", "").strip()
|
||
ocr_info = ocr_by_invoice.get(inv_num)
|
||
|
||
if not ocr_info:
|
||
for key, val in ocr_by_invoice.items():
|
||
if inv_num in key or key in inv_num:
|
||
ocr_info = val
|
||
break
|
||
|
||
if ocr_info:
|
||
matched += 1
|
||
|
||
if not row.get("人员姓名", "").strip() and ocr_info["人员姓名"]:
|
||
row["人员姓名"] = ocr_info["人员姓名"]
|
||
updated += 1
|
||
if not row.get("刷卡日期", "").strip() and ocr_info["刷卡日期"]:
|
||
row["刷卡日期"] = ocr_info["刷卡日期"]
|
||
updated += 1
|
||
if not row.get("刷卡金额", "").strip() and ocr_info["刷卡金额"]:
|
||
row["刷卡金额"] = ocr_info["刷卡金额"]
|
||
updated += 1
|
||
|
||
log.info(f"OCR 完成: 匹配 {matched}/{len(rows)} 行，更新 {updated} 个字段")
|
||
|
||
return rows |