def on_prepare(self):
"""
准备阶段，完成如下事件：
1. 加载vgg16风格损失网络
2. 加载风格损失图片
"""
# 获取输入参数
args_parser = ArgsParser(self.config_loader.callback_args)
# 获取风格图像位置
style_path = args_parser.get("style_image")
# 加载风格图像
self.style_image = train_tool.read_image(
style_path,
self.config_loader.image_width,
self.config_loader.image_height,
change_scale=False)
print("load style image {}....".format(style_path))
self.style_image = self.style_image / 255
print(self.style_image)
# 加载vgg16网络，若网络不存在，会直接经行下载
# 该网络使用image-net训练，不加载全连接层
vgg_16 = tf.keras.applications.vgg16.VGG16(weights='imagenet',
include_top=False)
print("load VGG 16....")
# 激活层id
activate_id = [3, 6, 10, 14, 18]
self.activate_list = []
# 生成各个激活层
for id in activate_id:
model = tf.keras.Model(inputs=vgg_16.input,
outputs=vgg_16.layers[id].output)
#需要将模型设置成不可训练，（若使用keras编译，则必须设置，
# 但使用tensorflow自定义训练，个人认为只要不把trainable_variable传递过去就不会进行训练（尚未验证））
model.trainable = False
self.activate_list.append(model)
print("VGG16 activation layers done....")
#由于风格图像给定，每一层特征就给定，可以提前计算
#直接计算gram举证
self.style_feature_list = []
for activate in self.activate_list:
self.style_feature_list.append(gram(activate(self.style_image)))
print("Gram feature done....")
# 执行父类方法
super(StyleContainer, self).on_prepare()
def on_prepare(self):
# 载入数据集
self.prepare_image_dataset()
# 创建模型
self.model = pixel.UNet(
input_shape=[128, 128, 6],
high_performance_enable=self.config_loader.high_performance)
self.hide_model = ExtractInvisible([128, 128, 6])
self.hide_model.summary()
self.reveal_model = pixel.UNet(
input_shape=[128, 128, 6],
high_performance_enable=self.config_loader.high_performance)
# 设置优化器
optimizer = tf.keras.optimizers.Adam(1e-4)
# 注册模型与优化器
self.register_model_and_optimizer(
optimizer, {
"model": self.model,
"hide_model": self.hide_model,
"reveal_model": self.reveal_model
}, "optimizer")
# self.register_model_and_optimizer(
# optimizer, {"model": self.model}, "optimizer")
print("Initial model and optimizer....")
#风格
style_image = train_tool.read_image("./style.jpg",
128,
128,
change_scale=False)
style_image = tf.cast(style_image, tf.float32)
style_image = style_image / 127.5 - 1
style_batch = np.zeros([self.config_loader.batch_size, 128, 128, 3])
for i in range(self.config_loader.batch_size):
style_batch[i, :, :, :] = style_image[0]
self.style_batch = style_batch
# 注册需要记录的损失名称
self.register_display_loss(["loss", "s", "h1", "h2", "r"])
# 调用父类
super(HideContainer, self).on_prepare()
def on_prepare(self):
# 准备数据集
self.on_prepare_dataset()
# 创建生成网络
self.generator = UNet(
input_shape=self.config_loader.config["dataset"]["image_size"]
["value"],
high_performance_enable=self.config_loader.high_performance)
print("Initial generator....")
self.extractor = ExtractInvisible()
print("Initial extractor....")
# 创建生成优化器
optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
print("Initial optimizer....")
self.register_model_and_optimizer(optimizer, {
"generator": self.generator,
"extractor": self.extractor
}, "opt")
print("register generator and optimizer....")
self.register_display_loss(["total_loss", "gen_loss", "wm_loss"])
self.wm = train_tool.read_image(
"C:\\Users\\happy\\Desktop\\Evil\\Eidolon-Tensorflow\\wm.png",
128,
128,
change_scale=True)
self.noise = np.zeros([1, 128, 128, 3])
super(WMContainer, self).on_prepare()
def eval_key_sensitive(key_path,
input_path,
target_path,
model_path,
visual_result_path=None,
key_bits=1024):
"""
评估密钥敏感性
"""
# load model
model = model_use.EncoderDecoder(model_path, key_enable=True)
print("initial model....")
# load key in numpy
secret_key = np.load(key_path)
# 把他变成长条是为了方便修改
secret_key = np.reshape(secret_key, [1024, 1])
print("load model....")
# input_image = train_tool.read_image(
# in, 32, 32, change_scale=True)
# print("load input....")
target_image = train_tool.read_image(target_path,
32,
32,
change_scale=True)
print("load target....")
result_list = []
# 修改秘钥， 测试改变不同位数下的敏感性
for i in range(0, key_bits + 1):
# copy key
temp_key = secret_key.copy()
# 修改位数，从头开始修改
for j in range(0, i):
temp_key[j, 0] = -temp_key[j, 0]
temp_key = np.reshape(temp_key, [1, 32, 32, 1])
# save temp key
np.save("~temp_key.npy", temp_key)
# decode
decode_image = model.decode_from_image(encode_image_path=input_path,
secret_key_path="~temp_key.npy")
# eval
decode_image = decode_image * 2 - 1
decode_image = tf.reshape(decode_image, [1, 32, 32, 3])
psnr = eval_util.evaluate(decode_image,
target_image,
psnr_enable=True,
ssim_enable=False,
ber_enable=False)
result_list.append(psnr["psnr"])
print("\r" + "eval {}..., psnr={}".format(i, psnr["psnr"]),
end="",
flush=True)
print()
# delete temp key file
os.remove("~temp_key.npy")
return result_list
def eval_all(data_path,
model_path,
visual_result_dir=None,
watermark_path=None,
wm_width=64,
wm_height=64,
watermark_binary=False,
attack_test_func=None):
"""
评估函数, 评估模型的PSNR与SSIM, 若存在水印，则评估水印的PSNR或BER
@since 2019.11.27
@author anomymity
:param model_path: 模型路径，包括生成网络和提取网络，名称默认为generator.h5和extractor.h5
"""
# create empty config, this is used when loading data
configs = config.ConfigLoader()
# load data
data_loader = loader.ImageLoader(data_path, is_training=False)
dataset = data_loader.load(configs)
print("load data....")
# load wateramrk
watermark_enable = False
wm_target = None
if watermark_path != None:
watermark_enable = True
"""
@author anomymity
@update 2019.11.28
change_scale没有指定，因此读取的水印为[0,1], 与要求[-1,1]不符，使得后续计算误码率有误。现已修复。
"""
wm_target = train_tool.read_image(watermark_path,
wm_width,
wm_height,
binary=watermark_binary,
change_scale=True)
print("load watermark....")
generator_path = os.path.join(model_path, "generator.h5")
extractor_path = None
if watermark_enable == True:
extractor_path = os.path.join(model_path, "extractor.h5")
# load model
print("load model....")
model = model_use.GeneratorModel(generator_path=generator_path,
wm_width=wm_width,
wm_height=wm_height,
watermark_enable=watermark_enable,
extractor_path=extractor_path,
binary=watermark_binary)
if visual_result_dir != None and os.path.exists(
visual_result_dir) == False:
os.mkdir(visual_result_dir)
# start eval
print("start eval....")
"""
The result set is a key-value dir,
which will save 1. mean_value, 2. value_list(the mean value is calculated bu this)
"""
result_set = {}
# the value list will save all the results by the image
value_list = []
# default value
image_mean_psnr = 0
image_mean_ssim = 0
wm_mean_error = 0
# image num
count = 0
# for each
for input_image, ground_truth in dataset:
# result each
result_each = {}
# generate
output_tensor, wm_tensor, wm_feature = model.generate(
input_image, attack_test_func=attack_test_func)
# eval image
image_result_each = eval_util.evaluate(output_tensor,
ground_truth,
psnr_enable=True,
ssim_enable=True,
ber_enable=False)
# save results
result_each["image_psnr"] = image_result_each["psnr"]
result_each["image_ssim"] = image_result_each["ssim"]
# caluclate total value
image_mean_psnr = image_mean_psnr + result_each["image_psnr"]
image_mean_ssim = image_mean_ssim + result_each["image_ssim"]
if watermark_path != None:
# eval watermark
wm_result_each = eval_util.evaluate(
wm_tensor,
wm_target,
psnr_enable=(not watermark_binary),
ssim_enable=False,
ber_enable=watermark_binary)
# calcualte watermark error
if watermark_binary == True:
result_each["wm_ber"] = wm_result_each["ber"]
wm_mean_error = wm_mean_error + result_each["wm_ber"]
else:
result_each["wm_psnr"] = wm_result_each["psnr"]
wm_mean_error = wm_mean_error + result_each["wm_psnr"]
# append
value_list.append(result_each)
# save visual results
if visual_result_dir != None:
# basic
image_list = [input_image, ground_truth, output_tensor]
title_list = ["IN", "GT", "PR"]
# wm
if watermark_path != None:
image_list.append(wm_tensor)
image_list.append(wm_feature)
title_list.append("WM")
title_list.append("WF")
# save image
train_tool.save_images(image_list, title_list, visual_result_dir,
count + 1)
# one image test finished
count = count + 1
# this print will flush at the same place
print("\r" + "testing image {} ...".format(count), end='', flush=True)
# change line now
print("")
# calculate image mean value
mean_value = {}
image_mean_psnr = image_mean_psnr / count
mean_value["psnr"] = image_mean_psnr
image_mean_ssim = image_mean_ssim / count
mean_value["ssim"] = image_mean_ssim
# mean watermark error
if watermark_path != None:
wm_mean_error = wm_mean_error / count
if watermark_binary == True:
mean_value["wm_ber"] = wm_mean_error
else:
mean_value["wm_psnr"] = wm_mean_error
# save all
result_set["mean_value"] = mean_value
result_set["value_list"] = value_list
# genrate report
eval_report = "The evaluating report comes here:\n'image psnr' = {}, 'image ssim' = {}".format(
image_mean_psnr, image_mean_ssim)
if watermark_path != None:
eval_report_wm = ", 'watermark {}' = {}"
if watermark_binary == True:
wm_format = "ber"
else:
wm_format = "psnr"
eval_report_wm = eval_report_wm.format("wm_" + wm_format,
mean_value["wm_" + wm_format])
eval_report = eval_report + eval_report_wm
# print(eval_report)
return result_set, eval_report
def on_prepare(self):
"""
准备阶段，完成以下事宜：
1. 创建水印提取网络
2. 加载解码器
3. 调用父类的方法，该方法完成以下事宜：
1. 加载数据集
2. 创建网络与优化器
3. 将网络与优化器注册到父类中，以便自动保存
4. 调用父类on_prepare
"""
"""
This will be invokde in initial process
"""
print("This container is for training the invisible watermark.")
# lambdas in loss
self.lambda_wm_positive = self.config_loader.lambda_array[0]
self.lambda_wm_negitive = self.config_loader.lambda_array[1]
# lambdas print
print("lp={}, ln={}".format(
self.lambda_wm_positive, self.lambda_wm_negitive))
# 解析所有配置
args_parser = ArgsParser(self.config_loader.callback_args)
# 解码器
self.decoder_path = args_parser.get("decoder")
# 水印
self.wm_path = args_parser.get("wm_path")
# 水印大小
self.wm_width = int(args_parser.get("wm_width"))
# training noise attack
self.noise_attack = args_parser.get("noise") in ("True", "true")
print("noise attack:{}".format(self.noise_attack))
# 训练随机裁剪增加水印裁剪攻击能力
self.crop_attack = args_parser.get("crop") in ("True", "true")
print("crop attack:{}".format(self.crop_attack))
# init extract network
print("initial WM extract network")
self.extractor = Extractor(
self.config_loader.config["dataset"]["image_size"]["value"])
# 注册
self.model_map["extractor"] = self.extractor
# load watermark
# 保留老版本的编码方法
if self.decoder_path != "new":
self.watermark_target = train_tool.read_image(
self.wm_path, self.config_loader.image_width, self.config_loader.image_height, change_scale=True)
else:
self.watermark_target = encode_watermark_from_image(
self.wm_path, self.config_loader.image_width, self.config_loader.image_height)
print("load watermark successfully...")
# create negitive. if no watermark, a 1 matrix will be out
self.negitive_target = tf.zeros(
shape=[1, self.config_loader.image_width, self.config_loader.image_height, 3])*1
print("create negative watermark successfully...")
# the pretrained encoder-decoder model
# @update 2019.11.27
# @author anomymity
# 修复相对路径bug,否则无法载入模型
# 新版本不需要解码器预训练
if self.decoder_path != "new":
self.decoder_model = EncoderDecoder(self.decoder_path)
print("load decoder successfully...")
# 调用父类
super(WMContainer, self).on_prepare()