AndroidRemoteController/python/onnxocr/onnxocr/ocr_images_pdfs.py

# logic.py
import sys
import os
# 添加父目录到sys.path，便于导入onnxocr包
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from onnxocr.onnx_paddleocr import ONNXPaddleOcr, sav2Img
import cv2
from typing import List, Callable
from pathlib import Path
import time
import numpy as np

# 尝试导入pdf2image用于PDF转图片
try:
    from pdf2image import convert_from_path
except ImportError:
    convert_from_path = None

# 尝试导入pymupdf用于PDF转图片
try:
    import fitz  # pymupdf
    def pdf_to_images(pdf_path, dpi=200):
        """
        使用pymupdf将PDF每一页转为图片（numpy数组）
        """
        doc = fitz.open(pdf_path)
        images = []
        for page in doc:
            pix = page.get_pixmap(dpi=dpi)
            img = np.frombuffer(pix.samples, dtype=np.uint8)
            img = img.reshape((pix.height, pix.width, pix.n))
            if pix.n == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            images.append(img)
        return images
except ImportError:
    pdf_to_images = None

class OCRLogic:
    """
    OCR 业务逻辑主类，支持批量图片/PDF识别，多线程加速，模型热切换等
    """
    def __init__(self, status_callback: Callable[[str], None]):
        """
        初始化，传入状态回调函数用于UI进度提示
        """
        self.status_callback = status_callback
        # 默认初始化OCR模型
        self.model = ONNXPaddleOcr(use_angle_cls=True, use_gpu=False)

    def run(self, files: List[str], save_txt: bool, merge_txt: bool, output_img: bool = False, file_time_callback=None, pdf_progress_callback=None, max_workers: int = 4):
        """
        批量图片/PDF识别主入口，支持多线程加速
        files: 待识别文件路径列表
        save_txt: 是否保存txt
        merge_txt: 是否合并为一个txt
        output_img: 是否输出带框图片
        file_time_callback: 单文件识别耗时回调
        pdf_progress_callback: PDF页进度回调
        max_workers: 最大线程数，默认4
        """
        import concurrent.futures
        start_time = time.time()
        all_text = [None] * len(files)  # 用于顺序合并结果
        def process_one(idx_file):
            idx, file = idx_file
            ext = os.path.splitext(file)[1].lower()
            self.status_callback(f"正在处理: {os.path.basename(file)} ({idx+1}/{len(files)})")
            t0 = time.time()
            text = ""
            if ext == ".pdf":
                # PDF转图片后识别
                if pdf_to_images is None:
                    raise RuntimeError("未安装pymupdf库，无法处理PDF文件。请先安装pymupdf。")
                images = pdf_to_images(file, dpi=300)
                text = self._ocr_images(images, file, save_txt, merge_txt, output_img=output_img, is_pdf=True, pdf_progress_callback=pdf_progress_callback, max_workers=max_workers)
            else:
                # 普通图片识别，兼容中文路径
                try:
                    if file.lower().endswith('.bmp'):
                        img = cv2.imdecode(np.fromfile(file, dtype=np.uint8), cv2.IMREAD_COLOR)
                    else:
                        with open(file, 'rb') as fimg:
                            img_array = np.frombuffer(fimg.read(), np.uint8)
                        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                except Exception as e:
                    self.status_callback(f"图片读取失败: {file}，错误: {e}")
                    if file_time_callback:
                        file_time_callback(idx, 0)
                    return (idx, "")
                if img is None:
                    self.status_callback(f"文件无法读取或不是有效图片: {file}")
                    if file_time_callback:
                        file_time_callback(idx, 0)
                    return (idx, "")
                text = self._ocr_image(img, file, save_txt, output_img=output_img)
            t1 = time.time()
            if file_time_callback:
                file_time_callback(idx, t1-t0)
            self.status_callback(f"{os.path.basename(file)} 识别用时: {t1-t0:.2f} 秒")
            if len(files) > 1:
                avg = (t1 - start_time) / (idx + 1)
                self.status_callback(f"已完成 {idx+1}/{len(files)}，平均单张用时: {avg:.2f} 秒")
            return (idx, text)
        # 多线程处理所有文件，结果按索引回填，保证顺序
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(process_one, (idx, file)) for idx, file in enumerate(files)]
            for future in concurrent.futures.as_completed(futures):
                idx, text = future.result()
                all_text[idx] = text
        # 合并写入txt
        if save_txt and merge_txt and len(files) > 1:
            out_dir = self._get_output_dir(files[0])
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            out_txt = os.path.join(out_dir, f"merged_ocr_{timestamp}.txt")
            with open(out_txt, "w", encoding="utf-8") as f:
                for text in all_text:
                    if text:
                        f.write(text)
                        f.write("\n\n")
        elapsed = time.time() - start_time
        if files:
            out_dir = self._get_output_dir(files[0])
            self.status_callback(f"识别完成，总耗时：{elapsed:.2f}秒，文件保存在：{out_dir}")
        else:
            self.status_callback(f"识别完成，总耗时：{elapsed:.2f}秒")

    def _ocr_images(self, images, pdf_path, save_txt, merge_txt, output_img=False, is_pdf=False, pdf_progress_callback=None, max_workers: int = 4):
        """
        PDF转图片后，批量图片识别，支持多线程加速
        images: PDF每页图片（numpy数组）
        pdf_path: 原PDF路径
        save_txt: 是否保存txt
        merge_txt: 是否合并txt（未用）
        output_img: 是否输出带框图片
        pdf_progress_callback: 页进度回调
        max_workers: 最大线程数，默认4
        """
        import concurrent.futures
        out_dir = self._get_output_dir(pdf_path)
        pdf_text = [None] * len(images)
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        total = len(images)
        def process_page(i_img):
            i, img = i_img
            img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            result = self.model.ocr(img_cv)
            if output_img:
                out_img_path = os.path.join(out_dir, f"{Path(pdf_path).stem}_page{i+1}_ocr.jpg")
                sav2Img(img_cv, result, name=out_img_path)
            page_text = self._result_to_text(result)
            return (i, page_text)
        # 多线程识别每一页，结果按页码顺序合并
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(process_page, (i, img)) for i, img in enumerate(images)]
            for future in concurrent.futures.as_completed(futures):
                i, page_text = future.result()
                pdf_text[i] = page_text
                if pdf_progress_callback:
                    pdf_progress_callback(i + 1, total)
        if save_txt:
            txt_path = os.path.join(out_dir, f"{Path(pdf_path).stem}_ocr_{timestamp}.txt")
            with open(txt_path, "w", encoding="utf-8") as f:
                f.write("\n\n".join(pdf_text))
        return "\n\n".join(pdf_text)

    def _ocr_image(self, img, img_path, save_txt, output_img=False):
        """
        单张图片OCR识别，支持保存txt和输出带框图片
        """
        out_dir = self._get_output_dir(img_path)
        result = self.model.ocr(img)
        if output_img:
            out_img_path = os.path.join(out_dir, f"{Path(img_path).stem}_ocr.jpg")
            sav2Img(img, result, name=out_img_path)
        text = self._result_to_text(result)
        if save_txt:
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            txt_path = os.path.join(out_dir, f"{Path(img_path).stem}_ocr_{timestamp}.txt")
            with open(txt_path, "w", encoding="utf-8") as f:
                f.write(text)
        return text

    def _result_to_text(self, result):
        """
        将OCR识别结果结构化为纯文本，兼容只检测无识别内容的情况
        """
        # 健壮性检查，防止result为空或结构异常
        if not result or not isinstance(result, list) or not result[0] or not isinstance(result[0], list):
            return "[未检测到内容]"
        lines = []
        for box in result[0]:
            # 兼容只检测无识别内容的情况
            if isinstance(box, list) and len(box) == 2 and isinstance(box[1], (list, tuple)) and len(box[1]) >= 1:
                lines.append(str(box[1][0]))
            elif isinstance(box, list) and (isinstance(box[0], (list, tuple)) or isinstance(box[0], float)):
                # 只有检测框，无识别内容
                lines.append("[未识别] " + str(box))
            else:
                lines.append(str(box))
        return "\n".join(lines)

    def _get_output_dir(self, file_path):
        """
        获取输出目录，自动创建
        """
        base_dir = os.path.dirname(file_path)
        out_dir = os.path.join(base_dir, "Output_OCR")
        os.makedirs(out_dir, exist_ok=True)
        return out_dir

    def set_model(self, model_name, use_gpu=False):
        """
        切换OCR模型，支持多模型热切换，所有模型统一用ppocrv5字典
        use_gpu: 是否启用GPU
        """
        import os
        import tkinter.messagebox as messagebox
        base_model_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "onnxocr", "models"))
        model_map = {
            "PP-OCRv5": "ppocrv5",
            "PP-OCRv4": "ppocrv4",
            "ch_ppocr_server_v2.0": "ch_ppocr_server_v2.0"
        }
        model_dir = model_map.get(model_name, "ppocrv5")
        model_path = os.path.join(base_model_dir, model_dir)
        det_model_dir = os.path.join(model_path, "det", "det.onnx")
        cls_model_dir = os.path.join(model_path, "cls", "cls.onnx")
        rec_char_dict_path = os.path.join(base_model_dir, "ppocrv5", "ppocrv5_dict.txt")
        rec_model_dir = os.path.join(model_path, "rec", "rec.onnx") if os.path.exists(os.path.join(model_path, "rec", "rec.onnx")) else None
        ocr_kwargs = dict(
            use_angle_cls=True,
            use_gpu=use_gpu,  # 关键：传递GPU参数
            det_model_dir=det_model_dir,
            cls_model_dir=cls_model_dir,
            rec_char_dict_path=rec_char_dict_path
        )
        if rec_model_dir and os.path.exists(rec_model_dir):
            ocr_kwargs["rec_model_dir"] = rec_model_dir
        try:
            self.model = ONNXPaddleOcr(**ocr_kwargs)
            if use_gpu:
                try:
                    import onnxruntime as ort
                    providers = self.model.session.get_providers() if hasattr(self.model, 'session') else []
                    if not any('CUDA' in p for p in providers):
                        msg = ("未检测到可用GPU，已自动切换为CPU推理。请检查CUDA/cuDNN环境配置。")
                        if hasattr(self, 'ui_ref') and hasattr(self.ui_ref, 'update_gpu_status'):
                            self.ui_ref.update_gpu_status(msg)
                        if hasattr(self, 'status_callback'):
                            self.status_callback("[警告] 未检测到可用GPU，已切换为CPU推理。请检查CUDA/cuDNN环境配置。")
                except Exception:
                    msg = ("检测GPU状态时发生异常，可能未正确安装CUDA/cuDNN或onnxruntime-gpu。已自动切换为CPU推理。")
                    if hasattr(self, 'ui_ref') and hasattr(self.ui_ref, 'update_gpu_status'):
                        self.ui_ref.update_gpu_status(msg)
                    if hasattr(self, 'status_callback'):
                        self.status_callback("[警告] GPU检测异常，已切换为CPU推理。请检查CUDA/cuDNN环境配置。")
        except Exception as e:
            if use_gpu:
                msg = f"GPU初始化失败，已自动切换为CPU。请检查CUDA/cuDNN环境配置。错误信息: {e}"
                if hasattr(self, 'ui_ref') and hasattr(self.ui_ref, 'update_gpu_status'):
                    self.ui_ref.update_gpu_status(msg)
                if hasattr(self, 'status_callback'):
                    self.status_callback("[警告] GPU初始化失败，已切换为CPU推理。请检查CUDA/cuDNN环境配置。")
                ocr_kwargs["use_gpu"] = False
                self.model = ONNXPaddleOcr(**ocr_kwargs)
            else:
                raise