def test_test_model_single_call():
from vgg import vgg19, vgg19_rev
import os.path as osp
import tensorlayer as tl
VGG19_PARTIAL_WEIGHTS_PATH = 'pretrained_models/predefined_vgg19_endwith(conv4_1)_weights.h5'
DEC_BEST_WEIGHTS_PATH = 'pretrained_models/dec_best_weights.h5'
CONTENT_DATA_PATH = './test_images/content'
STYLE_DATA_PATH = './test_images/style'
test_content_filenames = ['brad_pitt_01.jpg']
test_style_filenames = ['cat.jpg']
TEST_INPUT_CONSTRAINTED_SIZE = 800
TEST_OUTPUT_PATH = './test_images/output'
tl.logging.set_verbosity(tl.logging.DEBUG)
enc_net = vgg19(pretrained=False, end_with='conv4_1')
enc_net.load_weights(VGG19_PARTIAL_WEIGHTS_PATH, in_order=False)
dec_net = vgg19_rev(pretrained=False, batch_norm=False, input_depth=512)
dec_net.load_weights(DEC_BEST_WEIGHTS_PATH, skip=True)
i = 0 # only test 1 pair of input
test_content = utils.imread(
osp.join(CONTENT_DATA_PATH, test_content_filenames[i]))
test_style = utils.imread(
osp.join(STYLE_DATA_PATH, test_style_filenames[i]))
# import cv2
# test_content = cv2.cvtColor(test_content, cv2.COLOR_BGR2RGB) # <- moved to utils.imread
# test_style = cv2.cvtColor(test_style, cv2.COLOR_BGR2RGB) # <- moved to utils.imread
content_features = enc_net(test_content, is_train=False)
style_features = enc_net(test_style, is_train=False)
target_features = utils.AdaIN(content_features, style_features, alpha=1)
del content_features, style_features
generated = dec_net(target_features, is_train=False)
import tensorflow as tf
if isinstance(generated, tf.Tensor):
if generated.dtype == tf.float32:
generated = tf.cast(generated, tf.uint8)
generated = generated[0].numpy()
saved_path = f"{osp.splitext(test_style_filenames[i])[0]}+{osp.splitext(test_content_filenames[i])[0]}"
saved_path = osp.join(TEST_OUTPUT_PATH, f"{saved_path}.jpg")
# generated = cv2.cvtColor(generated, cv2.COLOR_RGB2BGR) # <- moved to utils.imsave
utils.imsave(saved_path, generated)
tl.logging.info(f"saved_path = {saved_path}")
tl.logging.info(f"generated.shape = {generated.shape}")
def test_test_arbitrary_sized_inputs():
from vgg import vgg19, vgg19_rev
import os.path as osp
import tensorlayer as tl
DEC_LATEST_WEIGHTS_PATH = 'pretrained_models/dec_latest_weights.h5'
STYLE_LAYERS = ('conv1_1', 'conv2_1', 'conv3_1', 'conv4_1') # for Encoders
CONTENT_DATA_PATH = './dataset/content_samples' # COCO_train_2014/'
STYLE_DATA_PATH = './dataset/style_samples' # wiki_all_images/'
test_content_filenames = ['000000532397.jpg'
] #, '000000048289.jpg', '000000526781.jpg']
test_style_filenames = ['53154.jpg'] #, '2821.jpg', '216.jpg']
TEST_INPUT_CONSTRAINTED_SIZE = 800
TEMP_IMAGE_PATH = './temp_images/'
tl.logging.set_verbosity(tl.logging.DEBUG)
enc_net = vgg19(pretrained=True, end_with='conv4_1')
# NOTE: batch_norm=True will lower quality of the generated image = need retrain
dec_net = vgg19_rev(pretrained=False, batch_norm=False, input_depth=512)
if osp.exists(DEC_LATEST_WEIGHTS_PATH):
dec_net.load_weights(DEC_LATEST_WEIGHTS_PATH, skip=True)
enc_net.eval()
dec_net.eval()
for epoch in range(1): # for test generator validity
# Note: generator need reset for reuse
test_inputs_gen = utils.single_inputs_generator(
list(zip(test_content_filenames, test_style_filenames)),
CONTENT_DATA_PATH, STYLE_DATA_PATH, TEST_INPUT_CONSTRAINTED_SIZE)
for i, (test_content, test_style) in enumerate(test_inputs_gen):
# shape=[1, w, h, c], so as to feed arbitrary sized test images one by one
content_features = enc_net(test_content)
style_features = enc_net(test_style, )
target_features = utils.AdaIN(content_features,
style_features,
alpha=1)
del content_features, style_features
generated_images = dec_net(target_features)
paired_name = f"{osp.splitext(test_style_filenames[i])[0]}+{osp.splitext(test_content_filenames[i])[0]}"
utils.imsave(
osp.join(TEMP_IMAGE_PATH,
f"temp_{paired_name}_epoch{epoch}.jpg"),
generated_images[0].numpy())
def forward(self, inputs, training=None, alpha=1):
"""
:param inputs: [content_batch, style_batch], both have shape as [batch_size, w, h, c]
:param training:
:param alpha:
:return:
"""
# TL1to2: preprocessing and reverse -> vgg forward() will handle it
# # switch RGB to BGR
# content = tf.reverse(content_input, axis=[-1])
# style = tf.reverse(style_input, axis=[-1])
# # preprocess image
# content = Encoder.preprocess(content_input)
# style = Encoder.preprocess(style_input)
content, style = inputs
# encode image
# we should initial global variables before restore model
content_features = self.enc_net(content)
style_features = self.enc_net(style)
# pass the encoded images to AdaIN # IMPROVE: try alpha gradients
target_features = utils.AdaIN(content_features,
style_features,
alpha=alpha)
# decode target features back to image
generated_img = self.dec_net(target_features)
# # deprocess image
# generated_img = Encoder.reverse_preprocess(generated_img)
# # switch BGR back to RGB
# generated_img = tf.reverse(generated_img, axis=[-1])
# # clip to 0..255
# generated_img = tf.clip_by_value(generated_img, 0.0, 255.0)
return generated_img
name='content_input')
style_input = tf.placeholder(tf.float32,
shape=INPUT_SHAPE,
name='style_input')
# switch RGB to BGR
content = tf.reverse(content_input, axis=[-1])
style = tf.reverse(style_input, axis=[-1])
# preprocess image
content = encoder.preprocess(content)
style = encoder.preprocess(style)
content_enc_net = encoder.encode(content, 'content/')
style_enc_net = encoder.encode(style, 'style/')
adain_features = utils.AdaIN(content_enc_net.outputs,
style_enc_net.outputs)
stylied_dec_net = decoder.decode(adain_features, 'stylized_dec/')
# add the mean values back
stylied_image = encoder.deprocess(stylied_dec_net.outputs)
# switch BGR back to RGB
stylied_image = tf.reverse(stylied_image, axis=[-1])
# clip to 0..255
stylied_image = tf.clip_by_value(stylied_image, 0.0, 255.0)
# switch RGB to BGR
stylied_image = tf.reverse(stylied_image, axis=[-1])
def forward(self, inputs: list, training=None, alpha=1):
"""
:param inputs: [content_batch, style_batch], both have shape as [batch_size, w, h, c]
:param training:
:param alpha:
:return:
"""
# TL1to2: preprocessing and reverse -> vgg forward() will handle it
# # switch RGB to BGR
# content = tf.reverse(content_input, axis=[-1])
# style = tf.reverse(style_input, axis=[-1])
# # preprocess image
# content = Encoder.preprocess(content_input)
# style = Encoder.preprocess(style_input)
content, style = inputs
# 1.encode image: get content features and style features (i.e. intermediate style features)
c_content_features = self.enc_net(content)
s_content_features, s_style_feats_in_layers = self.enc_net(style, observed_layer_names=STYLE_LAYERS)
# 2.pass the encoded content and style to AdaIN
target_features = utils.AdaIN(c_content_features, s_content_features, alpha=alpha)
# 3.decode target features back to generate an image
generated_images = self.dec_net(target_features)
# # de-preprocess image
# generated_images = Encoder.reverse_preprocess(generated_images)
# # switch BGR back to RGB
# generated_images = tf.reverse(generated_images, axis=[-1])
# # clip to 0..255
# generated_images = tf.clip_by_value(generated_images, 0.0, 255.0)
# 4.compute content features and style features of the generated image
g_content_features, g_style_feats_in_layers = self.enc_net(generated_images, observed_layer_names=STYLE_LAYERS)
tl.logging.info(
f"c_c_feat:{c_content_features.shape}, s_c_feat:{s_content_features.shape}, "
f"t_feat:{target_features.shape}, g:{generated_images.shape}, g_c_feat:{g_content_features.shape}")
# 5.compute losses
self.content_loss = tf.reduce_sum(
tf.reduce_mean(tf.square(g_content_features - target_features), axis=[1, 2]))
style_layer_loss = []
for idx, layer_name in enumerate(STYLE_LAYERS):
# TL1to2: tl.layers.get_layers_with_name -> observe intermediate outputs through model.__call__
# s_style_feat = tl.layers.get_layers_with_name(self.enc_s_net, 'style/' + layer, True)[0]
# g_style_feat = tl.layers.get_layers_with_name(self.enc_net, 'stylized_enc/' + layer, True)[0]
s_style_feat = s_style_feats_in_layers[idx]
g_style_feat = g_style_feats_in_layers[idx]
mean_s, var_s = tf.nn.moments(s_style_feat, [1, 2])
mean_g, var_g = tf.nn.moments(g_style_feat, [1, 2])
sigma_s = tf.sqrt(var_s + EPSILON)
sigma_g = tf.sqrt(var_g + EPSILON)
l2_mean = tf.reduce_sum(tf.square(mean_g - mean_s))
l2_sigma = tf.reduce_sum(tf.square(sigma_g - sigma_s))
style_layer_loss.append(l2_mean + l2_sigma)
self.style_loss = tf.reduce_sum(style_layer_loss)
self.loss = self.content_loss + self.style_weight * self.style_loss # IMPROVE: tune STYLE_WEIGHT
return generated_images
name='style_input')
# switch RGB to BGR
content = tf.reverse(content_input, axis=[-1])
style = tf.reverse(style_input, axis=[-1])
# preprocess image
content = encoder.preprocess(content)
style = encoder.preprocess(style)
# encode image
# we should initial global variables before restore model
enc_c_net = encoder.encode(content, 'content/')
enc_s_net = encoder.encode(style, 'style/')
# pass the encoded images to AdaIN
target_features = utils.AdaIN(enc_c_net.outputs, enc_s_net.outputs)
# decode target features back to image
dec_net = decoder.decode(target_features, prefix="decoder/")
generated_img = dec_net.outputs
# deprocess image
generated_img = encoder.deprocess(generated_img)
# switch BGR back to RGB
generated_img = tf.reverse(generated_img, axis=[-1])
# clip to 0..255
generated_img = tf.clip_by_value(generated_img, 0.0, 255.0)
def main(*args, **kwargs):
INFO('--- Experiment begins: {}---------------'.format(Config.ExperimentName))
config_experiment = ConfigSerializer.load(Path.ExperimentConfigAbs)
init_logging(VerbosityLevel.DEBUG)
from helpers.tf_helper import async_preload_gpu_devices
async_preload_gpu_devices()
# TODO: explain terms or key concepts in comments, ref: overlook.vsd
from modules.data.data_manager import DataManager
x_train, y_train = None, None
x_test, y_test = None, None
data = None
from modules.models.model_manager import ModelManager
model = None
if config_experiment.train.enabled:
INFO('--- Training begins ---------')
config_data: Params = config_experiment.data_set.data
data = DataManager.load_data(config_data.signature, **config_data)
INFO(f"data loaded: {dump_iterable_data(data)}")
if not isinstance(data, tuple) or any(not isinstance(_, tuple) for _ in data):
raise ValueError("data loaded must be in type of ((x_train, y_train), (x_test, y_test))")
(x_train, y_train), (x_test, y_test) = data # unpack: tuple of 4 np.ndarrays
# e.g.((60000,28,28), (60000,)), ((10000,28,28), (10000,))
config_model: Params = config_experiment.model_set.model_base
model = ModelManager.load_model(config_model.signature, **config_model)
model = ModelManager.model_train(model, data=(x_train, y_train), **config_experiment.train)
eval_metrics = ModelManager.model_evaluate(model, data=(x_test, y_test))
else:
INFO('--- Training was disabled ---------')
if config_experiment.predict.enabled:
INFO('--- Prediction begins ---------')
predictions = None
while True:
if 'meta_info' not in vars():
meta_info = {} # retrieve meta info from DataManager
if x_test is None or y_test is None: # not config_experiment.train.enabled
data_key = config_experiment.predict.data_inputs.__str__()
config_data_test: Params = config_experiment.data_set[data_key]
# test signature "ui_web_files", need to keep compatibility with other type of data
data = DataManager.load_data(config_data_test.signature,
meta_info=meta_info, **config_data_test)
INFO(f"data loaded: {dump_iterable_data(data)}")
import tensorflow as tf
from helpers.tf_helper import is_tfdataset
if isinstance(data, tf.data.Dataset):
from helpers.tf_helper import tf_obj_to_np_array
if type(data.element_spec) is tuple:
# x_test = data.map(lambda x, y: x)
# y_test = data.map(lambda x, y: y)
# IMPROVE: unzip the ZipDataset by dataset.map(lambda). Any tf API for unzip?
x_test = tf_obj_to_np_array(data.map(lambda x, y: x))
y_test = tf_obj_to_np_array(data.map(lambda x, y: y))
else:
# data = data.batch(1) # TODO: read config `batch_size` in model_train()
# data = data.prefetch(1)
x_test, y_test = tf_obj_to_np_array(data), None
enc, dec = None, None
if enc is None or dec is None: # not config_experiment.train.enabled
config_model_enc: Params = config_experiment.model_set.model_enc
config_model_dec: Params = config_experiment.model_set.model_dec
if not config_model_enc.is_defined() and config_model_dec.is_defined():
raise ValueError('Config error: `model_trained` node is not defined')
enc = ModelManager.load_model(config_model_enc.signature, **config_model_enc)
dec = ModelManager.load_model(config_model_dec.signature, **config_model_dec)
generated = None
generated_saved = False
webapp = ensure_web_app()
@webapp.on_task_query(namespace="main::model_predict", onetime=False)
def handle_task_query(task_id):
nonlocal generated
handler_result = {}
if not generated_saved:
handler_result.update({'status': 'processing'})
else:
# return abspath to webapp
handler_result.update({'status': 'finished', 'result': generated_save_path})
generated = None
return handler_result
c_feature = ModelManager.model_predict(enc, x_test[0], **config_experiment.predict_enc)
s_feature = ModelManager.model_predict(enc, x_test[1], **config_experiment.predict_enc)
target_features = utils.AdaIN(c_feature, s_feature, alpha=1)
generated = ModelManager.model_predict(dec, target_features, **config_experiment.predict_dec)
import tensorflow as tf
if isinstance(generated, tf.Tensor):
if generated.dtype == tf.float32:
generated = tf.cast(generated, tf.uint8)
generated = generated.numpy()
# show_image_mat(generated[0])
import cv2
generated_save_path = osp.join(Path.ExperimentFolderAbs, tmp_filename_by_time('jpg'))
save_image_mat(generated[0], generated_save_path)
generated_saved = True
# INFO(f"generated: {', '.join([str(_) for _ in generated])}")
else:
INFO('--- Prediction was disabled ---------')
INFO('--- Experiment ends: {}---------------'.format(Config.ExperimentName))