mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-05-20 20:04:40 +02:00
test-user
e5d8970add
- CLI with visible, invisible, all, metadata, and batch commands - Gemini watermark removal via reverse alpha blending - Invisible watermark removal via diffusion regeneration (SynthID, TreeRing) - AI metadata stripping (EXIF, PNG text, C2PA) - Face protection (YOLO/Haar) and analog humanizer - 137 tests covering all CLI modes and core engines - Ruff and Pyright clean
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
import numpy as np
|
|
|
|
from remove_ai_watermarks.face_protector import FaceProtector
|
|
|
|
|
|
def test_face_protector_initialization():
|
|
# Will fallback to Haar cascade if ultralytics is missing
|
|
fp = FaceProtector(use_yolo=False)
|
|
assert fp.use_yolo is False
|
|
assert fp.haar_cascade is not None
|
|
|
|
|
|
def test_face_protector_lifecycle():
|
|
fp = FaceProtector(use_yolo=False)
|
|
|
|
# Create dummy black image
|
|
img = np.zeros((200, 200, 3), dtype=np.uint8)
|
|
|
|
# Since it's a black image, haar cascade should find 0 faces
|
|
faces = fp.extract_faces(img)
|
|
assert isinstance(faces, list)
|
|
assert len(faces) == 0
|
|
|
|
# Restoring 0 faces should result in strictly equal image
|
|
restored = fp.restore_faces(img, faces)
|
|
assert np.array_equal(img, restored)
|
|
|
|
|
|
def test_face_protector_restore_bypass_on_size_mismatch():
|
|
fp = FaceProtector(use_yolo=False)
|
|
img_small = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
|
|
# Manually mock a face that is OUT OF BOUNDS for img_small
|
|
mock_bbox = (80, 80, 130, 130)
|
|
mock_crop = np.ones((50, 50, 3), dtype=np.uint8) * 255
|
|
mock_faces = [(mock_bbox, mock_crop)]
|
|
|
|
# Attempt to restore onto an image too small for this box
|
|
restored = fp.restore_faces(img_small, mock_faces)
|
|
|
|
# Should safely skip restoring and not crash
|
|
assert np.array_equal(restored, img_small)
|
|
|
|
|
|
def test_face_protector_restore_blending():
|
|
fp = FaceProtector(use_yolo=False)
|
|
# Background is black
|
|
img_target = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
|
|
# Face crop is white
|
|
mock_bbox = (25, 25, 75, 75)
|
|
mock_crop = np.ones((50, 50, 3), dtype=np.uint8) * 255
|
|
mock_faces = [(mock_bbox, mock_crop)]
|
|
|
|
restored = fp.restore_faces(img_target, mock_faces)
|
|
|
|
# The center of the face should be perfectly white (255)
|
|
assert restored[50, 50, 0] >= 254
|
|
# The corner of the target should remain perfectly black (0)
|
|
assert restored[0, 0, 0] == 0
|
|
# We should have a blending gradient between them due to the gaussian blur mask
|
|
# For example, around (30, 30) or similar
|
|
assert 0 <= restored[28, 28, 0] <= 255
|