qiurui
cbfe4a23dc
- Pipeline: analyze(), redact(), analyze_image() methods - API: /analyze, /redact, /analyze/frame, /analyze/frame/base64 endpoints - Detectors: regex, NER, face (RKNN NPU) - Privacy frame route added for KVM-Privacy Hub integration
143 lines
4.9 KiB
Python
143 lines
4.9 KiB
Python
# tests/test_parsers.py
|
|
import pytest
|
|
from pathlib import Path
|
|
from info_privacy.parsers.pdf_parser import PdfParser
|
|
|
|
@pytest.fixture
|
|
def pdf_path():
|
|
p = Path("tmp/sample_text.pdf")
|
|
if not p.exists():
|
|
pytest.skip("需先运行生成脚本")
|
|
return str(p)
|
|
|
|
def test_pdf_has_text_layer(pdf_path):
|
|
parser = PdfParser()
|
|
meta, blocks = parser.parse(pdf_path)
|
|
assert meta["has_text_layer"] is True
|
|
assert meta["format"] == "pdf"
|
|
assert len(blocks) > 0
|
|
|
|
def test_pdf_blocks_have_text(pdf_path):
|
|
parser = PdfParser()
|
|
_, blocks = parser.parse(pdf_path)
|
|
all_text = " ".join(b.text for b in blocks)
|
|
assert "身份证" in all_text or "张伟" in all_text
|
|
|
|
def test_pdf_blocks_layer_is_text(pdf_path):
|
|
parser = PdfParser()
|
|
_, blocks = parser.parse(pdf_path)
|
|
assert all(b.layer == "text" for b in blocks)
|
|
|
|
def test_pdf_blocks_have_bbox(pdf_path):
|
|
parser = PdfParser()
|
|
_, blocks = parser.parse(pdf_path)
|
|
for b in blocks:
|
|
assert len(b.bbox) == 4
|
|
assert b.bbox[2] > b.bbox[0] # x2 > x1
|
|
|
|
|
|
# Office 解析器测试
|
|
import docx as _docx
|
|
from info_privacy.parsers.office_parser import OfficeParser
|
|
|
|
@pytest.fixture
|
|
def sample_docx(tmp_path):
|
|
doc = _docx.Document()
|
|
doc.add_paragraph("申请人：李明，身份证：310101198501011234")
|
|
doc.add_paragraph("手机：13999998888")
|
|
p = str(tmp_path / "sample.docx")
|
|
doc.save(p)
|
|
return p
|
|
|
|
def test_office_parser_docx(sample_docx):
|
|
parser = OfficeParser()
|
|
meta, blocks = parser.parse(sample_docx)
|
|
assert meta["format"] == "docx"
|
|
assert meta["has_text_layer"] is True
|
|
assert meta["redact_strategy"] == "text_replace"
|
|
all_text = " ".join(b.text for b in blocks)
|
|
assert "李明" in all_text or "身份证" in all_text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 图像 OCR 解析器测试（Task 7）
|
|
# ---------------------------------------------------------------------------
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
FACE_MODEL = "/data/rockchip/paddle_ocr/models/rknn/PP-OCRv4/det/ppocrv4_det_ch_int8.rknn"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_image(tmp_path):
|
|
img = np.ones((200, 600, 3), dtype=np.uint8) * 255
|
|
cv2.putText(img, "ID: 110101199001011234", (10, 50),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
|
|
p = str(tmp_path / "sample.png")
|
|
cv2.imwrite(p, img)
|
|
return p
|
|
|
|
|
|
def _rknn_available() -> bool:
|
|
"""检查 RKNN 运行时是否可用（仅 RK3588 设备上可用）。"""
|
|
try:
|
|
import importlib.util
|
|
return importlib.util.find_spec("rknn") is not None or importlib.util.find_spec("rknnlite") is not None
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def test_image_parser_returns_blocks(sample_image):
|
|
det_model = "/data/rockchip/paddle_ocr/models/rknn/PP-OCRv4/det/ppocrv4_det_ch_int8.rknn"
|
|
rec_model = "/data/rockchip/paddle_ocr/models/rknn/PP-OCRv4/rec/ppocrv4_rec_ch_fp16.rknn"
|
|
dict_file = "/data/rockchip/paddle_ocr/data/dicts/ppocr_keys_v1.txt"
|
|
if not os.path.exists(det_model):
|
|
pytest.skip("RKNN模型不存在，跳过")
|
|
if not _rknn_available():
|
|
pytest.skip("rknn/rknnlite 不可用（非 RK3588 设备），跳过")
|
|
from info_privacy.parsers.image_parser import ImageParser
|
|
parser = ImageParser(det_model_path=det_model, rec_model_path=rec_model, dict_path=dict_file)
|
|
meta, blocks = parser.parse(sample_image)
|
|
assert meta["format"] == "image"
|
|
assert meta["has_text_layer"] is False
|
|
assert meta["redact_strategy"] == "image_mask"
|
|
assert isinstance(blocks, list)
|
|
|
|
|
|
def test_image_parser_block_layer(sample_image):
|
|
det_model = "/data/rockchip/paddle_ocr/models/rknn/PP-OCRv4/det/ppocrv4_det_ch_int8.rknn"
|
|
if not os.path.exists(det_model):
|
|
pytest.skip("RKNN模型不存在")
|
|
if not _rknn_available():
|
|
pytest.skip("rknn/rknnlite 不可用（非 RK3588 设备），跳过")
|
|
rec_model = "/data/rockchip/paddle_ocr/models/rknn/PP-OCRv4/rec/ppocrv4_rec_ch_fp16.rknn"
|
|
dict_file = "/data/rockchip/paddle_ocr/data/dicts/ppocr_keys_v1.txt"
|
|
from info_privacy.parsers.image_parser import ImageParser
|
|
parser = ImageParser(det_model_path=det_model, rec_model_path=rec_model, dict_path=dict_file)
|
|
_, blocks = parser.parse(sample_image)
|
|
for b in blocks:
|
|
assert b.layer == "image"
|
|
|
|
|
|
# Parser Factory 测试
|
|
from info_privacy.parsers.parser_factory import ParserFactory
|
|
|
|
def test_factory_routes_pdf(pdf_path):
|
|
factory = ParserFactory()
|
|
parser = factory.get_parser(pdf_path)
|
|
from info_privacy.parsers.pdf_parser import PdfParser
|
|
assert isinstance(parser, PdfParser)
|
|
|
|
def test_factory_routes_docx(sample_docx, tmp_path):
|
|
factory = ParserFactory()
|
|
parser = factory.get_parser(sample_docx)
|
|
from info_privacy.parsers.office_parser import OfficeParser
|
|
assert isinstance(parser, OfficeParser)
|
|
|
|
def test_factory_unsupported_raises():
|
|
factory = ParserFactory()
|
|
with pytest.raises(ValueError, match="不支持"):
|
|
factory.get_parser("file.xyz")
|