Spaces:
Running
Running
| # Copyright 2018 The TensorFlow Authors All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Utilities for artificially damaging segmentation masks.""" | |
| import numpy as np | |
| from scipy.ndimage import interpolation | |
| from skimage import morphology | |
| from skimage import transform | |
| import tensorflow as tf | |
| def damage_masks(labels, shift=True, scale=True, rotate=True, dilate=True): | |
| """Damages segmentation masks by random transformations. | |
| Args: | |
| labels: Int32 labels tensor of shape (height, width, 1). | |
| shift: Boolean, whether to damage the masks by shifting. | |
| scale: Boolean, whether to damage the masks by scaling. | |
| rotate: Boolean, whether to damage the masks by rotation. | |
| dilate: Boolean, whether to damage the masks by dilation. | |
| Returns: | |
| The damaged version of labels. | |
| """ | |
| def _damage_masks_np(labels_): | |
| return damage_masks_np(labels_, shift, scale, rotate, dilate) | |
| damaged_masks = tf.py_func(_damage_masks_np, [labels], tf.int32, | |
| name='damage_masks') | |
| damaged_masks.set_shape(labels.get_shape()) | |
| return damaged_masks | |
| def damage_masks_np(labels, shift=True, scale=True, rotate=True, dilate=True): | |
| """Performs the actual mask damaging in numpy. | |
| Args: | |
| labels: Int32 numpy array of shape (height, width, 1). | |
| shift: Boolean, whether to damage the masks by shifting. | |
| scale: Boolean, whether to damage the masks by scaling. | |
| rotate: Boolean, whether to damage the masks by rotation. | |
| dilate: Boolean, whether to damage the masks by dilation. | |
| Returns: | |
| The damaged version of labels. | |
| """ | |
| unique_labels = np.unique(labels) | |
| unique_labels = np.setdiff1d(unique_labels, [0]) | |
| # Shuffle to get random depth ordering when combining together. | |
| np.random.shuffle(unique_labels) | |
| damaged_labels = np.zeros_like(labels) | |
| for l in unique_labels: | |
| obj_mask = (labels == l) | |
| damaged_obj_mask = _damage_single_object_mask(obj_mask, shift, scale, | |
| rotate, dilate) | |
| damaged_labels[damaged_obj_mask] = l | |
| return damaged_labels | |
| def _damage_single_object_mask(mask, shift, scale, rotate, dilate): | |
| """Performs mask damaging in numpy for a single object. | |
| Args: | |
| mask: Boolean numpy array of shape(height, width, 1). | |
| shift: Boolean, whether to damage the masks by shifting. | |
| scale: Boolean, whether to damage the masks by scaling. | |
| rotate: Boolean, whether to damage the masks by rotation. | |
| dilate: Boolean, whether to damage the masks by dilation. | |
| Returns: | |
| The damaged version of mask. | |
| """ | |
| # For now we just do shifting and scaling. Better would be Affine or thin | |
| # spline plate transformations. | |
| if shift: | |
| mask = _shift_mask(mask) | |
| if scale: | |
| mask = _scale_mask(mask) | |
| if rotate: | |
| mask = _rotate_mask(mask) | |
| if dilate: | |
| mask = _dilate_mask(mask) | |
| return mask | |
| def _shift_mask(mask, max_shift_factor=0.05): | |
| """Damages a mask for a single object by randomly shifting it in numpy. | |
| Args: | |
| mask: Boolean numpy array of shape(height, width, 1). | |
| max_shift_factor: Float scalar, the maximum factor for random shifting. | |
| Returns: | |
| The shifted version of mask. | |
| """ | |
| nzy, nzx, _ = mask.nonzero() | |
| h = nzy.max() - nzy.min() | |
| w = nzx.max() - nzx.min() | |
| size = np.sqrt(h * w) | |
| offset = np.random.uniform(-size * max_shift_factor, size * max_shift_factor, | |
| 2) | |
| shifted_mask = interpolation.shift(np.squeeze(mask, axis=2), | |
| offset, order=0).astype('bool')[..., | |
| np.newaxis] | |
| return shifted_mask | |
| def _scale_mask(mask, scale_amount=0.025): | |
| """Damages a mask for a single object by randomly scaling it in numpy. | |
| Args: | |
| mask: Boolean numpy array of shape(height, width, 1). | |
| scale_amount: Float scalar, the maximum factor for random scaling. | |
| Returns: | |
| The scaled version of mask. | |
| """ | |
| nzy, nzx, _ = mask.nonzero() | |
| cy = 0.5 * (nzy.max() - nzy.min()) | |
| cx = 0.5 * (nzx.max() - nzx.min()) | |
| scale_factor = np.random.uniform(1.0 - scale_amount, 1.0 + scale_amount) | |
| shift = transform.SimilarityTransform(translation=[-cx, -cy]) | |
| inv_shift = transform.SimilarityTransform(translation=[cx, cy]) | |
| s = transform.SimilarityTransform(scale=[scale_factor, scale_factor]) | |
| m = (shift + (s + inv_shift)).inverse | |
| scaled_mask = transform.warp(mask, m) > 0.5 | |
| return scaled_mask | |
| def _rotate_mask(mask, max_rot_degrees=3.0): | |
| """Damages a mask for a single object by randomly rotating it in numpy. | |
| Args: | |
| mask: Boolean numpy array of shape(height, width, 1). | |
| max_rot_degrees: Float scalar, the maximum number of degrees to rotate. | |
| Returns: | |
| The scaled version of mask. | |
| """ | |
| cy = 0.5 * mask.shape[0] | |
| cx = 0.5 * mask.shape[1] | |
| rot_degrees = np.random.uniform(-max_rot_degrees, max_rot_degrees) | |
| shift = transform.SimilarityTransform(translation=[-cx, -cy]) | |
| inv_shift = transform.SimilarityTransform(translation=[cx, cy]) | |
| r = transform.SimilarityTransform(rotation=np.deg2rad(rot_degrees)) | |
| m = (shift + (r + inv_shift)).inverse | |
| scaled_mask = transform.warp(mask, m) > 0.5 | |
| return scaled_mask | |
| def _dilate_mask(mask, dilation_radius=5): | |
| """Damages a mask for a single object by dilating it in numpy. | |
| Args: | |
| mask: Boolean numpy array of shape(height, width, 1). | |
| dilation_radius: Integer, the radius of the used disk structure element. | |
| Returns: | |
| The dilated version of mask. | |
| """ | |
| disk = morphology.disk(dilation_radius, dtype=np.bool) | |
| dilated_mask = morphology.binary_dilation( | |
| np.squeeze(mask, axis=2), selem=disk)[..., np.newaxis] | |
| return dilated_mask | |