Coverage for markdiffusion / detection / robin / robin_detection.py: 100.00%

1# Copyright 2025 THU-BPM MarkDiffusion. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16import torch 

17from markdiffusion.detection.base import BaseDetector 

18from scipy.stats import ncx2 

19from torch.nn import functional as F 

20 

21class ROBINDetector(BaseDetector): 

22 

23 def __init__(self, 

24 watermarking_mask: torch.Tensor, 

25 gt_patch: torch.Tensor, 

26 threshold: float, 

27 device: torch.device, 

28 threshold_p_value: float = 0.01, 

29 threshold_cosine_similarity: float = 0.5): 

30 super().__init__(threshold, device) 

31 self.watermarking_mask = watermarking_mask 

32 self.gt_patch = gt_patch 

33 self.threshold_p_value = threshold_p_value 

34 self.threshold_cosine_similarity = threshold_cosine_similarity 

35 

36 def eval_watermark(self, 

37 reversed_latents: torch.Tensor, 

38 reference_latents: torch.Tensor = None, 

39 detector_type: str = "l1_distance") -> float: 

40 reversed_latents_fft = torch.fft.fftshift(torch.fft.fft2(reversed_latents), dim=(-1, -2)) 

41 

42 # Resize mask and gt_patch if dimensions don't match 

43 if self.watermarking_mask.shape[-1] != reversed_latents.shape[-1]: 

44 target_size = reversed_latents.shape[-1] 

45 

46 # Resize mask (nearest neighbor for boolean mask) 

47 mask_float = self.watermarking_mask.float() 

48 mask_resized = F.interpolate(mask_float, size=(target_size, target_size), mode='nearest') 

49 current_mask = mask_resized.bool() 

50 

51 # Resize gt_patch (bilinear for continuous values) 

52 # gt_patch is complex, so we need to handle real and imag parts separately 

53 gt_real = self.gt_patch.real 

54 gt_imag = self.gt_patch.imag 

55 

56 gt_real_resized = F.interpolate(gt_real, size=(target_size, target_size), mode='bilinear', align_corners=False) 

57 gt_imag_resized = F.interpolate(gt_imag, size=(target_size, target_size), mode='bilinear', align_corners=False) 

58 

59 current_gt_patch = torch.complex(gt_real_resized, gt_imag_resized) 

60 else: 

61 current_mask = self.watermarking_mask 

62 current_gt_patch = self.gt_patch 

63 

64 if detector_type == 'l1_distance': 

65 target_patch = current_gt_patch 

66 l1_distance = torch.abs(reversed_latents_fft[current_mask] - target_patch[current_mask]).mean().item() 

67 return { 

68 'is_watermarked': bool(l1_distance < self.threshold), 

69 'l1_distance': l1_distance 

70 } 

71 elif detector_type == 'p_value': 

72 reversed_latents_fft_wm_area = reversed_latents_fft[current_mask].flatten() 

73 target_patch = current_gt_patch[current_mask].flatten() 

74 target_patch = torch.concatenate([target_patch.real, target_patch.imag]) 

75 reversed_latents_fft_wm_area = torch.concatenate([reversed_latents_fft_wm_area.real, reversed_latents_fft_wm_area.imag]) 

76 sigma_ = reversed_latents_fft_wm_area.std() 

77 lambda_ = (target_patch ** 2 / sigma_ ** 2).sum().item() 

78 x = (((reversed_latents_fft_wm_area - target_patch) / sigma_) ** 2).sum().item() 

79 p = ncx2.cdf(x=x, df=len(target_patch), nc=lambda_) 

80 return { 

81 'is_watermarked': bool(p < self.threshold_p_value), 

82 'p_value': p 

83 } 

84 elif detector_type == 'cosine_similarity': 

85 reversed_latents_fft_wm_area = reversed_latents_fft[current_mask].flatten() 

86 target_patch = current_gt_patch[current_mask].flatten() 

87 cosine_similarity = F.cosine_similarity(reversed_latents_fft_wm_area.real, target_patch.real, dim=0) 

88 return { 

89 'is_watermarked': bool(cosine_similarity > self.threshold_cosine_similarity), 

90 'cosine_similarity': cosine_similarity 

91 } 

92 else: 

93 raise ValueError(f"Tree Ring's watermark detector type {self.detector_type} not supported") 

94