diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index ec3c616d..f39dfe1d 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -26,6 +26,7 @@ jobs: python examples/example_no_writing.py python examples/example_str.py python examples/example_str_multi.py + python -m unittest discover -s tests # pytest --cov . # - name: Upload coverage reports to Codecov with GitHub Action # uses: codecov/codecov-action@v3 diff --git a/blind_watermark/blind_watermark.py b/blind_watermark/blind_watermark.py index eea10ed1..2daf7908 100644 --- a/blind_watermark/blind_watermark.py +++ b/blind_watermark/blind_watermark.py @@ -11,6 +11,14 @@ from .version import bw_notes +def _bits_to_text(bits): + byte = ''.join(str((i >= 0.5) * 1) for i in bits) + hex_str = hex(int(byte, base=2))[2:] + if len(hex_str) % 2: + hex_str = '0' + hex_str + return bytes.fromhex(hex_str).decode('utf-8', errors='replace') + + class WaterMark: def __init__(self, password_wm=1, password_img=1, block_shape=(4, 4), mode='common', processes=None): bw_notes.print_notes() @@ -102,7 +110,6 @@ def extract(self, filename=None, embed_img=None, wm_shape=None, out_wm_name=None wm = 255 * wm.reshape(wm_shape[0], wm_shape[1]) cv2.imwrite(out_wm_name, wm) elif mode == 'str': - byte = ''.join(str((i >= 0.5) * 1) for i in wm) - wm = bytes.fromhex(hex(int(byte, base=2))[2:]).decode('utf-8', errors='replace') + wm = _bits_to_text(wm) return wm diff --git a/tests/test_extract_text_decode.py b/tests/test_extract_text_decode.py new file mode 100644 index 00000000..88a33807 --- /dev/null +++ b/tests/test_extract_text_decode.py @@ -0,0 +1,20 @@ +import unittest + +import numpy as np + +from blind_watermark.blind_watermark import WaterMark + + +class ExtractTextDecodeTest(unittest.TestCase): + def test_extract_str_pads_odd_length_hex(self): + bwm = WaterMark(password_img=1, password_wm=1) + bwm.extract_decrypt = lambda wm_avg: wm_avg + bwm.bwm_core.extract_with_kmeans = lambda img, wm_shape: np.array([1, 1, 1, 1]) + + wm = bwm.extract(embed_img=np.zeros((4, 4, 3)), wm_shape=4, mode='str') + + self.assertEqual(wm, '\x0f') + + +if __name__ == '__main__': + unittest.main()