Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow_addons/image/cutout_ops.py: 18%

1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Cutout op""" 

16 

17import tensorflow as tf 

18from tensorflow_addons.utils.types import TensorLike, Number 

19 

20 

21@tf.function 

22def _norm_params(mask_size, offset=None): 

23 tf.assert_equal( 

24 tf.reduce_any(mask_size % 2 != 0), 

25 False, 

26 "mask_size should be divisible by 2", 

27 ) 

28 if tf.rank(mask_size) == 0: 

29 mask_size = tf.stack([mask_size, mask_size]) 

30 if offset is not None and tf.rank(offset) == 1: 

31 offset = tf.expand_dims(offset, 0) 

32 return mask_size, offset 

33 

34 

35@tf.function 

36def _random_center(mask_dim_length, image_dim_length, batch_size, seed): 

37 if mask_dim_length >= image_dim_length: 

38 return tf.tile([image_dim_length // 2], [batch_size]) 

39 half_mask_dim_length = mask_dim_length // 2 

40 return tf.random.uniform( 

41 shape=[batch_size], 

42 minval=half_mask_dim_length, 

43 maxval=image_dim_length - half_mask_dim_length, 

44 dtype=tf.int32, 

45 seed=seed, 

46 ) 

47 

48 

49def random_cutout( 

50 images: TensorLike, 

51 mask_size: TensorLike, 

52 constant_values: Number = 0, 

53 seed: Number = None, 

54) -> tf.Tensor: 

55 """Apply [cutout](https://arxiv.org/abs/1708.04552) to images with random offset. 

56 

57 This operation applies a `(mask_height x mask_width)` mask of zeros to 

58 a random location within `images`. The pixel values filled in will be of 

59 the value `constant_values`. The location where the mask will be applied is 

60 randomly chosen uniformly over the whole images. 

61 

62 Args: 

63 images: A tensor of shape `(batch_size, height, width, channels)` (NHWC). 

64 mask_size: Specifies how big the zero mask that will be generated is that 

65 is applied to the images. The mask will be of size 

66 `(mask_height x mask_width)`. Note: mask_size should be divisible by 2. 

67 constant_values: What pixel value to fill in the images in the area that has 

68 the cutout mask applied to it. 

69 seed: A Python integer. Used in combination with `tf.random.set_seed` to 

70 create a reproducible sequence of tensors across multiple calls. 

71 Returns: 

72 A `Tensor` of the same shape and dtype as `images`. 

73 Raises: 

74 InvalidArgumentError: if `mask_size` can't be divisible by 2. 

75 """ 

76 images = tf.convert_to_tensor(images) 

77 mask_size = tf.convert_to_tensor(mask_size) 

78 

79 image_dynamic_shape = tf.shape(images) 

80 batch_size, image_height, image_width = ( 

81 image_dynamic_shape[0], 

82 image_dynamic_shape[1], 

83 image_dynamic_shape[2], 

84 ) 

85 

86 mask_size, _ = _norm_params(mask_size, offset=None) 

87 

88 cutout_center_height = _random_center(mask_size[0], image_height, batch_size, seed) 

89 cutout_center_width = _random_center(mask_size[1], image_width, batch_size, seed) 

90 

91 offset = tf.transpose([cutout_center_height, cutout_center_width], [1, 0]) 

92 return cutout(images, mask_size, offset, constant_values) 

93 

94 

95def cutout( 

96 images: TensorLike, 

97 mask_size: TensorLike, 

98 offset: TensorLike = (0, 0), 

99 constant_values: Number = 0, 

100) -> tf.Tensor: 

101 """Apply [cutout](https://arxiv.org/abs/1708.04552) to images. 

102 

103 This operation applies a `(mask_height x mask_width)` mask of zeros to 

104 a location within `images` specified by the offset. 

105 The pixel values filled in will be of the value `constant_values`. 

106 The location where the mask will be applied is randomly 

107 chosen uniformly over the whole images. 

108 

109 Args: 

110 images: A tensor of shape `(batch_size, height, width, channels)` (NHWC). 

111 mask_size: Specifies how big the zero mask that will be generated is that 

112 is applied to the images. The mask will be of size 

113 `(mask_height x mask_width)`. Note: mask_size should be divisible by 2. 

114 offset: A tuple of `(height, width)` or `(batch_size, 2)` 

115 constant_values: What pixel value to fill in the images in the area that has 

116 the cutout mask applied to it. 

117 Returns: 

118 A `Tensor` of the same shape and dtype as `images`. 

119 Raises: 

120 InvalidArgumentError: if `mask_size` can't be divisible by 2. 

121 """ 

122 with tf.name_scope("cutout"): 

123 images = tf.convert_to_tensor(images) 

124 mask_size = tf.convert_to_tensor(mask_size) 

125 offset = tf.convert_to_tensor(offset) 

126 

127 image_static_shape = images.shape 

128 image_dynamic_shape = tf.shape(images) 

129 image_height, image_width, channels = ( 

130 image_dynamic_shape[1], 

131 image_dynamic_shape[2], 

132 image_dynamic_shape[3], 

133 ) 

134 

135 mask_size, offset = _norm_params(mask_size, offset) 

136 mask_size = mask_size // 2 

137 

138 cutout_center_heights = offset[:, 0] 

139 cutout_center_widths = offset[:, 1] 

140 

141 lower_pads = tf.maximum(0, cutout_center_heights - mask_size[0]) 

142 upper_pads = tf.maximum(0, image_height - cutout_center_heights - mask_size[0]) 

143 left_pads = tf.maximum(0, cutout_center_widths - mask_size[1]) 

144 right_pads = tf.maximum(0, image_width - cutout_center_widths - mask_size[1]) 

145 

146 cutout_shape = tf.transpose( 

147 [ 

148 image_height - (lower_pads + upper_pads), 

149 image_width - (left_pads + right_pads), 

150 ], 

151 [1, 0], 

152 ) 

153 padding_dims = tf.stack( 

154 [ 

155 tf.stack([lower_pads, upper_pads], axis=1), 

156 tf.stack([left_pads, right_pads], axis=1), 

157 ], 

158 axis=1, 

159 ) 

160 mask = tf.map_fn( 

161 _generate_masks, 

162 [cutout_shape, padding_dims], 

163 fn_output_signature=tf.TensorSpec( 

164 shape=image_static_shape[1:-1], dtype=tf.bool 

165 ), 

166 ) 

167 mask = tf.expand_dims(mask, -1) 

168 mask = tf.tile(mask, [1, 1, 1, channels]) 

169 

170 images = tf.where( 

171 mask, 

172 images, 

173 tf.cast(constant_values, dtype=images.dtype), 

174 ) 

175 images.set_shape(image_static_shape) 

176 return images 

177 

178 

179def _generate_masks(args): 

180 cutout_shape, padding_dims = args 

181 mask = tf.pad( 

182 tf.zeros(cutout_shape, dtype=tf.bool), 

183 padding_dims, 

184 constant_values=True, 

185 ) 

186 return mask