| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131 |
- # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- 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
|