PaddleX/zhch/test_single_model.py

from paddlex import create_model
import time
from pathlib import Path

# 定义paddlex模型名称列表
MODEL_LIST = [
    # OCR文本检测模型
    {"model_name": "PP-OCRv5_mobile_det", "description": "轻量级OCR文本检测模型，适用于移动端部署"},
    {"model_name": "PP-OCRv5_server_det", "description": "PP-OCRv5_rec 是新一代文本识别模型。该模型致力于以单一模型高效、精准地支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字等复杂文本场景的识别。在保持识别效果的同时，兼顾推理速度和模型鲁棒性，为各种场景下的文档理解提供高效、精准的技术支撑。"},
    
    # OCR文本识别模型
    {"model_name": "PP-OCRv5_mobile_rec", "description": "轻量级OCR文本识别模型，适用于移动端部署"},
    {"model_name": "PP-OCRv5_server_rec", "description": "服务端OCR文本识别模型，高精度识别"},

    # 版面区域检测模型
    {"model_name": "PP-DocLayout_plus-L", "description": "版面检测模型，包含20个常见的类别：文档标题、段落标题、文本、页码、摘要、目录、参考文献、脚注、页眉、页脚、算法、公式、公式编号、图像、表格、图和表标题（图标题、表格标题和图表标题）、印章、图表、侧栏文本和参考文献内容"},
    {"model_name": "PP-DocBlockLayout", "description": "文档图像版面子模块检测，包含1个 版面区域 类别，能检测多栏的报纸、杂志的每个子文章的文本区域"},
    
    # 表格分类模型
    {"model_name": "PP-LCNet_x1_0_table_cls", "description": "wired_table, wireless_table"},
    
    # 表格识别模型
    {"model_name": "SLANet_plus", "description": "SLANet_plus 是百度飞桨视觉团队自研的表格结构识别模型 SLANet 的增强版。相较于 SLANet，SLANet_plus 对无线表、复杂表格的识别能力得到了大幅提升，并降低了模型对表格定位准确性的敏感度，即使表格定位出现偏移，也能够较准确地进行识别。"},
    {"model_name": "SLANeXt_wired", "description": "SLANeXt 系列是百度飞桨视觉团队自研的新一代表格结构识别模型。相较于 SLANet 和 SLANet_plus，SLANeXt 专注于对表格结构进行识别，并且对有线表格(wired)和无线表格(wireless)的识别分别训练了专用的权重，对各类型表格的识别能力都得到了明显提高，特别是对有线表格的识别能力得到了大幅提升。"},
    {"model_name": "SLANeXt_wireless", "description": "SLANeXt 系列是百度飞桨视觉团队自研的新一代表格结构识别模型。相较于 SLANet 和 SLANet_plus，SLANeXt 专注于对表格结构进行识别，并且对有线表格(wired)和无线表格(wireless)的识别分别训练了专用的权重，对各类型表格的识别能力都得到了明显提高，特别是对无线表格的识别能力得到了大幅提升。"},
    
    # 公式识别模型
    {"model_name": "PP-FormulaNet_plus-L", "description": "负责将图像中的数学公式转换为可编辑的文本或计算机可识别的格式。该模块的性能直接影响到整个OCR系统的准确性和效率。公式识别模块通常会输出数学公式的 LaTeX 或 MathML 代码"},
    
    # 文档图像方向分类模型
    {"model_name": "PP-LCNet_x1_0_doc_ori", "description": "基于PP-LCNet_x1_0的文档图像分类模型，含有四个类别，即0度，90度，180度，270度"},

    # 文本图像矫正模型
    {"model_name": "UVDoc", "description": "针对图像进行几何变换，以纠正图像中的文档扭曲、倾斜、透视变形等问题，以供后续的文本识别进行更加准确"},

    # 印章检测模型
    {"model_name": "PP-OCRv4_mobile_seal_det", "description": "PP-OCRv4的移动端印章文本检测模型，效率更高，适合在端侧部署"},
    {"model_name": "PP-OCRv4_server_seal_det", "description": "PP-OCRv4的服务端印章文本检测模型，精度更高，适合在较好的服务器上部署"},

]

def test_single_model(
    model_name: str, 
    # input_path: str = "sample_data/300674-母公司现金流量表-扫描.png", 
    input_path: str, 
    output_path: str = "./sample_data/output/"):
    """
    Test single model for layout detection.
    """
    start_time = time.time()
    print(f"\nTesting model: {model_name}")
    # Create the model
    model = create_model(model_name=model_name)

    # 参数通常用于目标检测模型（如 DetPredictor）中，用于重叠框过滤。
    # 检查模型是否支持 layout_nms 参数
    predict_kwargs = {}
    if hasattr(model._predictor, 'layout_nms'):
        predict_kwargs['layout_nms'] = True

    # 特殊处理 Doc VLM 模型（如 PP-Chart2Table）
    if model_name in ["PP-Chart2Table", "PP-DocBee-2B", "PP-DocBee-7B", "PP-DocBee2-3B"]:
        # Doc VLM 模型需要字典格式的输入
        input_data = {
            "image": input_path,
            "query": "请将图表转换为表格格式"  # 或其他适合的查询
        }
        output = model.predict(input_data, batch_size=1, **predict_kwargs)
    else:
        # 其他模型使用标准输入格式
        output = model.predict(input_path, batch_size=1, **predict_kwargs)

    for res in output:
        res.print()
        res.save_all(save_path=output_path)  # Save all results to the specified path

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Total time taken for {model_name}: {elapsed_time:.2f} seconds")

if __name__ == "__main__":
    # Specify the model name
    # 循环，如何指定模型名称，直到quit
    while True:
        model_name = input("请输入模型名称（或输入'quit'退出）：")
        if model_name.lower() == "quit":
            break
        output_path = Path(f"./sample_data/single_model_output/{model_name}/")
        output_path.mkdir(parents=True, exist_ok=True)

        test_single_model(model_name=model_name, 
                        #   input_path="sample_data/300674-母公司现金流量表-扫描.png", 
                          input_path="/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/data_PPStructureV3_Results/B用户_扫描流水/B用户_扫描流水_page_002.png", 
                          output_path=output_path.as_posix())
        print("\n" + "="*50 + "\n")