functions.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. # Copyright (c) 2021 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. import shapely.ops
  17. from shapely.geometry import Polygon, MultiPolygon, GeometryCollection
  18. import copy
  19. def normalize(im, mean, std, min_value=[0, 0, 0], max_value=[255, 255, 255]):
  20. # Rescaling (min-max normalization)
  21. range_value = np.asarray(
  22. [1. / (max_value[i] - min_value[i]) for i in range(len(max_value))],
  23. dtype=np.float32)
  24. im = (im - np.asarray(min_value, dtype=np.float32)) * range_value
  25. # Standardization (Z-score Normalization)
  26. im -= mean
  27. im /= std
  28. return im
  29. def permute(im, to_bgr=False):
  30. im = np.swapaxes(im, 1, 2)
  31. im = np.swapaxes(im, 1, 0)
  32. if to_bgr:
  33. im = im[[2, 1, 0], :, :]
  34. return im
  35. def center_crop(im, crop_size=224):
  36. height, width = im.shape[:2]
  37. w_start = (width - crop_size) // 2
  38. h_start = (height - crop_size) // 2
  39. w_end = w_start + crop_size
  40. h_end = h_start + crop_size
  41. im = im[h_start:h_end, w_start:w_end, :]
  42. return im
  43. def horizontal_flip(im):
  44. if len(im.shape) == 3:
  45. im = im[:, ::-1, :]
  46. elif len(im.shape) == 2:
  47. im = im[:, ::-1]
  48. return im
  49. def vertical_flip(im):
  50. if len(im.shape) == 3:
  51. im = im[::-1, :, :]
  52. elif len(im.shape) == 2:
  53. im = im[::-1, :]
  54. return im
  55. def rgb2bgr(im):
  56. return im[:, :, ::-1]
  57. def is_poly(poly):
  58. assert isinstance(poly, (list, dict)), \
  59. "Invalid poly type: {}".format(type(poly))
  60. return isinstance(poly, list)
  61. def horizontal_flip_poly(poly, width):
  62. flipped_poly = np.array(poly)
  63. flipped_poly[0::2] = width - np.array(poly[0::2])
  64. return flipped_poly.tolist()
  65. def horizontal_flip_rle(rle, height, width):
  66. import pycocotools.mask as mask_util
  67. if 'counts' in rle and type(rle['counts']) == list:
  68. rle = mask_util.frPyObjects(rle, height, width)
  69. mask = mask_util.decode(rle)
  70. mask = mask[:, ::-1]
  71. rle = mask_util.encode(np.array(mask, order='F', dtype=np.uint8))
  72. return rle
  73. def vertical_flip_poly(poly, height):
  74. flipped_poly = np.array(poly)
  75. flipped_poly[1::2] = height - np.array(poly[1::2])
  76. return flipped_poly.tolist()
  77. def vertical_flip_rle(rle, height, width):
  78. import pycocotools.mask as mask_util
  79. if 'counts' in rle and type(rle['counts']) == list:
  80. rle = mask_util.frPyObjects(rle, height, width)
  81. mask = mask_util.decode(rle)
  82. mask = mask[::-1, :]
  83. rle = mask_util.encode(np.array(mask, order='F', dtype=np.uint8))
  84. return rle
  85. def crop_poly(segm, crop):
  86. xmin, ymin, xmax, ymax = crop
  87. crop_coord = [xmin, ymin, xmin, ymax, xmax, ymax, xmax, ymin]
  88. crop_p = np.array(crop_coord).reshape(4, 2)
  89. crop_p = Polygon(crop_p)
  90. crop_segm = list()
  91. for poly in segm:
  92. poly = np.array(poly).reshape(len(poly) // 2, 2)
  93. polygon = Polygon(poly)
  94. if not polygon.is_valid:
  95. exterior = polygon.exterior
  96. multi_lines = exterior.intersection(exterior)
  97. polygons = shapely.ops.polygonize(multi_lines)
  98. polygon = MultiPolygon(polygons)
  99. multi_polygon = list()
  100. if isinstance(polygon, MultiPolygon):
  101. multi_polygon = copy.deepcopy(polygon)
  102. else:
  103. multi_polygon.append(copy.deepcopy(polygon))
  104. for per_polygon in multi_polygon:
  105. inter = per_polygon.intersection(crop_p)
  106. if not inter:
  107. continue
  108. if isinstance(inter, (MultiPolygon, GeometryCollection)):
  109. for part in inter:
  110. if not isinstance(part, Polygon):
  111. continue
  112. part = np.squeeze(
  113. np.array(part.exterior.coords[:-1]).reshape(1, -1))
  114. part[0::2] -= xmin
  115. part[1::2] -= ymin
  116. crop_segm.append(part.tolist())
  117. elif isinstance(inter, Polygon):
  118. crop_poly = np.squeeze(
  119. np.array(inter.exterior.coords[:-1]).reshape(1, -1))
  120. crop_poly[0::2] -= xmin
  121. crop_poly[1::2] -= ymin
  122. crop_segm.append(crop_poly.tolist())
  123. else:
  124. continue
  125. return crop_segm
  126. def crop_rle(rle, crop, height, width):
  127. import pycocotools.mask as mask_util
  128. if 'counts' in rle and type(rle['counts']) == list:
  129. rle = mask_util.frPyObjects(rle, height, width)
  130. mask = mask_util.decode(rle)
  131. mask = mask[crop[1]:crop[3], crop[0]:crop[2]]
  132. rle = mask_util.encode(np.array(mask, order='F', dtype=np.uint8))
  133. return rle
  134. def expand_poly(poly, x, y):
  135. expanded_poly = np.array(poly)
  136. expanded_poly[0::2] += x
  137. expanded_poly[1::2] += y
  138. return expanded_poly.tolist()
  139. def expand_rle(rle, x, y, height, width, h, w):
  140. import pycocotools.mask as mask_util
  141. if 'counts' in rle and type(rle['counts']) == list:
  142. rle = mask_util.frPyObjects(rle, height, width)
  143. mask = mask_util.decode(rle)
  144. expanded_mask = np.full((h, w), 0).astype(mask.dtype)
  145. expanded_mask[y:y + height, x:x + width] = mask
  146. rle = mask_util.encode(np.array(expanded_mask, order='F', dtype=np.uint8))
  147. return rle
  148. def resize_poly(poly, im_scale_x, im_scale_y):
  149. resized_poly = np.array(poly, dtype=np.float32)
  150. resized_poly[0::2] *= im_scale_x
  151. resized_poly[1::2] *= im_scale_y
  152. return resized_poly.tolist()
  153. def resize_rle(rle, im_h, im_w, im_scale_x, im_scale_y, interp):
  154. import pycocotools.mask as mask_util
  155. if 'counts' in rle and type(rle['counts']) == list:
  156. rle = mask_util.frPyObjects(rle, im_h, im_w)
  157. mask = mask_util.decode(rle)
  158. mask = cv2.resize(
  159. mask, None, None, fx=im_scale_x, fy=im_scale_y, interpolation=interp)
  160. rle = mask_util.encode(np.array(mask, order='F', dtype=np.uint8))
  161. return rle