# # import re # # import PyPDF2 # # import tempfile # # from utils.ocr_engine import OcrEngine # # # # # # # 假设您的OCR函数名为 `ocr_extract`，并接受图片文件路径作为参数 # # # 请将此函数替换为您实际的OCR实现 # # # # def clean_page_content(text, common_header): # # # 首先删除抬头公共部分 # # if common_header: # 确保有公共抬头才进行替换 # # for header_line in common_header.split('\n'): # # if header_line.strip(): # 只处理非空行 # # # 替换首次出现的完整行 # # text = re.sub(r'^' + re.escape(header_line.strip()) + r'\n?', '', text, count=1) # # # # # 删除页码 eg:89/129 这个代码分三步走可以把89/129完全删除 # # text = re.sub(r'^\s*\d+\s*(?=\D)', '', text) # 删除开头的页码，仅当紧跟非数字字符时 # # text = re.sub(r'\s+\d+\s*$', '', text) # 删除结尾的页码 # # text = re.sub(r'\s*\/\s*\d+\s*', '', text) # 删除形如 /129 的页码 # # text = re.sub(r'\s*[—-]\s*\d+\s*[—-]\s*', '', text) # 删除形如 '—2—' 或 '-2-' 的页码 # # return text # # # # # # def extract_common_header(pdf_path): # # from PyPDF2 import PdfReader # # # # def get_headers(pdf_document, start_page, pages_to_read): # # headers = [] # # for i in range(start_page, min(start_page + pages_to_read, len(pdf_document.pages))): # # page = pdf_document.pages[i] # # text = page.extract_text() or "" # # if text: # # # 只取每页的前三行，去除前后的空白字符 # # first_lines = [line.strip() for line in text.strip().split('\n')[:3]] # # headers.append(first_lines) # # return headers # # # # def find_common_headers(headers): # # if not headers: # # return [] # # # # # 使用 zip 对齐所有页的对应行 # # common_headers = [] # # for lines in zip(*headers): # # # 检查所有行是否完全相同 # # if all(line == lines[0] for line in lines[1:]): # # common_headers.append(lines[0]) # # return common_headers # # # # pdf_document = PdfReader(pdf_path) # # total_pages = len(pdf_document.pages) # # # # # 定义两个提取策略 # # strategies = [] # # if total_pages >= 3: # # # 策略1：中间的3页 # # middle_page = total_pages // 2 # # start_page = max(0, middle_page - 1) # # strategies.append((start_page, 3)) # # elif total_pages == 2: # # # 策略1：2页 # # strategies.append((0, 2)) # # else: # # # 策略1：1页 # # strategies.append((0, 1)) # # # # # 策略2：前三页 # # if total_pages >= 3: # # strategies.append((0, 3)) # # elif total_pages == 2: # # strategies.append((0, 2)) # # elif total_pages == 1: # # strategies.append((0, 1)) # # # # common_headers = [] # # # # for idx, (start, count) in enumerate(strategies): # # headers = get_headers(pdf_document, start, count) # # if len(headers) < 2: # # continue # 需要至少2页来比较 # # # # current_common = find_common_headers(headers) # # if current_common: # # common_headers = current_common # # break # 找到共同部分后退出 # # # 如果没有找到，继续下一个策略 # # # # return '\n'.join(common_headers) # # # # # # def extract_images_from_page(reader, page): # # images = [] # # try: # # for img in page.images: # # xref = img['xref'] # # image = reader.extract_image(xref) # # image_bytes = image['image'] # # image_ext = image['ext'] # # images.append({'data': image_bytes, 'ext': image_ext}) # # except Exception as e: # # print(f"提取第{reader.pages.index(page) + 1}页图片时出错: {e}") # # return images # # # # # # def extract_text_by_page(file_path): # # common_header = extract_common_header(file_path) # # # print(f"公共抬头：{common_header}") # # # print("--------------------正文开始-------------------") # # result = "" # # with open(file_path, 'rb') as file: # # reader = PyPDF2.PdfReader(file) # # num_pages = len(reader.pages) # # # print(f"Total pages: {num_pages}") # # for page_num in range(num_pages): # # page = reader.pages[page_num] # # text = page.extract_text() or "" # # # # # 提取图片并进行OCR # # images = extract_images_from_page(reader, page) # # ocr_text = "" # # for image in images: # # image_data = image['data'] # # image_ext = image['ext'] # # try: # # with tempfile.NamedTemporaryFile(delete=True, suffix='.' + image_ext) as temp_image: # # temp_image.write(image_data) # # temp_image.flush() # # # 调用OCR函数 # # ocr_result = OcrEngine.recognize_text_from_image(temp_image.name) # # ocr_text += ocr_result + "\n" # # except Exception as e: # # print(f"处理第{page_num + 1}页图片时出错: {e}") # # # # # 清理文本 # # cleaned_text = clean_page_content(text, common_header) # # # 合并OCR文本 # # if ocr_text.strip(): # # cleaned_text += "\n" + ocr_text # # result += cleaned_text # # return result # # # # # # def extract_text_json_by_page(file_path): # # common_header = extract_common_header(file_path) # # # print(f"公共抬头：{common_header}") # # # print("--------------------正文开始-------------------") # # result = {} # # with open(file_path, 'rb') as file: # # reader = PyPDF2.PdfReader(file) # # num_pages = len(reader.pages) # # # print(f"Total pages: {num_pages}") # # for page_num in range(num_pages): # # page = reader.pages[page_num] # # text = page.extract_text() or "" # # # # # 提取图片并进行OCR # # images = extract_images_from_page(reader, page) # # ocr_text = "" # # for image in images: # # image_data = image['data'] # # image_ext = image['ext'] # # try: # # with tempfile.NamedTemporaryFile(delete=True, suffix='.' + image_ext) as temp_image: # # temp_image.write(image_data) # # temp_image.flush() # # # 调用OCR函数 # # ocr_result = OcrEngine.recognize_text_from_image(temp_image.name) # # ocr_text += ocr_result + "\n" # # except Exception as e: # # print(f"处理第{page_num + 1}页图片时出错: {e}") # # # # # 清理文本 # # cleaned_text = clean_page_content(text, common_header) # # # 合并OCR文本 # # if ocr_text.strip(): # # cleaned_text += "\n" + ocr_text # # result[str(page_num + 1)] = cleaned_text # # return result # # # # # # if __name__ == '__main__': # # pdf_path = "C:/test/iDS-TCE900神捕电警抓拍单元产品介绍.pdf" # # res = extract_text_json_by_page(pdf_path) # # print(res) # import json # import os # import re # import shutil # import uuid # # import PyPDF2 # import fitz # PyMuPDF # import tempfile # from utils.ocr_engine import OcrEngine # # # # 假设您的OCR函数名为 `ocr_extract`，并接受图片文件路径作为参数 # def ocr_extract(image_path): # # 示例：调用您的OCR脚本并返回识别的文本 # # 例如： # # return your_ocr_function(image_path) # # 将图片保存到tmp目录 # # shutil.copy(image_path, f'tmp/{uuid.uuid4()}.png') # # return "OCR提取的文本" # 替换为实际的OCR结果 # return OcrEngine.recognize_text_from_image(image_path) # # # def clean_page_content(text, common_header): # # 首先删除抬头公共部分 # if common_header: # 确保有公共抬头才进行替换 # for header_line in common_header.split('\n'): # if header_line.strip(): # 只处理非空行 # # 替换首次出现的完整行 # text = re.sub(r'^' + re.escape(header_line.strip()) + r'\n?', '', text, count=1) # # # 删除页码 eg:89/129 这个代码分三步走可以把89/129完全删除 # text = re.sub(r'^\s*\d+\s*(?=\D)', '', text) # 删除开头的页码，仅当紧跟非数字字符时 # text = re.sub(r'\s+\d+\s*$', '', text) # 删除结尾的页码 # text = re.sub(r'\s*\/\s*\d+\s*', '', text) # 删除形如 /129 的页码 # text = re.sub(r'\s*[—-]\s*\d+\s*[—-]\s*', '', text) # 删除形如 '—2—' 或 '-2-' 的页码 # return text # # # def extract_common_header(pdf_path): # from PyPDF2 import PdfReader # # def get_headers(pdf_document, start_page, pages_to_read): # headers = [] # for i in range(start_page, min(start_page + pages_to_read, len(pdf_document.pages))): # page = pdf_document.pages[i] # text = page.extract_text() or "" # if text: # # 只取每页的前三行，去除前后的空白字符 # first_lines = [line.strip() for line in text.strip().split('\n')[:3]] # headers.append(first_lines) # return headers # # def find_common_headers(headers): # if not headers: # return [] # # # 使用 zip 对齐所有页的对应行 # common_headers = [] # for lines in zip(*headers): # # 检查所有行是否完全相同 # if all(line == lines[0] for line in lines[1:]): # common_headers.append(lines[0]) # return common_headers # # pdf_document = PdfReader(pdf_path) # total_pages = len(pdf_document.pages) # # # 定义两个提取策略 # strategies = [] # if total_pages >= 3: # # 策略1：中间的3页 # middle_page = total_pages // 2 # start_page = max(0, middle_page - 1) # strategies.append((start_page, 3)) # elif total_pages == 2: # # 策略1：2页 # strategies.append((0, 2)) # else: # # 策略1：1页 # strategies.append((0, 1)) # # # 策略2：前三页 # if total_pages >= 3: # strategies.append((0, 3)) # elif total_pages == 2: # strategies.append((0, 2)) # elif total_pages == 1: # strategies.append((0, 1)) # # common_headers = [] # # for idx, (start, count) in enumerate(strategies): # headers = get_headers(pdf_document, start, count) # if len(headers) < 2: # continue # 需要至少2页来比较 # # current_common = find_common_headers(headers) # if current_common: # common_headers = current_common # break # 找到共同部分后退出 # # 如果没有找到，继续下一个策略 # # return '\n'.join(common_headers) # # # def extract_images_with_pymupdf(pdf_path): # images = [] # try: # doc = fitz.open(pdf_path) # for page_num in range(len(doc)): # page = doc.load_page(page_num) # image_list = page.get_images(full=True) # for img in image_list: # xref = img[0] # base_image = doc.extract_image(xref) # image_bytes = base_image['image'] # image_ext = base_image.get('ext', 'png') # # image_width = base_image.get['xres'] # # image_height = base_image.get['yres'] # # images.append({'data': image_bytes, 'ext': image_ext, 'page_num': page_num + 1, 'width': image_width, 'height': image_height}). # images.append({'data': image_bytes, 'ext': image_ext, 'page_num': page_num + 1}) # # # 确保输出目录存在 # if not os.path.exists("output_dir"): # os.makedirs("output_dir") # print(f"创建输出目录: output_dir") # # 构建图片文件名，例如: page_1_img.png # image_filename = f"page_{page_num}_img.{image_ext}" # image_path = os.path.join("output_dir", image_filename) # try: # with open(image_path, 'wb') as img_file: # img_file.write(image_bytes) # print(f"保存图片: {image_path}") # except Exception as e: # print(f"保存图片 {image_filename} 时出错: {e}") # # except Exception as e: # print(f"提取图片时出错: {e}") # return images # # # def extract_text_json_by_page(file_path): # common_header = extract_common_header(file_path) # # print(f"公共抬头：{common_header}") # # print("--------------------正文开始-------------------") # result = {} # # 如果已经有保存的ocr结果,直接读取 # # TODO 待修改 # if os.path.exists(file_path + ".json"): # with open(file_path + ".json", 'r', encoding='utf-8') as f: # result = json.load(f) # return result # images = extract_images_with_pymupdf(file_path) # # filtered_images = filter_images(images) # ocr_text_dict = {} # for img in images: # try: # with tempfile.NamedTemporaryFile(delete=False, suffix='.' + img['ext']) as temp_image: # temp_image.write(img['data']) # temp_image.flush() # # 调用OCR函数 # ocr_result = ocr_extract(temp_image.name) # if img['page_num'] in ocr_text_dict: # ocr_text_dict[img['page_num']] += ocr_result + "\n" # else: # ocr_text_dict[img['page_num']] = ocr_result + "\n" # except Exception as e: # print(f"OCR处理失败: {e}") # finally: # try: # os.remove(temp_image.name) # except Exception as e: # print(f"删除临时文件失败: {e}") # # with open(file_path, 'rb') as file: # reader = PyPDF2.PdfReader(file) # num_pages = len(reader.pages) # # print(f"Total pages: {num_pages}") # for page_num in range(num_pages): # page = reader.pages[page_num] # text = page.extract_text() or "" # # 清理文本 # cleaned_text = clean_page_content(text, common_header) # # 合并OCR文本 # if (page_num + 1) in ocr_text_dict and ocr_text_dict[page_num + 1].strip(): # cleaned_text += "\n" + ocr_text_dict[page_num + 1] # result[str(page_num + 1)] = cleaned_text # print("pdf预处理完成") # return result # # # def filter_images(images, min_width=200, min_height=200, text_threshold=5): # """ # 过滤图像，保留可能包含文本的图像。 # :param images: 图像列表，每个图像包含 data, ext, page_num, width, height # :param min_width: 最小宽度 # :param min_height: 最小高度 # :param text_threshold: 检测到的文本字符数阈值 # :return: 过滤后的图像列表 # """ # filtered = [] # for img in images: # # 基于尺寸过滤 # if img['width'] < min_width or img['height'] < min_height: # continue # # 基于文本检测过滤 # try: # # with tempfile.NamedTemporaryFile(delete=True, suffix='.' + img['ext']) as temp_image: # # temp_image.write(img['data']) # # temp_image.flush() # # # 使用OCR快速检测文本 # # ocr_result = ocr_extract(temp_image.name) # # if len(ocr_result.strip()) >= text_threshold: # # img['ocr_text'] = ocr_result # # filtered.append(img) # filtered.append(img) # except Exception as e: # print(f"过滤图像时出错: {e}") # # return filtered # # # if __name__ == '__main__': # pdf_path = "C:/test/发言系统11.pdf" # res = extract_text_json_by_page(pdf_path) # print(res) import os import re import json import tempfile import shutil import uuid from PyPDF2 import PdfReader import fitz # PyMuPDF from utils.ocr_engine import OcrEngine # 请确保OcrEngine已经正确导入 from utils.local_ocr import LocalOCR local_ocr = LocalOCR() def ocr_extract(image_path): # 调用您的OCR引擎来识别图像中的文本 # return OcrEngine.recognize_text_from_image(image_path) # 调用本地ocr return local_ocr.run(image_path) def clean_page_content(text, common_header): # 删除公共抬头 if common_header: for header_line in common_header.split('\n'): if header_line.strip(): text = re.sub(r'^' + re.escape(header_line.strip()) + r'\n?', '', text, count=1) # 删除页码 text = re.sub(r'^\s*\d+\s*(?=\D)', '', text) # 开头页码 text = re.sub(r'\s+\d+\s*$', '', text) # 结尾页码 text = re.sub(r'\s*\/\s*\d+\s*', '', text) # 形如 /129 text = re.sub(r'\s*[—-]\s*\d+\s*[—-]\s*', '', text) # 形如 '—2—' 或 '-2-' return text def extract_common_header(pdf_path): def get_headers(pdf_document, start_page, pages_to_read): headers = [] for i in range(start_page, min(start_page + pages_to_read, len(pdf_document.pages))): page = pdf_document.pages[i] text = page.extract_text() or "" if text: first_lines = [line.strip() for line in text.strip().split('\n')[:3]] headers.append(first_lines) return headers def find_common_headers(headers): if not headers: return [] common_headers = [] for lines in zip(*headers): if all(line == lines[0] for line in lines[1:]): common_headers.append(lines[0]) return common_headers pdf_document = PdfReader(pdf_path) total_pages = len(pdf_document.pages) strategies = [] if total_pages >= 3: middle_page = total_pages // 2 start_page = max(0, middle_page - 1) strategies.append((start_page, 3)) elif total_pages == 2: strategies.append((0, 2)) else: strategies.append((0, 1)) if total_pages >= 3: strategies.append((0, 3)) elif total_pages == 2: strategies.append((0, 2)) elif total_pages == 1: strategies.append((0, 1)) common_headers = [] for start, count in strategies: headers = get_headers(pdf_document, start, count) if len(headers) < 2: continue current_common = find_common_headers(headers) if current_common: common_headers = current_common break return '\n'.join(common_headers) def extract_images_from_page(pdf_path, page_num): images = [] try: doc = fitz.open(pdf_path) if page_num < 0 or page_num >= len(doc): print(f"页码 {page_num + 1} 超出范围") return images page = doc.load_page(page_num) image_list = page.get_images(full=True) for img in image_list: xref = img[0] base_image = doc.extract_image(xref) image_bytes = base_image['image'] image_ext = base_image.get('ext', 'png') images.append({'data': image_bytes, 'ext': image_ext, 'page_num': page_num + 1}) except Exception as e: print(f"提取图片时出错: {e}") return images def extract_text_json_by_page(file_path, text_threshold=10): """ 提取PDF每页的文本，若文本量低于text_threshold，则提取图片并OCR。 :param file_path: PDF文件路径 :param text_threshold: 文本量阈值（字符数） :return: 字典，键为页码，值为对应的文本 """ common_header = extract_common_header(file_path) result = {} # 如果已经有保存的OCR结果,直接读取 if os.path.exists(file_path + ".json"): with open(file_path + ".json", 'r', encoding='utf-8') as f: result = json.load(f) return result try: reader = PdfReader(file_path) num_pages = len(reader.pages) for page_num in range(num_pages): page = reader.pages[page_num] text = page.extract_text() or "" cleaned_text = clean_page_content(text, common_header) # 检查文本量是否低于阈值 if len(cleaned_text.strip()) < text_threshold: print(f"第 {page_num + 1} 页文本量低，开始提取图片并OCR") images = extract_images_from_page(file_path, page_num) for img in images: try: with tempfile.NamedTemporaryFile(delete=False, suffix='.' + img['ext']) as temp_image: temp_image.write(img['data']) temp_image.flush() # 调用OCR函数 ocr_result = ocr_extract(temp_image.name) if ocr_result.strip(): cleaned_text += "\n" + ocr_result except Exception as e: print(f"OCR处理失败: {e}") finally: try: os.remove(temp_image.name) except Exception as e: print(f"删除临时文件失败: {e}") result[str(page_num + 1)] = cleaned_text except Exception as e: print(f"处理PDF时出错: {e}") # 保存结果到JSON文件 try: with open(file_path + ".json", 'w', encoding='utf-8') as f: json.dump(result, f, ensure_ascii=False, indent=4) except Exception as e: print(f"保存结果到JSON文件时出错: {e}") print("PDF预处理完成") return result if __name__ == '__main__': pdf_path = "C:/test/所投主要产品检测报告.pdf" res = extract_text_json_by_page(pdf_path) print(json.dumps(res, ensure_ascii=False, indent=4))