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, sub_idx):
path, bboxes, labels = self.annotation[sub_idx]
image = read_image(path)
bboxes, labels = np.array(bboxes), np.array(labels)
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 = preprocess_image(image, (self._image_size, self._image_size), bboxes)
labels = augment.onehot(labels, self.num_classes, self.label_smoothing)
return image, bboxes, labels
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)
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, weights = 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, weights = augment.random_crop_and_zoom(image, bboxes, labels, weights,
# (self._image_size, self._image_size))
image, bboxes, labels, weights = augment.bbox_filter(
image, bboxes, labels, weights)
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