mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 12:07:13 +02:00
Halve decoder memory with strip processing and a raveled Haar pass
Each Haar pass is one flat pywt.downcoef call over a raveled strip instead of pywt.dwt(..., axis=1)[0], and the plane is walked in strips so no full-plane float64 intermediate exists. Exact only while the last axis is even, so _approximation raises on an odd width rather than returning wrong bits, and TestRaveledHaarPass pins both that raise and the downcoef/dwt equivalence a pywt upgrade could take away. Drops the block constructor knob: the fold chains are written for 4, nothing ever passed another value, and a knob that silently decodes wrong is worse than no knob. Peak RSS 111 MB to 21 MB on a 4.3 MP image; the decoder itself 0.011s to 0.007s, which is only 0.4% of identify() now that it is under 2% of the run. Output bits and detector verdicts over 200 sampled data/ images, two synthesized carriers and eight degenerate shapes are byte-identical to the pre-vectorization decoder. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
37789e02f6
commit
4855586834
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user