def build_model(content_input_,
style_input_,
trainable,
is_training,
reuse=None,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=None,
style_weights=None,
total_variation_weight=None):
"""The image stylize function.
Args:
content_input_: Tensor. Batch of content input images.
style_input_: Tensor. Batch of style input images.
trainable: bool. Should the parameters be marked as trainable?
is_training: bool. Is it training phase or not?
reuse: bool. Whether to reuse model parameters. Defaults to False.
inception_end_point: string. Specifies the endpoint to construct the
inception_v3 network up to. This network is used for style prediction.
style_prediction_bottleneck: int. Specifies the bottleneck size in the
number of parameters of the style embedding.
adds_losses: wheather or not to add objectives to the model.
content_weights: dict mapping layer names to their associated content loss
weight. Keys that are missing from the dict won't have their content
loss computed.
style_weights: dict mapping layer names to their associated style loss
weight. Keys that are missing from the dict won't have their style
loss computed.
total_variation_weight: float. Coefficient for the total variation part of
the loss.
Returns:
Tensor for the output of the transformer network, Tensor for the total loss,
dict mapping loss names to losses, Tensor for the bottleneck activations of
the style prediction network.
"""
# Gets scope name and shape of the activations of transformer network which
# will be used to apply style.
[activation_names,
activation_depths] = transformer_model.style_normalization_activations()
# Defines the style prediction network.
style_params, bottleneck_feat = style_prediction(
style_input_,
activation_names,
activation_depths,
is_training=is_training,
trainable=trainable,
inception_end_point=inception_end_point,
style_prediction_bottleneck=style_prediction_bottleneck,
reuse=reuse)
# Defines the style transformer network.
stylized_images = transformer_model.transform(
content_input_,
normalizer_fn=ops.conditional_style_norm,
reuse=reuse,
trainable=trainable,
is_training=is_training,
normalizer_params={'style_params': style_params})
# Adds losses.
loss_dict = {}
total_loss = []
if adds_losses:
total_loss, loss_dict = losses.total_loss(
content_input_,
style_input_,
stylized_images,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=total_variation_weight)
return stylized_images, total_loss, loss_dict, bottleneck_feat
def build_mobilenet_model(content_input_,
style_input_,
mobilenet_trainable=True,
style_params_trainable=False,
transformer_trainable=False,
reuse=None,
mobilenet_end_point='layer_19',
transformer_alpha=0.25,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=None,
style_weights=None,
total_variation_weight=None):
"""The image stylize function using a MobileNetV2 instead of InceptionV3.
Args:
content_input_: Tensor. Batch of content input images.
style_input_: Tensor. Batch of style input images.
mobilenet_trainable: bool. Should the MobileNet parameters be trainable?
style_params_trainable: bool. Should the style parameters be trainable?
transformer_trainable: bool. Should the style transfer network be
trainable?
reuse: bool. Whether to reuse model parameters. Defaults to False.
mobilenet_end_point: string. Specifies the endpoint to construct the
MobileNetV2 network up to. This network is used for style prediction.
transformer_alpha: float. Width multiplier used to reduce the number of
filters in the model and slim it down.
style_prediction_bottleneck: int. Specifies the bottleneck size in the
number of parameters of the style embedding.
adds_losses: wheather or not to add objectives to the model.
content_weights: dict mapping layer names to their associated content loss
weight. Keys that are missing from the dict won't have their content
loss computed.
style_weights: dict mapping layer names to their associated style loss
weight. Keys that are missing from the dict won't have their style
loss computed.
total_variation_weight: float. Coefficient for the total variation part of
the loss.
Returns:
Tensor for the output of the transformer network, Tensor for the total loss,
dict mapping loss names to losses, Tensor for the bottleneck activations of
the style prediction network.
"""
[activation_names,
activation_depths] = transformer_model.style_normalization_activations(
alpha=transformer_alpha)
# Defines the style prediction network.
style_params, bottleneck_feat = style_prediction_mobilenet(
style_input_,
activation_names,
activation_depths,
mobilenet_end_point=mobilenet_end_point,
mobilenet_trainable=mobilenet_trainable,
style_params_trainable=style_params_trainable,
style_prediction_bottleneck=style_prediction_bottleneck,
reuse=reuse)
# Defines the style transformer network
stylized_images = transformer_model.transform(
content_input_,
alpha=transformer_alpha,
normalizer_fn=ops.conditional_style_norm,
reuse=reuse,
trainable=transformer_trainable,
is_training=transformer_trainable,
normalizer_params={'style_params': style_params})
# Adds losses
loss_dict = {}
total_loss = []
if adds_losses:
total_loss, loss_dict = losses.total_loss(
content_input_,
style_input_,
stylized_images,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=total_variation_weight)
return stylized_images, total_loss, loss_dict, bottleneck_feat
def build_model(content_input_,
style_input_,
trainable,
is_training,
reuse=None,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=None,
style_weights=None,
total_variation_weight=None):
"""The image stylize function.
Args:
content_input_: Tensor. Batch of content input images.
style_input_: Tensor. Batch of style input images.
trainable: bool. Should the parameters be marked as trainable?
is_training: bool. Is it training phase or not?
reuse: bool. Whether to reuse model parameters. Defaults to False.
inception_end_point: string. Specifies the endpoint to construct the
inception_v3 network up to. This network is used for style prediction.
style_prediction_bottleneck: int. Specifies the bottleneck size in the
number of parameters of the style embedding.
adds_losses: wheather or not to add objectives to the model.
content_weights: dict mapping layer names to their associated content loss
weight. Keys that are missing from the dict won't have their content
loss computed.
style_weights: dict mapping layer names to their associated style loss
weight. Keys that are missing from the dict won't have their style
loss computed.
total_variation_weight: float. Coefficient for the total variation part of
the loss.
Returns:
Tensor for the output of the transformer network, Tensor for the total loss,
dict mapping loss names to losses, Tensor for the bottleneck activations of
the style prediction network.
"""
# Gets scope name and shape of the activations of transformer network which
# will be used to apply style.
[activation_names,
activation_depths] = transformer_model.style_normalization_activations()
# Defines the style prediction network.
style_params, bottleneck_feat = style_prediction(
style_input_,
activation_names,
activation_depths,
is_training=is_training,
trainable=trainable,
inception_end_point=inception_end_point,
style_prediction_bottleneck=style_prediction_bottleneck,
reuse=reuse)
# Defines the style transformer network.
stylized_images = transformer_model.transform(
content_input_,
normalizer_fn=ops.conditional_style_norm,
reuse=reuse,
trainable=trainable,
is_training=is_training,
normalizer_params={'style_params': style_params})
# Adds losses.
loss_dict = {}
total_loss = []
if adds_losses:
total_loss, loss_dict = losses.total_loss(
content_input_,
style_input_,
stylized_images,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=total_variation_weight)
return stylized_images, total_loss, loss_dict, bottleneck_feat
def build_mobilenet_model(content_input_,
style_input_,
mobilenet_trainable=True,
style_params_trainable=False,
transformer_trainable=False,
reuse=None,
mobilenet_end_point='layer_19',
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=None,
style_weights=None,
total_variation_weight=None):
"""The image stylize function using a MobileNetV2 instead of InceptionV3.
Args:
content_input_: Tensor. Batch of content input images.
style_input_: Tensor. Batch of style input images.
mobilenet_trainable: bool. Should the MobileNet parameters be trainable?
style_params_trainable: bool. Should the style parameters be trainable?
transformer_trainable: bool. Should the style transfer network be
trainable?
reuse: bool. Whether to reuse model parameters. Defaults to False.
mobilenet_end_point: string. Specifies the endpoint to construct the
MobileNetV2 network up to. This network is used for style prediction.
style_prediction_bottleneck: int. Specifies the bottleneck size in the
number of parameters of the style embedding.
adds_losses: wheather or not to add objectives to the model.
content_weights: dict mapping layer names to their associated content loss
weight. Keys that are missing from the dict won't have their content
loss computed.
style_weights: dict mapping layer names to their associated style loss
weight. Keys that are missing from the dict won't have their style
loss computed.
total_variation_weight: float. Coefficient for the total variation part of
the loss.
Returns:
Tensor for the output of the transformer network, Tensor for the total loss,
dict mapping loss names to losses, Tensor for the bottleneck activations of
the style prediction network.
"""
[activation_names,
activation_depths] = transformer_model.style_normalization_activations()
# Defines the style prediction network.
style_params, bottleneck_feat = style_prediction_mobilenet(
style_input_,
activation_names,
activation_depths,
mobilenet_end_point=mobilenet_end_point,
mobilenet_trainable=mobilenet_trainable,
style_params_trainable=style_params_trainable,
style_prediction_bottleneck=style_prediction_bottleneck,
reuse=reuse
)
# Defines the style transformer network
stylized_images = transformer_model.transform(
content_input_,
normalizer_fn=ops.conditional_style_norm,
reuse=reuse,
trainable=transformer_trainable,
is_training=transformer_trainable,
normalizer_params={'style_params': style_params}
)
# Adds losses
loss_dict = {}
total_loss = []
if adds_losses:
total_loss, loss_dict = losses.total_loss(
content_input_,
style_input_,
stylized_images,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=total_variation_weight
)
return stylized_images, total_loss, loss_dict, bottleneck_feat
