PaddleX/zhch/ppstructurev3_single_client.py

"""PDF转图像后通过API统一处理"""
import json
import time
import os
import traceback
import argparse
import sys
import warnings
import base64
from pathlib import Path
from typing import List, Dict, Any, Union
import requests
from tqdm import tqdm

from dotenv import load_dotenv
load_dotenv(override=True)

from utils import (
    get_image_files_from_dir,
    get_image_files_from_list,
    get_image_files_from_csv,
    collect_pid_files,
    load_images_from_pdf,
    normalize_financial_numbers,
    normalize_markdown_table
)

def convert_pdf_to_images(pdf_file: str, output_dir: str | None = None, dpi: int = 200) -> List[str]:
    """
    将PDF转换为图像文件
    
    Args:
        pdf_file: PDF文件路径
        output_dir: 输出目录
        dpi: 图像分辨率
        
    Returns:
        生成的图像文件路径列表
    """
    pdf_path = Path(pdf_file)
    if not pdf_path.exists() or pdf_path.suffix.lower() != '.pdf':
        print(f"❌ Invalid PDF file: {pdf_path}")
        return []

    # 如果没有指定输出目录，使用PDF同名目录
    if output_dir is None:
        output_path = pdf_path.parent / f"{pdf_path.stem}"
    else:
        output_path = Path(output_dir) / f"{pdf_path.stem}"
    output_path = output_path.resolve()
    output_path.mkdir(parents=True, exist_ok=True)

    try:
        # 使用doc_utils中的函数加载PDF图像
        images = load_images_from_pdf(str(pdf_path), dpi=dpi)
        
        image_paths = []
        for i, image in enumerate(images):
            # 生成图像文件名
            image_filename = f"{pdf_path.stem}_page_{i+1:03d}.png"
            image_path = output_path / image_filename

            # 保存图像
            image.save(str(image_path))
            image_paths.append(str(image_path))
            
        print(f"✅ Converted {len(images)} pages from {pdf_path.name} to images")
        return image_paths
        
    except Exception as e:
        print(f"❌ Error converting PDF {pdf_path}: {e}")
        traceback.print_exc()
        return []

def get_input_files(args) -> List[str]:
    """
    获取输入文件列表，统一处理PDF和图像文件
    
    Args:
        args: 命令行参数
        
    Returns:
        处理后的图像文件路径列表
    """
    input_files = []
    
    # 获取原始输入文件
    if args.input_csv:
        raw_files = get_image_files_from_csv(args.input_csv, "fail")
    elif args.input_file_list:
        raw_files = get_image_files_from_list(args.input_file_list)
    elif args.input_file:
        raw_files = [Path(args.input_file).resolve()]
    else:
        input_dir = Path(args.input_dir).resolve()
        if not input_dir.exists():
            print(f"❌ Input directory does not exist: {input_dir}")
            return []
        
        # 获取所有支持的文件（图像和PDF）
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        pdf_extensions = ['.pdf']
        
        raw_files = []
        for ext in image_extensions + pdf_extensions:
            raw_files.extend(list(input_dir.glob(f"*{ext}")))
            raw_files.extend(list(input_dir.glob(f"*{ext.upper()}")))
        
        raw_files = [str(f) for f in raw_files]
    
    # 分别处理PDF和图像文件
    pdf_count = 0
    image_count = 0
    
    for file_path in raw_files:
        file_path = Path(file_path)
        
        if file_path.suffix.lower() == '.pdf':
            # 转换PDF为图像
            print(f"📄 Processing PDF: {file_path.name}")
            pdf_images = convert_pdf_to_images(
                str(file_path), 
                args.output_dir,
                dpi=args.pdf_dpi
            )
            input_files.extend(pdf_images)
            pdf_count += 1
        else:
            # 直接添加图像文件
            if file_path.exists():
                input_files.append(str(file_path))
                image_count += 1
    
    print(f"📊 Input summary:")
    print(f"  PDF files processed: {pdf_count}")
    print(f"  Image files found: {image_count}")
    print(f"  Total image files to process: {len(input_files)}")
    
    return input_files

def convert_api_result_to_json(api_result: Dict[str, Any], 
                              input_image_path: str, 
                              output_dir: str, 
                              filename: str,
                              normalize_numbers: bool = True) -> tuple[str, Dict[str, Any]]:
    """
    将API返回结果转换为标准JSON格式，并支持数字标准化
    """
    # 获取主要数据
    layout_parsing_results = api_result.get('layoutParsingResults', [])
    
    if not layout_parsing_results:
        print("⚠️ Warning: No layoutParsingResults found in API response")
        return {}
    
    # 取第一个结果（通常只有一个）
    main_result = layout_parsing_results[0]
    pruned_result = main_result.get('prunedResult', {})
    
    # 构造标准格式的JSON
    converted_json = {
        "input_path": input_image_path,
        "page_index": None,
        "model_settings": pruned_result.get('model_settings', {}),
        "parsing_res_list": pruned_result.get('parsing_res_list', []),
        "doc_preprocessor_res": {
            "input_path": None,
            "page_index": None,
            "model_settings": pruned_result.get('doc_preprocessor_res', {}).get('model_settings', {}),
            "angle": pruned_result.get('doc_preprocessor_res', {}).get('angle', 0)
        },
        "layout_det_res": {
            "input_path": None,
            "page_index": None,
            "boxes": pruned_result.get('layout_det_res', {}).get('boxes', [])
        },
        "overall_ocr_res": {
            "input_path": None,
            "page_index": None,
            "model_settings": pruned_result.get('overall_ocr_res', {}).get('model_settings', {}),
            "dt_polys": pruned_result.get('overall_ocr_res', {}).get('dt_polys', []),
            "text_det_params": pruned_result.get('overall_ocr_res', {}).get('text_det_params', {}),
            "text_type": pruned_result.get('overall_ocr_res', {}).get('text_type', 'general'),
            "textline_orientation_angles": pruned_result.get('overall_ocr_res', {}).get('textline_orientation_angles', []),
            "text_rec_score_thresh": pruned_result.get('overall_ocr_res', {}).get('text_rec_score_thresh', 0.0),
            "return_word_box": pruned_result.get('overall_ocr_res', {}).get('return_word_box', False),
            "rec_texts": pruned_result.get('overall_ocr_res', {}).get('rec_texts', []),
            "rec_scores": pruned_result.get('overall_ocr_res', {}).get('rec_scores', []),
            "rec_polys": pruned_result.get('overall_ocr_res', {}).get('rec_polys', []),
            "rec_boxes": pruned_result.get('overall_ocr_res', {}).get('rec_boxes', [])
        },
        "table_res_list": pruned_result.get('table_res_list', [])
    }
    
    # 数字标准化处理
    original_json = converted_json.copy()
    changes_count = 0
    
    if normalize_numbers:
        # 1. 标准化 parsing_res_list 中的文本内容
        for item in converted_json.get('parsing_res_list', []):
            if 'block_content' in item:
                original_content = item['block_content']
                normalized_content = original_content
                # 根据block_label类型选择标准化方法
                if item.get('block_label') == 'table':
                    normalized_content = normalize_markdown_table(original_content)
                # else:
                #     normalized_content = normalize_financial_numbers(original_content)
                
                if original_content != normalized_content:
                    item['block_content'] = normalized_content
                    changes_count += len([1 for o, n in zip(original_content, normalized_content) if o != n])
        
        # 2. 标准化 table_res_list 中的HTML表格
        for table_item in converted_json.get('table_res_list', []):
            if 'pred_html' in table_item:
                original_html = table_item['pred_html']
                normalized_html = normalize_markdown_table(original_html)
                
                if original_html != normalized_html:
                    table_item['pred_html'] = normalized_html
                    changes_count += len([1 for o, n in zip(original_html, normalized_html) if o != n])
        
        # 3. 标准化 overall_ocr_res 中的识别文本
        # ocr_res = converted_json.get('overall_ocr_res', {})
        # if 'rec_texts' in ocr_res:
        #     original_texts = ocr_res['rec_texts'][:]
        #     normalized_texts = []
            
        #     for text in original_texts:
        #         normalized_text = normalize_financial_numbers(text)
        #         normalized_texts.append(normalized_text)
        #         if text != normalized_text:
        #             changes_count += len([1 for o, n in zip(text, normalized_text) if o != n])
            
        #     ocr_res['rec_texts'] = normalized_texts
        
        # 添加标准化处理信息
        converted_json['processing_info'] = {
            "normalize_numbers": normalize_numbers,
            "changes_applied": changes_count > 0,
            "character_changes_count": changes_count
        }
        
        # if changes_count > 0:
        #     print(f"🔧 已标准化 {changes_count} 个字符（全角→半角）")
    else:
        converted_json['processing_info'] = {
            "normalize_numbers": False,
            "changes_applied": False,
            "character_changes_count": 0
        }
    
    # 保存JSON文件
    output_path = Path(output_dir).resolve() / f"{filename}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(converted_json, f, ensure_ascii=False, indent=2)
    
    # 如果启用了标准化且有变化，保存原始版本用于对比
    if normalize_numbers and changes_count > 0:
        original_output_path = output_path.parent / f"{output_path.stem}_original.json"
        with open(original_output_path, 'w', encoding='utf-8') as f:
            json.dump(original_json, f, ensure_ascii=False, indent=2)
    
    return str(output_path), converted_json

def save_markdown_content(api_result: Dict[str, Any], output_dir: str, 
                         filename: str, normalize_numbers: bool = True) -> str:
    """
    保存Markdown内容，支持数字标准化
    """
    layout_parsing_results = api_result.get('layoutParsingResults', [])
    if not layout_parsing_results:
        return ""
    
    main_result = layout_parsing_results[0]
    markdown_data = main_result.get('markdown', {})
    
    output_path = Path(output_dir).resolve()
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 保存Markdown文本
    markdown_text = markdown_data.get('text', '')
    
    # 数字标准化处理
    if normalize_numbers and markdown_text:
        original_markdown_text = markdown_text
        markdown_text = normalize_markdown_table(markdown_text)
        
        changes_count = len([1 for o, n in zip(original_markdown_text, markdown_text) if o != n])
        # if changes_count > 0:
        #     print(f"🔧 Markdown中已标准化 {changes_count} 个字符（全角→半角）")
    
    md_file_path = output_path / f"{filename}.md"
    with open(md_file_path, 'w', encoding='utf-8') as f:
        f.write(markdown_text)
    
    # 如果启用了标准化且有变化，保存原始版本用于对比
    if normalize_numbers and changes_count > 0:
        original_output_path = output_path.parent / f"{output_path.stem}_original.json"
        with open(original_output_path, 'w', encoding='utf-8') as f:
            f.write(original_markdown_text)

    return str(md_file_path)

def call_api_for_image(image_path: str, api_url: str, timeout: int = 300) -> Dict[str, Any]:
    """
    为单个图像调用API
    
    Args:
        image_path: 图像文件路径
        api_url: API URL
        timeout: 超时时间（秒）
        
    Returns:
        API返回结果
    """
    try:
        # 对本地图像进行Base64编码
        with open(image_path, "rb") as file:
            image_bytes = file.read()
            image_data = base64.b64encode(image_bytes).decode("ascii")

        payload = {
            "file": image_data,
            "fileType": 1,
        }

        # 调用API
        response = requests.post(api_url, json=payload, timeout=timeout)
        response.raise_for_status()
        
        return response.json()["result"]
        
    except requests.exceptions.Timeout:
        raise Exception(f"API调用超时 ({timeout}秒)")
    except requests.exceptions.RequestException as e:
        raise Exception(f"API调用失败: {e}")
    except KeyError:
        raise Exception("API返回格式错误，缺少'result'字段")
    except Exception as e:
        raise Exception(f"处理图像时发生错误: {e}")

def process_images_via_api(image_paths: List[str],
                          api_url: str,
                          output_dir: str = "./output",
                          normalize_numbers: bool = True,
                          timeout: int = 300) -> List[Dict[str, Any]]:
    """
    通过API统一处理图像文件
    
    Args:
        image_paths: 图像路径列表
        api_url: API URL
        output_dir: 输出目录
        normalize_numbers: 是否标准化数字格式
        timeout: API调用超时时间
        
    Returns:
        处理结果列表
    """
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    all_results = []
    total_images = len(image_paths)
    
    print(f"Processing {total_images} images via API")
    
    # 使用tqdm显示进度
    with tqdm(total=total_images, desc="Processing images", unit="img", 
              bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]') as pbar:
        
        # 逐个处理图像
        for img_path in image_paths:
            start_time = time.time()
            
            try:
                # 调用API处理图像
                api_result = call_api_for_image(img_path, api_url, timeout)
                processing_time = time.time() - start_time
                
                # 处理API返回结果
                input_path = Path(img_path)
                
                # 生成输出文件名
                output_filename = input_path.stem
                
                # 转换并保存标准JSON格式
                json_output_path, converted_json = convert_api_result_to_json(
                    api_result, 
                    str(input_path), 
                    output_dir,
                    output_filename,
                    normalize_numbers=normalize_numbers
                )
                
                # 保存Markdown内容
                md_output_path = save_markdown_content(
                    api_result, 
                    output_dir, 
                    output_filename,
                    normalize_numbers=normalize_numbers
                )
                
                # 记录处理结果
                all_results.append({
                    "image_path": str(input_path),
                    "processing_time": processing_time,
                    "success": True,
                    "api_url": api_url,
                    "output_json": json_output_path,
                    "output_md": md_output_path,
                    "is_pdf_page": "_page_" in input_path.name,  # 标记是否为PDF页面
                    "processing_info": converted_json.get('processing_info', {})
                })
                
                # 更新进度条
                success_count = sum(1 for r in all_results if r.get('success', False))
                
                pbar.update(1)
                pbar.set_postfix({
                    'time': f"{processing_time:.2f}s",
                    'success': f"{success_count}/{len(all_results)}",
                    'rate': f"{success_count/len(all_results)*100:.1f}%"
                })
                
            except Exception as e:
                print(f"Error processing {Path(img_path).name}: {e}", file=sys.stderr)
                import traceback
                traceback.print_exc()
                
                # 添加错误结果
                all_results.append({
                    "image_path": str(img_path),
                    "processing_time": 0,
                    "success": False,
                    "api_url": api_url,
                    "error": str(e),
                    "is_pdf_page": "_page_" in Path(img_path).name
                })
                pbar.update(1)
    
    return all_results

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="PaddleX PP-StructureV3 API Client - Unified PDF/Image Processor")
    
    # 参数定义
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument("--input_file", type=str, help="Input file (supports both PDF and image file)")
    input_group.add_argument("--input_dir", type=str, help="Input directory (supports both PDF and image files)")
    input_group.add_argument("--input_file_list", type=str, help="Input file list (one file per line)")
    input_group.add_argument("--input_csv", type=str, help="Input CSV file with image_path and status columns")

    parser.add_argument("--output_dir", type=str, required=True, help="Output directory")
    parser.add_argument("--api_url", type=str, default="http://localhost:8080/layout-parsing", help="API URL")
    parser.add_argument("--pdf_dpi", type=int, default=200, help="DPI for PDF to image conversion")
    parser.add_argument("--timeout", type=int, default=300, help="API timeout in seconds")
    parser.add_argument("--no-normalize", action="store_true", help="禁用数字标准化")
    parser.add_argument("--test_mode", action="store_true", help="Test mode (process only 20 files)")
    parser.add_argument("--collect_results", type=str, help="收集处理结果到指定CSV文件")

    args = parser.parse_args()
    
    normalize_numbers = not args.no_normalize
    
    try:
        # 获取并预处理输入文件
        print("🔄 Preprocessing input files...")
        input_files = get_input_files(args)
        
        if not input_files:
            print("❌ No input files found or processed")
            return 1
        
        if args.test_mode:
            input_files = input_files[:20]
            print(f"Test mode: processing only {len(input_files)} images")
        
        print(f"🌐 Using API: {args.api_url}")
        print(f"🔧 数字标准化: {'启用' if normalize_numbers else '禁用'}")
        print(f"⏱️ Timeout: {args.timeout} seconds")
        
        # 开始处理
        start_time = time.time()
        results = process_images_via_api(
            input_files,
            args.api_url,
            args.output_dir,
            normalize_numbers=normalize_numbers,
            timeout=args.timeout
        )
        total_time = time.time() - start_time
        
        # 统计结果
        success_count = sum(1 for r in results if r.get('success', False))
        error_count = len(results) - success_count
        pdf_page_count = sum(1 for r in results if r.get('is_pdf_page', False))
        total_changes = sum(r.get('processing_info', {}).get('character_changes_count', 0) for r in results if 'processing_info' in r)
        
        print(f"\n" + "="*60)
        print(f"✅ API Processing completed!")
        print(f"📊 Statistics:")
        print(f"  Total files processed: {len(input_files)}")
        print(f"  PDF pages processed: {pdf_page_count}")
        print(f"  Regular images processed: {len(input_files) - pdf_page_count}")
        print(f"  Successful: {success_count}")
        print(f"  Failed: {error_count}")
        if len(input_files) > 0:
            print(f"  Success rate: {success_count / len(input_files) * 100:.2f}%")
        if normalize_numbers:
            print(f"  总标准化字符数: {total_changes}")
        print(f"⏱️ Performance:")
        print(f"  Total time: {total_time:.2f} seconds")
        if total_time > 0:
            print(f"  Throughput: {len(input_files) / total_time:.2f} files/second")
            print(f"  Avg time per file: {total_time / len(input_files):.2f} seconds")
        
        # 保存结果统计
        stats = {
            "total_files": len(input_files),
            "pdf_pages": pdf_page_count,
            "regular_images": len(input_files) - pdf_page_count,
            "success_count": success_count,
            "error_count": error_count,
            "success_rate": success_count / len(input_files) if len(input_files) > 0 else 0,
            "total_time": total_time,
            "throughput": len(input_files) / total_time if total_time > 0 else 0,
            "avg_time_per_file": total_time / len(input_files) if len(input_files) > 0 else 0,
            "api_url": args.api_url,
            "pdf_dpi": args.pdf_dpi,
            "normalize_numbers": normalize_numbers,
            "total_character_changes": total_changes,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        # 保存最终结果
        output_file_name = Path(args.output_dir).name
        output_file = os.path.join(args.output_dir, f"{output_file_name}_api_results.json")
        final_results = {
            "stats": stats,
            "results": results
        }
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(final_results, f, ensure_ascii=False, indent=2)
        
        print(f"💾 Results saved to: {output_file}")

        # 如果没有收集结果的路径，使用缺省文件名，和output_dir同一路径
        if not args.collect_results:
            output_file_processed = Path(args.output_dir) / f"processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv"
        else:
            output_file_processed = Path(args.collect_results).resolve()
        
        processed_files = collect_pid_files(output_file)
        with open(output_file_processed, 'w', encoding='utf-8') as f:
            f.write("image_path,status\n")
            for file_path, status in processed_files:
                f.write(f"{file_path},{status}\n")
        print(f"💾 Processed files saved to: {output_file_processed}")

        return 0
        
    except Exception as e:
        print(f"❌ Processing failed: {e}", file=sys.stderr)
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    print(f"🚀 启动PP-StructureV3 API客户端...")
    print(f"🔧 环境变量检查: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
    
    if len(sys.argv) == 1:
        # 如果没有命令行参数，使用默认配置运行
        print("ℹ️  No command line arguments provided. Running with default configuration...")
        
        # 默认配置
        default_config = {
            "input_dir": "../../OmniDocBench/OpenDataLab___OmniDocBench/images",
            "output_dir": "./OmniDocBench_API_Results",
            "api_url": "http://10.192.72.11:8111/layout-parsing",
            "timeout": "300",
            "collect_results": f"./OmniDocBench_API_Results/processed_files_{time.strftime('%Y%m%d_%H%M%S')}.csv",
        }
        
        # 构造参数
        sys.argv = [sys.argv[0]]
        for key, value in default_config.items():
            sys.argv.extend([f"--{key}", str(value)])
        
        # sys.argv.append("--no-normalize")

        # 测试模式
        # sys.argv.append("--test_mode")
    
    sys.exit(main())