subocr/ocr_engine.py at main · ezatgz/subocr · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import sys
import os
import json
from paddleocr import PaddleOCR
from PySide6.QtCore import QThread, Signal
import traceback


def get_resource_path(relative_path):
    """获取资源文件的绝对路径，支持开发环境和打包后的exe环境"""
    try:
        # PyInstaller创建临时文件夹，并将路径存储在_MEIPASS中
        base_path = sys._MEIPASS
    except Exception:
        # 如果不是打包环境，使用当前脚本目录
        base_path = os.path.dirname(os.path.abspath(__file__))

    return os.path.join(base_path, relative_path)


# 从配置文件读取text_det_limit_side_len参数
def load_text_det_limit():
    try:
        if os.path.exists('settings.json'):
            with open('settings.json', 'r', encoding='utf-8') as f:
                config = json.load(f)
                return config.get('text_det_limit_side_len', 1280)
    except Exception as e:
        pass  # 静默处理配置文件读取失败
    return 1280

# PaddleOCR 模型全局加载一次，避免重复加载
# 首次运行会下载模型，请耐心等待
try:
    default_text_det_limit = load_text_det_limit()

    # 获取模型文件的正确路径
    det_model_dir = get_resource_path("ocrmodels/PP-OCRv5_mobile_det")
    rec_model_dir = get_resource_path("ocrmodels/PP-OCRv5_mobile_rec")

    ocr_engine = PaddleOCR(
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
        text_detection_model_name="PP-OCRv5_mobile_det",
        text_recognition_model_name="PP-OCRv5_mobile_rec",
        lang="cn",
        #lang="chinese_cht",
        #lang="en"
        cpu_threads=8,
        text_det_limit_side_len=default_text_det_limit,
        text_detection_model_dir=det_model_dir,
        text_recognition_model_dir=rec_model_dir,
        text_det_limit_type="max"
        # 移除enable_mkldnn=False以测试是否能正常运行
    )

except Exception as e:
    ocr_engine = None


class OcrThread(QThread):
    """
    用于在后台执行 OCR 任务的线程
    """
    ocr_finished_signal = Signal(str, str)  # filepath, ocr_text

    def __init__(self, image_path, output_directory="", text_det_limit_side_len=1920):
        super().__init__()
        self.image_path = image_path
        self.output_directory = output_directory
        self.text_det_limit_side_len = text_det_limit_side_len

    def run(self):
        if not ocr_engine:
            self.ocr_finished_signal.emit(self.image_path, "OCR引擎初始化失败。")
            return
        try:
            # 根据demo.py中的示例，使用input参数
            # result = ocr_engine.predict(input=self.image_path)
            result = ocr_engine.predict(input=self.image_path, text_det_limit_side_len=self.text_det_limit_side_len, text_det_limit_type="max")
            ocr_text = ""
            # 根据demo.py中的示例处理结果
            if isinstance(result, list) and len(result) > 0:
                # 取第一个元素（即你的完整 JSON 字典）
                ocr_result_dict = result[0]
                # 提取 rec_texts
                rec_texts = ocr_result_dict.get("rec_texts", [])
                ocr_text = "\n".join(rec_texts)

                # 将OCR结果保存为txt文件
                txt_filename = os.path.splitext(os.path.basename(self.image_path))[0] + ".txt"
                if self.output_directory:
                    # 如果指定了输出目录，则保存到输出目录
                    txt_file_path = os.path.join(self.output_directory, txt_filename)
                    # 确保输出目录存在
                    os.makedirs(self.output_directory, exist_ok=True)
                else:
                    # 否则保存到原文件相同目录
                    txt_file_path = os.path.splitext(self.image_path)[0] + ".txt"

                # 只有当ocr_text不为空时才保存
                if ocr_text:
                    try:
                        with open(txt_file_path, 'w', encoding='utf-8') as f:
                            f.write(ocr_text)
                    except Exception as e:
                        pass  # 静默处理保存失败

            self.ocr_finished_signal.emit(self.image_path, ocr_text)
        except Exception as e:
            self.ocr_finished_signal.emit(self.image_path, f"OCR 识别失败: {e}")