-
-
Notifications
You must be signed in to change notification settings - Fork 683
/
example_str.py
166 lines (126 loc) · 7.75 KB
/
example_str.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# embed string
import numpy as np
from blind_watermark import WaterMark
from blind_watermark import att
from blind_watermark.recover import estimate_crop_parameters, recover_crop
import cv2
import os
os.chdir(os.path.dirname(__file__))
bwm = WaterMark(password_img=1, password_wm=1)
bwm.read_img('pic/ori_img.jpeg')
wm = '@guofei9987 开源万岁！'
bwm.read_wm(wm, mode='str')
bwm.embed('output/embedded.png')
len_wm = len(bwm.wm_bit) # 解水印需要用到长度
print('Put down the length of wm_bit {len_wm}'.format(len_wm=len_wm))
ori_img_shape = cv2.imread('pic/ori_img.jpeg').shape[:2] # 抗攻击有时需要知道原图的shape
h, w = ori_img_shape
# %% 解水印
bwm1 = WaterMark(password_img=1, password_wm=1)
wm_extract = bwm1.extract('output/embedded.png', wm_shape=len_wm, mode='str')
print("不攻击的提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %%截屏攻击1 = 裁剪攻击 + 缩放攻击 + 知道攻击参数（之后按照参数还原）
loc_r = ((0.1, 0.1), (0.5, 0.5))
scale = 0.7
x1, y1, x2, y2 = int(w * loc_r[0][0]), int(h * loc_r[0][1]), int(w * loc_r[1][0]), int(h * loc_r[1][1])
# 截屏攻击
att.cut_att3(input_filename='output/embedded.png', output_file_name='output/截屏攻击1.png',
loc=(x1, y1, x2, y2), scale=scale)
recover_crop(template_file='output/截屏攻击1.png', output_file_name='output/截屏攻击1_还原.png',
loc=(x1, y1, x2, y2), image_o_shape=ori_img_shape)
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/截屏攻击1_还原.png', wm_shape=len_wm, mode='str')
print("截屏攻击，知道攻击参数。提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %% 截屏攻击2 = 剪切攻击 + 缩放攻击 + 不知道攻击参数（因此需要 estimate_crop_parameters 来推测攻击参数）
loc_r = ((0.1, 0.1), (0.7, 0.6))
scale = 0.7
x1, y1, x2, y2 = int(w * loc_r[0][0]), int(h * loc_r[0][1]), int(w * loc_r[1][0]), int(h * loc_r[1][1])
print(f'Crop attack\'s real parameters: x1={x1},y1={y1},x2={x2},y2={y2}')
att.cut_att3(input_filename='output/embedded.png', output_file_name='output/截屏攻击2.png',
loc=(x1, y1, x2, y2), scale=scale)
# estimate crop attack parameters:
(x1, y1, x2, y2), image_o_shape, score, scale_infer = estimate_crop_parameters(original_file='output/embedded.png',
template_file='output/截屏攻击2.png',
scale=(0.5, 2), search_num=200)
print(f'Crop att estimate parameters: x1={x1},y1={y1},x2={x2},y2={y2}, scale_infer = {scale_infer}. score={score}')
# recover from attack:
recover_crop(template_file='output/截屏攻击2.png', output_file_name='output/截屏攻击2_还原.png',
loc=(x1, y1, x2, y2), image_o_shape=image_o_shape)
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/截屏攻击2_还原.png', wm_shape=len_wm, mode='str')
print("截屏攻击，不知道攻击参数。提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %%裁剪攻击1 = 裁剪 + 不做缩放 + 知道攻击参数
loc_r = ((0.1, 0.2), (0.5, 0.5))
x1, y1, x2, y2 = int(w * loc_r[0][0]), int(h * loc_r[0][1]), int(w * loc_r[1][0]), int(h * loc_r[1][1])
att.cut_att3(input_filename='output/embedded.png', output_file_name='output/随机裁剪攻击.png',
loc=(x1, y1, x2, y2), scale=None)
# recover from attack:
recover_crop(template_file='output/随机裁剪攻击.png', output_file_name='output/随机裁剪攻击_还原.png',
loc=(x1, y1, x2, y2), image_o_shape=image_o_shape)
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/随机裁剪攻击_还原.png', wm_shape=len_wm, mode='str')
print("裁剪攻击，知道攻击参数。提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %% 裁剪攻击2 = 裁剪 + 不做缩放 + 不知道攻击参数
loc_r = ((0.1, 0.1), (0.5, 0.4))
x1, y1, x2, y2 = int(w * loc_r[0][0]), int(h * loc_r[0][1]), int(w * loc_r[1][0]), int(h * loc_r[1][1])
att.cut_att3(input_filename='output/embedded.png', output_file_name='output/随机裁剪攻击2.png',
loc=(x1, y1, x2, y2), scale=None)
print(f'Cut attack\'s real parameters: x1={x1},y1={y1},x2={x2},y2={y2}')
# estimate crop attack parameters:
(x1, y1, x2, y2), image_o_shape, score, scale_infer = estimate_crop_parameters(original_file='output/embedded.png',
template_file='output/随机裁剪攻击2.png',
scale=(1, 1), search_num=None)
print(f'Cut attack\'s estimate parameters: x1={x1},y1={y1},x2={x2},y2={y2}. score={score}')
# recover from attack:
recover_crop(template_file='output/随机裁剪攻击2.png', output_file_name='output/随机裁剪攻击2_还原.png',
loc=(x1, y1, x2, y2), image_o_shape=image_o_shape)
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/随机裁剪攻击2_还原.png', wm_shape=len_wm, mode='str')
print("裁剪攻击，不知道攻击参数。提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %%椒盐攻击
ratio = 0.05
att.salt_pepper_att(input_filename='output/embedded.png', output_file_name='output/椒盐攻击.png', ratio=ratio)
# ratio是椒盐概率
# 提取
wm_extract = bwm1.extract('output/椒盐攻击.png', wm_shape=len_wm, mode='str')
print(f"椒盐攻击ratio={ratio}后的提取结果：", wm_extract)
assert np.all(wm == wm_extract), '提取水印和原水印不一致'
# %%旋转攻击
angle = 60
att.rot_att(input_filename='output/embedded.png', output_file_name='output/旋转攻击.png', angle=angle)
att.rot_att(input_filename='output/旋转攻击.png', output_file_name='output/旋转攻击_还原.png', angle=-angle)
# 提取水印
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/旋转攻击_还原.png', wm_shape=len_wm, mode='str')
print(f"旋转攻击angle={angle}后的提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %%遮挡攻击
n = 60
att.shelter_att(input_filename='output/embedded.png', output_file_name='output/多遮挡攻击.png', ratio=0.1, n=n)
# 提取
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/多遮挡攻击.png', wm_shape=len_wm, mode='str')
print(f"遮挡攻击{n}次后的提取结果：", wm_extract)
assert wm == wm_extract, '提取水印和原水印不一致'
# %%缩放攻击
att.resize_att(input_filename='output/embedded.png', output_file_name='output/缩放攻击.png', out_shape=(400, 300))
att.resize_att(input_filename='output/缩放攻击.png', output_file_name='output/缩放攻击_还原.png',
out_shape=ori_img_shape[::-1])
# out_shape 是分辨率，需要颠倒一下
bwm1 = WaterMark(password_wm=1, password_img=1)
wm_extract = bwm1.extract('output/缩放攻击_还原.png', wm_shape=len_wm, mode='str')
print("缩放攻击后的提取结果：", wm_extract)
assert np.all(wm == wm_extract), '提取水印和原水印不一致'
# %%亮度攻击
att.bright_att(input_filename='output/embedded.png', output_file_name='output/亮度攻击.png', ratio=0.9)
att.bright_att(input_filename='output/亮度攻击.png', output_file_name='output/亮度攻击_还原.png', ratio=1.1)
wm_extract = bwm1.extract('output/亮度攻击_还原.png', wm_shape=len_wm, mode='str')
print("亮度攻击后的提取结果：", wm_extract)
assert np.all(wm == wm_extract), '提取水印和原水印不一致'