mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 15:36:01 +02:00
113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
"""Regression tests for motion-compensated visible-video fill."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from remove_ai_watermarks.video_temporal import stabilize_filled_frame
|
|
|
|
|
|
def _translated_pair() -> tuple[np.ndarray, np.ndarray]:
|
|
rng = np.random.default_rng(7)
|
|
previous = rng.integers(0, 256, (96, 128, 3), dtype=np.uint8)
|
|
previous = cv2.GaussianBlur(previous, (5, 5), 0)
|
|
current = cv2.warpAffine(
|
|
previous,
|
|
np.float32(((1, 0, 2), (0, 1, 1))),
|
|
(128, 96),
|
|
borderMode=cv2.BORDER_REFLECT,
|
|
)
|
|
return previous, current
|
|
|
|
|
|
def test_motion_aligned_prior_reduces_independent_fill_error() -> None:
|
|
previous, current = _translated_pair()
|
|
mask = np.zeros(current.shape[:2], dtype=np.uint8)
|
|
mask[30:66, 45:85] = 255
|
|
rng = np.random.default_rng(11)
|
|
current_fill = current.copy()
|
|
noise = rng.normal(0, 18, (36, 40, 3))
|
|
current_fill[30:66, 45:85] = np.clip(
|
|
current_fill[30:66, 45:85].astype(np.float32) + noise,
|
|
0,
|
|
255,
|
|
).astype(np.uint8)
|
|
|
|
stabilized = stabilize_filled_frame(
|
|
previous,
|
|
previous,
|
|
mask,
|
|
current,
|
|
current_fill,
|
|
mask,
|
|
)
|
|
|
|
hole = mask > 0
|
|
before = float(np.mean((current_fill[hole].astype(np.float32) - current[hole]) ** 2))
|
|
after = float(np.mean((stabilized[hole].astype(np.float32) - current[hole]) ** 2))
|
|
assert after < before * 0.5
|
|
assert np.array_equal(stabilized[~hole], current_fill[~hole])
|
|
|
|
|
|
def test_owned_fill_can_be_stabilized_without_full_frame_copy() -> None:
|
|
previous, current = _translated_pair()
|
|
mask = np.zeros(current.shape[:2], dtype=np.uint8)
|
|
mask[30:66, 45:85] = 255
|
|
current_fill = current.copy()
|
|
current_fill[mask > 0] = 127
|
|
|
|
stabilized = stabilize_filled_frame(
|
|
previous,
|
|
previous,
|
|
mask,
|
|
current,
|
|
current_fill,
|
|
mask,
|
|
copy=False,
|
|
)
|
|
|
|
assert stabilized is current_fill
|
|
|
|
|
|
def test_scene_cut_keeps_independent_current_fill() -> None:
|
|
previous, _current = _translated_pair()
|
|
rng = np.random.default_rng(13)
|
|
current = rng.integers(0, 256, previous.shape, dtype=np.uint8)
|
|
mask = np.zeros(current.shape[:2], dtype=np.uint8)
|
|
mask[30:66, 45:85] = 255
|
|
current_fill = current.copy()
|
|
current_fill[mask > 0] = 0
|
|
|
|
stabilized = stabilize_filled_frame(
|
|
previous,
|
|
previous,
|
|
mask,
|
|
current,
|
|
current_fill,
|
|
mask,
|
|
)
|
|
|
|
assert np.array_equal(stabilized, current_fill)
|
|
|
|
|
|
def test_disjoint_prior_mask_cannot_reintroduce_old_mark() -> None:
|
|
previous, current = _translated_pair()
|
|
previous_mask = np.zeros(current.shape[:2], dtype=np.uint8)
|
|
previous_mask[8:24, 8:24] = 255
|
|
current_mask = np.zeros(current.shape[:2], dtype=np.uint8)
|
|
current_mask[60:76, 96:112] = 255
|
|
current_fill = current.copy()
|
|
current_fill[current_mask > 0] = 127
|
|
|
|
stabilized = stabilize_filled_frame(
|
|
previous,
|
|
previous,
|
|
previous_mask,
|
|
current,
|
|
current_fill,
|
|
current_mask,
|
|
)
|
|
|
|
assert np.array_equal(stabilized, current_fill)
|