Coverage for visualize / robin / robin_visualizer.py: 96.55%

1from typing import Optional 

2import torch 

3import matplotlib.pyplot as plt 

4from matplotlib.axes import Axes 

5import numpy as np 

6from visualize.base import BaseVisualizer 

7from visualize.data_for_visualization import DataForVisualization 

8 

9class ROBINVisualizer(BaseVisualizer): 

10 """ROBIN watermark visualization class""" 

11 

12 def __init__(self, data_for_visualization: DataForVisualization, dpi: int = 300, watermarking_step: int = -1): 

13 super().__init__(data_for_visualization, dpi, watermarking_step) 

14 # ROBIN uses a specific watermarking step 

15 if hasattr(self.data, 'watermarking_step'): 

16 self.watermarking_step = self.data.watermarking_step 

17 else: 

18 raise ValueError("watermarking_step is required for ROBIN visualization") 

19 

20 def draw_pattern_fft(self, 

21 title: str = None, 

22 cmap: str = "viridis", 

23 use_color_bar: bool = True, 

24 vmin: Optional[float] = None, 

25 vmax: Optional[float] = None, 

26 ax: Optional[Axes] = None, 

27 **kwargs) -> Axes: 

28 """ 

29 Draw FFT visualization with original watermark pattern, with all 0 background. 

30  

31 Parameters: 

32 title (str): The title of the plot. If None, includes watermarking step info. 

33 cmap (str): The colormap to use. 

34 use_color_bar (bool): Whether to display the colorbar. 

35 ax (Axes): The axes to plot on. 

36  

37 Returns: 

38 Axes: The plotted axes. 

39 """ 

40 # Use custom title with watermarking step if not provided 

41 if title is None: 

42 title = f"ROBIN FFT with Watermark Area (Step {self.watermarking_step})" 

43 

44 orig_latent = self.data.optimized_watermark[0, self.data.w_channel].cpu() 

45 watermarking_mask = self.data.watermarking_mask[0, self.data.w_channel].cpu() 

46 

47 fft_data = torch.from_numpy(self._fft_transform(orig_latent)) 

48 fft_vis = torch.zeros_like(fft_data) 

49 fft_vis[watermarking_mask] = fft_data[watermarking_mask] 

50 

51 im = ax.imshow(np.abs(fft_vis.cpu().numpy()), cmap=cmap, vmin=vmin, vmax=vmax, **kwargs) 

52 if title != "": 

53 ax.set_title(title) 

54 if use_color_bar: 

55 ax.figure.colorbar(im, ax=ax) 

56 ax.axis('off') 

57 

58 return ax 

59 

60 def draw_inverted_pattern_fft(self, 

61 step: Optional[int] = None, 

62 title: str = None, 

63 cmap: str = "viridis", 

64 use_color_bar: bool = True, 

65 vmin: Optional[float] = None, 

66 vmax: Optional[float] = None, 

67 ax: Optional[Axes] = None, 

68 **kwargs) -> Axes: 

69 """ 

70 Draw FFT visualization with inverted pattern, with all 0 background. 

71  

72 Parameters: 

73 step (Optional[int]): The timestep of the inverted latents. If None, uses ROBIN's specific step. 

74 title (str): The title of the plot. If None, includes watermarking step info. 

75 cmap (str): The colormap to use. 

76 use_color_bar (bool): Whether to display the colorbar. 

77 ax (Axes): The axes to plot on. 

78  

79 Returns: 

80 Axes: The plotted axes. 

81 """ 

82 # For ROBIN, we need to use the specific watermarking step 

83 if step is None: 

84 # Calculate the actual step index for ROBIN 

85 # ROBIN uses: num_steps_to_use - 1 - self.config.watermarking_step 

86 num_steps = len(self.data.reversed_latents) 

87 actual_step = num_steps - 1 - self.watermarking_step 

88 inverted_latent = self.data.reversed_latents[actual_step][0, self.data.w_channel] 

89 else: 

90 inverted_latent = self.data.reversed_latents[step][0, self.data.w_channel] 

91 

92 # Use custom title with watermarking step if not provided 

93 if title is None: 

94 title = f"ROBIN FFT with Inverted Watermark Area (Step {self.watermarking_step})" 

95 

96 watermarking_mask = self.data.watermarking_mask[0, self.data.w_channel].cpu() 

97 

98 fft_data = torch.from_numpy(self._fft_transform(inverted_latent)) 

99 fft_vis = torch.zeros_like(fft_data).to(fft_data.device) 

100 fft_vis[watermarking_mask] = fft_data[watermarking_mask] 

101 

102 im = ax.imshow(np.abs(fft_vis.cpu().numpy()), cmap=cmap, vmin=vmin, vmax=vmax, **kwargs) 

103 if title != "": 

104 ax.set_title(title) 

105 if use_color_bar: 

106 ax.figure.colorbar(im, ax=ax) 

107 ax.axis('off') 

108 

109 return ax 

110 

111 def draw_optimized_watermark(self, 

112 title: str = None, 

113 cmap: str = "viridis", 

114 use_color_bar: bool = True, 

115 vmin: Optional[float] = None, 

116 vmax: Optional[float] = None, 

117 ax: Optional[Axes] = None, 

118 **kwargs) -> Axes: 

119 """ 

120 Draw the optimized watermark pattern (ROBIN-specific). 

121  

122 Parameters: 

123 title (str): The title of the plot. If None, includes watermarking step info. 

124 cmap (str): The colormap to use. 

125 use_color_bar (bool): Whether to display the colorbar. 

126 ax (Axes): The axes to plot on. 

127  

128 Returns: 

129 Axes: The plotted axes. 

130 """ 

131 # Use custom title with watermarking step if not provided 

132 if title is None: 

133 title = f"ROBIN Optimized Watermark (Step {self.watermarking_step})" 

134 

135 optimized_watermark = self.data.optimized_watermark[0, self.data.w_channel].cpu() 

136 

137 im = ax.imshow(np.abs(optimized_watermark.numpy()), cmap=cmap, vmin=vmin, vmax=vmax, **kwargs) 

138 if title != "": 

139 ax.set_title(title) 

140 if use_color_bar: 

141 ax.figure.colorbar(im, ax=ax) 

142 ax.axis('off') 

143 

144 return ax