Spaces:
Sleeping
Sleeping
File size: 5,024 Bytes
749745d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
"""Mixup detection dataset wrapper."""
from __future__ import absolute_import
import numpy as np
import torch
import torch.utils.data as data
class MixupDetection(data.Dataset):
"""Detection dataset wrapper that performs mixup for normal dataset.
Parameters
----------
dataset : mx.gluon.data.Dataset
Gluon dataset object.
mixup : callable random generator, e.g. np.random.uniform
A random mixup ratio sampler, preferably a random generator from numpy.random
A random float will be sampled each time with mixup(*args).
Use None to disable.
*args : list
Additional arguments for mixup random sampler.
"""
def __init__(self, dataset, mixup=None, preproc=None, *args):
super().__init__(dataset.input_dim)
self._dataset = dataset
self.preproc = preproc
self._mixup = mixup
self._mixup_args = args
def set_mixup(self, mixup=None, *args):
"""Set mixup random sampler, use None to disable.
Parameters
----------
mixup : callable random generator, e.g. np.random.uniform
A random mixup ratio sampler, preferably a random generator from numpy.random
A random float will be sampled each time with mixup(*args)
*args : list
Additional arguments for mixup random sampler.
"""
self._mixup = mixup
self._mixup_args = args
def __len__(self):
return len(self._dataset)
@Dataset.resize_getitem
def __getitem__(self, idx):
self._dataset._input_dim = self.input_dim
# first image
img1, label1, _, _ = self._dataset.pull_item(idx)
lambd = 1
# draw a random lambda ratio from distribution
if self._mixup is not None:
lambd = max(0, min(1, self._mixup(*self._mixup_args)))
if lambd >= 1:
weights1 = np.ones((label1.shape[0], 1))
label1 = np.hstack((label1, weights1))
height, width, _ = img1.shape
img_info = (width, height)
if self.preproc is not None:
img_o, target_o = self.preproc(img1, label1, self.input_dim)
return img_o, target_o, img_info, idx
# second image
idx2 = int(np.random.choice(np.delete(np.arange(len(self)), idx)))
img2, label2, _, _ = self._dataset.pull_item(idx2)
# mixup two images
height = max(img1.shape[0], img2.shape[0])
width = max(img1.shape[1], img2.shape[1])
mix_img = np.zeros((height, width, 3), dtype=np.float32)
mix_img[: img1.shape[0], : img1.shape[1], :] = img1.astype(np.float32) * lambd
mix_img[: img2.shape[0], : img2.shape[1], :] += img2.astype(np.float32) * (1.0 - lambd)
mix_img = mix_img.astype(np.uint8)
y1 = np.hstack((label1, np.full((label1.shape[0], 1), lambd)))
y2 = np.hstack((label2, np.full((label2.shape[0], 1), 1.0 - lambd)))
mix_label = np.vstack((y1, y2))
if self.preproc is not None:
mix_img, padded_labels = self.preproc(mix_img, mix_label, self.input_dim)
img_info = (width, height)
return mix_img, padded_labels, img_info, idx
def pull_item(self, idx):
self._dataset._input_dim = self.input_dim
# first image
img1, label1, _, _ = self._dataset.pull_item(idx)
lambd = 1
# draw a random lambda ratio from distribution
if self._mixup is not None:
lambd = max(0, min(1, self._mixup(*self._mixup_args)))
if lambd >= 1:
weights1 = np.ones((label1.shape[0], 1))
label1 = np.hstack((label1, weights1))
height, width, _ = img1.shape
img_info = (width, height)
if self.preproc is not None:
img_o, target_o = self.preproc(img1, label1, self.input_dim)
return img_o, target_o, img_info, idx
# second image
idx2 = int(np.random.choice(np.delete(np.arange(len(self)), idx)))
img2, label2 = self._dataset.pull_item(idx2)
# mixup two images
height = max(img1.shape[0], img2.shape[0])
width = max(img1.shape[1], img2.shape[1])
mix_img = np.zeros((height, width, 3), dtype=np.float32)
mix_img[: img1.shape[0], : img1.shape[1], :] = img1.astype(np.float32) * lambd
mix_img[: img2.shape[0], : img2.shape[1], :] += img2.astype(np.float32) * (1.0 - lambd)
mix_img = mix_img.astype(np.uint8)
y1 = np.hstack((label1, np.full((label1.shape[0], 1), lambd)))
y2 = np.hstack((label2, np.full((label2.shape[0], 1), 1.0 - lambd)))
mix_label = np.vstack((y1, y2))
if self.preproc is not None:
mix_img, padded_labels = self.preproc(mix_img, mix_label, self.input_dim)
img_info = (width, height)
return mix_img, padded_labels, img_info, idx
|