-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
138 lines (111 loc) · 5.54 KB
/
main.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
from time import time
from datetime import datetime
from config import *
from attacks import *
from tools import *
from detection_failedfouriertransform import extract_watermark, similarity, wpsnr
from embedment_failedfouriertransform import embed_watermark
from roc_failedfouriertransform import read_model
def evaluate_model(params, order_of_execution):
original_img, watermarked_image_name, watermark, attacks, watermarked_image, alpha, level, subband = params
model_name = '_'.join(watermarked_image_name.split('_')[1:])
print(model_name)
_, _, threshold, tpr, _ = read_model(model_name)
successful = 0
wpsnr_tot = 0
model_stats = {}
print("Threshold: ", threshold)
for attack in attacks:
attacked_img, _ = do_attacks(watermarked_image,attack)
extracted_watermark = None
extracted_watermark = extract_watermark(original_img, attacked_img, alpha, level, subband)
sim = similarity(watermark, extracted_watermark)
_wpsnr = wpsnr(original_img, attacked_img)
if sim < threshold and _wpsnr > 35:
successful += 1
wpsnr_tot += _wpsnr
print("The following attack (%s) was succesfull with a similarity of %f and a wpsnr of %f" % (describe_attacks(attack), sim, _wpsnr))
model_stats['Succesful Attacks'] = successful
model_stats['Unsuccesful Attacks'] = RUNS_PER_IMAGE - successful
model_stats['WPSNR Original-Watermarked'] = wpsnr(original_img, watermarked_image)
score = 0
if successful > 0:
score += attacked_wpsnr_to_mark(wpsnr_tot / successful)
else:
score += 6 + 2
score += wpsnr_to_mark(model_stats['WPSNR Original-Watermarked'])
model_stats['Score'] = score * tpr
return order_of_execution, model_name, model_stats
def multiproc_embed_watermark(params, order_of_execution):
original_img, img_name, watermark, alpha, level, subband = params
watermarked_img = embed_watermark(original_img, watermark, alpha, level, subband)
return order_of_execution, original_img, img_name, alpha, level, subband, watermarked_img
def main():
# Get N_IMAGES_LIMIT random images
images = import_images('images/'+TEAM_NAME+'/original/',N_IMAGES_LIMIT,True)
# Read watermark
watermark = np.load("failedfouriertransform.npy").reshape((MARK_SIZE, MARK_SIZE))
attacks = []
# Get list of attacks
for _ in range(0, RUNS_PER_IMAGE):
attacks.append(get_random_attacks(randint(1, MAX_N_ATTACKS)))
watermarked_images = {}
print('Welcome to multiprocessing city')
print('Embedding...')
work = []
for image in images:
original_img, img_name = image
for alpha in range(18,25):
for level in [2]:
for subband in [["LL"]]:
work.append((original_img, img_name, watermark, alpha, level, subband))
results = multiprocessed_workload(multiproc_embed_watermark,work)
for result in results:
original_img, img_name, alpha, level, subband, watermarked_img = result
params = '_'.join([str(alpha),str(level),'-'.join(subband)])
watermarked_images[img_name + '_' + params] = (original_img, img_name, watermarked_img, alpha, level, subband)
print("Let the Hunger Games begin!")
work = []
for watermarked_image_name in watermarked_images:
original_img, img_name, watermarked_image, alpha, level, subband = watermarked_images[watermarked_image_name]
work.append((original_img, watermarked_image_name, watermark, attacks, watermarked_image, alpha, level, subband))
results = multiprocessed_workload(evaluate_model,work)
all_models = {}
# Merge results from models on N_IMAGES_LIMIT different images
# Let's say we've run model x on two different images, we want to know the average WPSNR that image had with certain embedding method
# For example the average how many attacks the image with the watermark embedded a certain way did survive and so on
for model in results:
model_name, model = model
if model_name in all_models:
all_models[model_name]['AVG WPSNR Original-Watermarked'] += round((model['WPSNR Original-Watermarked'] / N_IMAGES_LIMIT),2)
all_models[model_name]['AVG Score'] += round((model['Score'] / N_IMAGES_LIMIT),2)
all_models[model_name]['AVG Succesful Attacks'] += round((model['Succesful Attacks'] / N_IMAGES_LIMIT),2)
all_models[model_name]['AVG Unsuccesful Attacks'] += round((model['Unsuccesful Attacks'] / N_IMAGES_LIMIT),2)
else:
new_model = {}
new_model['AVG WPSNR Original-Watermarked'] = round((model['WPSNR Original-Watermarked'] / N_IMAGES_LIMIT),2)
new_model['AVG Score'] = round((model['Score'] / N_IMAGES_LIMIT),2)
new_model['AVG Succesful Attacks'] = round((model['Succesful Attacks'] / N_IMAGES_LIMIT),2)
new_model['AVG Unsuccesful Attacks'] = round((model['Unsuccesful Attacks'] / N_IMAGES_LIMIT),2)
all_models[model_name] = new_model
print("Models sorted by score:")
lst_models = [(model_name, all_models[model_name]) for model_name in all_models]
lst_models.sort(key=lambda x: x[1]['AVG Score'], reverse=True)
for model_name, model in lst_models:
print(model_name.ljust(10), model)
best_model_name, _ = lst_models[0]
_, _, threshold, _, _ = read_model(best_model_name)
alpha, level, subband = best_model_name.split('_')
subband = subband.split('-')
update_parameters('detection_failedfouriertransform.py', ALPHA = alpha, DWT_LEVEL = level, SUBBANDS = subband, DETECTION_THRESHOLD = threshold, MARK_SIZE = MARK_SIZE)
for image in images:
original_img, img_name = image
print(img_name, best_model_name)
print(watermarked_images[img_name + '_' + best_model_name])
_, _, watermarked_img, _, _, _ = watermarked_images[img_name + '_' + best_model_name]
save_image(watermarked_img, img_name, "watermarked", TEAM_NAME)
if __name__ == '__main__':
st = time()
main()
et = time()
print(et-st)