Skip to content

Instantly share code, notes, and snippets.

@ozgurshn
Last active April 10, 2020 13:55
Show Gist options
  • Save ozgurshn/6b2c0682c75fd4b2c4e376585b4811a1 to your computer and use it in GitHub Desktop.
Save ozgurshn/6b2c0682c75fd4b2c4e376585b4811a1 to your computer and use it in GitHub Desktop.
add_reflective_padding_and_crop add layer manually in Core ML
##https://heartbeat.fritz.ai/advanced-tips-for-core-ml-77c9e0231a9
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment