mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-20 08:25:14 +00:00
refine repo structure
This commit is contained in:
nttstar
@@ -18,97 +18,103 @@ from PIL import Image
|
||||
from config import config
|
||||
from skimage import transform as tf
|
||||
|
||||
|
||||
class FaceSegIter(DataIter):
|
||||
def __init__(self, batch_size,
|
||||
per_batch_size = 0,
|
||||
path_imgrec = None,
|
||||
aug_level = 0,
|
||||
force_mirror = False,
|
||||
exf = 1,
|
||||
use_coherent = 0,
|
||||
args = None,
|
||||
data_name = "data",
|
||||
label_name = "softmax_label"):
|
||||
self.aug_level = aug_level
|
||||
self.force_mirror = force_mirror
|
||||
self.use_coherent = use_coherent
|
||||
self.exf = exf
|
||||
self.batch_size = batch_size
|
||||
self.per_batch_size = per_batch_size
|
||||
self.data_name = data_name
|
||||
self.label_name = label_name
|
||||
assert path_imgrec
|
||||
logging.info('loading recordio %s...',
|
||||
path_imgrec)
|
||||
path_imgidx = path_imgrec[0:-4]+".idx"
|
||||
self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
|
||||
self.oseq = list(self.imgrec.keys)
|
||||
print('train size', len(self.oseq))
|
||||
self.cur = 0
|
||||
self.reset()
|
||||
self.data_shape = (3, config.input_img_size, config.input_img_size)
|
||||
self.num_classes = config.num_classes
|
||||
self.input_img_size = config.input_img_size
|
||||
#self.label_classes = self.num_classes
|
||||
if config.losstype=='heatmap':
|
||||
if aug_level>0:
|
||||
self.output_label_size = config.output_label_size
|
||||
self.label_shape = (self.num_classes, self.output_label_size, self.output_label_size)
|
||||
def __init__(self,
|
||||
batch_size,
|
||||
per_batch_size=0,
|
||||
path_imgrec=None,
|
||||
aug_level=0,
|
||||
force_mirror=False,
|
||||
exf=1,
|
||||
use_coherent=0,
|
||||
args=None,
|
||||
data_name="data",
|
||||
label_name="softmax_label"):
|
||||
self.aug_level = aug_level
|
||||
self.force_mirror = force_mirror
|
||||
self.use_coherent = use_coherent
|
||||
self.exf = exf
|
||||
self.batch_size = batch_size
|
||||
self.per_batch_size = per_batch_size
|
||||
self.data_name = data_name
|
||||
self.label_name = label_name
|
||||
assert path_imgrec
|
||||
logging.info('loading recordio %s...', path_imgrec)
|
||||
path_imgidx = path_imgrec[0:-4] + ".idx"
|
||||
self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec,
|
||||
'r') # pylint: disable=redefined-variable-type
|
||||
self.oseq = list(self.imgrec.keys)
|
||||
print('train size', len(self.oseq))
|
||||
self.cur = 0
|
||||
self.reset()
|
||||
self.data_shape = (3, config.input_img_size, config.input_img_size)
|
||||
self.num_classes = config.num_classes
|
||||
self.input_img_size = config.input_img_size
|
||||
#self.label_classes = self.num_classes
|
||||
if config.losstype == 'heatmap':
|
||||
if aug_level > 0:
|
||||
self.output_label_size = config.output_label_size
|
||||
self.label_shape = (self.num_classes, self.output_label_size,
|
||||
self.output_label_size)
|
||||
else:
|
||||
self.output_label_size = self.input_img_size
|
||||
#self.label_shape = (self.num_classes, 2)
|
||||
self.label_shape = (self.num_classes, self.output_label_size,
|
||||
self.output_label_size)
|
||||
else:
|
||||
self.output_label_size = self.input_img_size
|
||||
#self.label_shape = (self.num_classes, 2)
|
||||
self.label_shape = (self.num_classes, self.output_label_size, self.output_label_size)
|
||||
else:
|
||||
if aug_level>0:
|
||||
self.output_label_size = config.output_label_size
|
||||
self.label_shape = (self.num_classes, 2)
|
||||
else:
|
||||
self.output_label_size = self.input_img_size
|
||||
#self.label_shape = (self.num_classes, 2)
|
||||
self.label_shape = (self.num_classes, 2)
|
||||
self.provide_data = [(data_name, (batch_size,) + self.data_shape)]
|
||||
self.provide_label = [(label_name, (batch_size,) + self.label_shape)]
|
||||
self.img_num = 0
|
||||
self.invalid_num = 0
|
||||
self.mode = 1
|
||||
self.vis = 0
|
||||
self.stats = [0,0]
|
||||
self.flip_order = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
|
||||
26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, 34, 33, 32, 31,
|
||||
45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50, 49, 48,
|
||||
59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66, 65]
|
||||
#self.mirror_set = [
|
||||
# (22,23),
|
||||
# (21,24),
|
||||
# (20,25),
|
||||
# (19,26),
|
||||
# (18,27),
|
||||
# (40,43),
|
||||
# (39,44),
|
||||
# (38,45),
|
||||
# (37,46),
|
||||
# (42,47),
|
||||
# (41,48),
|
||||
# (33,35),
|
||||
# (32,36),
|
||||
# (51,53),
|
||||
# (50,54),
|
||||
# (62,64),
|
||||
# (61,65),
|
||||
# (49,55),
|
||||
# (49,55),
|
||||
# (68,66),
|
||||
# (60,56),
|
||||
# (59,57),
|
||||
# (1,17),
|
||||
# (2,16),
|
||||
# (3,15),
|
||||
# (4,14),
|
||||
# (5,13),
|
||||
# (6,12),
|
||||
# (7,11),
|
||||
# (8,10),
|
||||
# ]
|
||||
if aug_level > 0:
|
||||
self.output_label_size = config.output_label_size
|
||||
self.label_shape = (self.num_classes, 2)
|
||||
else:
|
||||
self.output_label_size = self.input_img_size
|
||||
#self.label_shape = (self.num_classes, 2)
|
||||
self.label_shape = (self.num_classes, 2)
|
||||
self.provide_data = [(data_name, (batch_size, ) + self.data_shape)]
|
||||
self.provide_label = [(label_name, (batch_size, ) + self.label_shape)]
|
||||
self.img_num = 0
|
||||
self.invalid_num = 0
|
||||
self.mode = 1
|
||||
self.vis = 0
|
||||
self.stats = [0, 0]
|
||||
self.flip_order = [
|
||||
16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 26, 25,
|
||||
24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, 34, 33, 32, 31,
|
||||
45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50,
|
||||
49, 48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66, 65
|
||||
]
|
||||
#self.mirror_set = [
|
||||
# (22,23),
|
||||
# (21,24),
|
||||
# (20,25),
|
||||
# (19,26),
|
||||
# (18,27),
|
||||
# (40,43),
|
||||
# (39,44),
|
||||
# (38,45),
|
||||
# (37,46),
|
||||
# (42,47),
|
||||
# (41,48),
|
||||
# (33,35),
|
||||
# (32,36),
|
||||
# (51,53),
|
||||
# (50,54),
|
||||
# (62,64),
|
||||
# (61,65),
|
||||
# (49,55),
|
||||
# (49,55),
|
||||
# (68,66),
|
||||
# (60,56),
|
||||
# (59,57),
|
||||
# (1,17),
|
||||
# (2,16),
|
||||
# (3,15),
|
||||
# (4,14),
|
||||
# (5,13),
|
||||
# (6,12),
|
||||
# (7,11),
|
||||
# (8,10),
|
||||
# ]
|
||||
|
||||
def get_data_shape(self):
|
||||
return self.data_shape
|
||||
@@ -118,181 +124,190 @@ class FaceSegIter(DataIter):
|
||||
|
||||
def get_shape_dict(self):
|
||||
D = {}
|
||||
for (k,v) in self.provide_data:
|
||||
for (k, v) in self.provide_data:
|
||||
D[k] = v
|
||||
for (k,v) in self.provide_label:
|
||||
for (k, v) in self.provide_label:
|
||||
D[k] = v
|
||||
return D
|
||||
|
||||
def get_label_names(self):
|
||||
D = []
|
||||
for (k,v) in self.provide_label:
|
||||
for (k, v) in self.provide_label:
|
||||
D.append(k)
|
||||
return D
|
||||
|
||||
def reset(self):
|
||||
#print('reset')
|
||||
if self.aug_level==0:
|
||||
self.seq = self.oseq
|
||||
else:
|
||||
self.seq = []
|
||||
for _ in range(self.exf):
|
||||
_seq = self.oseq[:]
|
||||
random.shuffle(_seq)
|
||||
self.seq += _seq
|
||||
print('train size after reset', len(self.seq))
|
||||
self.cur = 0
|
||||
#print('reset')
|
||||
if self.aug_level == 0:
|
||||
self.seq = self.oseq
|
||||
else:
|
||||
self.seq = []
|
||||
for _ in range(self.exf):
|
||||
_seq = self.oseq[:]
|
||||
random.shuffle(_seq)
|
||||
self.seq += _seq
|
||||
print('train size after reset', len(self.seq))
|
||||
self.cur = 0
|
||||
|
||||
def next_sample(self):
|
||||
"""Helper function for reading in next sample."""
|
||||
if self.cur >= len(self.seq):
|
||||
raise StopIteration
|
||||
idx = self.seq[self.cur]
|
||||
self.cur += 1
|
||||
s = self.imgrec.read_idx(idx)
|
||||
header, img = recordio.unpack(s)
|
||||
img = mx.image.imdecode(img).asnumpy()
|
||||
hlabel = np.array(header.label).reshape( (self.num_classes,2) )
|
||||
if not config.label_xfirst:
|
||||
hlabel = hlabel[:,::-1] #convert to X/W first
|
||||
annot = {'scale': config.base_scale}
|
||||
"""Helper function for reading in next sample."""
|
||||
if self.cur >= len(self.seq):
|
||||
raise StopIteration
|
||||
idx = self.seq[self.cur]
|
||||
self.cur += 1
|
||||
s = self.imgrec.read_idx(idx)
|
||||
header, img = recordio.unpack(s)
|
||||
img = mx.image.imdecode(img).asnumpy()
|
||||
hlabel = np.array(header.label).reshape((self.num_classes, 2))
|
||||
if not config.label_xfirst:
|
||||
hlabel = hlabel[:, ::-1] #convert to X/W first
|
||||
annot = {'scale': config.base_scale}
|
||||
|
||||
#ul = np.array( (50000,50000), dtype=np.int32)
|
||||
#br = np.array( (0,0), dtype=np.int32)
|
||||
#for i in range(hlabel.shape[0]):
|
||||
# h = int(hlabel[i][0])
|
||||
# w = int(hlabel[i][1])
|
||||
# key = np.array((h,w))
|
||||
# ul = np.minimum(key, ul)
|
||||
# br = np.maximum(key, br)
|
||||
#ul = np.array( (50000,50000), dtype=np.int32)
|
||||
#br = np.array( (0,0), dtype=np.int32)
|
||||
#for i in range(hlabel.shape[0]):
|
||||
# h = int(hlabel[i][0])
|
||||
# w = int(hlabel[i][1])
|
||||
# key = np.array((h,w))
|
||||
# ul = np.minimum(key, ul)
|
||||
# br = np.maximum(key, br)
|
||||
|
||||
return img, hlabel, annot
|
||||
return img, hlabel, annot
|
||||
|
||||
def get_flip(self, data, label):
|
||||
data_flip = np.zeros_like(data)
|
||||
label_flip = np.zeros_like(label)
|
||||
for k in range(data_flip.shape[2]):
|
||||
data_flip[:,:,k] = np.fliplr(data[:,:,k])
|
||||
for k in range(label_flip.shape[0]):
|
||||
label_flip[k,:] = np.fliplr(label[k,:])
|
||||
#print(label[0,:].shape)
|
||||
label_flip = label_flip[self.flip_order,:]
|
||||
return data_flip, label_flip
|
||||
data_flip = np.zeros_like(data)
|
||||
label_flip = np.zeros_like(label)
|
||||
for k in range(data_flip.shape[2]):
|
||||
data_flip[:, :, k] = np.fliplr(data[:, :, k])
|
||||
for k in range(label_flip.shape[0]):
|
||||
label_flip[k, :] = np.fliplr(label[k, :])
|
||||
#print(label[0,:].shape)
|
||||
label_flip = label_flip[self.flip_order, :]
|
||||
return data_flip, label_flip
|
||||
|
||||
def get_data(self, data, label, annot):
|
||||
if self.vis:
|
||||
self.img_num+=1
|
||||
#if self.img_num<=self.vis:
|
||||
# filename = './vis/raw_%d.jpg' % (self.img_num)
|
||||
# print('save', filename)
|
||||
# draw = data.copy()
|
||||
# for i in range(label.shape[0]):
|
||||
# cv2.circle(draw, (label[i][1], label[i][0]), 1, (0, 0, 255), 2)
|
||||
# scipy.misc.imsave(filename, draw)
|
||||
if self.vis:
|
||||
self.img_num += 1
|
||||
#if self.img_num<=self.vis:
|
||||
# filename = './vis/raw_%d.jpg' % (self.img_num)
|
||||
# print('save', filename)
|
||||
# draw = data.copy()
|
||||
# for i in range(label.shape[0]):
|
||||
# cv2.circle(draw, (label[i][1], label[i][0]), 1, (0, 0, 255), 2)
|
||||
# scipy.misc.imsave(filename, draw)
|
||||
|
||||
rotate = 0
|
||||
#scale = 1.0
|
||||
if 'scale' in annot:
|
||||
scale = annot['scale']
|
||||
else:
|
||||
scale = max(data.shape[0], data.shape[1])
|
||||
if 'center' in annot:
|
||||
center = annot['center']
|
||||
else:
|
||||
center = np.array( (data.shape[1]/2, data.shape[0]/2) )
|
||||
max_retry = 3
|
||||
if self.aug_level==0: #validation mode
|
||||
max_retry = 6
|
||||
retry = 0
|
||||
found = False
|
||||
base_scale = scale
|
||||
while retry<max_retry:
|
||||
retry+=1
|
||||
succ = True
|
||||
_scale = base_scale
|
||||
if self.aug_level>0:
|
||||
rotate = np.random.randint(-40, 40)
|
||||
scale_config = 0.2
|
||||
#rotate = 0
|
||||
#scale_config = 0.0
|
||||
scale_ratio = min(1+scale_config, max(1-scale_config, (np.random.randn() * scale_config) + 1))
|
||||
_scale = int(base_scale * scale_ratio)
|
||||
#translate = np.random.randint(-5, 5, size=(2,))
|
||||
#center += translate
|
||||
data_out, trans = img_helper.transform(data, center, self.input_img_size, _scale, rotate)
|
||||
#data_out = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
|
||||
label_out = np.zeros(self.label_shape, dtype=np.float32)
|
||||
#print('out shapes', data_out.shape, label_out.shape)
|
||||
for i in range(label.shape[0]):
|
||||
pt = label[i].copy()
|
||||
#pt = pt[::-1]
|
||||
npt = img_helper.transform_pt(pt, trans)
|
||||
if npt[0]>=data_out.shape[1] or npt[1]>=data_out.shape[0] or npt[0]<0 or npt[1]<0:
|
||||
succ = False
|
||||
#print('err npt', npt)
|
||||
break
|
||||
if config.losstype=='heatmap':
|
||||
pt_scale = float(self.output_label_size)/self.input_img_size
|
||||
npt *= pt_scale
|
||||
npt = npt.astype(np.int32)
|
||||
img_helper.gaussian(label_out[i], npt, config.gaussian)
|
||||
else:
|
||||
label_out[i] = (npt/self.input_img_size)
|
||||
#print('before gaussian', label_out[i].shape, pt.shape)
|
||||
#trans = img_helper.transform(pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate)
|
||||
#print(trans.shape)
|
||||
#if not img_helper.gaussian(label_out[i], trans, _g):
|
||||
# succ = False
|
||||
# break
|
||||
if not succ:
|
||||
if self.aug_level==0:
|
||||
base_scale+=20
|
||||
continue
|
||||
|
||||
flip_data_out = None
|
||||
flip_label_out = None
|
||||
if config.net_coherent:
|
||||
flip_data_out, flip_label_out = self.get_flip(data_out, label_out)
|
||||
elif ((self.aug_level>0 and np.random.rand() < 0.5) or self.force_mirror): #flip aug
|
||||
flip_data_out, flip_label_out = self.get_flip(data_out, label_out)
|
||||
data_out, label_out = flip_data_out, flip_label_out
|
||||
rotate = 0
|
||||
#scale = 1.0
|
||||
if 'scale' in annot:
|
||||
scale = annot['scale']
|
||||
else:
|
||||
scale = max(data.shape[0], data.shape[1])
|
||||
if 'center' in annot:
|
||||
center = annot['center']
|
||||
else:
|
||||
center = np.array((data.shape[1] / 2, data.shape[0] / 2))
|
||||
max_retry = 3
|
||||
if self.aug_level == 0: #validation mode
|
||||
max_retry = 6
|
||||
retry = 0
|
||||
found = False
|
||||
base_scale = scale
|
||||
while retry < max_retry:
|
||||
retry += 1
|
||||
succ = True
|
||||
_scale = base_scale
|
||||
if self.aug_level > 0:
|
||||
rotate = np.random.randint(-40, 40)
|
||||
scale_config = 0.2
|
||||
#rotate = 0
|
||||
#scale_config = 0.0
|
||||
scale_ratio = min(
|
||||
1 + scale_config,
|
||||
max(1 - scale_config,
|
||||
(np.random.randn() * scale_config) + 1))
|
||||
_scale = int(base_scale * scale_ratio)
|
||||
#translate = np.random.randint(-5, 5, size=(2,))
|
||||
#center += translate
|
||||
data_out, trans = img_helper.transform(data, center,
|
||||
self.input_img_size, _scale,
|
||||
rotate)
|
||||
#data_out = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
|
||||
label_out = np.zeros(self.label_shape, dtype=np.float32)
|
||||
#print('out shapes', data_out.shape, label_out.shape)
|
||||
for i in range(label.shape[0]):
|
||||
pt = label[i].copy()
|
||||
#pt = pt[::-1]
|
||||
npt = img_helper.transform_pt(pt, trans)
|
||||
if npt[0] >= data_out.shape[1] or npt[1] >= data_out.shape[
|
||||
0] or npt[0] < 0 or npt[1] < 0:
|
||||
succ = False
|
||||
#print('err npt', npt)
|
||||
break
|
||||
if config.losstype == 'heatmap':
|
||||
pt_scale = float(
|
||||
self.output_label_size) / self.input_img_size
|
||||
npt *= pt_scale
|
||||
npt = npt.astype(np.int32)
|
||||
img_helper.gaussian(label_out[i], npt, config.gaussian)
|
||||
else:
|
||||
label_out[i] = (npt / self.input_img_size)
|
||||
#print('before gaussian', label_out[i].shape, pt.shape)
|
||||
#trans = img_helper.transform(pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate)
|
||||
#print(trans.shape)
|
||||
#if not img_helper.gaussian(label_out[i], trans, _g):
|
||||
# succ = False
|
||||
# break
|
||||
if not succ:
|
||||
if self.aug_level == 0:
|
||||
base_scale += 20
|
||||
continue
|
||||
|
||||
found = True
|
||||
break
|
||||
flip_data_out = None
|
||||
flip_label_out = None
|
||||
if config.net_coherent:
|
||||
flip_data_out, flip_label_out = self.get_flip(
|
||||
data_out, label_out)
|
||||
elif ((self.aug_level > 0 and np.random.rand() < 0.5)
|
||||
or self.force_mirror): #flip aug
|
||||
flip_data_out, flip_label_out = self.get_flip(
|
||||
data_out, label_out)
|
||||
data_out, label_out = flip_data_out, flip_label_out
|
||||
|
||||
found = True
|
||||
break
|
||||
|
||||
#self.stats[0]+=1
|
||||
if not found:
|
||||
#self.stats[1]+=1
|
||||
#print('find aug error', retry)
|
||||
#print(self.stats)
|
||||
#print('!!!ERR')
|
||||
return None
|
||||
#print('found with scale', _scale, rotate)
|
||||
#self.stats[0]+=1
|
||||
if not found:
|
||||
#self.stats[1]+=1
|
||||
#print('find aug error', retry)
|
||||
#print(self.stats)
|
||||
#print('!!!ERR')
|
||||
return None
|
||||
#print('found with scale', _scale, rotate)
|
||||
|
||||
if self.vis > 0 and self.img_num <= self.vis:
|
||||
print('crop', data.shape, center, _scale, rotate, data_out.shape)
|
||||
filename = './vis/cropped_%d.jpg' % (self.img_num)
|
||||
print('save', filename)
|
||||
draw = data_out.copy()
|
||||
alabel = label_out.copy()
|
||||
for i in range(label.shape[0]):
|
||||
a = cv2.resize(alabel[i],
|
||||
(self.input_img_size, self.input_img_size))
|
||||
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
|
||||
cv2.circle(draw, (ind[1], ind[0]), 1, (0, 0, 255), 2)
|
||||
scipy.misc.imsave(filename, draw)
|
||||
filename = './vis/raw_%d.jpg' % (self.img_num)
|
||||
scipy.misc.imsave(filename, data)
|
||||
|
||||
if self.vis>0 and self.img_num<=self.vis:
|
||||
print('crop', data.shape, center, _scale, rotate, data_out.shape)
|
||||
filename = './vis/cropped_%d.jpg' % (self.img_num)
|
||||
print('save', filename)
|
||||
draw = data_out.copy()
|
||||
alabel = label_out.copy()
|
||||
for i in range(label.shape[0]):
|
||||
a = cv2.resize(alabel[i], (self.input_img_size, self.input_img_size))
|
||||
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
|
||||
cv2.circle(draw, (ind[1], ind[0]), 1, (0, 0, 255), 2)
|
||||
scipy.misc.imsave(filename, draw)
|
||||
filename = './vis/raw_%d.jpg' % (self.img_num)
|
||||
scipy.misc.imsave(filename, data)
|
||||
|
||||
return data_out, label_out, flip_data_out,flip_label_out
|
||||
return data_out, label_out, flip_data_out, flip_label_out
|
||||
|
||||
def next(self):
|
||||
"""Returns the next batch of data."""
|
||||
#print('next')
|
||||
batch_size = self.batch_size
|
||||
batch_data = nd.empty((batch_size,)+self.data_shape)
|
||||
batch_label = nd.empty((batch_size,)+self.label_shape)
|
||||
batch_data = nd.empty((batch_size, ) + self.data_shape)
|
||||
batch_label = nd.empty((batch_size, ) + self.label_shape)
|
||||
i = 0
|
||||
#self.cutoff = random.randint(800,1280)
|
||||
try:
|
||||
@@ -301,7 +316,7 @@ class FaceSegIter(DataIter):
|
||||
data, label, annot = self.next_sample()
|
||||
R = self.get_data(data, label, annot)
|
||||
if R is None:
|
||||
continue
|
||||
continue
|
||||
data_out, label_out, flip_data_out, flip_label_out = R
|
||||
if not self.use_coherent:
|
||||
data = nd.array(data_out)
|
||||
@@ -323,18 +338,17 @@ class FaceSegIter(DataIter):
|
||||
batch_data[i][:] = data
|
||||
batch_label[i][:] = label
|
||||
#i+=1
|
||||
j = i+self.per_batch_size//2
|
||||
j = i + self.per_batch_size // 2
|
||||
batch_data[j][:] = data2
|
||||
batch_label[j][:] = label2
|
||||
i += 1
|
||||
if j%self.per_batch_size==self.per_batch_size-1:
|
||||
i = j+1
|
||||
if j % self.per_batch_size == self.per_batch_size - 1:
|
||||
i = j + 1
|
||||
except StopIteration:
|
||||
if i<batch_size:
|
||||
if i < batch_size:
|
||||
raise StopIteration
|
||||
|
||||
#return {self.data_name : batch_data,
|
||||
# self.label_name : batch_label}
|
||||
#print(batch_data.shape, batch_label.shape)
|
||||
return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user