class TripletSamplingDataFetcher(Process):
def __init__(self, queue, sketch_data, image_data, sketch_label, image_label, triplets, queue_paras, phase):
"""Setup the TripletSamplingDataLayer."""
super(TripletSamplingDataFetcher, self).__init__()
# mean = mean
self.im_size = queue_paras['im_size']
self.cp_size = queue_paras['cp_size']
self.chns = queue_paras['chns']
self.mean = queue_paras['mean']
self.batchsize = queue_paras['batch_size']
self.num_epoch = queue_paras['num_epochs']
self._queue = queue
self._phase = phase
self.sketch_transformer = Transformer(self.cp_size, self.chns, self.mean, self._phase == "TRAIN")
if sketch_label is None or image_label is None:
self.has_label = False
else:
self.has_label = True
self.anc_bm = MemoryBlockManager(sketch_data, sketch_label, self.has_label)
self.pos_neg_bm = MemoryBlockManager(image_data, image_label, self.has_label)
self.triplets = triplets
def get_next_batch_data(self):
anc_batch = []; pos_batch = []; neg_batch = []
# sampling
anc_inds, pos_inds, neg_inds = zip(*sample_triplets(self.triplets, self.batchsize, self.num_epoch).next())
#pos_inds, neg_inds = sample_triplets_trueMatch(anc_inds, self.triplets)
# fetch data
for (anc_id, pos_id, neg_id) in zip(anc_inds, pos_inds, neg_inds):
anc_batch.append(self.anc_bm.get_sample(anc_id).reshape((self.im_size, self.im_size, self.chns)))
pos_batch.append(self.pos_neg_bm.get_sample(pos_id).reshape((self.im_size, self.im_size, self.chns)))
neg_batch.append(self.pos_neg_bm.get_sample(neg_id).reshape((self.im_size, self.im_size, self.chns)))
# apply transform
anc_batch = self.sketch_transformer.transform_all(anc_batch)
pos_batch = self.sketch_transformer.transform_all(pos_batch)
neg_batch = self.sketch_transformer.transform_all(neg_batch)
self._queue.put((anc_batch, pos_batch, neg_batch))
def get_next_batch_data_label(self):
anc_batch = []; pos_batch = []; neg_batch = []
# sampling
anc_inds, pos_inds, neg_inds = zip(*sample_triplets(self.triplets, self.batchsize, self.num_epoch).next())
#pos_inds, neg_inds = sample_triplets_trueMatch(anc_inds, self.triplets)
# fetch data
for (anc_id, pos_id, neg_id) in zip(anc_inds, pos_inds, neg_inds):
anc_batch.append(self.anc_bm.get_sample(anc_id))
pos_batch.append(self.pos_neg_bm.get_sample(pos_id))
neg_batch.append(self.pos_neg_bm.get_sample(neg_id))
# apply transform
anc_batch_data, anc_batch_label = self.sketch_transformer.transform_all_with_label(anc_batch)
pos_batch_data, pos_batch_label = self.sketch_transformer.transform_all_with_label(pos_batch)
neg_batch_data, neg_batch_label = self.sketch_transformer.transform_all_with_label(neg_batch)
self._queue.put((anc_batch_data, pos_batch_data, neg_batch_data, anc_batch_label, pos_batch_label, neg_batch_label))
def run(self):
print ('TripletSamplingDataFetcher started')
if self.has_label:
while True:
self.get_next_batch_data_label()
else:
while True:
self.get_next_batch_data()
