Merge pull request #28 from 8baby8/main

Modifying Ocr Scripts
This commit is contained in:
chaoke 2024-02-29 12:27:20 +08:00 committed by GitHub
commit 5193ad656b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

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