data_extractor.py

"""
数据提取器
从OCR结果中提取订单关键信息
"""
import re
from datetime import datetime


class DataExtractor:
    def __init__(self):
        # 定义提取规则
        self.patterns = {
            'order_no': [
                r'订单[号编号]?[：:]?\s*([A-Za-z0-9\-]+)',
                r'单号[：:]?\s*([A-Za-z0-9\-]+)',
                r'Order[No\.]?[：:]?\s*([A-Za-z0-9\-]+)',
                r'PO[：:]?\s*([A-Za-z0-9\-]+)',
            ],
            'customer': [
                r'客户[名称]?[：:]?\s*([^\n]+)',
                r'买方[：:]?\s*([^\n]+)',
                r'采购[方单位]?[：:]?\s*([^\n]+)',
                r'Customer[：:]?\s*([^\n]+)',
            ],
            'amount': [
                r'金额[：:]?\s*[¥￥]?\s*([\d,]+\.?\d*)',
                r'总价[：:]?\s*[¥￥]?\s*([\d,]+\.?\d*)',
                r'合计[：:]?\s*[¥￥]?\s*([\d,]+\.?\d*)',
                r'Amount[：:]?\s*\$?\s*([\d,]+\.?\d*)',
            ],
            'date': [
                r'日期[：:]?\s*(\d{4}[-/年]\d{1,2}[-/月]\d{1,2})',
                r'下单时间[：:]?\s*(\d{4}[-/年]\d{1,2}[-/月]\d{1,2})',
                r'Date[：:]?\s*(\d{4}[-/]\d{1,2}[-/]\d{1,2})',
            ],
            'product': [
                r'品[名][：:]?\s*([^\n]+)',
                r'货物[名称]?[：:]?\s*([^\n]+)',
                r'Product[：:]?\s*([^\n]+)',
            ],
            'quantity': [
                r'数量[：:]?\s*(\d+)',
                r'Qty[：:]?\s*(\d+)',
            ]
        }

    def extract(self, ocr_results):
        """
        从OCR结果中提取结构化数据
        :param ocr_results: OCR识别结果列表
        :return: 提取的订单信息字典
        """
        # 合并所有文本
        full_text = ' '.join([item['text'] for item in ocr_results])

        extracted = {
            'order_no': '',
            'customer': '',
            'amount': '',
            'date': '',
            'product': '',
            'quantity': '',
            'raw_text': full_text,
            'extract_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }

        # 按规则提取
        for field, patterns in self.patterns.items():
            for pattern in patterns:
                match = re.search(pattern, full_text, re.IGNORECASE)
                if match:
                    value = match.group(1).strip()
                    # 清理金额中的逗号
                    if field == 'amount':
                        value = value.replace(',', '')
                    extracted[field] = value
                    break

        return extracted

    def extract_batch(self, batch_results):
        """
        批量提取
        :param batch_results: {图片路径: OCR结果}
        :return: [提取的数据列表]
        """
        all_data = []
        for image_path, ocr_result in batch_results.items():
            if isinstance(ocr_result, dict) and 'error' in ocr_result:
                all_data.append({
                    'image_path': image_path,
                    'error': ocr_result['error']
                })
            else:
                data = self.extract(ocr_result)
                data['image_path'] = image_path
                all_data.append(data)
        return all_data