Spaces:

jun2024
/

fc

Paused

fc / apps /evidence_sorting /services /classifier.py

jun2024

Upload 2241 files

b4f2238 verified 7 months ago

7.49 kB

"""OCR 关键词分类服务 - 将图片分为对账单/出库单/收款凭证/其他"""

from __future__ import annotations

import base64
import logging
import re
from dataclasses import dataclass, field
from importlib import import_module
from typing import Any

logger = logging.getLogger("apps.evidence_sorting")

# 文件类型常量
TYPE_STATEMENT = "statement"  # 对账单
TYPE_DELIVERY = "delivery"  # 出库单/出仓单
TYPE_RECEIPT = "receipt"  # 收款凭证
TYPE_OTHER = "other"  # 其他

# 关键词权重表
_KEYWORDS: dict[str, list[str]] = {
    TYPE_STATEMENT: ["对账单", "对帐单", "对账", "月结", "月度汇总"],
    TYPE_DELIVERY: [
        "出库单",
        "出仓单",
        "出货单",
        "发货单",
        "送货单",
        "承运单",
        "提货单",
        "出库",
        "出仓",
    ],
    TYPE_RECEIPT: [
        "收款",
        "转账",
        "付款",
        "汇款",
        "收到",
        "到账",
        "银行回单",
        "电子回单",
        "交易成功",
        "支付成功",
        "中国农业银行",
        "中国工商银行",
        "中国建设银行",
        "中国银行",
        "招商银行",
        "微信支付",
        "支付宝",
    ],
}


@dataclass
class ClassifiedImage:
    """分类后的图片"""

    filename: str
    category: str  # statement / delivery / receipt / other
    ocr_text: str
    date: str | None = None  # YYYYMMDD
    amount: str | None = None  # 如 "65500"
    signed: bool | None = None  # 对账单是否签名
    confidence: float = 0.0
    image_data: str = ""  # base64
    rotation: int = 0


@dataclass
class ClassifyResult:
    """分类结果"""

    images: list[ClassifiedImage] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)


class ClassifierService:
    """图片 OCR + 关键词分类"""

    def __init__(self) -> None:
        self._ocr_service: Any | None = None

    def _get_ocr_service(self) -> Any:
        if self._ocr_service is None:
            module = import_module("apps.image_rotation.services.orientation.service")
            orientation_service_cls = module.OrientationDetectionService
            self._ocr_service = orientation_service_cls()
        return self._ocr_service

    def classify_images(self, images: list[dict[str, Any]]) -> ClassifyResult:
        """
        批量 OCR + 分类
        Args:
            images: [{"filename": str, "data": str(base64)}]
        Returns:
            ClassifyResult
        """
        result = ClassifyResult()
        ocr_svc = self._get_ocr_service()

        for idx, img in enumerate(images):
            filename: str = img.get("filename", f"image_{idx}")
            data: str = img.get("data", "")
            try:
                # 去掉 data:image/xxx;base64, 前缀
                raw_data = data.split(",", 1)[-1] if "," in data else data
                image_bytes = base64.b64decode(raw_data)

                # OCR + 方向检测
                ocr_result = ocr_svc.detect_orientation_with_text(image_bytes)
                ocr_text: str = ocr_result.get("ocr_text", "")
                rotation: int = ocr_result.get("rotation", 0)
                confidence: float = ocr_result.get("confidence", 0.0)

                # 关键词分类
                category = self._classify_by_keywords(ocr_text, filename)

                # 提取日期和金额
                date = self._extract_date(ocr_text)
                amount = self._extract_amount(ocr_text)

                # 对账单检测签名
                signed: bool | None = None
                if category == TYPE_STATEMENT:
                    signed = self._detect_signed(ocr_text)

                result.images.append(
                    ClassifiedImage(
                        filename=filename,
                        category=category,
                        ocr_text=ocr_text,
                        date=date,
                        amount=amount,
                        signed=signed,
                        confidence=confidence,
                        image_data=data,
                        rotation=rotation,
                    )
                )
            except Exception as e:
                logger.warning("分类失败: %s - %s", filename, e)
                result.errors.append(f"{filename}: {e!s}")
                result.images.append(
                    ClassifiedImage(
                        filename=filename,
                        category=TYPE_OTHER,
                        ocr_text="",
                        image_data=data,
                    )
                )

        logger.info(
            "分类完成: %d 张图片, %d 个错误",
            len(result.images),
            len(result.errors),
        )
        return result

    def _classify_by_keywords(self, text: str, filename: str = "") -> str:
        """关键词匹配分类，返回匹配数最多的类型"""
        if not text:
            # OCR 识别不出文字的 PNG 文件，大概率是对账单截图
            if filename.lower().endswith(".png"):
                return TYPE_STATEMENT
            return TYPE_OTHER

        scores: dict[str, int] = dict.fromkeys(_KEYWORDS, 0)
        for category, keywords in _KEYWORDS.items():
            for kw in keywords:
                if kw in text:
                    scores[category] += 1

        best = max(scores, key=lambda k: scores[k])
        if scores[best] == 0:
            return TYPE_OTHER
        return best

    def _extract_date(self, text: str) -> str | None:
        """从 OCR 文本提取日期 → YYYYMMDD"""
        # 匹配 YYYY年MM月DD日 / YYYY-MM-DD / YYYY/MM/DD
        patterns = [
            r"(\d{4})\s*年\s*(\d{1,2})\s*月\s*(\d{1,2})\s*日",
            r"(\d{4})[-/](\d{1,2})[-/](\d{1,2})",
        ]
        for pat in patterns:
            m = re.search(pat, text)
            if m:
                y, mo, d = m.group(1), m.group(2).zfill(2), m.group(3).zfill(2)
                return f"{y}{mo}{d}"
        return None

    def _extract_amount(self, text: str) -> str | None:
        """从 OCR 文本提取金额"""
        # 匹配 ¥123,456.78 / 123456.78元 / 金额:123456 等
        patterns = [
            r"[¥￥]\s*([\d,]+\.?\d*)",
            r"([\d,]+\.?\d*)\s*元",
            r"金额[：:]\s*([\d,]+\.?\d*)",
            r"合计[：:]\s*([\d,]+\.?\d*)",
            r"总计[：:]\s*([\d,]+\.?\d*)",
        ]
        amounts: list[float] = []
        for pat in patterns:
            for m in re.finditer(pat, text):
                try:
                    val = float(m.group(1).replace(",", ""))
                    amounts.append(val)
                except ValueError:
                    pass
        if amounts:
            # 返回最大金额
            best = max(amounts)
            # 保留小数（如果有）
            if best == int(best):
                return str(int(best))
            return str(best)
        return None

    def _detect_signed(self, text: str) -> bool:
        """检测对账单是否有签名痕迹（通过 OCR 文本中的签名相关关键词）"""
        sign_keywords = ["签名", "签章", "盖章", "确认", "签字"]
        # 如果 OCR 文本中有大量手写体识别（置信度低的文字），可能有签名
        # 简单策略：检查关键词
        return any(kw in text for kw in sign_keywords)

	"""OCR 关键词分类服务 - 将图片分为对账单/出库单/收款凭证/其他"""

	from __future__ import annotations

	import base64
	import logging
	import re
	from dataclasses import dataclass, field
	from importlib import import_module
	from typing import Any

	logger = logging.getLogger("apps.evidence_sorting")

	# 文件类型常量
	TYPE_STATEMENT = "statement" # 对账单
	TYPE_DELIVERY = "delivery" # 出库单/出仓单
	TYPE_RECEIPT = "receipt" # 收款凭证
	TYPE_OTHER = "other" # 其他

	# 关键词权重表
	_KEYWORDS: dict[str, list[str]] = {
	TYPE_STATEMENT: ["对账单", "对帐单", "对账", "月结", "月度汇总"],
	TYPE_DELIVERY: [
	"出库单",
	"出仓单",
	"出货单",
	"发货单",
	"送货单",
	"承运单",
	"提货单",
	"出库",
	"出仓",
	],
	TYPE_RECEIPT: [
	"收款",
	"转账",
	"付款",
	"汇款",
	"收到",
	"到账",
	"银行回单",
	"电子回单",
	"交易成功",
	"支付成功",
	"中国农业银行",
	"中国工商银行",
	"中国建设银行",
	"中国银行",
	"招商银行",
	"微信支付",
	"支付宝",
	],
	}


	@dataclass
	class ClassifiedImage:
	"""分类后的图片"""

	filename: str
	category: str # statement / delivery / receipt / other
	ocr_text: str
	date: str \| None = None # YYYYMMDD
	amount: str \| None = None # 如 "65500"
	signed: bool \| None = None # 对账单是否签名
	confidence: float = 0.0
	image_data: str = "" # base64
	rotation: int = 0


	@dataclass
	class ClassifyResult:
	"""分类结果"""

	images: list[ClassifiedImage] = field(default_factory=list)
	errors: list[str] = field(default_factory=list)


	class ClassifierService:
	"""图片 OCR + 关键词分类"""

	def __init__(self) -> None:
	self._ocr_service: Any \| None = None

	def _get_ocr_service(self) -> Any:
	if self._ocr_service is None:
	module = import_module("apps.image_rotation.services.orientation.service")
	orientation_service_cls = module.OrientationDetectionService
	self._ocr_service = orientation_service_cls()
	return self._ocr_service

	def classify_images(self, images: list[dict[str, Any]]) -> ClassifyResult:
	"""
	批量 OCR + 分类

	Args:
	images: [{"filename": str, "data": str(base64)}]

	Returns:
	ClassifyResult
	"""
	result = ClassifyResult()
	ocr_svc = self._get_ocr_service()

	for idx, img in enumerate(images):
	filename: str = img.get("filename", f"image_{idx}")
	data: str = img.get("data", "")
	try:
	# 去掉 data:image/xxx;base64, 前缀
	raw_data = data.split(",", 1)[-1] if "," in data else data
	image_bytes = base64.b64decode(raw_data)

	# OCR + 方向检测
	ocr_result = ocr_svc.detect_orientation_with_text(image_bytes)
	ocr_text: str = ocr_result.get("ocr_text", "")
	rotation: int = ocr_result.get("rotation", 0)
	confidence: float = ocr_result.get("confidence", 0.0)

	# 关键词分类
	category = self._classify_by_keywords(ocr_text, filename)

	# 提取日期和金额
	date = self._extract_date(ocr_text)
	amount = self._extract_amount(ocr_text)

	# 对账单检测签名
	signed: bool \| None = None
	if category == TYPE_STATEMENT:
	signed = self._detect_signed(ocr_text)

	result.images.append(
	ClassifiedImage(
	filename=filename,
	category=category,
	ocr_text=ocr_text,
	date=date,
	amount=amount,
	signed=signed,
	confidence=confidence,
	image_data=data,
	rotation=rotation,
	)
	)
	except Exception as e:
	logger.warning("分类失败: %s - %s", filename, e)
	result.errors.append(f"{filename}: {e!s}")
	result.images.append(
	ClassifiedImage(
	filename=filename,
	category=TYPE_OTHER,
	ocr_text="",
	image_data=data,
	)
	)

	logger.info(
	"分类完成: %d 张图片, %d 个错误",
	len(result.images),
	len(result.errors),
	)
	return result

	def _classify_by_keywords(self, text: str, filename: str = "") -> str:
	"""关键词匹配分类，返回匹配数最多的类型"""
	if not text:
	# OCR 识别不出文字的 PNG 文件，大概率是对账单截图
	if filename.lower().endswith(".png"):
	return TYPE_STATEMENT
	return TYPE_OTHER

	scores: dict[str, int] = dict.fromkeys(_KEYWORDS, 0)
	for category, keywords in _KEYWORDS.items():
	for kw in keywords:
	if kw in text:
	scores[category] += 1

	best = max(scores, key=lambda k: scores[k])
	if scores[best] == 0:
	return TYPE_OTHER
	return best

	def _extract_date(self, text: str) -> str \| None:
	"""从 OCR 文本提取日期 → YYYYMMDD"""
	# 匹配 YYYY年MM月DD日 / YYYY-MM-DD / YYYY/MM/DD
	patterns = [
	r"(\d{4})\s年\s(\d{1,2})\s月\s(\d{1,2})\s*日",
	r"(\d{4})[-/](\d{1,2})[-/](\d{1,2})",
	]
	for pat in patterns:
	m = re.search(pat, text)
	if m:
	y, mo, d = m.group(1), m.group(2).zfill(2), m.group(3).zfill(2)
	return f"{y}{mo}{d}"
	return None

	def _extract_amount(self, text: str) -> str \| None:
	"""从 OCR 文本提取金额"""
	# 匹配 ¥123,456.78 / 123456.78元 / 金额:123456 等
	patterns = [
	r"[¥￥]\s([\d,]+\.?\d)",
	r"([\d,]+\.?\d)\s元",
	r"金额[：:]\s([\d,]+\.?\d)",
	r"合计[：:]\s([\d,]+\.?\d)",
	r"总计[：:]\s([\d,]+\.?\d)",
	]
	amounts: list[float] = []
	for pat in patterns:
	for m in re.finditer(pat, text):
	try:
	val = float(m.group(1).replace(",", ""))
	amounts.append(val)
	except ValueError:
	pass
	if amounts:
	# 返回最大金额
	best = max(amounts)
	# 保留小数（如果有）
	if best == int(best):
	return str(int(best))
	return str(best)
	return None

	def _detect_signed(self, text: str) -> bool:
	"""检测对账单是否有签名痕迹（通过 OCR 文本中的签名相关关键词）"""
	sign_keywords = ["签名", "签章", "盖章", "确认", "签字"]
	# 如果 OCR 文本中有大量手写体识别（置信度低的文字），可能有签名
	# 简单策略：检查关键词
	return any(kw in text for kw in sign_keywords)