优化PDF文本查找功能,支持列表类型查询,新增预处理选项以提高模糊匹配准确性,修复多个匹配结果的处理逻辑
This commit is contained in:
@@ -1,7 +1,22 @@
|
|||||||
import fitz # pymupdf
|
import fitz # pymupdf
|
||||||
import regex # 支持多行正则
|
import regex # 支持多行正则
|
||||||
from rapidfuzz import fuzz
|
from rapidfuzz import fuzz
|
||||||
|
import re
|
||||||
|
def normalize_text(text):
|
||||||
|
"""标准化文本,移除多余空白字符"""
|
||||||
|
# 将换行符、制表符等替换为空格,然后合并多个空格为一个
|
||||||
|
import re
|
||||||
|
normalized = re.sub(r'\s+', ' ', text.strip())
|
||||||
|
return normalized
|
||||||
|
|
||||||
|
|
||||||
|
def clean_text_for_fuzzy_match(text):
|
||||||
|
"""清理文本用于模糊匹配,移除特殊字符,只保留字母数字和空格"""
|
||||||
|
# 移除标点符号和特殊字符,只保留字母、数字、中文字符和空格
|
||||||
|
cleaned = re.sub(r'[^\w\s\u4e00-\u9fff]', '', text)
|
||||||
|
# 标准化空白字符
|
||||||
|
cleaned = re.sub(r'\s+', ' ', cleaned.strip())
|
||||||
|
return cleaned
|
||||||
def _merge_lines(lines):
|
def _merge_lines(lines):
|
||||||
"""
|
"""
|
||||||
把多行文本合并成一段,同时记录每行 bbox 的并集。
|
把多行文本合并成一段,同时记录每行 bbox 的并集。
|
||||||
@@ -43,16 +58,25 @@ def _collect_lines(page):
|
|||||||
return lines
|
return lines
|
||||||
|
|
||||||
def find_text_in_pdf(pdf_path,
|
def find_text_in_pdf(pdf_path,
|
||||||
query,
|
query, # 修改为支持list类型
|
||||||
use_regex=False,
|
use_regex=False,
|
||||||
threshold=80, # rapidfuzz 默认 0~100
|
threshold=80, # rapidfuzz 默认 0~100
|
||||||
page_range=None): # 例如 (1,5) 只搜 1-4 页
|
page_range=None,
|
||||||
|
preprocess=True): # 添加预处理选项
|
||||||
"""
|
"""
|
||||||
高级查找函数
|
高级查找函数
|
||||||
query: 正则表达式字符串 或 普通字符串
|
query: 正则表达式字符串 或 普通字符串,或它们的列表
|
||||||
|
preprocess: 是否对文本进行预处理以提高模糊匹配准确性
|
||||||
返回: list[dict] 每个 dict 含 page, bbox, matched_text
|
返回: list[dict] 每个 dict 含 page, bbox, matched_text
|
||||||
"""
|
"""
|
||||||
results = []
|
# 处理单个查询字符串的情况
|
||||||
|
if isinstance(query, str):
|
||||||
|
queries = [query]
|
||||||
|
else:
|
||||||
|
queries = query # 假设已经是列表
|
||||||
|
# 初始化结果列表
|
||||||
|
batch_results = [[] for _ in queries]
|
||||||
|
|
||||||
doc = fitz.open(pdf_path)
|
doc = fitz.open(pdf_path)
|
||||||
pages = range(len(doc)) if page_range is None else range(page_range[0]-1, page_range[1])
|
pages = range(len(doc)) if page_range is None else range(page_range[0]-1, page_range[1])
|
||||||
|
|
||||||
@@ -63,58 +87,88 @@ def find_text_in_pdf(pdf_path,
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
full_text, _ = _merge_lines(lines) # 整页纯文本
|
full_text, _ = _merge_lines(lines) # 整页纯文本
|
||||||
positions = [] # 记录匹配区间在 full_text 中的起止字符索引
|
|
||||||
|
|
||||||
if use_regex:
|
# 如果启用预处理,则对整页文本进行预处理
|
||||||
# regex 支持 (?s) 使 . 匹配换行
|
processed_full_text = full_text
|
||||||
pattern = regex.compile(query)
|
if preprocess and not use_regex:
|
||||||
for match in pattern.finditer(full_text):
|
processed_full_text = clean_text_for_fuzzy_match(full_text)
|
||||||
positions.append((match.start(), match.end(), match.group()))
|
|
||||||
else:
|
|
||||||
# 模糊匹配:滑动窗口(整页 vs 查询)
|
|
||||||
# 修改:支持多个匹配结果并计算相似度分数
|
|
||||||
potential_matches = []
|
|
||||||
# 使用不同的方法获取多个可能的匹配
|
|
||||||
for i in range(len(full_text) - len(query) + 1):
|
|
||||||
if i < 0:
|
|
||||||
continue
|
|
||||||
window_text = full_text[i:i + len(query)]
|
|
||||||
if window_text.strip(): # 只处理非空文本
|
|
||||||
score = fuzz.partial_ratio(query, window_text)
|
|
||||||
if score >= threshold:
|
|
||||||
potential_matches.append((i, i + len(query), window_text, score))
|
|
||||||
|
|
||||||
# 如果找到了潜在匹配,按分数排序并只取最高分的匹配
|
# 一次性计算所有查询的匹配结果
|
||||||
if potential_matches:
|
for idx ,q in enumerate(queries):
|
||||||
# 按分数降序排序
|
positions = [] # 记录匹配区间在 full_text 中的起止字符索引
|
||||||
potential_matches.sort(key=lambda x: x[3], reverse=True)
|
results = []
|
||||||
# 只取分数最高的匹配
|
if use_regex:
|
||||||
best_match = potential_matches[0]
|
# regex 支持 (?s) 使 . 匹配换行
|
||||||
positions.append((best_match[0], best_match[1], best_match[2]))
|
pattern = regex.compile(q)
|
||||||
|
for match in pattern.finditer(full_text):
|
||||||
|
positions.append((match.start(), match.end(), match.group()))
|
||||||
|
else:
|
||||||
|
# 模糊匹配:滑动窗口(整页 vs 查询)
|
||||||
|
# 修改:支持多个匹配结果并计算相似度分数
|
||||||
|
potential_matches = []
|
||||||
|
query_text = q
|
||||||
|
# 如果启用预处理,则对查询文本也进行预处理
|
||||||
|
if preprocess:
|
||||||
|
query_text = clean_text_for_fuzzy_match(q)
|
||||||
|
score = fuzz.partial_ratio(processed_full_text, query_text, score_cutoff=threshold)
|
||||||
|
if score >= threshold:
|
||||||
|
# 这里简单返回整页;如需精确定位,可再做二次对齐
|
||||||
|
positions.append((0, len(full_text), full_text))
|
||||||
|
|
||||||
# 将字符区间映射回行
|
# query_len = len(query_text)
|
||||||
for start, end, matched_text in positions:
|
# text_len = len(processed_full_text)
|
||||||
# 计算每一行在 full_text 中的起止字符偏移
|
|
||||||
offset = 0
|
# # 优化:只在合理范围内进行滑动窗口匹配
|
||||||
matched_lines = []
|
# # 添加早期终止机制,一旦找到足够高的匹配就停止搜索
|
||||||
for text, bbox in lines:
|
# best_score = 0
|
||||||
line_start = offset
|
# for i in range(text_len - query_len + 1):
|
||||||
line_end = offset + len(text)
|
# window_text = processed_full_text[i:i + query_len]
|
||||||
# 检查该行是否与匹配区间有重叠 - 更严格的条件
|
# # 优化:只处理非空文本
|
||||||
if line_start < end and line_end > start:
|
# if window_text.strip():
|
||||||
matched_lines.append((text, bbox))
|
# # 优化:使用更快速的相似度计算方法
|
||||||
# 修正:正确计算偏移量,包括换行符
|
# score = fuzz.partial_ratio(query_text, window_text)
|
||||||
offset += len(text) + 1 # 加上换行符的长度
|
# if score >= threshold:
|
||||||
# 修正:只有当确实匹配到文本时才添加结果
|
# # 优化:记录当前最佳分数
|
||||||
if matched_lines:
|
# if score > best_score:
|
||||||
_, merged_bbox = _merge_lines(matched_lines)
|
# best_score = score
|
||||||
results.append({
|
# potential_matches.append((i, i + query_len, window_text, score))
|
||||||
"page": p + 1,
|
# # 优化:如果找到非常高分的匹配,可以提前终止
|
||||||
"bbox": merged_bbox,
|
# if score >= 95: # 如果匹配度已经很高,可以提前结束
|
||||||
"matched_text": matched_text
|
# break
|
||||||
})
|
|
||||||
|
# 如果找到了潜在匹配,按分数排序并只取最高分的匹配
|
||||||
|
# if potential_matches:
|
||||||
|
# # 按分数降序排序
|
||||||
|
# potential_matches.sort(key=lambda x: x[3], reverse=True)
|
||||||
|
# # 只取分数最高的匹配
|
||||||
|
# best_match = potential_matches[0]
|
||||||
|
# positions.append((best_match[0], best_match[1], best_match[2]))
|
||||||
|
|
||||||
|
# 将字符区间映射回行
|
||||||
|
for start, end, matched_text in positions:
|
||||||
|
# 计算每一行在 full_text 中的起止字符偏移
|
||||||
|
offset = 0
|
||||||
|
matched_lines = []
|
||||||
|
for text, bbox in lines:
|
||||||
|
line_start = offset
|
||||||
|
line_end = offset + len(text)
|
||||||
|
# 检查该行是否与匹配区间有重叠 - 更严格的条件
|
||||||
|
if line_start < end and line_end > start:
|
||||||
|
matched_lines.append((text, bbox))
|
||||||
|
# 修正:正确计算偏移量,包括换行符
|
||||||
|
offset += len(text) + 1 # 加上换行符的长度
|
||||||
|
# 修正:只有当确实匹配到文本时才添加结果
|
||||||
|
if matched_lines:
|
||||||
|
_, merged_bbox = _merge_lines(matched_lines)
|
||||||
|
results.append({
|
||||||
|
"page": p + 1,
|
||||||
|
"bbox": merged_bbox,
|
||||||
|
"matched_text": matched_text
|
||||||
|
})
|
||||||
|
if results:
|
||||||
|
batch_results[idx].append(results)
|
||||||
doc.close()
|
doc.close()
|
||||||
return results
|
return batch_results
|
||||||
|
|
||||||
def highlight_matches(pdf_path, matches, output_path="highlighted.pdf"):
|
def highlight_matches(pdf_path, matches, output_path="highlighted.pdf"):
|
||||||
"""
|
"""
|
||||||
@@ -161,8 +215,19 @@ if __name__ == "__main__":
|
|||||||
threshold=75
|
threshold=75
|
||||||
|
|
||||||
)
|
)
|
||||||
for match in matches:
|
# 修改:正确处理二维列表结果
|
||||||
print(f"第 {match['page']} 页 匹配: {match['matched_text'][:50]}... 位置: {match['bbox']}")
|
print(matches)
|
||||||
|
print("------------------")
|
||||||
|
for idx,query_matches in enumerate(matches):
|
||||||
|
for m_item in query_matches:
|
||||||
|
highlight_matches(pdf_path, m_item, "example_highlighted.pdf")
|
||||||
|
for m in m_item:
|
||||||
|
# 输出匹配结果
|
||||||
|
#print(m)
|
||||||
|
|
||||||
|
print(f"第 {m['page']} 页 匹配: {m['matched_text'][:50]}... 位置: {m['bbox']}")
|
||||||
|
|
||||||
# 2. 高亮并保存
|
# 2. 高亮并保存
|
||||||
highlight_matches(pdf_path, matches, "example_highlighted.pdf")
|
# 修改:展平二维列表用于高亮
|
||||||
|
# flattened_matches = [match for query_matches in matches for match in query_matches]
|
||||||
|
# highlight_matches(pdf_path, flattened_matches, "example_highlighted.pdf")
|
||||||
|
Reference in New Issue
Block a user