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