# pip install pix2text
import logging
import os

from pix2text import Pix2Text

logger = logging.getLogger(__name__)


# 识别图片中的文字及公式
def pix2text_ocr(OUTPUT_DIR, image_path):
    # 检查目标文件是否已存在
    md_file = os.path.join(OUTPUT_DIR, 'ocr_result.md')
    if os.path.exists(md_file):
        logger.info(f"OCR结果文件已存在，直接返回: {md_file}")
        return md_file
    
    # 文件不存在时执行OCR识别
    logger.info(f"开始OCR识别: {image_path}")
    p2t = Pix2Text.from_config()
    outs = p2t.recognize_text_formula(image_path, resized_shape=1024, return_text=True)
    outs = outs.replace('\sb', '')

    # 确保输出目录存在
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    # 保存成文件
    with open(md_file, 'w', encoding='utf-8') as f:
        f.write(outs)
    
    logger.info(f"OCR结果已保存到: {md_file}")
    return md_file