def on_prepare(self):
# 载入数据
self.prepare_image_dataset()
#获取图像尺寸
image_size=self.config_loader.config["dataset"]["image_size"]["value"]
# 创建模型
# self.model = basic.make_Conv_models(self.config_loader.config["dataset"]["image_size"]["value"])
self.encoder = pixel.UNet([image_size[0],image_size[1],image_size[2]*2], high_performance_enable=self.config_loader.high_performance)
self.decoder = pixel.UNet(image_size, high_performance_enable=self.config_loader.high_performance)
# 设置优化器
optimizer = tf.keras.optimizers.Adam(1e-4)
# 注册模型与优化器
self.register_model_and_optimizer(
optimizer, {"encoder": self.encoder,"decoder":self.decoder}, "optimizer")
print("Initial model and optimizer....")
# 注册需要记录的损失名称
self.register_display_loss(["encoder loss","decoder loss","task loss"])
self.mask=np.ones([1,128,128,3])
self.mask[0:20,0:20]=0
self.mask=tf.cast(self.mask, tf.float32)
# 调用父类
super(HideEncoderContainer, self).on_prepare()
def on_prepare(self):
# 载入数据
self.prepare_image_dataset(
load_function=loader.load_single_crop_no_rescale)
#获取图像尺寸
image_size = self.config_loader.config["dataset"]["image_size"][
"value"]
# 创建模型
# self.model = basic.make_Conv_models(self.config_loader.config["dataset"]["image_size"]["value"])
self.encoder = pixel.UNet(
[image_size[0], image_size[1], image_size[2] * 2],
high_performance_enable=self.config_loader.high_performance)
self.decoder = pixel.UNet(
image_size,
high_performance_enable=self.config_loader.high_performance)
# 设置优化器
optimizer = tf.keras.optimizers.Adam(1e-4)
# 注册模型与优化器
self.register_model_and_optimizer(optimizer, {
"encoder": self.encoder,
"decoder": self.decoder
}, "optimizer")
print("Initial model and optimizer....")
# 注册需要记录的损失名称
self.register_display_loss(["encoder loss", "decoder loss"])
# 调用父类
super(StegoContainer, 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()