WEIUI 的安装方式为Docker 安装才可以使用


title: 高精度 OCR 文字提取
author: kmsbox.com
version: 2.1
"""

import os
import fitz
import pytesseract
from PIL import Image, ImageOps, ImageEnhance, ImageFilter
import glob
import time


class Tools:
    def __init__(self):
        # 容器内 Tesseract 默认路径已由系统环境配置
        pass

    def extract_text(self, __user__: dict = None) -> str:
        """
        提取最新上传图片或 PDF 中的文字。已通过隐藏类型注解修复 Pydantic 架构解析错误。
        """
        # 1. 定义 Docker 内部上传目录
        target_dir = "/app/backend/data/uploads"
        if not os.path.exists(target_dir):
            return f"错误：容器内未找到上传目录 {target_dir}"

        # 2. 扫描并获取最新文件
        files = glob.glob(os.path.join(target_dir, "*.*"))
        if not files:
            return "错误：上传目录为空，请确认图片已成功上传并在输入框上方显示。"

        latest_file = max(files, key=os.path.getmtime)
        file_name = os.path.basename(latest_file)

        try:
            # 3. PDF 处理逻辑
            if latest_file.lower().endswith(".pdf"):
                doc = fitz.open(latest_file)
                full_text = ""
                for page in doc:
                    # 提高 PDF 渲染清晰度
                    pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    full_text += self._execute_ocr_logic(img)
                doc.close()
                return f"【PDF 识别结果 - {file_name}】\n\n{full_text}"

            # 4. 图片处理逻辑
            else:
                img = Image.open(latest_file)
                text = self._execute_ocr_logic(img)

                # 识别完成后，建议在对话中引导 AI 进行语义纠错
                return f"【图片识别结果 - {file_name}】\n\n{text}"

        except Exception as e:
            return f"工具执行出错。文件名: {file_name}, 错误详情: {str(e)}"

    def _execute_ocr_logic(self, img_obj) -> str:
        """
        私有方法：执行核心 OCR。
        注意：此处不写参数类型注解（如 img_obj: Image），以绕过 Pydantic 的自动 Schema 检查。
        """
        try:
            # --- 图像预处理增强 ---
            # A. 转为 RGB 并大幅度放大（3倍），增加中文字符像素密度
            img = img_obj.convert("RGB")
            w, h = img.size
            img = img.resize((w * 3, h * 3), Image.Resampling.LANCZOS)

            # B. 边缘锐化处理
            img = img.filter(ImageFilter.SHARPEN)

            # C. 极致对比度增强，过滤背景杂色
            img = ImageEnhance.Contrast(img).enhance(2.8)
            img = ImageEnhance.Brightness(img).enhance(1.1)

            # D. 二值化处理（将图片转为纯黑白，阈值设为 150）
            img = img.convert("L").point(lambda x: 0 if x < 150 else 255, "1")

            # --- 调用 Tesseract 引擎 ---
            # --psm 6: 假设为单一字体块（最适合股票/行情表格）
            # lang='chi_sim+eng': 中英文混合识别
            custom_config = r"--oem 3 --psm 6"
            result = pytesseract.image_to_string(
                img, lang="chi_sim+eng", config=custom_config
            )

            return result
        except Exception as e:
            return f"(预处理/OCR 阶段失败: {str(e)})"