def __get_default_coco(self, img, anns, num_objs): boxes = [] if self.num_keypoints > 0: kpts = [] for k in range(num_objs): ann = anns[k] bbox = self._coco_box_to_bbox(ann['bbox']) boxes.append(BoundingBox(*bbox)) if self.num_keypoints > 0: if 'keypoints' not in ann: ann['keypoints'] = np.zeros((3 * self.num_keypoints, )) kpt = [ Keypoint(*x) for x in np.array(ann['keypoints']).reshape(-1, 3)[:, :2] ] kpts.extend(kpt) bbs = BoundingBoxesOnImage(boxes, shape=img.shape) if self.num_keypoints > 0: kpts = KeypointsOnImage(kpts, shape=img.shape) if self.augmentation is not None: if self.num_keypoints > 0: img_aug, bbs_aug, kpts_aug = self.augmentation( image=img, bounding_boxes=bbs, keypoints=kpts) else: img_aug, bbs_aug = self.augmentation(image=img, bounding_boxes=bbs) else: if self.num_keypoints > 0: kpts_aug = kpts.copy() img_aug, bbs_aug = np.copy(img), bbs.copy() if self.num_keypoints > 0: img_aug, bbs_aug, kpts_aug = self.resize(image=img_aug, bounding_boxes=bbs_aug, keypoints=kpts_aug) else: img_aug, bbs_aug = self.resize(image=img_aug, bounding_boxes=bbs_aug) img = (img_aug.astype(np.float32) / 255.) inp = (img - self.mean) / self.std inp = inp.transpose(2, 0, 1) output_h = self.input_size[1] // self.down_ratio output_w = self.input_size[0] // self.down_ratio num_classes = self.num_classes hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32) wh = np.zeros((self.max_detections, 2), dtype=np.float32) reg = np.zeros((self.max_detections, 2), dtype=np.float32) ind = np.zeros((self.max_detections), dtype=np.int64) reg_mask = np.zeros((self.max_detections), dtype=np.uint8) gt_det = np.zeros((self.max_detections, 6), dtype=np.float32) gt_areas = np.zeros((self.max_detections), dtype=np.float32) if self.num_keypoints > 0: kp = np.zeros((self.max_detections, self.num_keypoints * 2), dtype=np.float32) gt_kp = np.zeros((self.max_detections, self.num_keypoints, 2), dtype=np.float32) kp_reg_mask = np.zeros( (self.max_detections, self.num_keypoints * 2), dtype=np.uint8) bbs_aug, kpts_aug = self.resize_out(bounding_boxes=bbs_aug, keypoints=kpts_aug) else: bbs_aug = self.resize_out(bounding_boxes=bbs_aug) for k in range(num_objs): ann = anns[k] bbox_aug = bbs_aug[k].clip_out_of_image((output_w, output_h)) bbox = np.array( [bbox_aug.x1, bbox_aug.y1, bbox_aug.x2, bbox_aug.y2]) cls_id = int(self.cat_mapping[ann['category_id']]) bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, output_w - 1) bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, output_h - 1) h, w = bbox[3] - bbox[1], bbox[2] - bbox[0] if h > 0 and w > 0: radius = gaussian_radius((np.ceil(h), np.ceil(w))) radius = max(0, int(radius)) ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32) ct_int = ct.astype(np.int32) draw_gaussian(hm[cls_id], ct_int, radius) wh[k] = 1. * w, 1. * h ind[k] = ct_int[1] * output_w + ct_int[0] reg[k] = ct - ct_int reg_mask[k] = 1 gt_det[k] = ([ ct[0] - w / 2, ct[1] - h / 2, ct[0] + w / 2, ct[1] + h / 2, 1, cls_id ]) if self.num_keypoints > 0: valid = np.array(ann["keypoints"]).reshape(-1, 3)[:, -1] for i, p in enumerate( kpts_aug[k * self.num_keypoints:k * self.num_keypoints + self.num_keypoints]): kp[k][i * 2] = p.x - ct_int[0] kp[k][i * 2 + 1] = p.y - ct_int[1] is_valid = valid[i] == 2 and not p.is_out_of_image( (output_w, output_w)) kp_reg_mask[k, i * 2] = int(is_valid) kp_reg_mask[k, i * 2 + 1] = int(is_valid) gt_kp[k][i] = p.x, p.y if "area" not in ann: gt_areas[k] = w * h else: gt_areas[k] = ann["area"] del bbs del bbs_aug del img_aug gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else np.zeros( (1, 6), dtype=np.float32) ret = { 'input': inp, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh, 'reg': reg, 'gt_dets': gt_det, 'gt_areas': gt_areas, } if self.num_keypoints > 0: ret['kps'] = kp ret['gt_kps'] = gt_kp ret['kp_reg_mask'] = kp_reg_mask del kpts_aug return ret
def __get_rotated_coco(self, img, anns, num_objs): kpts = [] kpts_tmp = [] for k in range(num_objs): ann = anns[k] ann_rotated = get_annotation_with_angle(ann) ann_rotated[4] = ann_rotated[4] rot = rotate_bbox(*ann_rotated) kpts.extend([Keypoint(*x) for x in rot]) if self.num_keypoints > 0: if 'keypoints' not in ann: ann['keypoints'] = np.zeros((3 * self.num_keypoints, )) kpt = [ Keypoint(*x) for x in np.array(ann['keypoints']).reshape(-1, 3)[:, :2] ] kpts_tmp.extend(kpt) idx_boxes = len(kpts) if self.num_keypoints > 0: kpts.extend(kpts_tmp) kpts = KeypointsOnImage(kpts, shape=img.shape) if self.augmentation is not None: img_aug, kpts_aug = self.augmentation(image=img, keypoints=kpts) else: img_aug, kpts_aug = np.copy(img), kpts.copy() img_aug, kpts_aug = self.resize(image=img_aug, keypoints=kpts_aug) img = (img_aug.astype(np.float32) / 255.) inp = (img - self.mean) / self.std inp = inp.transpose(2, 0, 1) output_h = self.input_size[1] // self.down_ratio output_w = self.input_size[0] // self.down_ratio num_classes = self.num_classes hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32) wh = np.zeros((self.max_detections, 3), dtype=np.float32) reg = np.zeros((self.max_detections, 2), dtype=np.float32) ind = np.zeros((self.max_detections), dtype=np.int64) reg_mask = np.zeros((self.max_detections), dtype=np.uint8) gt_det = np.zeros( (self.max_detections, 7 if self.use_rotated_boxes else 6), dtype=np.float32) gt_areas = np.zeros((self.max_detections), dtype=np.float32) if self.num_keypoints > 0: kp = np.zeros((self.max_detections, self.num_keypoints * 2), dtype=np.float32) gt_kp = np.zeros((self.max_detections, self.num_keypoints, 2), dtype=np.float32) kp_reg_mask = np.zeros( (self.max_detections, self.num_keypoints * 2), dtype=np.uint8) kpts_aug = self.resize_out(keypoints=kpts_aug) box_kpts_aug, kpts_aug = kpts_aug[:idx_boxes], kpts_aug[idx_boxes:] assert num_objs == len(box_kpts_aug) // 4 for k in range(num_objs): ann = anns[k] points = [] for p in box_kpts_aug[k * 4:k * 4 + 4]: box_kp = list( (np.clip(p.x, 0, output_w - 1), np.clip(p.y, 0, output_h - 1))) points.append(box_kp) points = np.array(points).astype(np.float32) cv_ct, cv_wh, cv_angle = cv2.minAreaRect(points) if cv_wh[0] == 0 or cv_wh[1] == 0: continue cx, cy, w, h, angle = get_annotation_with_angle({ 'rbbox': np.array([cv_ct[0], cv_ct[1], cv_wh[0], cv_wh[1], cv_angle]) }) ct = np.array((cx, cy)) cls_id = int(self.cat_mapping[ann['category_id']]) if h > 0 and w > 0: radius = gaussian_radius((np.ceil(h), np.ceil(w))) radius = max(0, int(radius)) ct_int = ct.astype(np.int32) draw_gaussian(hm[cls_id], ct_int, radius) wh[k] = w, h, angle ind[k] = ct_int[1] * output_w + ct_int[0] reg[k] = ct - ct_int reg_mask[k] = 1 gt_det[k] = ([ct[0], ct[1], w, h, angle, 1, cls_id]) if self.num_keypoints > 0: valid = np.array(ann["keypoints"]).reshape(-1, 3)[:, -1] for i, p in enumerate( kpts_aug[k * self.num_keypoints:k * self.num_keypoints + self.num_keypoints]): kp[k][i * 2] = p.x - ct_int[0] kp[k][i * 2 + 1] = p.y - ct_int[1] is_valid = valid[i] == 2 and not p.is_out_of_image( (output_w, output_w)) kp_reg_mask[k, i * 2] = int(is_valid) kp_reg_mask[k, i * 2 + 1] = int(is_valid) gt_kp[k][i] = p.x, p.y if "area" not in ann: gt_areas[k] = w * h else: gt_areas[k] = ann["area"] del box_kpts_aug del img_aug gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else np.zeros( (1, 7), dtype=np.float32) ret = { 'input': inp, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh, 'reg': reg, 'gt_dets': gt_det, 'gt_areas': gt_areas, } if self.num_keypoints > 0: ret['kps'] = kp ret['gt_kps'] = gt_kp ret['kp_reg_mask'] = kp_reg_mask del kpts_aug return ret
def __get_rotated_coco(self, img, anns, num_objs): kpts = [] for k in range(num_objs): ann = get_annotation_with_angle(anns[k]) ann[4] = np.radians(ann[4]) rot = rotate_bbox(*ann) kpts.extend([Keypoint(*x) for x in rot]) kpts = KeypointsOnImage(kpts, shape=img.shape) if self.augmentation is not None: img_aug, kpts_aug = self.augmentation(image=img, keypoints=kpts) else: img_aug, kpts_aug = np.copy(img), kpts.copy() img_aug, kpts_aug = self.resize(image=img_aug, keypoints=kpts_aug) img = (img_aug.astype(np.float32) / 255.) inp = (img - self.mean) / self.std inp = inp.transpose(2, 0, 1) output_h = self.input_size[1] // self.down_ratio output_w = self.input_size[0] // self.down_ratio num_classes = self.num_classes hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32) wh = np.zeros((self.max_detections, 3), dtype=np.float32) dense_wh = np.zeros((2, output_h, output_w), dtype=np.float32) reg = np.zeros((self.max_detections, 2), dtype=np.float32) ind = np.zeros((self.max_detections), dtype=np.int64) reg_mask = np.zeros((self.max_detections), dtype=np.uint8) gt_det = np.zeros( (self.max_detections, 7 if self.use_rotated_boxes else 6), dtype=np.float32) gt_areas = np.zeros((self.max_detections), dtype=np.float32) kpts_aug = self.resize_out(keypoints=kpts_aug) assert num_objs == len(kpts_aug) // 4 for k in range(num_objs): ann = anns[k] points = [] for p in kpts_aug[k * 4:k * 4 + 4]: kp = list( (np.clip(p.x, 0, output_w - 1), np.clip(p.y, 0, output_h - 1))) points.append(kp) points = np.array(points).astype(np.float32) cv_ct, cv_wh, cv_angle = cv2.minAreaRect(points) if cv_wh[0] == 0 or cv_wh[1] == 0: continue cx, cy, w, h, angle = get_annotation_with_angle({ 'rbbox': np.array([cv_ct[0], cv_ct[1], cv_wh[0], cv_wh[1], cv_angle]) }) ct = np.array((cx, cy)) cls_id = int(self.cat_mapping[ann['category_id']]) if h > 0 and w > 0: angle = np.radians(angle) radius = gaussian_radius((np.ceil(h), np.ceil(w))) radius = max(0, int(radius)) ct_int = ct.astype(np.int32) draw_gaussian(hm[cls_id], ct_int, radius) wh[k] = w, h, angle ind[k] = ct_int[1] * output_w + ct_int[0] reg[k] = ct - ct_int reg_mask[k] = 1 gt_det[k] = ([ct[0], ct[1], w, h, angle, 1, cls_id]) if "area" not in ann: gt_areas[k] = w * h else: gt_areas[k] = ann["area"] del kpts del kpts_aug del img_aug gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else np.zeros( (1, 7), dtype=np.float32) ret = { 'input': inp, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh, 'reg': reg, 'gt_dets': gt_det, 'gt_areas': gt_areas, } return ret