def __init__(self, params, result):
self.result = result
#load data
if params['source'][-4:] == '.mat':
self.img_data, self.img_labels = mat2py_imdb(params['source'])
self.img_data = np.rollaxis(self.img_data,2) #swap dims to have Num x height x width
elif os.path.isdir(params['source']): #assume it is lmdb
lmdb_ = lmdbpy()
self.img_data, self.img_labels = lmdb_.read(params['source']) #use read_encoded if the img is encoded
self.img_data = self.img_data.squeeze()
else:
assert 0, 'Invalid format for source {}\n\
need either lmdb or .mat data'.format(params['source'])
if params['mean_file']==0:
self.img_mean = 0
elif params['mean_file'][-4:] == '.mat':
self.img_mean = mat2py_mean(params['mean_file'])
elif params['mean_file'][-12:] == '.binaryproto':
self.img_mean = biproto2py(params['mean_file']).squeeze()
else:
assert 0, 'Invalid format for mean_file {}'.format(params['mean_file'])
self.indexlist = range(len(self.img_labels))
shuffle(self.indexlist)
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'], rot = params['rot'])
print "\nBatchLoader initialized with {} images".format(
len(self.img_labels))
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.classes_per_batch = params['classes_per_batch']
self.img_lmdb = lmdbs(params['img_source'])
if params['skt_source'].endswith('.pkl'):
self.skt_lmdb = svgs(params['skt_source'])
else:
self.skt_lmdb = lmdbs(params['skt_source'])
self.img_labels = self.img_lmdb.get_label_list()
self.skt_labels = self.skt_lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.num_classes = len(set(self.skt_labels))
assert len(self.skt_labels)%self.num_classes==0, \
'Unequal sketch training samples for each class'
self.skt_per_class = len(self.skt_labels) / self.num_classes
if 'hard_pos' in params:
self.hard_sel = 'hard_pos'
self.hard_sel_file = params['hard_pos']
elif 'hard_neg' in params:
self.hard_sel = 'hard_neg'
self.hard_sel_file = params['hard_neg']
elif 'hard_pn' in params:
self.hard_sel = 'hard_pn'
self.hard_sel_file = params['hard_pn']
else:
assert False, 'Hard selection must be on'
self.watchChange = WatchDog(
self.hard_sel_file) #check if file has been updated
self.hardsel_tab = np.load(self.hard_sel_file)
self.hard_pos = self.hardsel_tab['pos']
self.hard_neg = self.hardsel_tab['neg']
self.img_labels_dict, self.classes = vec2dic(self.img_labels)
self.NSKTS = len(self.skt_labels)
self.indexlist = range(self.NSKTS)
self.shuffle_keeping_min_classes_per_batch()
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=params['shape'],
scale=params['scale'],
rot=params['rot'])
if 'verbose' not in params:
print "BatchLoader initialized with {} sketches, {} images of {} classes".format(
len(self.skt_labels), len(self.img_labels), self.num_classes)
print('Hard selection: {}'.format(self.hard_sel))
#create threadpools for parallel augmentation
self.pool = ThreadPool() #4
def __init__(self, params):
#load data
self.batch_size = params['batch_size']
self.outshape = params['shape']
if params['source'][-4:] == '.mat':
self.img_data, self.img_labels = mat2py_imdb(params['source'])
self.img_data = self.img_data.transpose(
2, 0, 1) #swap dims to have Num x height x width
elif os.path.isdir(params['source']): #assume it is lmdb
lmdb_ = lmdbpy()
self.img_data, self.img_labels = lmdb_.read(
params['source']) #use read_encoded if the img is encoded
self.img_data = self.img_data.squeeze()
else:
assert 0, 'Invalid format for source {}\n\
need either lmdb or .mat data'.format(params['source'])
label_ids = list(set(self.img_labels))
NCATS = len(label_ids)
if label_ids[0] != 0 or label_ids[-1] != NCATS - 1:
print 'Your data labels are not [0:{}]. Converting label ...'.format(
NCATS - 1)
self.img_labels = [
label_ids.index(label) for label in self.img_labels
]
if params['mean_file'] == 0:
self.img_mean = 0
elif params['mean_file'][-4:] == '.mat':
self.img_mean = mat2py_mean(params['mean_file'])
elif params['mean_file'][-12:] == '.binaryproto':
self.img_mean = biproto2py(params['mean_file']).squeeze()
else:
assert 0, 'Invalid format for mean_file {}'.format(
params['mean_file'])
self.indexlist = range(len(self.img_labels))
shuffle(self.indexlist)
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=params['shape'],
scale=params['scale'],
rot=params['rot'])
print "\nBatchLoader initialized with {} images".format(
len(self.img_labels))
self.pool = ThreadPool() #4
class BatchLoader(object):
"""
This class abstracts away the loading of images.
Images can either be loaded singly, or in a batch. The latter is used for
the asyncronous data layer to preload batches while other processing is
performed.
"""
def __init__(self, params, result):
self.result = result
#load data
if params['source'][-4:] == '.mat':
self.img_data, self.img_labels = mat2py_imdb(params['source'])
self.img_data = np.rollaxis(self.img_data,2) #swap dims to have Num x height x width
elif os.path.isdir(params['source']): #assume it is lmdb
lmdb_ = lmdbpy()
self.img_data, self.img_labels = lmdb_.read(params['source']) #use read_encoded if the img is encoded
self.img_data = self.img_data.squeeze()
else:
assert 0, 'Invalid format for source {}\n\
need either lmdb or .mat data'.format(params['source'])
if params['mean_file']==0:
self.img_mean = 0
elif params['mean_file'][-4:] == '.mat':
self.img_mean = mat2py_mean(params['mean_file'])
elif params['mean_file'][-12:] == '.binaryproto':
self.img_mean = biproto2py(params['mean_file']).squeeze()
else:
assert 0, 'Invalid format for mean_file {}'.format(params['mean_file'])
self.indexlist = range(len(self.img_labels))
shuffle(self.indexlist)
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'], rot = params['rot'])
print "\nBatchLoader initialized with {} images".format(
len(self.img_labels))
def load_next_image(self):
"""
Load the next image in a batch.
"""
# Did we finish an epoch?
if self._cur == len(self.indexlist):
self._cur = 0
shuffle(self.indexlist)
# Load a datum
index = self.indexlist[self._cur] # Get the index
img = self.img_data[index]
label = self.img_labels[index]
self._cur += 1
return (self.img_augment.augment(img), label)
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.lmdb = lmdbs(params['source'])
self.labels = self.lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.NIMGS = len(self.labels)
self.num_batches = int(np.ceil(self.NIMGS/float(self.batch_size)))
self._cur = 0 # current batch
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'])
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.lmdb = lmdbs(params['source'])
self.labels = self.lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.NIMGS = len(self.labels)
assert self.NIMGS%self.batch_size==0,'NIMGS {} not dividible by batchsize {}'.format(
self.NIMGS,self.batch_size)
self.num_batches = self.NIMGS/self.batch_size
self._cur = 0 # current batch
self.labels_tab = self.labels.reshape((self.num_batches,self.batch_size))
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'])
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.classes_per_batch = params['classes_per_batch']
self.img_lmdb = lmdbs(params['img_source'])
if params['skt_source'].endswith('.pkl'):
self.skt_lmdb = svgs(params['skt_source'])
else:
self.skt_lmdb = lmdbs(params['skt_source'])
self.img_labels = self.img_lmdb.get_label_list()
self.skt_labels = self.skt_lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.num_classes = len(set(self.skt_labels))
assert len(self.skt_labels)%self.num_classes==0, \
'Unequal sketch training samples for each class'
self.skt_per_class = len(self.skt_labels) / self.num_classes
self.hard_sel = 0
self.img_labels_dict, self.classes = vec2dic(self.img_labels)
self.indexlist = range(len(self.skt_labels))
self.shuffle_keeping_min_classes_per_batch()
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=params['shape'],
scale=params['scale'],
rot=params['rot'])
if 'verbose' not in params:
print "BatchLoader initialized with {} sketches, {} images of {} classes".format(
len(self.skt_labels), len(self.img_labels), self.num_classes)
#create threadpools for parallel augmentation
self.pool = ThreadPool() #4
class batchloader(object):
"""
load lmdb in mini-batch. Roll-over at the end of the epoch.
Requirement: NIMGS dividible to batch_size
Useful in datalayer of CNN training network.
"""
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.lmdb = lmdbs(params['source'])
self.labels = self.lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.NIMGS = len(self.labels)
assert self.NIMGS%self.batch_size==0,'NIMGS {} not dividible by batchsize {}'.format(
self.NIMGS,self.batch_size)
self.num_batches = self.NIMGS/self.batch_size
self._cur = 0 # current batch
self.labels_tab = self.labels.reshape((self.num_batches,self.batch_size))
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'])
#create threadpools for parallel augmentation
#self.pool = ThreadPool() #4
def get_info(self):
res = dict(NIMGS = self.NIMGS,batch_size = self.batch_size,
num_batches = self.num_batches,cur = self._cur)
return res
def get_labels(self):
return self.labels
def IsEpochEnded(self):
return self._cur == self.num_batches
def load_next_batch(self):
if self._cur == self.num_batches:
self._cur = 0
batch_lst = np.arange(self.batch_size) + self._cur * self.batch_size
chunk = self.img_augment.augment_deploy(self.lmdb.get_data(batch_lst))
labels = self.labels_tab[self._cur]
self._cur +=1
return chunk,labels
class batchloader2(object):
"""
batchloader in general.
load lmdb in mini-batch. Roll-over at the end of the epoch.
DIFFERENT: does not require NIMGS dividible to batch_size; hence the last batch
may have different size.
Useful in CNN deploy.
"""
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.lmdb = lmdbs(params['source'])
self.labels = self.lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.NIMGS = len(self.labels)
self.num_batches = int(np.ceil(self.NIMGS/float(self.batch_size)))
self._cur = 0 # current batch
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,shape=params['shape'],
scale = params['scale'])
#create threadpools for parallel augmentation
#self.pool = ThreadPool() #4
def get_info(self):
res = dict(NIMGS = self.NIMGS,batch_size = self.batch_size,
num_batches = self.num_batches,cur = self._cur)
return res
def get_labels(self):
return self.labels
def IsEpochEnded(self):
return self._cur == self.num_batches
def load_next_batch(self):
if self._cur == self.num_batches:
self._cur = 0
batch_lst = np.arange(self._cur * self.batch_size,min([self.NIMGS,(self._cur+1) * self.batch_size]))
chunk = self.img_augment.augment_deploy(self.lmdb.get_data(batch_lst))
labels = self.labels[batch_lst]
self._cur +=1
return chunk,labels
def __init__(self, params):
self.batch_size = params['batch_size']
self.img_shape = params['shape']
self.classes_per_batch = params['classes_per_batch']
self.img_lmdb = lmdbs(params['img_source'])
if params['skt_source'].endswith('.pkl'):
self.skt_lmdb = svgs(params['skt_source'])
else:
self.skt_lmdb = lmdbs(params['skt_source'])
self.img_labels = self.img_lmdb.get_label_list()
self.skt_labels = self.skt_lmdb.get_label_list()
label_ids = list(set(self.img_labels))
NCATS = len(label_ids)
if label_ids[0] != 0 or label_ids[-1] != NCATS - 1:
if 'verbose' not in params:
print 'Your data labels are not [0:{}]. Converting label ...'.format(
NCATS - 1)
self.img_labels = [
label_ids.index(label) for label in self.img_labels
]
self.skt_labels = [
label_ids.index(label) for label in self.skt_labels
]
self.img_mean = biproto2py(params['mean_file']).squeeze()
#self.skt_mean = biproto2py(params['skt_mean']).squeeze()
self.num_classes = len(set(self.skt_labels))
assert self.num_classes == NCATS, 'XX!!Sketch & image datasets unequal #categories'
assert len(self.skt_labels)%self.num_classes==0, \
'Unequal sketch training samples for each class'
self.skt_per_class = len(self.skt_labels) / self.num_classes
if 'hard_pos' in params:
self.hard_sel = 1
self.hard_pos = np.load(params['hard_pos'])['pos']
elif 'hard_neg' in params:
self.hard_sel = 2
self.hard_neg = np.load(params['hard_neg'])['neg']
elif 'hard_pn' in params:
self.hard_sel = 3
tmp = np.load(params['hard_pn'])
self.hard_pos = tmp['pos']
self.hard_neg = tmp['neg']
else: #hard selection turn off
self.hard_sel = 0
#self.img_labels_dict, self.classes = vec2dic(self.img_labels)
self.indexlist = range(len(self.skt_labels))
self.indexlist_img = range(len(self.img_labels))
#self.shuffle_keeping_min_classes_per_batch()
shuffle(self.indexlist)
shuffle(self.indexlist_img)
self._cur = 0 # current image
self._cur_img = 0
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=self.img_shape,
scale=params['scale'],
rot=params['rot'])
print "BatchLoader initialized with {} sketches, {} images of {} classes".format(
len(self.skt_labels), len(self.img_labels), self.num_classes)
#create threadpools for parallel augmentation
self.pool = ThreadPool() #4
class BatchLoader(object):
"""
This class abstracts away the loading of images.
Images can either be loaded singly, or in a batch. The latter is used for
the asyncronous data layer to preload batches while other processing is
performed.
"""
def __init__(self, params):
self.batch_size = params['batch_size']
self.img_shape = params['shape']
self.classes_per_batch = params['classes_per_batch']
self.img_lmdb = lmdbs(params['img_source'])
if params['skt_source'].endswith('.pkl'):
self.skt_lmdb = svgs(params['skt_source'])
else:
self.skt_lmdb = lmdbs(params['skt_source'])
self.img_labels = self.img_lmdb.get_label_list()
self.skt_labels = self.skt_lmdb.get_label_list()
label_ids = list(set(self.img_labels))
NCATS = len(label_ids)
if label_ids[0] != 0 or label_ids[-1] != NCATS - 1:
if 'verbose' not in params:
print 'Your data labels are not [0:{}]. Converting label ...'.format(
NCATS - 1)
self.img_labels = [
label_ids.index(label) for label in self.img_labels
]
self.skt_labels = [
label_ids.index(label) for label in self.skt_labels
]
self.img_mean = biproto2py(params['mean_file']).squeeze()
#self.skt_mean = biproto2py(params['skt_mean']).squeeze()
self.num_classes = len(set(self.skt_labels))
assert self.num_classes == NCATS, 'XX!!Sketch & image datasets unequal #categories'
assert len(self.skt_labels)%self.num_classes==0, \
'Unequal sketch training samples for each class'
self.skt_per_class = len(self.skt_labels) / self.num_classes
if 'hard_pos' in params:
self.hard_sel = 1
self.hard_pos = np.load(params['hard_pos'])['pos']
elif 'hard_neg' in params:
self.hard_sel = 2
self.hard_neg = np.load(params['hard_neg'])['neg']
elif 'hard_pn' in params:
self.hard_sel = 3
tmp = np.load(params['hard_pn'])
self.hard_pos = tmp['pos']
self.hard_neg = tmp['neg']
else: #hard selection turn off
self.hard_sel = 0
#self.img_labels_dict, self.classes = vec2dic(self.img_labels)
self.indexlist = range(len(self.skt_labels))
self.indexlist_img = range(len(self.img_labels))
#self.shuffle_keeping_min_classes_per_batch()
shuffle(self.indexlist)
shuffle(self.indexlist_img)
self._cur = 0 # current image
self._cur_img = 0
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=self.img_shape,
scale=params['scale'],
rot=params['rot'])
print "BatchLoader initialized with {} sketches, {} images of {} classes".format(
len(self.skt_labels), len(self.img_labels), self.num_classes)
#create threadpools for parallel augmentation
self.pool = ThreadPool() #4
def load_next_pair(self, l):
"""
Load the next pair in a batch.
"""
# Did we finish an epoch?
l.acquire() #5
if self._cur == len(self.indexlist):
self._cur = 0
#self.shuffle_keeping_min_classes_per_batch()
shuffle(self.indexlist)
if self._cur_img == len(self.indexlist_img):
self._cur_img = 0
#self.shuffle_keeping_min_classes_per_batch()
shuffle(self.indexlist_img)
# Load a sketch
index = self.indexlist[self._cur] # Get the sketch index
index_img = self.indexlist_img[self._cur_img]
self._cur += 1
self._cur_img += 1
l.release() #6
skt = self.skt_lmdb.get_datum(index).squeeze()
label = self.skt_labels[index]
img_i = self.img_lmdb.get_datum(index_img).squeeze()
label_i = self.img_labels[index_img]
#==============================================================================
# #randomly select paired image
# diff_label = np.random.choice(2)
# label_i = label
# if self.hard_sel == 0: #hard selection turned off
# if diff_label: #paired image has different label
# while label_i == label:
# label_i = np.random.choice(self.classes)
# index_i = np.random.choice(self.img_labels_dict[str(label_i)])
# elif self.hard_sel == 1: #hard positive selection
# if diff_label:
# while label_i == label:
# label_i = np.random.choice(self.classes)
# index_i = np.random.choice(self.img_labels_dict[str(label_i)])
# else:
# index_i = np.random.choice(self.hard_pos[index])
# elif self.hard_sel == 2: #hard neg
# if diff_label:
# index_i = np.random.choice(self.hard_neg[index])
# else:
# index_i = np.random.choice(self.img_labels_dict[str(label_i)])
# else: #hard pos and neg
# if diff_label:
# index_i = np.random.choice(self.hard_neg[index])
# else:
# index_i = np.random.choice(self.hard_pos[index])
#
# img_i = self.img_lmdb.get_image_deprocess(index_i)
#==============================================================================
res = dict(sketch=self.img_augment.augment(skt),
image=self.img_augment.augment(img_i),
label_s=label,
label_i=label_i)
return res
def load_next_batch(self):
res = {}
#7
lock = Lock()
threads = [self.pool.apply_async(self.load_next_pair,(lock,)) for \
i in range (self.batch_size)]
thread_res = [thread.get() for thread in threads]
res['data_s'] = np.asarray([tri['sketch']
for tri in thread_res])[:, None, :, :]
res['data_i'] = np.asarray([tri['image']
for tri in thread_res])[:, None, :, :]
res['label_s'] = np.asarray([tri['label_s'] for tri in thread_res],
dtype=np.float32)[:, None]
res['label_i'] = np.asarray([tri['label_i'] for tri in thread_res],
dtype=np.float32)[:, None]
return res
#==============================================================================
# res['data_s'] = np.zeros((self.batch_size,1,self.outshape[0],\
# self.outshape[1]),dtype = np.float32)
# res['data_i'] = np.zeros_like(res['data_a'],dtype=np.float32)
# res['label'] = np.zeros((self.batch_size,1),dtype = np.float32)
# for itt in range(self.batch_size):
# trp = self.load_next_pair(1)
# res['data_s'][itt,...] = trp['sketch']
# res['data_i'][itt,...] = trp['image']
# res['label'][itt,...] = trp['label']
# return res
#==============================================================================
def shuffle_keeping_min_classes_per_batch(self):
shuffle(self.indexlist)
# sort index list by class
# (using lambda to restrict sorting to just first element of tuple
sort_indexlist = sorted(self.indexlist,
key=lambda index: self.skt_labels[index])
# make it 2D based on classes
sort_indexlist = np.array(sort_indexlist, dtype=np.int64)
sort_indexlist = sort_indexlist.reshape(self.num_classes,
self.skt_per_class)
# permute the class positions and flat it (make it 1D again)
sort_indexlist = np.random.permutation(sort_indexlist)
sort_indexlist = sort_indexlist.reshape(self.num_classes *
self.skt_per_class)
temp_indexlist = np.array([], dtype=np.int64)
skt_per_batch_per_class = int(self.batch_size / self.classes_per_batch)
# apply dark magic
# (slices of classes sliced, appended together until list is over)
for k in range(0, self.skt_per_class, skt_per_batch_per_class):
temp_indexlist = np.append(temp_indexlist, [
sort_indexlist[(i * self.skt_per_class) + j]
for i in range(self.num_classes)
for j in range(k, k + skt_per_batch_per_class)
])
# convert back to list
self.indexlist = temp_indexlist.tolist()
def extract_cnn_feat(net_params, DB, OUT, layer=0, verbose=True):
"""
extract features from CNN
DB: lmdb data you want to extract feature
net_params: dictionary with keys "DEPLOY_PRO","data_mean",
"WEIGHTS","scale_factor", batch_size
OUT: save output in mat file
layer: 0 for last layer, -1: one before the last layer, -2: ...
"""
assert layer <= 0, 'layer should be a non-positive integer'
DEPLOY_PRO = net_params['DEPLOY_PRO']
WEIGHTS = net_params['WEIGHTS']
scale_factor = net_params['scale_factor']
data_mean = net_params['data_mean']
batch_size = net_params['batch_size']
net = caffe.Net(DEPLOY_PRO, WEIGHTS, caffe.TEST)
if verbose:
print 'Extracting cnn feats...'
print ' Model def: {}\n Weights: {}'.format(DEPLOY_PRO, WEIGHTS)
start_t = time.time()
db = lmdbs(DB)
labels = db.get_label_list()
NIMGS = labels.size
img_mean = biproto2py(data_mean)
inblob = net.inputs[0]
in_dim = net.blobs[inblob].data.shape[1:]
prep = SimpleAugment(mean=img_mean, shape=in_dim[-2:])
feat_l = net.blobs.keys()[layer - 1]
out_dim = net.blobs[feat_l].data.squeeze().shape[-1]
feats = np.zeros((NIMGS, out_dim), dtype=np.float32)
for i in xrange(0, NIMGS, batch_size):
batch = range(i, min(i + batch_size, NIMGS))
if verbose:
print('Processing sample #{} - {}'.format(batch[0], batch[-1]))
new_shape = (len(batch), ) + in_dim
net.blobs[inblob].reshape(*new_shape)
chunk = db.get_data(batch)
net.blobs[inblob].data[...] = prep.augment_deploy(chunk)
temp = net.forward()
feats[batch] = net.blobs[feat_l].data.squeeze()
#apply scale factor
feats *= scale_factor
if OUT.endswith('.mat'):
py2mat(dict(feats=feats, labels=labels), OUT)
elif OUT.endswith('.npz'):
np.savez(OUT, feats=feats, labels=labels)
else: #assume it is pickle
helps = helper()
helps.save(OUT, feats=feats, labels=labels)
net = None
if verbose:
end_t = time.time()
print 'Save features to {}.'.format(OUT)
print 'Time: {}\n'.format(timedelta(seconds=int(end_t - start_t)))
class BatchLoader(object):
"""
This class abstracts away the loading of images.
Images can either be loaded singly, or in a batch. The latter is used for
the asyncronous data layer to preload batches while other processing is
performed.
"""
def __init__(self, params):
#load data
self.batch_size = params['batch_size']
self.outshape = params['shape']
if params['source'][-4:] == '.mat':
self.img_data, self.img_labels = mat2py_imdb(params['source'])
self.img_data = self.img_data.transpose(
2, 0, 1) #swap dims to have Num x height x width
elif os.path.isdir(params['source']): #assume it is lmdb
lmdb_ = lmdbpy()
self.img_data, self.img_labels = lmdb_.read(
params['source']) #use read_encoded if the img is encoded
self.img_data = self.img_data.squeeze()
else:
assert 0, 'Invalid format for source {}\n\
need either lmdb or .mat data'.format(params['source'])
label_ids = list(set(self.img_labels))
NCATS = len(label_ids)
if label_ids[0] != 0 or label_ids[-1] != NCATS - 1:
print 'Your data labels are not [0:{}]. Converting label ...'.format(
NCATS - 1)
self.img_labels = [
label_ids.index(label) for label in self.img_labels
]
if params['mean_file'] == 0:
self.img_mean = 0
elif params['mean_file'][-4:] == '.mat':
self.img_mean = mat2py_mean(params['mean_file'])
elif params['mean_file'][-12:] == '.binaryproto':
self.img_mean = biproto2py(params['mean_file']).squeeze()
else:
assert 0, 'Invalid format for mean_file {}'.format(
params['mean_file'])
self.indexlist = range(len(self.img_labels))
shuffle(self.indexlist)
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=params['shape'],
scale=params['scale'],
rot=params['rot'])
print "\nBatchLoader initialized with {} images".format(
len(self.img_labels))
self.pool = ThreadPool() #4
def load_next_image(self, l):
"""
Load the next image in a batch.
"""
# Did we finish an epoch?
l.acquire() #5
if self._cur == len(self.indexlist):
self._cur = 0
shuffle(self.indexlist)
# Load a datum
index = self.indexlist[self._cur] # Get the index
self._cur += 1
l.release() #6
img = self.img_data[index]
label = self.img_labels[index]
return (self.img_augment.augment(img), label)
def load_next_batch(self):
res = {}
#7
lock = Lock()
threads = [self.pool.apply_async(self.load_next_image,(lock,)) for \
i in range (self.batch_size)]
thread_res = [thread.get() for thread in threads]
res['data'] = np.asarray([datum[0]
for datum in thread_res])[:, None, :, :]
res['label'] = np.asarray([datum[1] for datum in thread_res],
dtype=np.float32)
return res
class BatchLoader(object):
"""
This class abstracts away the loading of images.
Images can either be loaded singly, or in a batch. The latter is used for
the asyncronous data layer to preload batches while other processing is
performed.
"""
def __init__(self, params):
self.batch_size = params['batch_size']
self.outshape = params['shape']
self.classes_per_batch = params['classes_per_batch']
self.img_lmdb = lmdbs(params['img_source'])
if params['skt_source'].endswith('.pkl'):
self.skt_lmdb = svgs(params['skt_source'])
else:
self.skt_lmdb = lmdbs(params['skt_source'])
self.img_labels = self.img_lmdb.get_label_list()
self.skt_labels = self.skt_lmdb.get_label_list()
self.img_mean = biproto2py(params['mean_file']).squeeze()
self.num_classes = len(set(self.skt_labels))
assert len(self.skt_labels)%self.num_classes==0, \
'Unequal sketch training samples for each class'
self.skt_per_class = len(self.skt_labels) / self.num_classes
self.hard_sel = 0
self.img_labels_dict, self.classes = vec2dic(self.img_labels)
self.indexlist = range(len(self.skt_labels))
self.shuffle_keeping_min_classes_per_batch()
self._cur = 0 # current image
# this class does some simple data-manipulations
self.img_augment = SimpleAugment(mean=self.img_mean,
shape=params['shape'],
scale=params['scale'],
rot=params['rot'])
if 'verbose' not in params:
print "BatchLoader initialized with {} sketches, {} images of {} classes".format(
len(self.skt_labels), len(self.img_labels), self.num_classes)
#create threadpools for parallel augmentation
self.pool = ThreadPool() #4
def load_next_triplet(self, l):
"""
Load the next triplet in a batch.
"""
# Did we finish an epoch?
l.acquire() #5
if self._cur == len(self.indexlist):
self._cur = 0
self.shuffle_keeping_min_classes_per_batch()
# Load a sketch
index = self.indexlist[self._cur] # Get the sketch index
self._cur += 1
l.release() #6
skt = self.skt_lmdb.get_datum(index).squeeze()
label = self.skt_labels[index]
#randomly select pos and neg img
index_p = np.random.choice(self.img_labels_dict[str(label)])
label_n = label
while label_n == label:
label_n = np.random.choice(self.classes)
index_n = np.random.choice(self.img_labels_dict[str(label_n)])
img_p = self.img_lmdb.get_datum(index_p).squeeze()
img_n = self.img_lmdb.get_datum(index_n).squeeze()
res = dict(anchor=self.img_augment.augment(skt),
pos=self.img_augment.augment(img_p),
neg=self.img_augment.augment(img_n),
label_a=label,
label_n=label_n)
return res
def load_next_batch(self):
res = {}
#7
lock = Lock()
threads = [self.pool.apply_async(self.load_next_triplet,(lock,)) for \
i in range (self.batch_size)]
thread_res = [thread.get() for thread in threads]
res['data_a'] = np.asarray([tri['anchor']
for tri in thread_res])[:, None, :, :]
res['data_p'] = np.asarray([tri['pos']
for tri in thread_res])[:, None, :, :]
res['data_n'] = np.asarray([tri['neg']
for tri in thread_res])[:, None, :, :]
res['label_a'] = np.asarray([tri['label_a'] for tri in thread_res],
dtype=np.float32)
res['label_n'] = np.asarray([tri['label_n'] for tri in thread_res],
dtype=np.float32)
return res
#==============================================================================
# res['data_a'] = np.zeros((self.batch_size,1,self.outshape[0],\
# self.outshape[1]),dtype = np.float32)
# res['data_p'] = np.zeros_like(res['data_a'],dtype=np.float32)
# res['data_n'] = np.zeros_like(res['data_a'],dtype=np.float32)
# res['label'] = np.zeros((self.batch_size,3),dtype = np.float32)
# for itt in range(self.batch_size):
# trp = self.load_next_triplet(1)
# res['data_a'][itt,...] = trp['anchor']
# res['data_p'][itt,...] = trp['pos']
# res['data_n'][itt,...] = trp['neg']
# res['label'][itt,...] = trp['label']
# return res
#==============================================================================
def shuffle_keeping_min_classes_per_batch(self):
shuffle(self.indexlist)
# sort index list by class
# (using lambda to restrict sorting to just first element of tuple
sort_indexlist = sorted(self.indexlist,
key=lambda index: self.skt_labels[index])
# make it 2D based on classes
sort_indexlist = np.array(sort_indexlist, dtype=np.int64)
sort_indexlist = sort_indexlist.reshape(self.num_classes,
self.skt_per_class)
# permute the class positions and flat it (make it 1D again)
sort_indexlist = np.random.permutation(sort_indexlist)
sort_indexlist = sort_indexlist.reshape(self.num_classes *
self.skt_per_class)
temp_indexlist = np.array([], dtype=np.int64)
skt_per_batch_per_class = int(self.batch_size / self.classes_per_batch)
# apply dark magic
# (slices of classes sliced, appended together until list is over)
for k in range(0, self.skt_per_class, skt_per_batch_per_class):
temp_indexlist = np.append(temp_indexlist, [
sort_indexlist[(i * self.skt_per_class) + j]
for i in range(self.num_classes)
for j in range(k, k + skt_per_batch_per_class)
])
# convert back to list
self.indexlist = temp_indexlist.tolist()