imgaug_support.py 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. import numpy as np
  2. #def execute_imgaug(augmenter, im, bboxes=None, polygons=None, segment_map=None):
  3. # import imgaug.augmentables.polys as polys
  4. # import imgaug.augmentables.bbs as bbs
  5. #
  6. # # 预处理,将bboxes, polygons转换成imgaug格式
  7. # aug_bboxes = None
  8. # if bboxes is not None:
  9. # aug_bboxes = list()
  10. # for i in range(len(bboxes)):
  11. # x1 = bboxes[i, 0]-1
  12. # y1 = bboxes[i, 1]
  13. # x2 = bboxes[i, 2]
  14. # y2 = bboxes[i, 3]
  15. # aug_bboxes.append(bbs.BoundingBox(x1, y1, x2, y2))
  16. #
  17. # aug_polygons = None
  18. # lod_info = list()
  19. # if polygons is not None:
  20. # aug_polygons = list()
  21. # for i in range(len(polygons)):
  22. # num = len(polygons[i])
  23. # lod_info.append(num)
  24. # for j in range(num):
  25. # points = np.reshape(polygons[i][j], (-1, 2))
  26. # aug_polygons.append(polys.Polygon(points))
  27. #
  28. # aug_segment_map = None
  29. # if segment_map is not None:
  30. # if len(segment_map.shape) == 2:
  31. # h, w = segment_map.shape
  32. # aug_segment_map = np.reshape(segment_map, (1, h, w, 1))
  33. # elif len(segment_map.shape) == 3:
  34. # h, w, c = segment_map.shape
  35. # aug_segment_map = np.reshape(segment_map, (1, h, w, c))
  36. # else:
  37. # raise Exception("Only support 2-dimensions for 3-dimensions for segment_map")
  38. #
  39. # aug_im = im.astype('uint8')
  40. # aug_im, aug_bboxes, aug_polygons, aug_seg_map = augmenter.augment(image=aug_im, bounding_boxes=aug_bboxes,
  41. # polygons=aug_polygons,
  42. # segmentation_maps=aug_segment_map)
  43. #
  44. # if aug_polygons is not None:
  45. # assert len(aug_bboxes) == len(lod_info), "length of bboxes should be equal to the length of polygons"
  46. #
  47. # if aug_bboxes is not None:
  48. # # 裁剪掉在图像之外的bbox和polygon
  49. # for i in range(len(aug_bboxes)):
  50. # aug_bboxes[i] = aug_bboxes[i].clip_out_of_image(aug_im)
  51. # if aug_polygons is not None:
  52. # for i in range(len(aug_polygons)):
  53. # aug_polygons[i] = aug_polygons[i].clip_out_of_image(aug_im)
  54. #
  55. # # 过滤掉无效的bbox和polygon,并转换为训练数据格式
  56. # converted_bboxes = list()
  57. # converted_polygons = list()
  58. # poly_index = 0
  59. # for i in range(len(aug_bboxes)):
  60. # # 过滤width或height不足1像素的框
  61. # if aug_bboxes[i].width < 1 or aug_bboxes[i].height < 1:
  62. # continue
  63. # if aug_polygons is None:
  64. # converted_bboxes.append([aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2, aug_bboxes[i].y2])
  65. # continue
  66. #
  67. # # 如若有polygons,将会继续执行下面代码
  68. # polygons_this_box = list()
  69. # for ps in aug_polygons[poly_index:poly_index+lod_info[i]]:
  70. # if len(ps) == 0:
  71. # continue
  72. # for p in ps:
  73. # # 没有3个point的polygon被过滤
  74. # if len(p.exterior) < 3:
  75. # continue
  76. # polygons_this_box.append(p.exterior.flatten().tolist())
  77. # poly_index += lod_info[i]
  78. #
  79. # if len(polygons_this_box) == 0:
  80. # continue
  81. # converted_bboxes.append([aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2, aug_bboxes[i].y2])
  82. # converted_polygons.append(polygons_this_box)
  83. # if len(converted_bboxes) == 0:
  84. # # 若完成增强后,没有可用的目标,则仍然返回原图进行训练
  85. # aug_im = im
  86. # converted_bboxes = bboxes
  87. # converted_polygons = polygons
  88. #
  89. # result = [aug_im.astype('float32')]
  90. # if bboxes is not None:
  91. # result.append(np.array(converted_bboxes))
  92. # if polygons is not None:
  93. # result.append(converted_polygons)
  94. # if segment_map is not None:
  95. # n, h, w, c = aug_seg_map.shape
  96. # if len(segment_map.shape) == 2:
  97. # aug_seg_map = np.reshape(aug_seg_map, (h, w))
  98. # elif len(segment_map.shape) == 3:
  99. # aug_seg_map = np.reshape(aug_seg_map, (h, w, c))
  100. # result.append(aug_seg_map)
  101. # return result
  102. import numpy as np
  103. def execute_imgaug(augmenter, im, bboxes=None, polygons=None,
  104. segment_map=None):
  105. # 预处理,将bboxes, polygons转换成imgaug格式
  106. import imgaug.augmentables.polys as polys
  107. import imgaug.augmentables.bbs as bbs
  108. aug_im = im.astype('uint8')
  109. aug_bboxes = None
  110. if bboxes is not None:
  111. aug_bboxes = list()
  112. for i in range(len(bboxes)):
  113. x1 = bboxes[i, 0] - 1
  114. y1 = bboxes[i, 1]
  115. x2 = bboxes[i, 2]
  116. y2 = bboxes[i, 3]
  117. aug_bboxes.append(bbs.BoundingBox(x1, y1, x2, y2))
  118. aug_polygons = None
  119. lod_info = list()
  120. if polygons is not None:
  121. aug_polygons = list()
  122. for i in range(len(polygons)):
  123. num = len(polygons[i])
  124. lod_info.append(num)
  125. for j in range(num):
  126. points = np.reshape(polygons[i][j], (-1, 2))
  127. aug_polygons.append(polys.Polygon(points))
  128. aug_segment_map = None
  129. if segment_map is not None:
  130. if len(segment_map.shape) == 2:
  131. h, w = segment_map.shape
  132. aug_segment_map = np.reshape(segment_map, (1, h, w, 1))
  133. elif len(segment_map.shape) == 3:
  134. h, w, c = segment_map.shape
  135. aug_segment_map = np.reshape(segment_map, (1, h, w, c))
  136. else:
  137. raise Exception(
  138. "Only support 2-dimensions for 3-dimensions for segment_map")
  139. aug_im, aug_bboxes, aug_polygons, aug_seg_map = augmenter.augment(
  140. image=aug_im,
  141. bounding_boxes=aug_bboxes,
  142. polygons=aug_polygons,
  143. segmentation_maps=aug_segment_map)
  144. aug_im = aug_im.astype('float32')
  145. if aug_polygons is not None:
  146. assert len(aug_bboxes) == len(
  147. lod_info
  148. ), "Number of aug_bboxes should be equal to number of aug_polygons"
  149. if aug_bboxes is not None:
  150. # 裁剪掉在图像之外的bbox和polygon
  151. for i in range(len(aug_bboxes)):
  152. aug_bboxes[i] = aug_bboxes[i].clip_out_of_image(aug_im)
  153. if aug_polygons is not None:
  154. for i in range(len(aug_polygons)):
  155. aug_polygons[i] = aug_polygons[i].clip_out_of_image(aug_im)
  156. # 过滤掉无效的bbox和polygon,并转换为训练数据格式
  157. converted_bboxes = list()
  158. converted_polygons = list()
  159. poly_index = 0
  160. for i in range(len(aug_bboxes)):
  161. # 过滤width或height不足1像素的框
  162. if aug_bboxes[i].width < 1 or aug_bboxes[i].height < 1:
  163. continue
  164. if aug_polygons is None:
  165. converted_bboxes.append([
  166. aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2,
  167. aug_bboxes[i].y2
  168. ])
  169. continue
  170. # 如若有polygons,将会继续执行下面代码
  171. polygons_this_box = list()
  172. for ps in aug_polygons[poly_index:poly_index + lod_info[i]]:
  173. if len(ps) == 0:
  174. continue
  175. for p in ps:
  176. # 没有3个point的polygon被过滤
  177. if len(p.exterior) < 3:
  178. continue
  179. polygons_this_box.append(p.exterior.flatten().tolist())
  180. poly_index += lod_info[i]
  181. if len(polygons_this_box) == 0:
  182. continue
  183. converted_bboxes.append([
  184. aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2,
  185. aug_bboxes[i].y2
  186. ])
  187. converted_polygons.append(polygons_this_box)
  188. if len(converted_bboxes) == 0:
  189. aug_im = im
  190. converted_bboxes = bboxes
  191. converted_polygons = polygons
  192. result = [aug_im]
  193. if bboxes is not None:
  194. result.append(np.array(converted_bboxes))
  195. if polygons is not None:
  196. result.append(converted_polygons)
  197. if segment_map is not None:
  198. n, h, w, c = aug_seg_map.shape
  199. if len(segment_map.shape) == 2:
  200. aug_seg_map = np.reshape(aug_seg_map, (h, w))
  201. elif len(segment_map.shape) == 3:
  202. aug_seg_map = np.reshape(aug_seg_map, (h, w, c))
  203. result.append(aug_seg_map)
  204. return result