| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220 |
- 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
|