|
|
@@ -0,0 +1,220 @@
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+#def execute_imgaug(augmenter, im, bboxes=None, polygons=None, segment_map=None):
|
|
|
+# import imgaug.augmentables.polys as polys
|
|
|
+# import imgaug.augmentables.bbs as bbs
|
|
|
+#
|
|
|
+# # 预处理,将bboxes, polygons转换成imgaug格式
|
|
|
+# aug_bboxes = None
|
|
|
+# if bboxes is not None:
|
|
|
+# aug_bboxes = list()
|
|
|
+# for i in range(len(bboxes)):
|
|
|
+# x1 = bboxes[i, 0]-1
|
|
|
+# y1 = bboxes[i, 1]
|
|
|
+# x2 = bboxes[i, 2]
|
|
|
+# y2 = bboxes[i, 3]
|
|
|
+# aug_bboxes.append(bbs.BoundingBox(x1, y1, x2, y2))
|
|
|
+#
|
|
|
+# aug_polygons = None
|
|
|
+# lod_info = list()
|
|
|
+# if polygons is not None:
|
|
|
+# aug_polygons = list()
|
|
|
+# for i in range(len(polygons)):
|
|
|
+# num = len(polygons[i])
|
|
|
+# lod_info.append(num)
|
|
|
+# for j in range(num):
|
|
|
+# points = np.reshape(polygons[i][j], (-1, 2))
|
|
|
+# aug_polygons.append(polys.Polygon(points))
|
|
|
+#
|
|
|
+# aug_segment_map = None
|
|
|
+# if segment_map is not None:
|
|
|
+# if len(segment_map.shape) == 2:
|
|
|
+# h, w = segment_map.shape
|
|
|
+# aug_segment_map = np.reshape(segment_map, (1, h, w, 1))
|
|
|
+# elif len(segment_map.shape) == 3:
|
|
|
+# h, w, c = segment_map.shape
|
|
|
+# aug_segment_map = np.reshape(segment_map, (1, h, w, c))
|
|
|
+# else:
|
|
|
+# raise Exception("Only support 2-dimensions for 3-dimensions for segment_map")
|
|
|
+#
|
|
|
+# aug_im = im.astype('uint8')
|
|
|
+# aug_im, aug_bboxes, aug_polygons, aug_seg_map = augmenter.augment(image=aug_im, bounding_boxes=aug_bboxes,
|
|
|
+# polygons=aug_polygons,
|
|
|
+# segmentation_maps=aug_segment_map)
|
|
|
+#
|
|
|
+# if aug_polygons is not None:
|
|
|
+# assert len(aug_bboxes) == len(lod_info), "length of bboxes should be equal to the length of polygons"
|
|
|
+#
|
|
|
+# if aug_bboxes is not None:
|
|
|
+# # 裁剪掉在图像之外的bbox和polygon
|
|
|
+# for i in range(len(aug_bboxes)):
|
|
|
+# aug_bboxes[i] = aug_bboxes[i].clip_out_of_image(aug_im)
|
|
|
+# if aug_polygons is not None:
|
|
|
+# for i in range(len(aug_polygons)):
|
|
|
+# aug_polygons[i] = aug_polygons[i].clip_out_of_image(aug_im)
|
|
|
+#
|
|
|
+# # 过滤掉无效的bbox和polygon,并转换为训练数据格式
|
|
|
+# converted_bboxes = list()
|
|
|
+# converted_polygons = list()
|
|
|
+# poly_index = 0
|
|
|
+# for i in range(len(aug_bboxes)):
|
|
|
+# # 过滤width或height不足1像素的框
|
|
|
+# if aug_bboxes[i].width < 1 or aug_bboxes[i].height < 1:
|
|
|
+# continue
|
|
|
+# if aug_polygons is None:
|
|
|
+# converted_bboxes.append([aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2, aug_bboxes[i].y2])
|
|
|
+# continue
|
|
|
+#
|
|
|
+# # 如若有polygons,将会继续执行下面代码
|
|
|
+# polygons_this_box = list()
|
|
|
+# for ps in aug_polygons[poly_index:poly_index+lod_info[i]]:
|
|
|
+# if len(ps) == 0:
|
|
|
+# continue
|
|
|
+# for p in ps:
|
|
|
+# # 没有3个point的polygon被过滤
|
|
|
+# if len(p.exterior) < 3:
|
|
|
+# continue
|
|
|
+# polygons_this_box.append(p.exterior.flatten().tolist())
|
|
|
+# poly_index += lod_info[i]
|
|
|
+#
|
|
|
+# if len(polygons_this_box) == 0:
|
|
|
+# continue
|
|
|
+# converted_bboxes.append([aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2, aug_bboxes[i].y2])
|
|
|
+# converted_polygons.append(polygons_this_box)
|
|
|
+# if len(converted_bboxes) == 0:
|
|
|
+# # 若完成增强后,没有可用的目标,则仍然返回原图进行训练
|
|
|
+# aug_im = im
|
|
|
+# converted_bboxes = bboxes
|
|
|
+# converted_polygons = polygons
|
|
|
+#
|
|
|
+# result = [aug_im.astype('float32')]
|
|
|
+# if bboxes is not None:
|
|
|
+# result.append(np.array(converted_bboxes))
|
|
|
+# if polygons is not None:
|
|
|
+# result.append(converted_polygons)
|
|
|
+# if segment_map is not None:
|
|
|
+# n, h, w, c = aug_seg_map.shape
|
|
|
+# if len(segment_map.shape) == 2:
|
|
|
+# aug_seg_map = np.reshape(aug_seg_map, (h, w))
|
|
|
+# elif len(segment_map.shape) == 3:
|
|
|
+# aug_seg_map = np.reshape(aug_seg_map, (h, w, c))
|
|
|
+# result.append(aug_seg_map)
|
|
|
+# return result
|
|
|
+
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+
|
|
|
+def execute_imgaug(augmenter, im, bboxes=None, polygons=None,
|
|
|
+ segment_map=None):
|
|
|
+ # 预处理,将bboxes, polygons转换成imgaug格式
|
|
|
+ import imgaug.augmentables.polys as polys
|
|
|
+ import imgaug.augmentables.bbs as bbs
|
|
|
+
|
|
|
+ aug_im = im.astype('uint8')
|
|
|
+
|
|
|
+ aug_bboxes = None
|
|
|
+ if bboxes is not None:
|
|
|
+ aug_bboxes = list()
|
|
|
+ for i in range(len(bboxes)):
|
|
|
+ x1 = bboxes[i, 0] - 1
|
|
|
+ y1 = bboxes[i, 1]
|
|
|
+ x2 = bboxes[i, 2]
|
|
|
+ y2 = bboxes[i, 3]
|
|
|
+ aug_bboxes.append(bbs.BoundingBox(x1, y1, x2, y2))
|
|
|
+
|
|
|
+ aug_polygons = None
|
|
|
+ lod_info = list()
|
|
|
+ if polygons is not None:
|
|
|
+ aug_polygons = list()
|
|
|
+ for i in range(len(polygons)):
|
|
|
+ num = len(polygons[i])
|
|
|
+ lod_info.append(num)
|
|
|
+ for j in range(num):
|
|
|
+ points = np.reshape(polygons[i][j], (-1, 2))
|
|
|
+ aug_polygons.append(polys.Polygon(points))
|
|
|
+
|
|
|
+ aug_segment_map = None
|
|
|
+ if segment_map is not None:
|
|
|
+ if len(segment_map.shape) == 2:
|
|
|
+ h, w = segment_map.shape
|
|
|
+ aug_segment_map = np.reshape(segment_map, (1, h, w, 1))
|
|
|
+ elif len(segment_map.shape) == 3:
|
|
|
+ h, w, c = segment_map.shape
|
|
|
+ aug_segment_map = np.reshape(segment_map, (1, h, w, c))
|
|
|
+ else:
|
|
|
+ raise Exception(
|
|
|
+ "Only support 2-dimensions for 3-dimensions for segment_map")
|
|
|
+
|
|
|
+ aug_im, aug_bboxes, aug_polygons, aug_seg_map = augmenter.augment(
|
|
|
+ image=aug_im,
|
|
|
+ bounding_boxes=aug_bboxes,
|
|
|
+ polygons=aug_polygons,
|
|
|
+ segmentation_maps=aug_segment_map)
|
|
|
+
|
|
|
+ aug_im = aug_im.astype('float32')
|
|
|
+
|
|
|
+ if aug_polygons is not None:
|
|
|
+ assert len(aug_bboxes) == len(
|
|
|
+ lod_info
|
|
|
+ ), "Number of aug_bboxes should be equal to number of aug_polygons"
|
|
|
+
|
|
|
+ if aug_bboxes is not None:
|
|
|
+ # 裁剪掉在图像之外的bbox和polygon
|
|
|
+ for i in range(len(aug_bboxes)):
|
|
|
+ aug_bboxes[i] = aug_bboxes[i].clip_out_of_image(aug_im)
|
|
|
+ if aug_polygons is not None:
|
|
|
+ for i in range(len(aug_polygons)):
|
|
|
+ aug_polygons[i] = aug_polygons[i].clip_out_of_image(aug_im)
|
|
|
+
|
|
|
+ # 过滤掉无效的bbox和polygon,并转换为训练数据格式
|
|
|
+ converted_bboxes = list()
|
|
|
+ converted_polygons = list()
|
|
|
+ poly_index = 0
|
|
|
+ for i in range(len(aug_bboxes)):
|
|
|
+ # 过滤width或height不足1像素的框
|
|
|
+ if aug_bboxes[i].width < 1 or aug_bboxes[i].height < 1:
|
|
|
+ continue
|
|
|
+ if aug_polygons is None:
|
|
|
+ converted_bboxes.append([
|
|
|
+ aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2,
|
|
|
+ aug_bboxes[i].y2
|
|
|
+ ])
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 如若有polygons,将会继续执行下面代码
|
|
|
+ polygons_this_box = list()
|
|
|
+ for ps in aug_polygons[poly_index:poly_index + lod_info[i]]:
|
|
|
+ if len(ps) == 0:
|
|
|
+ continue
|
|
|
+ for p in ps:
|
|
|
+ # 没有3个point的polygon被过滤
|
|
|
+ if len(p.exterior) < 3:
|
|
|
+ continue
|
|
|
+ polygons_this_box.append(p.exterior.flatten().tolist())
|
|
|
+ poly_index += lod_info[i]
|
|
|
+
|
|
|
+ if len(polygons_this_box) == 0:
|
|
|
+ continue
|
|
|
+ converted_bboxes.append([
|
|
|
+ aug_bboxes[i].x1, aug_bboxes[i].y1, aug_bboxes[i].x2,
|
|
|
+ aug_bboxes[i].y2
|
|
|
+ ])
|
|
|
+ converted_polygons.append(polygons_this_box)
|
|
|
+ if len(converted_bboxes) == 0:
|
|
|
+ aug_im = im
|
|
|
+ converted_bboxes = bboxes
|
|
|
+ converted_polygons = polygons
|
|
|
+
|
|
|
+ result = [aug_im]
|
|
|
+ if bboxes is not None:
|
|
|
+ result.append(np.array(converted_bboxes))
|
|
|
+ if polygons is not None:
|
|
|
+ result.append(converted_polygons)
|
|
|
+ if segment_map is not None:
|
|
|
+ n, h, w, c = aug_seg_map.shape
|
|
|
+ if len(segment_map.shape) == 2:
|
|
|
+ aug_seg_map = np.reshape(aug_seg_map, (h, w))
|
|
|
+ elif len(segment_map.shape) == 3:
|
|
|
+ aug_seg_map = np.reshape(aug_seg_map, (h, w, c))
|
|
|
+ result.append(aug_seg_map)
|
|
|
+ return result
|