From c1971a3e8d9be7a89afdf8ca67d5bd8a586449fc Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sat, 20 Jun 2026 15:34:39 -0700 Subject: [PATCH] feat(invisible): region-targeted regeneration for AI-enhanced composites For AI-enhanced composites (digitalSourceType compositeWithTrainedAlgorithmicMedia, identify ai_source_kind == "enhanced"; roadmap P1#8): regenerate ONLY the AI region and preserve the real photo elsewhere, instead of regenerating the whole frame. - noai.tiling.feather_region_composite(base, regenerated, box, *, feather): pure, model-free compositor that blends the regenerated AI box back over the original with a feathered seam, leaving pixels OUTSIDE the box exactly equal to base. Fully unit-tested (outside-box exactness, interior == regenerated, hard paste at feather 0, monotonic seam ramp, dtype/grayscale/clamp/empty-box/shape-mismatch). - WatermarkRemover.remove_watermark(region=, region_feather=) and the module-level convenience function thread it through: the remover regenerates (or tiles) the frame, then composites only the AI box back over the original input. The box is caller-supplied -- a C2PA composite manifest carries no reliable machine-readable region, so none is fabricated. The no-model lossless region path stays region_eraser.erase. Co-Authored-By: Claude Opus 4.8 --- src/remove_ai_watermarks/noai/tiling.py | 53 ++++++++++++++ .../noai/watermark_remover.py | 33 +++++++++ tests/test_tiling.py | 70 +++++++++++++++++++ 3 files changed, 156 insertions(+) diff --git a/src/remove_ai_watermarks/noai/tiling.py b/src/remove_ai_watermarks/noai/tiling.py index 4eb6502..a858c91 100644 --- a/src/remove_ai_watermarks/noai/tiling.py +++ b/src/remove_ai_watermarks/noai/tiling.py @@ -100,6 +100,59 @@ def feather_weights(width: int, height: int, overlap: int) -> NDArray[Any]: return weights +def feather_region_composite( + base: NDArray[Any], + regenerated: NDArray[Any], + box: tuple[int, int, int, int], + *, + feather: int = 64, +) -> NDArray[Any]: + """Composite ``regenerated`` over ``base`` inside ``box`` only, feathering the seam. + + For AI-ENHANCED composites (digitalSourceType ``compositeWithTrainedAlgorithmicMedia``): + the diffusion remover regenerates the whole frame, but only the AI-composited + REGION should change -- the rest is a real photo that must be preserved. This + blends the regenerated pixels in over ``box = (x, y, w, h)`` with a separable + linear taper of ``feather`` px at the box edges, so the result equals ``base`` + EXACTLY outside the box and ramps smoothly (no hard seam) at the boundary. + + Pure and model-free (unit-tested): ``base`` and ``regenerated`` must be the same + shape (H x W, or H x W x C). The output preserves ``base``'s dtype. ``feather`` is + clamped to half the box on each axis, so a small region still tapers symmetrically; + ``feather=0`` is a hard-edged paste. + """ + import numpy as np + + if base.shape != regenerated.shape: + raise ValueError(f"shape mismatch: base {base.shape} vs regenerated {regenerated.shape}") + h, w = base.shape[:2] + x, y, bw, bh = box + x0, y0 = max(0, x), max(0, y) + x1, y1 = min(w, x + bw), min(h, y + bh) + out = base.copy() + if x1 <= x0 or y1 <= y0: + return out # empty / off-image box -> nothing regenerated + + def taper(n: int) -> NDArray[Any]: + win = np.ones(n, dtype=np.float32) + f = min(max(feather, 0), n // 2) + if f > 0: + ramp = (np.arange(f, dtype=np.float32) + 1.0) / (f + 1.0) # in (0, 1), 0 at the edge + win[:f] = ramp + win[n - f :] = ramp[::-1] + return win + + rh, rw = y1 - y0, x1 - x0 + wmap = np.outer(taper(rh), taper(rw)) # ~0 at the box edge, 1 in the interior + if base.ndim == 3: + wmap = wmap[:, :, None] + roi_base = base[y0:y1, x0:x1].astype(np.float32) + roi_gen = regenerated[y0:y1, x0:x1].astype(np.float32) + blended = roi_base * (1.0 - wmap) + roi_gen * wmap + out[y0:y1, x0:x1] = np.clip(blended, 0, 255).astype(base.dtype) + return out + + def run_tiled( generate_tile: Callable[[PILImage.Image], PILImage.Image], image: PILImage.Image, diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py index 3c4a2c6..fa38b9f 100644 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ b/src/remove_ai_watermarks/noai/watermark_remover.py @@ -566,6 +566,8 @@ class WatermarkRemover: tile: bool = False, tile_size: int = 1024, tile_overlap: int = 128, + region: tuple[int, int, int, int] | None = None, + region_feather: int = 64, ) -> Path: """Remove watermark from an image using regeneration attack. @@ -589,6 +591,15 @@ class WatermarkRemover: tile_size: Tile dimension in px (default 1024, SDXL's training size). tile_overlap: Overlap between adjacent tiles in px (default 128), feather- blended so there is no visible seam. + region: Restrict the regeneration to the AI-composited box ``(x, y, w, h)`` + and feather-composite it back over the ORIGINAL pixels everywhere else. + For AI-ENHANCED composites (digitalSourceType + ``compositeWithTrainedAlgorithmicMedia``, surfaced as + ``identify.ProvenanceReport.ai_source_kind == "enhanced"``): the real + photo outside the box is preserved exactly, only the AI region is + scrubbed. The box is supplied by the caller (a C2PA composite manifest + does not carry a reliable machine-readable region). None -> whole frame. + region_feather: Seam taper in px for ``region`` compositing (default 64). Returns: Path to the cleaned image. @@ -660,6 +671,22 @@ class WatermarkRemover: self._controlnet_pipeline = None cleaned_image = _generate() + # Region-targeted regeneration for AI-enhanced composites: keep the real photo + # outside the AI box pixel-exact, blend only the regenerated AI region back in. + if region is not None: + import numpy as np + + from remove_ai_watermarks.noai.tiling import feather_region_composite + + gen = cleaned_image.convert("RGB") + if gen.size != init_image.size: # a downscaled/tiled pass can resize + gen = gen.resize(init_image.size) + cleaned_image = gen + base_rgb = np.asarray(init_image) # original RGB, untouched outside the box + merged = feather_region_composite(base_rgb, np.asarray(gen), region, feather=region_feather) + cleaned_image = Image.fromarray(merged) + self._set_progress(f"Region-targeted regeneration: AI box {region}, real photo preserved") + self._set_progress(f"Regeneration complete ยท Output: {w}x{h}px {cleaned_image.mode}") output_path.parent.mkdir(parents=True, exist_ok=True) @@ -877,12 +904,17 @@ def remove_watermark( model_id: str | None = None, device: str | None = None, hf_token: str | None = None, + region: tuple[int, int, int, int] | None = None, ) -> Path: """Convenience function to remove watermark from an image. ``strength=None`` lets the profile pick its vendor-adaptive default (0.20 OpenAI / 0.30 Google / 0.30 unknown, from the C2PA SynthID proxy on the input; same ladder for the controlnet and sdxl pipelines). Pass a value to override. + + ``region=(x, y, w, h)`` restricts the regeneration to that box and preserves the + real photo elsewhere -- for AI-enhanced composites (see + ``WatermarkRemover.remove_watermark``). """ from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength @@ -892,4 +924,5 @@ def remove_watermark( output_path=output_path, strength=strength, vendor=vendor_for_strength(image_path), + region=region, ) diff --git a/tests/test_tiling.py b/tests/test_tiling.py index fbe440f..ab6df75 100644 --- a/tests/test_tiling.py +++ b/tests/test_tiling.py @@ -15,6 +15,7 @@ from PIL import Image from remove_ai_watermarks.noai.tiling import ( Tile, _axis_positions, + feather_region_composite, feather_weights, plan_tiles, run_tiled, @@ -138,3 +139,72 @@ class TestRunTiled: image = Image.new("RGB", (1500, 1100), (200, 100, 50)) out = run_tiled(generate, image, tile_size=1024, overlap=128) assert out.size == (1500, 1100) + + +class TestFeatherRegionComposite: + """Region-targeted compositing for AI-enhanced composites: only the AI box is + regenerated, the real photo outside it stays pixel-exact (roadmap P1#8).""" + + @staticmethod + def _frames(h=200, w=300): + base = np.full((h, w, 3), 80, np.uint8) + regenerated = np.full((h, w, 3), 200, np.uint8) + return base, regenerated + + def test_outside_box_is_pixel_exact(self): + base, regen = self._frames() + out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8) + # Far corners are well outside the box -> identical to base. + assert np.array_equal(out[:50, :80], base[:50, :80]) + assert np.array_equal(out[150:, 220:], base[150:, 220:]) + + def test_interior_equals_regenerated(self): + base, regen = self._frames() + out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8) + # Deep interior of the box (past the feather ramp) is fully regenerated. + assert np.array_equal(out[80:90, 130:150], regen[80:90, 130:150]) + + def test_hard_paste_when_no_feather(self): + base, regen = self._frames() + out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=0) + assert np.array_equal(out[60:110, 100:180], regen[60:110, 100:180]) + assert np.array_equal(out[:60], base[:60]) + + def test_seam_is_monotonic_ramp(self): + base, regen = self._frames() + out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=10).astype(np.float32) + # Along a horizontal line crossing the left edge, values rise from base(80) + # toward regenerated(200) monotonically through the feather band. + row = out[85, 100:115, 0] + assert row[0] < row[-1] + assert np.all(np.diff(row) >= -1e-3) + + def test_dtype_preserved(self): + base, regen = self._frames() + out = feather_region_composite(base, regen, (50, 50, 40, 40), feather=4) + assert out.dtype == base.dtype + + def test_grayscale_2d_supported(self): + base = np.full((100, 120), 30, np.uint8) + regen = np.full((100, 120), 220, np.uint8) + out = feather_region_composite(base, regen, (40, 30, 30, 30), feather=4) + assert out.shape == base.shape + assert np.array_equal(out[:30], base[:30]) + + def test_empty_or_offimage_box_returns_base(self): + base, regen = self._frames() + assert np.array_equal(feather_region_composite(base, regen, (0, 0, 0, 0)), base) + assert np.array_equal(feather_region_composite(base, regen, (500, 500, 40, 40)), base) + + def test_box_clamped_to_image_bounds(self): + base, regen = self._frames() + # Box overhangs the bottom-right; only the in-image part is composited. + out = feather_region_composite(base, regen, (280, 180, 60, 60), feather=0) + assert np.array_equal(out[180:, 280:], regen[180:, 280:]) + assert out.shape == base.shape + + def test_shape_mismatch_raises(self): + base, _ = self._frames(200, 300) + bad = np.full((100, 100, 3), 200, np.uint8) + with pytest.raises(ValueError, match="shape mismatch"): + feather_region_composite(base, bad, (10, 10, 20, 20))