# -*- coding: utf-8 -*-
# @Author:逐梦苍穹
# @Time: 2024/1/18 1:23
import os
import cv2
import fitz
import shutil
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from cnocr import CnOcr
from tqdm import tqdm
from PIL import Image
from docx import Documentdef mkdir(folder_path):os.makedirs(folder_path, exist_ok=True)def get_max_folder_number(folder_path):image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]return len(image_files)def crop_top_and_save(args):i, image_path, output_path, crop_percentage = argsimage = cv2.imread(image_path)height = image.shape[0]crop_pixels = int(height * crop_percentage)cropped_image = image[crop_pixels:, :]cv2.imwrite(output_path, cropped_image)print(f"Image {i + 1} cropped and saved: {output_path}")def startCutImage(input_folder, output_folder, crop_percentage, max_workers_input):mkdir(output_folder)max_folder_number = get_max_folder_number(input_folder)with ThreadPoolExecutor(max_workers=max_workers_input) as executor:args_list = [(i, os.path.join(input_folder, f'{i + 1}.jpg'), os.path.join(output_folder, f'{i + 1}.jpg'), crop_percentage)for i in range(max_folder_number)]executor.map(crop_top_and_save, args_list)def get_max_folder_number(folder_path):image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]return len(image_files)def delete_images(folder_path, start_index, end_index):for i in range(start_index, end_index + 1):imageStr = str(i) + ".jpg"file_path = os.path.join(folder_path, imageStr)os.remove(file_path)print(f"Deleted: {file_path}")def startDeleteImage(folder_path, end_index, back_start_index_input):max_folder_number = get_max_folder_number(folder_path)if max_folder_number > 0:start_index = 1back_start_index = max_folder_number - back_start_index_input + 1back_end_index = max_folder_numberdelete_images(folder_path, start_index, end_index)delete_images(folder_path, back_start_index, back_end_index)else:print("No valid image files found.")def convert_page_to_image(pdf_path, output_folder, page_number):pdf_document = fitz.open(pdf_path)page = pdf_document[page_number]image_list = page.get_images(full=True)for img_index, img_info in enumerate(image_list):image_index = img_info[0]base_image = pdf_document.extract_image(image_index)image_bytes = base_image["image"]image_filename = os.path.join(output_folder, f"{page_number + 1}.jpg")with open(image_filename, "wb") as image_file:image_file.write(image_bytes)pdf_document.close()def pdf_to_images(pdf_path, output_folder, num_workers):if not os.path.exists(output_folder):os.makedirs(output_folder)pdf_document = fitz.open(pdf_path)total_pages = pdf_document.page_countwith ThreadPoolExecutor(max_workers=num_workers) as executor:futures = [executor.submit(convert_page_to_image, pdf_path, output_folder, page_number)for page_number in range(total_pages)]for future in tqdm(futures, desc="Converting pages", total=total_pages):future.result()def get_max_folder_number(folder_path):image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]return len(image_files)def ocr_and_append(image_str, startList, event, lock):pagelist = []image_number = int(os.path.basename(image_str).split('.')[0])image = Image.open(image_str)ocr = CnOcr(debug=False)res = ocr.ocr(image)for j in range(len(res)):pagelist.append(res[j]['text'])with lock:startList[image_number - 11] = pagelistprint(f"Image {image_number} OCR completed.")event.set()def process_images_with_ocr(start, end, startList, event, max_workers_input, lock):with ThreadPoolExecutor(max_workers=max_workers_input) as executor:futures = [executor.submit(ocr_and_append, f'output-image-cut/{i}.jpg', startList, event, lock) for i intqdm(range(start, end + 1), desc="OCR Processing")]for future in as_completed(futures):future.result()def startOCR(file_folder, test_end_index, word_file_output_path, max_workers_input):max_folder_number = get_max_folder_number(file_folder)start = 1 + test_end_indexend = test_end_index + max_folder_numberstartList = [''] * max_folder_numberlock = threading.Lock()event = threading.Event()process_images_with_ocr(start, end, startList, event, max_workers_input, lock)event.wait()print(len(startList))for i in startList:print(len(i))print(startList)startWriteWordFile(word_file_output_path, startList)def startWriteWordFile(word_file_path, listPages):doc = Document()for i in listPages:write_list_to_word(word_file_path, i, doc)doc.add_page_break()doc.save(word_file_path)def write_list_to_word(word_file, my_list, doc):endList = []for item in my_list:item = item.strip()if len(item) < 10:if endList != None and len(endList) > 0:doc.add_paragraph("".join(endList))endList = []doc.add_paragraph(item)else:endList.append(item)if item == my_list[-1]:doc.add_paragraph("".join(endList))endList = []def main():try:crop_percentage = 0.0end_index = 0back_start_index_input = 0print("==========================")print("这是一个OCR识别程序，作者CSDN@逐梦苍穹，该程序不做商用请切记！")print("==========================")pdf_file_path = str(input("请输入PDF文件的绝对路径："))print("该PDF文档是否是书籍?")isBookPdf = int(input("(是，请输入1；否，请输入0):"))if isBookPdf == 1:print("输入要删除非正文的页面(封面、目录、结尾)")end_index = int(input("需要删除该文档的前几页: "))back_start_index_input = int(input("需要删除该文档的倒数前几页: "))print("询问是否需要截取图片区域?")isCutImage = int(input("(是，请输入1；否，请输入0):"))if isCutImage == 1:crop_percentage = float(input("请输入截取图片区域的纵向比例："))max_workers = int(input("请输入本次运行过程中想要使用的进程数："))if max_workers <= 0:max_workers = 1except Exception as e:print("参数有误,报错信息如下：", e)try:pdf_to_images(pdf_file_path, 'output-image', max_workers)if isCutImage == 1:startCutImage('output-image', 'output-image-cut', crop_percentage, max_workers)if isBookPdf == 1:startDeleteImage('output-image-cut', end_index, back_start_index_input)file_name = os.path.splitext(os.path.basename(pdf_file_path))[0]directory_path = os.path.dirname(pdf_file_path)word_file_path = directory_path + '\\' + file_name + '(OCR_XZL).docx'startOCR(r'output-image-cut', end_index, word_file_path, max_workers)except Exception as e:print("处理过程出错，报错信息如下：", e)try:delete_path_1 = r"output-image"delete_path_2 = r"output-image-cut"shutil.rmtree(delete_path_1)shutil.rmtree(delete_path_2)print(f"Folder '{delete_path_1}' successfully removed.")print(f"Folder '{delete_path_2}' successfully removed.")print("OCR处理完成！")except Exception as e:print("删除过程性文件出错，原因如下：", e)if __name__ == '__main__':main()