|
|
@@ -0,0 +1,35 @@
|
|
|
+"""WatermarkProcessor 与 preset 单测。"""
|
|
|
+import numpy as np
|
|
|
+import pytest
|
|
|
+
|
|
|
+from ocr_utils.watermark import WatermarkProcessor, merge_watermark_config, get_preset
|
|
|
+
|
|
|
+
|
|
|
+def test_merge_cell_preset_has_mask():
|
|
|
+ cfg = merge_watermark_config("cell", {"method": "masked_adaptive"})
|
|
|
+ assert cfg["method"] == "masked_adaptive"
|
|
|
+ assert "mask" in cfg
|
|
|
+ assert cfg["mask"]["hough_min_line_length"] < get_preset("page", "masked_adaptive")["mask"][
|
|
|
+ "hough_min_line_length"
|
|
|
+ ]
|
|
|
+
|
|
|
+
|
|
|
+def test_processor_disabled_passthrough():
|
|
|
+ proc = WatermarkProcessor({"enabled": False}, scope="cell")
|
|
|
+ img = np.ones((40, 80, 3), dtype=np.uint8) * 255
|
|
|
+ out, stages = proc.process(img)
|
|
|
+ assert stages == []
|
|
|
+ assert out.shape == img.shape
|
|
|
+
|
|
|
+
|
|
|
+def test_processor_cell_force_wm():
|
|
|
+ proc = WatermarkProcessor(
|
|
|
+ {"enabled": True, "method": "masked_adaptive"}, scope="cell"
|
|
|
+ )
|
|
|
+ img = np.ones((50, 100, 3), dtype=np.uint8) * 240
|
|
|
+ cv2 = pytest.importorskip("cv2")
|
|
|
+ for x in range(0, 100, 8):
|
|
|
+ cv2.line(img, (x, 0), (x + 30, 50), (180, 180, 180), 1)
|
|
|
+ out, stages = proc.process(img, force=True)
|
|
|
+ assert "wm" in stages
|
|
|
+ assert out.shape == img.shape
|