diff --git a/docs/module-internals.md b/docs/module-internals.md index cfb2fa3..89bce5a 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -595,18 +595,32 @@ available from a hand-rolled numpy Haar is unreachable rather than merely untaken: pywt's C convolution contracts into an FMA that numpy has no ufunc for, and `np.longdouble` is 64-bit on arm64 macOS. -Measured on a 1536x2816 image: the decoder went 0.280 s to 0.019 s, and a warm -`identify()` on the same file 1.757 s to 1.365 s (-22.3%), both arms timed in one -process. Verified by recording decoder output and detector verdict over 200 -sampled `data/` images plus two synthesized carriers before and after the change: -the record is byte-identical, as are five degenerate shapes (`1x65536` through -`8x8192`) that clear the caller's area check. +Each pass is one flat `pywt.downcoef` call over a raveled strip rather than +`pywt.dwt(..., axis=1)[0]`, and the plane is processed in strips of `_STRIP` +block-rows so no full-plane float64 intermediate is ever materialized. The strip +height is not a tuned value: 8 through 64 all scored inside each other's noise +with unstable ordering, and only "strips at all" versus whole-plane matters. -Two further optimizations were identified and deliberately not taken, because -each buys speed with a new precondition rather than with less work: -`pywt.downcoef("a", ...)` on a flattened plane (valid only while the last axis -stays even and C-contiguous), and a single-pass `np.abs(..., out=)` over the -block gather. Both would need their own ablation. +`_approximation`'s even-last-axis check is load-bearing, not defensive. Haar's +filter is length 2, so an even row length keeps every pair inside its own row; +on an odd width the pairs walk across row boundaries and the reshape still +succeeds whenever the total is even, which would be wrong bits with no +exception. Both call sites are even by construction today, and +`TestRaveledHaarPass` pins both halves -- the `downcoef`/`dwt` equivalence, +which a pywt upgrade could take away, and the raise on an odd width. + +Measured on a 1536x2816 image, all arms timed in one process: the decoder went +0.112 s (the original per-block Python loop) to 0.011 s vectorized to 0.007 s +with strips, and a warm `identify()` 1.757 s to 1.365 s on the first step and a +further 0.4% on the second. That last figure is the point at which this target +is finished: the decoder is now under 2% of `identify()`, so speed here has +stopped buying anything. What the strips buy is peak RSS in the stage, 111 MB to +21 MB on a 4.3 MP image, which is what matters on the memory-limited Space. + +Both steps were verified by recording decoder output and detector verdict over +200 sampled `data/` images plus two synthesized carriers before and after: the +record is byte-identical, as are seven degenerate shapes (`1x65536` through +`8x8192`, plus an odd width) that clear the caller's area check. Regression coverage: diff --git a/src/remove_ai_watermarks/dwt_dct.py b/src/remove_ai_watermarks/dwt_dct.py index 5f45898..ba6fe13 100644 --- a/src/remove_ai_watermarks/dwt_dct.py +++ b/src/remove_ai_watermarks/dwt_dct.py @@ -27,7 +27,39 @@ if TYPE_CHECKING: from numpy.typing import NDArray _DEFAULT_SCALES = (0, 36, 36) -_DEFAULT_BLOCK = 4 +# Fixed by the format being decoded, and the 4-way fold chains in `_frame_bits` +# are written for it. It was a constructor parameter while the block scan was a +# generic Python loop; nothing ever passed another value, and a knob that now +# silently returns wrong bits is worse than no knob. +_BLOCK = 4 + +# Block-rows of the approximation band handled per strip. Not a tuned value: +# every height measured landed inside the others' noise. What matters is strips +# at all rather than a full-plane intermediate, not this number. +_STRIP = 16 + + +def _approximation(rows: NDArray[Any]) -> NDArray[Any]: + """One Haar pass along the last axis, approximation band only. + + ``pywt.dwt(x, "haar", axis=1)[0]`` computes and allocates the detail band as + well, and dispatches per row. Flattening lets one ``downcoef`` call do the + whole plane, and it is the same numbers in the same order **only while the + last axis is even**: Haar's filter is length 2, so an even row length makes + every pair fall inside its own row with no boundary extension. On an odd + width the pairs walk across row boundaries and the reshape below still + succeeds whenever the total is even -- wrong bits, no exception, and the + upstream-parity test is `skipif`-gated. Hence the explicit check. + + The transform itself stays inside pywt however slow that looks. Its C + convolution contracts into an FMA that no numpy expression reproduces, and + the caller's threshold is ``peak % 36 > 18.0`` against values that are exact + multiples of 0.5, so a 1-ulp difference flips real bits. + """ + height, width = rows.shape + if width % 2: + raise RuntimeError(f"row length {width} is odd; the raveled Haar pass requires an even last axis") + return pywt.downcoef("a", rows.ravel(), "haar").reshape(height, width // 2) class _DecodeMaxDct: @@ -37,11 +69,9 @@ class _DecodeMaxDct: self, wm_lengths: tuple[int, ...], scales: tuple[int, int, int] = _DEFAULT_SCALES, - block: int = _DEFAULT_BLOCK, ) -> None: self._wm_lengths = wm_lengths self._scales = scales - self._block = block def decode(self, bgr: NDArray[Any]) -> dict[int, NDArray[Any]]: row, col, _channels = bgr.shape @@ -49,7 +79,7 @@ class _DecodeMaxDct: trimmed = yuv[: row // 4 * 4, : col // 4 * 4] per_channel = [ - self._frame_bits(self._approximation(trimmed, channel), self._scales[channel]) + self._plane_bits(trimmed, channel, self._scales[channel]) for channel in range(2) if self._scales[channel] > 0 ] @@ -66,43 +96,57 @@ class _DecodeMaxDct: decoded[wm_len] = sums * 255 > counts * 127 return decoded - @staticmethod - def _approximation(trimmed: NDArray[Any], channel: int) -> NDArray[Any]: - """The Haar approximation band, and only it. + def _plane_bits(self, trimmed: NDArray[Any], channel: int, scale: int) -> NDArray[Any]: + """Block bits for one colour plane, a strip of block-rows at a time. ``dwt2`` is ``dwtn``: it transforms along axis 0, then along axis 1 over both halves, and three of the four bands it returns are discarded here. - Two ``dwt`` calls keeping ``[0]`` skip that, and transposing between them - lets pywt walk a contiguous axis instead of a column. + Only the approximation is ever asked for, and transposing between the + two passes lets pywt walk a contiguous axis instead of a column -- that + access pattern, not the arithmetic saved, is where the time goes. - The result must stay bit-identical to ``dwt2``'s, which is why the - transform is left to pywt however slow that is: the caller's threshold is - ``peak % 36 > 18.0``, and for uint8 input the exact value is a multiple - of 0.5, so it lands exactly on the threshold often enough that a 1-ulp - difference flips real bits. + Strips mean no full-plane float64 intermediate is ever materialized, and + they are seam-free for the reason ``_approximation`` documents: output + row ``k`` reads input rows ``2k`` and ``2k + 1`` only. Strips start on + multiples of ``2 * _BLOCK``, so neither a pair nor a 4x4 block straddles + one. """ if trimmed.shape[0] == 0 or trimmed.shape[1] == 0: - # Reachable: a 1x65536 image clears the caller's area check. Left to - # dwt2 so the exception stays the one this module has always raised. + # Reachable: a 1x65536 image clears the caller's area check and + # trims to an empty plane. Left to dwt2 so the exception stays the + # one this module has always raised -- returning empty bits here + # instead would silently turn a raise into an all-false verdict. return pywt.dwt2(trimmed[:, :, channel], "haar")[0] - columns = cv2.transpose(cv2.extractChannel(trimmed, channel)) - along_rows = pywt.dwt(columns, "haar", axis=1)[0] - return pywt.dwt(cv2.transpose(along_rows), "haar", axis=1)[0] + rows = trimmed.shape[0] // (2 * _BLOCK) + cols = trimmed.shape[1] // (2 * _BLOCK) + if rows == 0 or cols == 0: + return np.zeros(0, dtype=np.float64) + + width = cols * 2 * _BLOCK + pieces: list[NDArray[Any]] = [] + for start in range(0, rows, _STRIP): + stop = min(start + _STRIP, rows) + strip = trimmed[start * 2 * _BLOCK : stop * 2 * _BLOCK, :width] + columns = cv2.transpose(cv2.extractChannel(strip, channel)) + band = _approximation(cv2.transpose(_approximation(columns))) + pieces.append(self._frame_bits(band, scale)) + return np.concatenate(pieces) def _frame_bits(self, frame: NDArray[Any], scale: int) -> NDArray[Any]: """One bit per 4x4 block, in row-major block order. Upstream's per-block loop, said to numpy once instead of to the - interpreter ~135k times per image. + interpreter ~135k times per image. Zeroing the DC term in the absolute + band says "ignore index 0" without materializing a ``(nblocks, 16)`` + copy, and the 4-way ``np.maximum`` chains reduce over contiguous rows. """ - block = self._block - rows = frame.shape[0] // block - cols = frame.shape[1] // block - if rows == 0 or cols == 0: - return np.zeros(0, dtype=np.float64) - aligned = frame[: rows * block, : cols * block] - blocks = aligned.reshape(rows, block, cols, block).swapaxes(1, 2) - peak = np.abs(blocks.reshape(rows * cols, block * block)[:, 1:]).max(axis=1) + rows = frame.shape[0] // _BLOCK + cols = frame.shape[1] // _BLOCK + band = np.abs(frame[: rows * _BLOCK, : cols * _BLOCK]) + band[::_BLOCK, ::_BLOCK] = 0.0 + folded = np.maximum(np.maximum(band[0::4], band[1::4]), np.maximum(band[2::4], band[3::4])) + folded = folded.reshape(rows * cols, _BLOCK) + peak = np.maximum(np.maximum(folded[:, 0], folded[:, 1]), np.maximum(folded[:, 2], folded[:, 3])) return ((peak % scale) > 0.5 * scale).astype(np.float64) diff --git a/tests/test_invisible_watermark.py b/tests/test_invisible_watermark.py index a516146..db0568c 100644 --- a/tests/test_invisible_watermark.py +++ b/tests/test_invisible_watermark.py @@ -59,6 +59,47 @@ class TestHelpers: assert _bytes_match_frac(b"abc", b"abcd") == 0.0 +class TestRaveledHaarPass: + """The precondition that makes the decoder's flat Haar pass legitimate. + + `_approximation` replaces `pywt.dwt(x, "haar", axis=1)[0]` with one + `downcoef` call over `x.ravel()`. That is exact only while the last axis is + even. Neither half of this is checked anywhere else: the equivalence is a + property of pywt's implementation that an upgrade could take away, and an + odd width produces wrong bits with no exception, since the reshape still + succeeds whenever the total length is even. + """ + + @pytest.mark.parametrize("shape", [(64, 64), (7, 128), (129, 2), (2, 2), (33, 400)]) + def test_matches_pywt_dwt_on_even_widths(self, shape: tuple[int, int]): + import pywt + + from remove_ai_watermarks.dwt_dct import _approximation + + rng = np.random.default_rng(0) + # uint8 is what the FIRST production pass receives -- `decode` builds its + # plane with cvtColor, and extractChannel and transpose preserve the + # dtype -- so a divergence in how downcoef coerces integers would be + # invisible to a float-only parametrization. + for array in ( + rng.random(shape), + rng.integers(0, 256, shape).astype(np.float64), + rng.integers(0, 256, shape).astype(np.uint8), + ): + expected = pywt.dwt(array, "haar", axis=1)[0] + got = _approximation(array) + assert got.shape == expected.shape + assert np.array_equal(got, expected), "downcoef diverged from dwt -- a pywt upgrade may have changed it" + + def test_odd_width_raises_instead_of_returning_wrong_bits(self): + from remove_ai_watermarks.dwt_dct import _approximation + + # 4x6 ravels to 24, an even total, so the reshape would happily produce + # a 4x3 array of numbers that pair across row boundaries. + with pytest.raises(RuntimeError, match="odd"): + _approximation(np.zeros((4, 6))[:, :5]) + + class TestDetect: def test_in_tree_decoder_matches_upstream(self, tmp_path: Path): from imwatermark import WatermarkDecoder