add_reflective_padding.py

def add_reflective_padding_and_crop(mlmodel, padding_size=20):
    """Add reflective padding and crop layers to remove edge artifcats.

    Because the convolution layers rely on 'same' padding, stylized images have
    a small ring of distortion around the outer edge. This can be eliminated
    with reflective padding on the input image. This method modifies the
    original MLModel spec to add a padding layer after the input and a crop
    layer before the output to remove the padding at the end.

    Args:
        mlmodel (coremltools.models.MLModel): an MLModel spec.
        padding_size (Optional, int): the number of pixels to pad.
    
    Returns:
        new_mlmodel (coremltools.models.MLModel): a new MLModel spec.
    """
    new_spec = mlmodel.get_spec()
    # Clear all the layers
    while new_spec.neuralNetwork.layers:
        new_spec.neuralNetwork.layers.pop()

    # Add a reflective padding layer first
    spec_layer = new_spec.neuralNetwork.layers.add()
    spec_layer.name = 'padding_1'
    spec_layer.input.append('image')
    spec_layer.output.append('padding_1_output')
    spec_layer_params = spec_layer.padding
    spec_layer_params.reflection.MergeFromString(b'')             
    height_border = spec_layer_params.paddingAmounts.borderAmounts.add()
    height_border.startEdgeSize = padding_size
    height_border.endEdgeSize = padding_size
    width_border = spec_layer_params.paddingAmounts.borderAmounts.add()
    width_border.startEdgeSize = padding_size
    width_border.endEdgeSize = padding_size

    # Add the rest of the layers
    for layer in mlmodel._spec.neuralNetwork.layers:
        spec_layer = new_spec.neuralNetwork.layers.add()
        spec_layer.MergeFrom(layer)

    # Crop the padding as a final layer.
    spec_layer = new_spec.neuralNetwork.layers.add()
    spec_layer.name = 'crop_1'
    spec_layer.input.append(new_spec.neuralNetwork.layers[-2].name + '_output')
    spec_layer.output.append('stylizedImage')
    spec_layer_params = spec_layer.crop          
    height_border = spec_layer_params.cropAmounts.borderAmounts.add()
    height_border.startEdgeSize = padding_size
    height_border.endEdgeSize = padding_size
    width_border = spec_layer_params.cropAmounts.borderAmounts.add()
    width_border.startEdgeSize = padding_size
    width_border.endEdgeSize = padding_size

    # Fix the inputs and outputs for the padding and crop layers
    new_spec.neuralNetwork.layers[1].input.pop()
    new_spec.neuralNetwork.layers[1].input.append('padding_1_output')

    new_spec.neuralNetwork.layers[-2].output.pop()
    new_spec.neuralNetwork.layers[-2].output.append(
        new_spec.neuralNetwork.layers[-2].name + '_output')

    return coremltools.models.MLModel(new_spec)