import easyocr

import os

import re

from PIL import Image

import cv2

import numpy as np

import logging

# 禁用不必要的日志记录

logging.disable(logging.WARNING)

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

def preprocess_image(image_path):

    """图像预处理以提高OCR准确率"""

    try:

        # 使用 OpenCV 读取图像

        img = cv2.imread(image_path)

        if img is None:

            # 如果 OpenCV 无法读取，使用 PIL

            pil_img = Image.open(image_path)

            img = np.array(pil_img)

            # 转换 BGR 格式（OpenCV 使用 BGR）

            if len(img.shape) == 3:

                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # 转换为灰度图

        if len(img.shape) == 3:

            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        else:

            gray = img

        # 提高对比度

        gray = cv2.convertScaleAbs(gray, alpha=1.5, beta=0)

        # 二值化处理

        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        return binary

    except Exception as e:

        print(f"图像预处理失败: {e}")

        # 如果预处理失败，返回原始图像

        return cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

def clean_filename(text):

    """清理文本，使其适合作为文件名"""

    # 移除 Windows 文件名中不允许的字符

    invalid_chars = '<>:"/\\|?*\''

    for char in invalid_chars:

        text = text.replace(char, '')

    # 移除换行符和制表符

    text = text.replace('\n', '').replace('\r', '').replace('\t', '')

    # 移除首尾空格

    return text.strip()

def extract_keyword_text(text, keyword):

    """从文本中提取包含关键字在内的20个字符，忽略所有特殊字符"""

    if not text or not keyword:

        return None

    # 查找关键字位置（不区分大小写）

    pattern = re.compile(re.escape(keyword), re.IGNORECASE)

    match = pattern.search(text)

    if not match:

        return None

    # 获取关键字开始的位置

    start_pos = match.start()

    # 从关键字开始的位置提取文本

    remaining_text = text[start_pos:]

    # 移除所有非字母数字字符（只保留字母和数字）

    # 使用正则表达式移除所有非字母数字字符

    cleaned_text = re.sub(r'[^A-Za-z0-9]', '', remaining_text)

    # 取前 20 个字符

    return cleaned_text[:20]

def main():

    # 设置模型路径（确保您已经将模型文件复制到这个位置）

    model_storage_directory = r'C:\Users\Administrator\.EasyOCR\model'

    # 检查模型文件是否存在

    if not os.path.exists(model_storage_directory):

        print(f"错误: 模型目录不存在: {model_storage_directory}")

        print("请先将模型文件复制到该目录")

        return

    # 初始化 EasyOCR 阅读器，指定模型路径

    try:

        print("正在初始化OCR引擎...")

        reader = easyocr.Reader(

            ['ch_sim', 'en'], 

            gpu=False,

            download_enabled=False,  # 禁用下载

            model_storage_directory=model_storage_directory

        )

        print("OCR引擎初始化完成")

    except Exception as e:

        print(f"OCR引擎初始化失败: {e}")

        print("请检查模型文件是否正确放置")

        return

    # 指定图片文件夹路径

    folder_path = input("请输入合同图片文件夹路径: ").strip().strip('"')

    # 确保文件夹存在

    if not os.path.exists(folder_path):

        print(f"错误: 文件夹不存在: {folder_path}")

        return

    # 获取用户输入的关键字

    keyword = input("请输入验收单号的关键字: ").strip()

    print(f"开始处理文件夹: {folder_path}")

    print(f"使用关键字: {keyword}")

    # 支持的图片格式

    image_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif')

    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(image_extensions)]

    if not image_files:

        print("未找到图片文件")

        return

    print(f"找到 {len(image_files)} 个图片文件")

    success_count = 0

    processed_count = 0

    for filename in image_files:

        file_path = os.path.join(folder_path, filename)

        print(f"\n正在处理 ({processed_count + 1}/{len(image_files)}): {filename}")

        processed_count += 1

        try:

            # 图像预处理

            processed_img = preprocess_image(file_path)

            # 使用 EasyOCR 进行识别

            print("正在进行OCR识别...")

            results = reader.readtext(processed_img, detail=0)

            text = '\n'.join(results)

            if text:

                print(f"识别到的文本:\n{text[:25]}...")  # 只显示前 25 个字符

            else:

                print("未识别到文本")

                continue

            # 提取关键字后的文本

            extracted_text = extract_keyword_text(text, keyword)

            if extracted_text:

                print(f"找到提取的文本: {extracted_text}")

                # 清理文本，确保它可以作为文件名

                clean_text = clean_filename(extracted_text)

                # 构建新文件名

                file_extension = os.path.splitext(filename)[1]

                new_filename = f"{clean_text}{file_extension}"

                new_file_path = os.path.join(folder_path, new_filename)

                # 检查是否已存在同名文件

                counter = 1

                while os.path.exists(new_file_path):

                    new_filename = f"{clean_text}_{counter}{file_extension}"

                    new_file_path = os.path.join(folder_path, new_filename)

                    counter += 1

                # 重命名文件

                os.rename(file_path, new_file_path)

                success_count += 1

                print(f"成功重命名: {filename} -> {new_filename}")

            else:

                print(f"未找到包含关键字 '{keyword}' 的文本")

        except Exception as e:

            print(f"处理 {filename} 时出错: {str(e)}")

    print(f"\n处理完成! 成功重命名 {success_count}/{processed_count} 个文件")

if __name__ == "__main__":

    main()