def __getitem__(self, idx): l_bound = idx * self.batch_size r_bound = (idx + 1) * self.batch_size if r_bound > len(self.annotation): r_bound = len(self.annotation) l_bound = r_bound - self.batch_size self._on_batch_start(idx) batch_image = np.zeros( (r_bound - l_bound, self._image_size, self._image_size, 3), dtype=np.float32) batch_label = [ np.zeros((r_bound - l_bound, size, size, len(mask_per_layer) * (5 + self.num_classes)), dtype=np.float32) for size, mask_per_layer in zip(self._grid_size, self.mask) ] for i, sub_idx in enumerate(range(l_bound, r_bound)): image, bboxes, labels = self._getitem(sub_idx) if self.mosaic: sub_idx = np.random.choice( np.delete(np.arange(self.num_anno), idx), 3, False) image2, bboxes2, labels2 = self._getitem(sub_idx[0]) image3, bboxes3, labels3 = self._getitem(sub_idx[1]) image4, bboxes4, labels4 = self._getitem(sub_idx[2]) image, bboxes, labels = augment.mosic(image, bboxes, labels, image2, bboxes2, labels2, image3, bboxes3, labels3, image4, bboxes4, labels4) if self.normal_method: image = augment.random_distort(image) image = augment.random_grayscale(image) image, bboxes = augment.random_flip_lr(image, bboxes) image, bboxes = augment.random_rotate(image, bboxes) image, bboxes, labels = augment.random_crop_and_zoom( image, bboxes, labels, (self._image_size, self._image_size)) image, bboxes, labels = augment.bbox_filter(image, bboxes, labels) labels = self._preprocess_true_boxes(bboxes, labels) batch_image[i] = image for j in range(len(self.mask)): batch_label[j][i, :, :, :] = labels[j][:, :, :] return batch_image, batch_label
def __getitem__(self, idx): l_bound = idx * self.batch_size r_bound = (idx + 1) * self.batch_size if r_bound > len(self.annotation): r_bound = len(self.annotation) l_bound = r_bound - self.batch_size self._on_batch_start(idx) batch_image = np.zeros( (r_bound - l_bound, self._image_size, self._image_size, 3), dtype=np.float32) batch_label = [ np.zeros((r_bound - l_bound, size, size, len(mask_per_layer) * (6 + self.num_classes)), dtype=np.float32) for size, mask_per_layer in zip(self._grid_size, self.mask) ] for i, sub_idx in enumerate(range(l_bound, r_bound)): image, bboxes, labels = self._getitem(sub_idx) weights = np.full(len(labels), 1.) # high level augment if self.mix_up and np.random.randint(2): sub_idx2 = np.random.choice( np.delete(np.arange(self.num_anno), sub_idx)) image2, bboxes2, labels2 = self._getitem(sub_idx2) image, bboxes, labels, weights = augment.mix_up( image, bboxes, labels, image2, bboxes2, labels2) elif self.cut_mix and np.random.randint(2): sub_idx2 = np.random.choice( np.delete(np.arange(self.num_anno), sub_idx)) image2, bboxes2, labels2 = self._getitem(sub_idx2) image, bboxes, labels, weights = augment.cut_mix( image, bboxes, labels, image2, bboxes2, labels2) elif self.mosaic and np.random.randint(2): sub_idx = np.random.choice( np.delete(np.arange(self.num_anno), idx), 3, False) image2, bboxes2, labels2 = self._getitem(sub_idx[0]) image3, bboxes3, labels3 = self._getitem(sub_idx[1]) image4, bboxes4, labels4 = self._getitem(sub_idx[2]) image, bboxes, labels, weights = augment.mosic( image, bboxes, labels, image2, bboxes2, labels2, image3, bboxes3, labels3, image4, bboxes4, labels4) bboxes = np.divide(bboxes, self._image_size) bboxes[..., [0, 2]] = np.clip(bboxes[..., [0, 2]], 0., 1.) bboxes[..., [1, 3]] = np.clip(bboxes[..., [1, 3]], 0., 1.) labels = self._preprocess_true_boxes(bboxes, labels, weights) batch_image[i] = image for j in range(len(self.mask)): batch_label[j][i, :, :, :] = labels[j][:, :, :] return batch_image, batch_label