import os
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_prim_skip_onednn"] = "1"
os.environ["FLAGS_enable_pir_api"] = "0"

import cv2
import re
from paddleocr import PaddleOCR

# 初始化 PaddleOCR（只需初始化一次，可复用）
ocr = PaddleOCR(use_angle_cls=True, lang='ch')

def extract_text_from_image(image_path):
    """使用 PaddleOCR 从图片中提取文本"""
    try:
        # 使用 PaddleOCR 进行识别
        result = ocr.ocr(image_path)

        # 提取所有文本
        text_list = []
        if result and result[0]:
            for line in result[0]:
                text_list.append(line[1][0])

        # 合并所有识别的文本
        full_text = '\n'.join(text_list)
        return full_text.strip()
    except Exception as e:
        print(f"OCR 识别错误: {str(e)}")
        return ""

def extract_id_card_info(image_path):
    """
    从身份证图片中提取姓名和身份证号
    返回: {'name': 姓名, 'id_number': 身份证号}
    """
    try:
        # 使用 PaddleOCR 识别身份证
        result = ocr.ocr(image_path)

        name = None
        id_number = None

        if result and result[0]:
            for line in result[0]:
                text = line[1][0]
                confidence = line[1][1]

                # 识别身份证号（18位或15位）
                id_match = re.search(r'\d{17}[\dxX]|\d{15}', text)
                if id_match and confidence > 0.8:
                    id_number = id_match.group(0).upper()

                # 识别姓名（通常在"姓名"关键字后面）
                if '姓名' in text:
                    # 获取姓名字段
                    name_match = re.search(r'姓名[:\s]*([^\s]+)', text)
                    if name_match:
                        name = name_match.group(1)
                # 或者直接识别中文姓名（2-4个汉字）
                elif not name and confidence > 0.9:
                    name_match = re.search(r'^[一-龥]{2,4}$', text)
                    if name_match:
                        name = name_match.group(0)

        return {
            'name': name,
            'id_number': id_number
        }
    except Exception as e:
        print(f"身份证识别错误: {str(e)}")
        return {
            'name': None,
            'id_number': None
        }

# 示例调用
if __name__ == "__main__":
    # 测试普通 OCR
    text = extract_text_from_image('path_to_image.png')
    print("提取的文本:", text)

    # 测试身份证识别
    id_info = extract_id_card_info('path_to_id_card.jpg')
    print("身份证信息:", id_info)