import time
import os
from paddleocr import PaddleOCR
from memory_profiler import profile

def main():
    start_load = time.time()
    ocr = PaddleOCR(
        use_doc_orientation_classify=False, 
        use_doc_unwarping=False, 
        use_textline_orientation=False, 
        device="cpu",
        text_detection_model_name="PP-OCRv5_mobile_det",
        text_detection_model_dir="./models/PP-OCRv5_mobile_det",
        text_recognition_model_name="PP-OCRv5_mobile_rec",
        text_recognition_model_dir="./models/PP-OCRv5_mobile_rec",
        cpu_threads=4,
        precision="fp16",
        enable_mkldnn=True,
        mkldnn_cache_capacity=10
    )
    end_load = time.time()
    print(f"模型加载时间: {end_load - start_load:.2f} 秒")

    # 指定图片目录
    image_dir = "./barcode"  # 替换为你的图片目录路径

    # 获取目录中的所有图片文件（支持常见格式）
    supported_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"]
    image_files = [
        os.path.join(image_dir, f) 
        for f in os.listdir(image_dir) 
        if os.path.splitext(f)[1].lower() in supported_extensions
    ]

    print(f"共找到 {len(image_files)} 张图片，开始 OCR 识别...")

    # 遍历每张图片并单独预测
      for img_path in image_files:
          print(f"\n正在处理: {img_path}")
          
          start_time = time.time()
          result = ocr.predict(img_path)  # 单张图片预测
          end_time = time.time()
          
          elapsed_time = end_time - start_time
          print(f"OCR 识别耗时: {elapsed_time:.2f} 秒")
          
          # # 保存结果（每张图片单独保存）
          # for res in result:
          #     base_name = os.path.splitext(os.path.basename(img_path))[0]
          #     res.save_to_img(f"output/{base_name}_result")  # 保存图片结果
          #     res.save_to_json(f"output/{base_name}_result")  # 保存 JSON 结果

    print("\n所有图片处理完成！")

if __name__ == "__main__":
    main()