# -*- coding: utf-8 -*-
# @Time : 2024/10/18 22:09
# @Author : 黄子寒
# @Email : 1064071566@qq.com
# @File : OCR.py
# @Project : EmoLLM
import cv2
from paddleocr import PaddleOCR
import os
import glob

# 初始化OCR模型
ocr = PaddleOCR(use_angle_cls=True, lang='ch')


image_dir = "output"
output_txt_dir = "output_txt"


if not os.path.exists(output_txt_dir):
    os.makedirs(output_txt_dir)

image_list = glob.glob(os.path.join(image_dir, "*.png"))

# 批量识别处理
for img_path in image_list:
    # 读取图像
    img = cv2.imread(img_path)

    # 使用OCR模型进行识别
    result = ocr.ocr(img)

    # 获取图像文件名（不带扩展名）
    img_name = os.path.splitext(os.path.basename(img_path))[0]

    # 将OCR结果整理为文本
    txt_file_path = os.path.join(output_txt_dir, f"{img_name}.txt")

    # 打开文件以写入OCR结果
    with open(txt_file_path, 'w', encoding='utf-8') as f:
        for line in result:
            for word_info in line:
                # 提取识别到的文本和其置信度
                word, confidence = word_info[1][0], word_info[1][1]

                f.write(f"{word}\n")

                print(f" {word}, C: {confidence}")

    print(f"{txt_file_path}")