|
|
@@ -13,36 +13,41 @@
|
|
|
# limitations under the License.
|
|
|
|
|
|
import numpy as np
|
|
|
+import copy
|
|
|
|
|
|
|
|
|
def execute_imgaug(augmenter, im, bboxes=None, polygons=None,
|
|
|
segment_map=None):
|
|
|
# 预处理,将bboxes, polygons转换成imgaug格式
|
|
|
- import imgaug.augmentables.polys as polys
|
|
|
+ import imgaug.augmentables.kps as kps
|
|
|
import imgaug.augmentables.bbs as bbs
|
|
|
|
|
|
aug_im = im.astype('uint8')
|
|
|
+ aug_im = augmenter.augment(image=aug_im)
|
|
|
+ return aug_im
|
|
|
|
|
|
+ # TODO imgaug的标注处理逻辑与paddlex已存的transform存在部分差异
|
|
|
+ # 目前仅支持对原图进行处理,因此只能使用pixlevel的imgaug增强操作
|
|
|
+ # 以下代码暂不会执行
|
|
|
aug_bboxes = None
|
|
|
if bboxes is not None:
|
|
|
aug_bboxes = list()
|
|
|
for i in range(len(bboxes)):
|
|
|
- x1 = bboxes[i, 0] - 1
|
|
|
+ x1 = bboxes[i, 0]
|
|
|
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()
|
|
|
+ aug_points = None
|
|
|
if polygons is not None:
|
|
|
- aug_polygons = list()
|
|
|
+ aug_points = 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))
|
|
|
+ tmp = np.reshape(polygons[i][j], (-1, 2))
|
|
|
+ for k in range(len(tmp)):
|
|
|
+ aug_points.append(kps.Keypoint(tmp[k, 0], tmp[k, 1]))
|
|
|
|
|
|
aug_segment_map = None
|
|
|
if segment_map is not None:
|
|
|
@@ -56,72 +61,47 @@ def execute_imgaug(augmenter, im, bboxes=None, polygons=None,
|
|
|
raise Exception(
|
|
|
"Only support 2-dimensions for 3-dimensions for segment_map")
|
|
|
|
|
|
- aug_im, aug_bboxes, aug_polygons, aug_seg_map = augmenter.augment(
|
|
|
+ unnormalized_batch = augmenter.augment(
|
|
|
image=aug_im,
|
|
|
bounding_boxes=aug_bboxes,
|
|
|
- polygons=aug_polygons,
|
|
|
- segmentation_maps=aug_segment_map)
|
|
|
+ keypoints=aug_points,
|
|
|
+ segmentation_maps=aug_segment_map,
|
|
|
+ return_batch=True)
|
|
|
+ aug_im = unnormalized_batch.images_aug[0]
|
|
|
+ aug_bboxes = unnormalized_batch.bounding_boxes_aug
|
|
|
+ aug_points = unnormalized_batch.keypoints_aug
|
|
|
+ aug_seg_map = unnormalized_batch.segmentation_maps_aug
|
|
|
|
|
|
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
|
|
|
+ aug_bboxes = converted_bboxes
|
|
|
+
|
|
|
+ aug_polygons = None
|
|
|
+ if aug_points is not None:
|
|
|
+ aug_polygons = copy.deepcopy(polygons)
|
|
|
+ idx = 0
|
|
|
+ for i in range(len(aug_polygons)):
|
|
|
+ num = len(aug_polygons[i])
|
|
|
+ for j in range(num):
|
|
|
+ num_points = len(aug_polygons[i][j]) // 2
|
|
|
+ for k in range(num_points):
|
|
|
+ aug_polygons[i][j][k * 2] = aug_points[idx].x
|
|
|
+ aug_polygons[i][j][k * 2 + 1] = aug_points[idx].y
|
|
|
+ idx += 1
|
|
|
|
|
|
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:
|
|
|
+ if aug_bboxes is not None:
|
|
|
+ result.append(np.array(aug_bboxes))
|
|
|
+ if aug_polygons is not None:
|
|
|
+ result.append(aug_polygons)
|
|
|
+ if aug_seg_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))
|