wehub-resource-sync
31 lines
957 B
Python
31 lines
957 B
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
|
|
import torch
|
|
from audiocraft.models.watermark import AudioSeal
|
|
from tests.common_utils.wav_utils import get_white_noise
|
|
|
|
|
|
class TestWatermarkModel:
|
|
|
|
def test_base(self):
|
|
sr = 16_000
|
|
duration = 1.0
|
|
wav = get_white_noise(1, int(sr * duration)).unsqueeze(0)
|
|
wm = AudioSeal.get_pretrained(name="base")
|
|
|
|
secret_message = torch.randint(0, 2, (1, 16), dtype=torch.int32)
|
|
watermarked_wav = wm(wav, message=secret_message, sample_rate=sr, alpha=0.8)
|
|
result = wm.detect_watermark(watermarked_wav)
|
|
|
|
detected = (
|
|
torch.count_nonzero(torch.gt(result[:, 1, :], 0.5)) / result.shape[-1]
|
|
)
|
|
detect_prob = detected.cpu().item() # type: ignore
|
|
|
|
assert detect_prob >= 0.0
|