from paddlex import create_pipeline
import time
from pathlib import Path

# input_path = "./sample_data/300674-母公司现金流量表-扫描.png"
input_path = "/Users/zhch158/workspace/data/流水分析/对公_招商银行图/table_recognition_v2_Results/对公_招商银行图/对公_招商银行图_page_001.png"

pipeline_path = "./my_config/table_recognition_v2.yaml"
pipeline_name = Path(pipeline_path).stem
output_path = Path(f"./sample_data/single_pipeline_output/{pipeline_name}/")

pipeline = create_pipeline(pipeline=pipeline_path)

# For Image
output = pipeline.predict(
    input=input_path,
    device="gpu",  # 或者 "gpu" 如果你有 GPU 支持
    use_doc_orientation_classify=True, # 开启文档方向分类
    use_doc_unwarping=False, # 开启文档去畸变
    # use_e2e_wireless_table_rec_model=True, # 开启端到端无线表格识别
    use_wireless_table_cells_trans_to_html=True, # 开启无线表格单元格转 HTML
)

# 可视化结果并保存 json 结果
for res in output:
    res.print() 
    # res.save_to_json(save_path="sample_data/output") 
    # res.save_to_markdown(save_path="sample_data/output") 
    output_path.mkdir(parents=True, exist_ok=True)
    res.save_all(save_path=output_path.as_posix())  # 保存所有结果到指定路径