functional.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  1. # Copyright (c) 2020 PaddlePaddle 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. import cv2
  15. import numpy as np
  16. from PIL import Image, ImageEnhance
  17. from scipy.ndimage.morphology import distance_transform_edt
  18. def normalize(im, mean, std):
  19. im = im.astype(np.float32, copy=False) / 255.0
  20. im -= mean
  21. im /= std
  22. return im
  23. def resize(im, target_size=608, interp=cv2.INTER_LINEAR):
  24. if isinstance(target_size, list) or isinstance(target_size, tuple):
  25. w = target_size[0]
  26. h = target_size[1]
  27. else:
  28. w = target_size
  29. h = target_size
  30. im = cv2.resize(im, (w, h), interpolation=interp)
  31. return im
  32. def resize_long(im, long_size=224, interpolation=cv2.INTER_LINEAR):
  33. value = max(im.shape[0], im.shape[1])
  34. scale = float(long_size) / float(value)
  35. resized_width = int(round(im.shape[1] * scale))
  36. resized_height = int(round(im.shape[0] * scale))
  37. im = cv2.resize(
  38. im, (resized_width, resized_height), interpolation=interpolation)
  39. return im
  40. def horizontal_flip(im):
  41. if len(im.shape) == 3:
  42. im = im[:, ::-1, :]
  43. elif len(im.shape) == 2:
  44. im = im[:, ::-1]
  45. return im
  46. def vertical_flip(im):
  47. if len(im.shape) == 3:
  48. im = im[::-1, :, :]
  49. elif len(im.shape) == 2:
  50. im = im[::-1, :]
  51. return im
  52. def brightness(im, brightness_lower, brightness_upper):
  53. brightness_delta = np.random.uniform(brightness_lower, brightness_upper)
  54. im = ImageEnhance.Brightness(im).enhance(brightness_delta)
  55. return im
  56. def contrast(im, contrast_lower, contrast_upper):
  57. contrast_delta = np.random.uniform(contrast_lower, contrast_upper)
  58. im = ImageEnhance.Contrast(im).enhance(contrast_delta)
  59. return im
  60. def saturation(im, saturation_lower, saturation_upper):
  61. saturation_delta = np.random.uniform(saturation_lower, saturation_upper)
  62. im = ImageEnhance.Color(im).enhance(saturation_delta)
  63. return im
  64. def hue(im, hue_lower, hue_upper):
  65. hue_delta = np.random.uniform(hue_lower, hue_upper)
  66. im = np.array(im.convert('HSV'))
  67. im[:, :, 0] = im[:, :, 0] + hue_delta
  68. im = Image.fromarray(im, mode='HSV').convert('RGB')
  69. return im
  70. def rotate(im, rotate_lower, rotate_upper):
  71. rotate_delta = np.random.uniform(rotate_lower, rotate_upper)
  72. im = im.rotate(int(rotate_delta))
  73. return im
  74. def mask_to_onehot(mask, num_classes):
  75. """
  76. Convert a mask (H, W) to onehot (K, H, W).
  77. Args:
  78. mask (np.ndarray): Label mask with shape (H, W)
  79. num_classes (int): Number of classes.
  80. Returns:
  81. np.ndarray: Onehot mask with shape(K, H, W).
  82. """
  83. _mask = [mask == i for i in range(num_classes)]
  84. _mask = np.array(_mask).astype(np.uint8)
  85. return _mask
  86. def onehot_to_binary_edge(mask, radius):
  87. """
  88. Convert a onehot mask (K, H, W) to a edge mask.
  89. Args:
  90. mask (np.ndarray): Onehot mask with shape (K, H, W)
  91. radius (int|float): Radius of edge.
  92. Returns:
  93. np.ndarray: Edge mask with shape(H, W).
  94. """
  95. if radius < 1:
  96. raise ValueError('`radius` should be greater than or equal to 1')
  97. num_classes = mask.shape[0]
  98. edge = np.zeros(mask.shape[1:])
  99. # pad borders
  100. mask = np.pad(mask, ((0, 0), (1, 1), (1, 1)),
  101. mode='constant',
  102. constant_values=0)
  103. for i in range(num_classes):
  104. dist = distance_transform_edt(mask[i, :]) + distance_transform_edt(
  105. 1.0 - mask[i, :])
  106. dist = dist[1:-1, 1:-1]
  107. dist[dist > radius] = 0
  108. edge += dist
  109. edge = np.expand_dims(edge, axis=0)
  110. edge = (edge > 0).astype(np.uint8)
  111. return edge
  112. def mask_to_binary_edge(mask, radius, num_classes):
  113. """
  114. Convert a segmentic segmentation mask (H, W) to a binary edge mask(H, W).
  115. Args:
  116. mask (np.ndarray): Label mask with shape (H, W)
  117. radius (int|float): Radius of edge.
  118. num_classes (int): Number of classes.
  119. Returns:
  120. np.ndarray: Edge mask with shape(H, W).
  121. """
  122. mask = mask.squeeze()
  123. onehot = mask_to_onehot(mask, num_classes)
  124. edge = onehot_to_binary_edge(onehot, radius)
  125. return edge