-
Notifications
You must be signed in to change notification settings - Fork 216
/
mixupdetection.py
125 lines (106 loc) · 4.8 KB
/
mixupdetection.py
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
"""Mixup detection dataset wrapper."""
from __future__ import absolute_import
import numpy as np
import torch
#from mxnet.gluon.data import Dataset
from .dataloading import Dataset
class MixupDetection(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. - 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. - 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. - 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. - 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