51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
# -*- 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}")
|