Watermark Detection

# From audioseal/models.py:355
class AudioSealDetector(torch.nn.Module):
    def __init__(
        self,
        encoder: SEANetEncoderKeepDimension,
        normalizer: Optional[NormalizationProcessor] = None,
        nbits: int = 0,
    ):
        super().__init__()
        last_layer = torch.nn.Conv1d(encoder.output_dim, 2 + nbits, 1)
        self.detector = torch.nn.Sequential(encoder, last_layer)
        self.nbits = nbits

last_layer = torch.nn.Conv1d(encoder.output_dim, 2 + nbits, 1)

# From audioseal/models.py:444
if self.normalizer is not None and not torch.jit.is_scripting():
    x = self.normalizer.loudness_normalization(x)

result = self.detector(x)  # Shape: (batch, 2+nbits, frames)

# From audioseal/models.py:452
# Softmax on first 2 channels for detection
result[:, :2, :] = torch.softmax(result[:, :2, :], dim=1)

# From audioseal/models.py:421
@torch.jit.export
def decode_message(self, result: torch.Tensor) -> torch.Tensor:
    """
    Decode the message from the watermark result (batch x nbits x frames)
    Returns: The message of size batch x nbits (probability of 1 for each bit)
    """
    decoded_message = result.mean(dim=-1)  # Average across all frames
    return torch.sigmoid(decoded_message)  # Convert to [0, 1] probabilities

# From audioseal/models.py:390
@torch.jit.export
def detect_watermark(
    self,
    x: torch.Tensor,
    sample_rate: Optional[int] = None,
    message_threshold: float = 0.5,
    detection_threshold: float = 0.5,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Returns:
        detect_prob: Probability of audio being watermarked (scalar per batch)
        message: Binary message tensor (batch x nbits)
    """
    result, message = self.forward(x, sample_rate=sample_rate)
    
    # Count frames above threshold
    detect_prob = (
        torch.count_nonzero(
            torch.gt(result[:, 1, :], detection_threshold), dim=-1
        ) / result.shape[-1]
    )
    
    # Convert message probabilities to binary
    message = torch.gt(message, message_threshold).int()
    
    return detect_prob, message

detection_threshold: float = 0.5  # Default

message_threshold: float = 0.5  # Default

from audioseal import AudioSeal

# Load detector
detector = AudioSeal.load_detector("audioseal_detector_16bits")
detector.eval()

# Detect watermark (high-level API)
detect_prob, message = detector.detect_watermark(audio)

print(f"Detection probability: {detect_prob.item():.2%}")
if detect_prob > 0.5:
    print(f"Watermarked! Message: {message}")
else:
    print("No watermark detected")

# Get per-frame probabilities
result, message = detector(audio)

# result shape: (batch, 2, frames)
# Extract watermark probability for each frame
wm_prob_per_frame = result[:, 1, :]  # Shape: (batch, frames)

# Find watermarked regions
import torch
watermarked_frames = torch.where(wm_prob_per_frame > 0.5)[1]

print(f"Watermark detected in {len(watermarked_frames)} frames")
print(f"Total frames: {wm_prob_per_frame.shape[1]}")

# More sensitive detection
detect_prob, message = detector.detect_watermark(
    audio,
    detection_threshold=0.3,  # Lower threshold
    message_threshold=0.5
)

# More conservative detection
detect_prob, message = detector.detect_watermark(
    audio,
    detection_threshold=0.7,  # Higher threshold
    message_threshold=0.5
)

# Detect watermarks in potentially edited audio
result, message = detector(edited_audio)
wm_prob = result[:, 1, :]  # Per-frame probabilities

# Find contiguous watermarked segments
from scipy.ndimage import label
watermarked_binary = (wm_prob[0] > 0.5).cpu().numpy()
segments, num_segments = label(watermarked_binary)

print(f"Found {num_segments} watermarked segments")

# Assuming 16kHz sample rate, 1 frame ≈ 1 sample
for i in range(1, num_segments + 1):
    segment_frames = np.where(segments == i)[0]
    start_time = segment_frames[0] / 16000
    end_time = segment_frames[-1] / 16000
    print(f"Segment {i}: {start_time:.2f}s - {end_time:.2f}s")