def __getitem__(self, index) -> dict:
res = super().__getitem__(index)
# 图像Resize至网络输入大小
na_img = res[HK.NATIVE_IMAGE]
img_resize, np_kps_resize, np_boxes_resize = self.rkr.resize(
na_img, res[HK.KEYPOINTS], res[HK.BOXES])
res[HK.RE_IMAGE] = img_resize
res[HK.RE_KEYPOINTS] = np_kps_resize
res[HK.RE_BOXES] = np_boxes_resize
if 'visual_debug' in self.kwargs and self.kwargs.get('visual_debug'):
img_draw = KeypointsOnImage.from_xy_array(res[HK.KEYPOINTS].reshape(-1, 2), shape=na_img.shape) \
.draw_on_image(na_img, size=5)
img_draw = BoundingBoxesOnImage.from_xyxy_array(res[HK.BOXES].reshape(-1, 4), shape=na_img.shape) \
.draw_on_image(img_draw, size=2)
res[HK.DEBUG_NATIVE_IMAGE] = img_draw
img_draw = KeypointsOnImage.from_xy_array(res[HK.RE_KEYPOINTS].reshape(-1, 2), shape=img_resize.shape) \
.draw_on_image(img_resize, size=5)
img_draw = BoundingBoxesOnImage.from_xyxy_array(res[HK.RE_BOXES].reshape(-1, 4), shape=img_resize.shape) \
.draw_on_image(img_draw, size=2)
res[HK.DEBUG_RE_IMAGE] = img_draw
return res
def aug(self, img, pts: np.ndarray, boxes: np.ndarray):
pts_shape = pts.shape
pts = pts.reshape((-1, 2))
boxes_shape = boxes.shape
boxes = boxes.reshape((-1, 4))
kps_on_image = KeypointsOnImage.from_xy_array(pts, shape=img.shape)
boxes_on_img = BoundingBoxesOnImage.from_xyxy_array(boxes,
shape=img.shape)
seq = iaa.Sequential([
iaa.Multiply((0.8, 1.2)), # change brightness
iaa.Affine(
rotate=(-5, 5),
scale=(0.9, 1.05),
translate_percent={
"x": (-0.1, 0.1),
"y": (-0.1, 0.1)
},
),
iaa.GaussianBlur(sigma=(0, 0.7)),
iaa.Sometimes(0.3, iaa.MotionBlur(k=(3, 7)))
])
det = seq.to_deterministic()
img_aug = det.augment_image(img)
kps_aug = det.augment_keypoints(kps_on_image)
boxes_aug = det.augment_bounding_boxes(boxes_on_img)
np_kps_aug = kps_aug.to_xy_array()
np_kps_aug = np_kps_aug.reshape(pts_shape)
np_boxes_aug = boxes_aug.to_xy_array()
np_boxes_aug = np_boxes_aug.reshape(boxes_shape)
return img_aug, np_kps_aug, np_boxes_aug
def resize(self, img, pts: np.ndarray, boxes: np.ndarray):
pts_shape = pts.shape
pts = pts.reshape((-1, 2))
boxes_shape = boxes.shape
boxes = boxes.reshape((-1, 4))
tw, th = self.target_size
ih, iw, ic = img.shape
kps_on_image = KeypointsOnImage.from_xy_array(pts, shape=img.shape)
boxes_on_img = BoundingBoxesOnImage.from_xyxy_array(boxes,
shape=img.shape)
seq = self.__aug_sequence((iw, ih), (tw, th))
det = seq.to_deterministic()
img_aug = det.augment_image(img)
kps_aug = det.augment_keypoints(kps_on_image)
boxes_aug = det.augment_bounding_boxes(boxes_on_img)
np_kps_aug = kps_aug.to_xy_array()
np_kps_aug = np_kps_aug.reshape(pts_shape)
np_boxes_aug = boxes_aug.to_xy_array()
np_boxes_aug = np_boxes_aug.reshape(boxes_shape)
return img_aug, np_kps_aug, np_boxes_aug
def __getitem__(self, index) -> dict:
res = super().__getitem__(index)
# 图像Resize至网络输入大小
img = res[HK.RE_IMAGE]
if self.is_train:
img_aug, np_kps_aug, np_boxes_aug = self.aug.aug(
img, res[HK.RE_KEYPOINTS], res[HK.RE_BOXES])
else:
img_aug, np_kps_aug, np_boxes_aug = img, res[HK.RE_KEYPOINTS], res[
HK.RE_BOXES]
res[HK.AUG_IMAGE] = img_aug
res[HK.AUG_KEYPOINTS] = np_kps_aug
res[HK.AUG_BOXES] = np_boxes_aug
if 'visual_debug' in self.kwargs and self.kwargs.get('visual_debug'):
img_draw = KeypointsOnImage.from_xy_array(res[HK.AUG_KEYPOINTS].reshape(-1, 2), shape=img_aug.shape)\
.draw_on_image(img_aug, size=5)
img_draw = BoundingBoxesOnImage.from_xyxy_array(res[HK.AUG_BOXES].reshape(-1, 4), shape=img_aug.shape)\
.draw_on_image(img_draw, size=2)
res[HK.DEBUG_AUG_IMAGE] = img_draw
return res