Network Module

This module holds the implementation of all the networks, their losses, and their components

class stransfer.network.ContentLoss(target)

Implementation of the content loss

Parameters

target (Tensor) – the target image we want to use to calculate the content loss

forward(input)

Calculate the distance between the input image and the target. This is saved to a local loss attribute.

Return type

Tensor

set_target(target)

This method allows us to set the target image of this content loss instance

class stransfer.network.FeatureReconstructionLoss(target)

Implementation of the feature reconstruction loss.

..note::

this loss is currently not used since it doesn’t seem to provide much improvement over the normal ContentLoss

forward(input)

Get the feature loss of input with respect to the target

Return type

Tensor

class stransfer.network.ImageTransformNet(style_image, batch_size=4)

This the implementation of the fast style transform, image transform network, as defined in:

Perceptual Losses for Real-Time Style Transfer and Super-Resolution

Parameters

• style_image (Tensor) – The image we want to use as as the style reference

• batch_size – the size of the batch

get_optimizer(optimizer=<class 'torch.optim.adam.Adam'>)

Get an initialized optimizer

get_total_variation_regularization_loss(transformed_image, regularization_factor=1e-06)

Calculate a regularization loss, which will tell us how ‘noisy’ is the current image. Penalize if it is very noisy. See: https://en.wikipedia.org/wiki/Total_variation_denoising#2D_signal_images

Parameters

• transformed_image (Tensor) – image for which we will get the loss

• regularization_factor – ‘weight’ to scale the loss

Return type

Tensor

Returns

the regularization loss

process_image(image_path, style_name='nsp', out_dir='results/')

Processes a given input image at image_path with a network pretrained on the style style_name.

Saves the processed image to out_dir

Parameters

• image_path (str) – path to the image we want to stylize

• style_name – name of the style we want to apply to the image. Note that a pretrained model with said style must exist in data/models/

• out_dir – directory were the stylized image will be saved

Return type

None

static_test(test_loader, loss_network, style_weight=100000, feature_weight=1)

Tests the performance of a fast style transfer network on still images

static_train(style_name='nsp', epochs=50, style_weight=100000, content_weight=1)

Trains a fast style transfer network for style transfer on still images.

class stransfer.network.ResidualBlock(in_channels, out_channels, kernel_size=3, stride=1)

forward(x)

Given a tensor input, pass it through the residual block, and return the output.

Return type

Tensor

class stransfer.network.StyleLoss(target)

Implementation of the style loss

Parameters

target (Tensor) – the tensor representing the style image we want to take as reference during training

forward(input)

Compare the gram matrix of the input with that of the target. By doing this we calculate the style loss, which is saved to a local loss attribute.

Return type

Tensor

gram_matrix(input)

Calculate the gram matrix for the input tensor

Return type

Tensor

set_target(target)

This method allows us to change the style target of the loss

class stransfer.network.StyleNetwork(style_image, content_image=None)

Implementation of the StyleNetwork as defined in

A Neural Algorithm of Artistic Style - Gatys (2015)

Parameters

• style_image (Tensor) – tensor of the image we want to use as a source for the style

• content_image (Optional[Tensor]) – tensor of the image we want to use as the source for the content

forward(input_image, content_image=None, style_image=None)

Given an input image pass it through all layers in the network

Parameters

• input_image (Tensor) – the image to pass through the network

• content_image – if specified then this will change the curret target for the content loss

• style_image – if specified then this will change the current target for the style loss

Return type

None

get_total_current_content_loss(weight=1)

Returns the sum of all the loss present in all content nodes

Return type

Tensor

get_total_current_feature_loss(weight=1)

Returns the sum of all the loss present in all content nodes

Return type

Tensor

get_total_current_style_loss(weight=1)

Returns the sum of all the loss present in all style nodes

Return type

Tensor

run_through_pieces(input_g, until=-1)

Runs ths input input_g through all the pieces of the network, or until the specified layer if until is not -1

Parameters

• input_g (Tensor) – the input to run through the network

• until – by default -1. If changed then it specified until which layer we want to run the input through

Return type

Tensor

Returns

the output of running the input through all the specified layers

train_gatys(style_image, content_image, steps=550, style_weight=100000, content_weight=1)

Creates a new image with the style of style_image and the content of content_image

Return type

Tensor

Returns

the converted image

class stransfer.network.VideoTransformNet(style_image, batch_size=4, fast_transfer_dict=None)

Implementation of the video transform net.

Parameters

• style_image (Tensor) – image we’ll use as style reference

• batch_size – size of the batch

• fast_transfer_dict – state dict from a pretrained ‘fast style network’. It allows us to start training the video model from this, which allows to bootstrap training. It is recommended to do this since the current video set is not very big.

get_temporal_loss(old_content, old_stylized, current_content, current_stylized, temporal_weight=1)

Calculates the temporal loss See https://github.com/tupini07/StyleTransfer/issues/5

Parameters

• old_content – tensor representing the content of the previous frame

• old_stylized – tensor representing the stylized previous frame

• current_content – tensor representing the content of the current frame

• current_stylized – tensor representing the stylized current frame

• temporal_weight – weight for the temporal loss

Return type

Tensor

Returns

the temporal loss

process_video(video_path, style_name='nsp', working_dir='workdir/', out_dir='results/', fps=24.0)

Applies style to a single video, using pretrained weights. Note that the weights must exist, if not an exception will be raised.

Parameters

• video_path (str) – the path of the video to stylize

• style_name – the name of the style to apply to the video. The weights for a video transform model using said style must exist in data/models/

• working_dir – directory where the transformed frames will be saved

• out_dir – directory where the final transformed video will be saved

• fps – the frames per second to use in the final video

video_train(style_name='nsp', epochs=50, temporal_weight=0.8, style_weight=100000, feature_weight=1, content_weight=1)

Trains the video network

Parameters

• style_name – the name of the style (used for saving and loading checkpoints)

• epochs – how many epochs should the training go through

• temporal_weight – the weight for the temporal loss

• style_weight – the weight for the style loss

• feature_weight – the weight for the feature loss

• content_weight – the weight for the content loss

Return type

None

Returns

stransfer.network.adaptive_torch_load(weights_path)

When loading saved weights, we check if need to map them to either cuda or cpu.

Parameters

weights_path (str) – paths of the weights to load

Returns

the loaded weights

stransfer.network.get_tensorboard_writer(path)

Creates a TensorBoard writer at the folder with path.

If path exists then it is deleted and recreated.

Parameters

path (str) – the path where our SummaryWriter will write

Return type

SummaryWriter

Returns

an initialized SummaryWriter