mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-30 19:21:37 +02:00
feat: GFPGAN face-identity restoration post-pass
Add an optional, commercial-safe face-restoration post-pass that recovers face identity the diffusion removal pass drifts (canny holds structure, not likeness) while still scrubbing the pixel watermark in the face regions. - face_restore.py: GFPGANer singleton (CPU unless CUDA), the basicsr torchvision.transforms.functional_tensor shim, and the pure feather _composite_faces helper (unit-tested without the model). GFPGAN re-synthesizes each face from a StyleGAN2 prior, so composited face pixels are GAN-generated (no watermark, no pixel-copy) -- oracle-clean at weight 0.5 with identity preserved. - InvisibleEngine.remove_watermark: restore_faces / restore_faces_weight, best-effort, auto-skips when the extra is absent or no face is detected. - CLI --restore-faces/--no-restore-faces + --restore-faces-weight on invisible/all/batch (on by default). - restore extra (gfpgan/facexlib/basicsr), numpy<2-pinned (scipy<1.18, numba<0.60) and kept out of `all`; basicsr needs Python <3.13 + setuptools<69 to build, so pin .python-version 3.12. Commercial-safe: GFPGAN Apache-2.0, RetinaFace MIT. The CodeFormer alternative is non-commercial and is not shipped. The earlier IP-Adapter FaceID layer was removed (footgun: needs high strength, corrupts faces at the low removal strength). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
d90d5d886a
commit
411ef16ec3
@@ -0,0 +1,85 @@
|
||||
"""Tests for the GFPGAN face-restoration post-pass.
|
||||
|
||||
The pure feather-composite helper is unit-tested without the model; the
|
||||
model-running paths are gated behind ``is_available()`` (a multi-hundred-MB
|
||||
download), matching the discipline used for the other ML-adjacent modules.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks import face_restore
|
||||
|
||||
|
||||
class TestIsAvailable:
|
||||
def test_returns_bool(self):
|
||||
assert isinstance(face_restore.is_available(), bool)
|
||||
|
||||
def test_reflects_dependencies(self):
|
||||
import importlib.util
|
||||
|
||||
expected = all(importlib.util.find_spec(m) is not None for m in ("gfpgan", "facexlib"))
|
||||
assert face_restore.is_available() is expected
|
||||
|
||||
|
||||
class TestCompositeFaces:
|
||||
"""Unit tests for the pure ``_composite_faces`` helper (cv2/numpy only)."""
|
||||
|
||||
def _base_and_restored(self, h: int = 100, w: int = 120):
|
||||
base = np.zeros((h, w, 3), dtype=np.uint8) # black
|
||||
restored = np.full((h, w, 3), 255, dtype=np.uint8) # white
|
||||
return base, restored
|
||||
|
||||
def test_output_shape_and_dtype(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)])
|
||||
assert out.shape == base.shape
|
||||
assert out.dtype == np.uint8
|
||||
|
||||
def test_box_region_pulls_toward_restored(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)])
|
||||
# Center of the box should be near the restored (white) value.
|
||||
cy, cx = 50, 60
|
||||
assert out[cy, cx].mean() > 200
|
||||
|
||||
def test_far_from_box_stays_base(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)], pad=2)
|
||||
# Top-left corner is far from the box and feather, so it stays black.
|
||||
assert out[0, 0].mean() < 5
|
||||
|
||||
def test_empty_boxes_returns_base_unchanged(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [])
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_box_fully_outside_is_skipped(self):
|
||||
base, restored = self._base_and_restored(h=100, w=120)
|
||||
# Box entirely beyond the right/bottom edge -> clipped to empty -> no-op.
|
||||
out = face_restore._composite_faces(base, restored, [(200.0, 200.0, 260.0, 260.0)], pad=0)
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_near_edge_box_clips_without_error(self):
|
||||
base, restored = self._base_and_restored(h=100, w=120)
|
||||
# Box reaching past the bottom-right corner must clip, not raise.
|
||||
out = face_restore._composite_faces(base, restored, [(100.0, 80.0, 130.0, 110.0)], pad=10)
|
||||
assert out.shape == base.shape
|
||||
# The clipped in-bounds region still pulls toward white.
|
||||
assert out[95, 115].mean() > 100
|
||||
|
||||
|
||||
@pytest.mark.skipif(not face_restore.is_available(), reason="requires the 'restore' extra (gfpgan/facexlib)")
|
||||
class TestRestoreFacesModel:
|
||||
"""Model-running smoke test, gated behind the optional extra."""
|
||||
|
||||
def test_no_faces_returns_cleaned_unchanged(self):
|
||||
# A flat gray image has no faces; restore_faces must return the cleaned
|
||||
# input unchanged (the no-op path).
|
||||
cleaned = np.full((128, 128, 3), 127, dtype=np.uint8)
|
||||
original = np.full((128, 128, 3), 127, dtype=np.uint8)
|
||||
out = face_restore.restore_faces(original, cleaned)
|
||||
assert out.shape == cleaned.shape
|
||||
assert np.array_equal(out, cleaned)
|
||||
Reference in New Issue
Block a user