online_test.py

import requests
from pathlib import Path
from config import INPUT_PATH,OUTPUT_PATH
import os


def ocr_pdf(pdf_path, server_url="http://localhost:8001", save_result=True):
    """
    对 PDF 文档进行 OCR 识别
    
    参数:
        pdf_path: PDF 文件路径
        server_url: OCR 服务地址
        save_result: 是否保存识别结果到文件
        
    返回:
        dict: 包含识别结果的字典
    """
    # 1. 检查文件
    pdf_file = Path(pdf_path)
    if not pdf_file.exists():
        raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
    
    if pdf_file.suffix.lower() != '.pdf':
        raise ValueError(f"不是 PDF 文件: {pdf_path}")
    
    # 2. 读取文件大小（用于显示）
    file_size_mb = pdf_file.stat().st_size / (1024 * 1024)
    
    print(f"文件名: {pdf_file.name}")
    print(f"文件大小: {file_size_mb:.2f} MB")
    print(f"开始处理...")
    
    # 3. 准备请求
    api_url = f"{server_url}/ocr"
    
    # 4. 发送请求
    with open(pdf_path, 'rb') as f:
        files = {'file': (pdf_file.name, f, 'application/pdf')}
        
        # 这里可以添加额外参数
        # data = {'enable_description': True}  # 启用图片描述（会增加处理时间）
        
        response = requests.post(api_url, files=files)
    
    # 5. 处理结果
    if response.status_code == 200:
        result = response.json()
        
        print(f"处理完成！")
        print(f"统计信息:")
        print(f"   - 总页数: {result['page_count']} 页")
        print(f"   - 处理耗时: {result['processing_time']:.2f} 秒")
        print(f"   - 平均速度: {result['processing_time'] / result['page_count']:.2f} 秒/页")
        
        # 6. 保存结果到文件
        if save_result:
            os.makedirs(OUTPUT_PATH, exist_ok=True)
            output_file = pdf_file.with_suffix('.md')
            file_path = os.path.join(OUTPUT_PATH, output_file)

            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(result['markdown'])
            print(f"结果已保存到: {output_file}")
        
        return result
    else:
        print(f"处理失败: {response.status_code}")
        print(f"错误信息: {response.text}")
        return None

# 使用示例
# 替换为你的实际 PDF 路径
#pdf_file = "./doc/DeepSeek_OCR_paper_layouts.pdf"
pdf_file = INPUT_PATH
# 调用 OCR 函数
result = ocr_pdf(pdf_file)

if result:
    # 显示部分识别结果
    print("\n" + "="*60)
    print("识别结果预览:")
    print("="*60)
    
    # 显示前 1000 个字符
    preview_text = result['markdown'][:1000]
    print(preview_text)
    
    if len(result['markdown']) > 1000:
        print("\n... (内容过长，已截断)")
        print(f"\n完整内容共 {len(result['markdown'])} 个字符")