2024-02-29 11:56:58 +08:00
|
|
|
|
import os
|
|
|
|
|
import sys
|
|
|
|
|
import glob
|
|
|
|
|
try:
|
|
|
|
|
import cv2
|
|
|
|
|
except :
|
|
|
|
|
os.system('pip install opencv-python')
|
2024-02-29 12:36:38 +08:00
|
|
|
|
import cv2
|
2024-02-29 11:56:58 +08:00
|
|
|
|
try :
|
|
|
|
|
from paddleocr import PaddleOCR , draw_ocr , download_with_progressbar
|
|
|
|
|
except:
|
|
|
|
|
os.system('pip install paddleocr')
|
2024-02-29 12:36:38 +08:00
|
|
|
|
from paddleocr import PaddleOCR , draw_ocr , download_with_progressbar
|
2024-02-29 11:56:58 +08:00
|
|
|
|
output_folder_path = 'res/'
|
|
|
|
|
if not os.path.exists(output_folder_path):
|
|
|
|
|
os.makedirs(output_folder_path)
|
|
|
|
|
|
|
|
|
|
def get_pdf_files_in_directory(directory_path):
|
|
|
|
|
# 确保路径存在
|
|
|
|
|
if os.path.exists(directory_path) and os.path.isdir(directory_path):
|
|
|
|
|
return glob.glob(os.path.join(directory_path, '**', '*.pdf'), recursive=True)
|
|
|
|
|
else:
|
|
|
|
|
return []
|
|
|
|
|
def ocr_pdf_folder(folder_path):
|
|
|
|
|
ocr = PaddleOCR ( use_angle_cls = True , lang = "ch" , page_num = 0 ) # 只需运行一次即可将模型下载并加载到内存中
|
|
|
|
|
print("ppocrv4 加载完毕!!!")
|
|
|
|
|
pdf_paths = get_pdf_files_in_directory(folder_path)
|
|
|
|
|
print(f"共检测到 {len(pdf_paths)} 个PDF文件")
|
|
|
|
|
# 打印所有PDF文件的路径
|
|
|
|
|
for pdf_path in pdf_paths:
|
|
|
|
|
print(f'正在处理文件:{pdf_path}')
|
|
|
|
|
|
|
|
|
|
result = ocr.ocr (pdf_path , cls = True )
|
|
|
|
|
for idx in range(len(result)):
|
|
|
|
|
res = result[idx]
|
|
|
|
|
for line in res :
|
|
|
|
|
print(line)
|
|
|
|
|
print(f'{pdf_path} 处理完毕')
|
|
|
|
|
ocr_result = ""
|
|
|
|
|
for idx in range(len(result)):
|
|
|
|
|
res = result[idx]
|
|
|
|
|
for line in res:
|
|
|
|
|
# print(line[1][0])
|
|
|
|
|
ocr_result = f"{ocr_result} {str(line[1][0])}"
|
|
|
|
|
|
|
|
|
|
filename = os.path.splitext(os.path.basename(pdf_path))[0]
|
|
|
|
|
|
|
|
|
|
# 构建TXT文件的完整路径
|
|
|
|
|
txt_path = os.path.join('res/', f'{filename}.txt')
|
|
|
|
|
|
|
|
|
|
# 将提取的文本写入TXT文件
|
|
|
|
|
with open(txt_path, 'w', encoding='utf-8') as txt_file:
|
|
|
|
|
txt_file.write(ocr_result)
|
|
|
|
|
|
|
|
|
|
print(f'生成的txt文档保存在{txt_path}')
|
|
|
|
|
# break
|
|
|
|
|
# print(ocr_result)
|
|
|
|
|
# with open('my_file.txt', 'a') as f:
|
|
|
|
|
# # 写入字符串
|
|
|
|
|
# f.write(ocr_result)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
|
if len(sys.argv) > 1:
|
|
|
|
|
# sys.argv[0] 是脚本名,sys.argv[1:] 是传递给脚本的参数列表
|
|
|
|
|
pdf_path = sys.argv[1]
|
|
|
|
|
print(f'需要处理的文件夹是:{pdf_path}')
|
|
|
|
|
ocr_pdf_folder(pdf_path)
|
|
|
|
|
|