tm-desensitize/tmkit/redactor.py

49 lines
1.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""坐标回写:在 PII bbox 处黑条遮盖,输出脱敏后 PDF bytes。
两种方案:
A) add_redact_annot + apply_redactions(官方推荐,真正删除文字内容)
B) draw_rect fill=black(只画黑条,文字仍在文本层,可被复制提取)
默认用 A(删得最干净);fill 参数可同时画出黑底。测试验证 A 能让
get_text 中 PII 消失。
"""
import io
from typing import Iterable, List
import fitz # PyMuPDF
def redact_pdf(pdf_bytes: bytes, hits: List[dict], method: str = "redact") -> bytes:
"""按 hits([{page, bbox}, ...])脱敏,返回新 PDF bytes。
method:
'redact' → add_redact_annot + apply_redactions(删除内容,推荐)
'draw' → 仅画黑矩形(文字仍在)
"""
by_page: dict = {}
for h in hits:
by_page.setdefault(h["page"], []).append(h["bbox"])
src = fitz.open(stream=pdf_bytes, filetype="pdf")
out = fitz.open()
for pno, page in enumerate(src):
boxes = by_page.get(pno, [])
if method == "redact":
for bbox in boxes:
page.add_redact_annot(bbox, fill=(0, 0, 0))
if boxes:
page.apply_redactions()
elif method == "draw":
for bbox in boxes:
page.draw_rect(bbox, color=None, fill=(0, 0, 0))
else:
raise ValueError(f"unknown method: {method}")
out.insert_pdf(src, from_page=pno, to_page=pno)
# 注意:redact 修改的是 src 中的页,insert_pdf 会拷贝修改后的页
buf = io.BytesIO()
out.save(buf)
src.close()
out.close()
return buf.getvalue()