# THIS FILE IS FOR EXPERIMENTS, USE image_iter.py FOR NORMAL IMAGE LOADING. from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import random import logging import sys import numbers import math import sklearn import datetime import numpy as np import cv2 import mxnet as mx from mxnet import ndarray as nd #from . import _ndarray_internal as _internal #from mxnet._ndarray_internal import _cvimresize as imresize #from ._ndarray_internal import _cvcopyMakeBorder as copyMakeBorder from mxnet import io from mxnet import recordio sys.path.append(os.path.join(os.path.dirname(__file__), 'common')) import face_preprocess logger = logging.getLogger() class FaceImageIter(io.DataIter): def __init__(self, batch_size, data_shape, path_imgrec = None, shuffle=False, aug_list=None, rand_mirror = False, cutoff = 0, ctx_num = 0, images_per_identity = 0, triplet_params = None, mx_model = None, data_name='data', label_name='softmax_label', **kwargs): super(FaceImageIter, self).__init__() assert path_imgrec assert shuffle logging.info('loading recordio %s...', path_imgrec) path_imgidx = path_imgrec[0:-4]+".idx" self.imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type s = self.imgrec.read_idx(0) header, _ = recordio.unpack(s) assert header.flag>0 print('header0 label', header.label) self.header0 = (int(header.label[0]), int(header.label[1])) #assert(header.flag==1) self.imgidx = range(1, int(header.label[0])) self.id2range = {} self.seq_identity = range(int(header.label[0]), int(header.label[1])) for identity in self.seq_identity: s = self.imgrec.read_idx(identity) header, _ = recordio.unpack(s) a,b = int(header.label[0]), int(header.label[1]) self.id2range[identity] = (a,b) print('id2range', len(self.id2range)) self.seq = self.imgidx print(len(self.seq)) self.check_data_shape(data_shape) self.provide_data = [(data_name, (batch_size,) + data_shape)] self.batch_size = batch_size self.data_shape = data_shape self.shuffle = shuffle self.image_size = '%d,%d'%(data_shape[1],data_shape[2]) self.rand_mirror = rand_mirror print('rand_mirror', rand_mirror) self.cutoff = cutoff #self.cast_aug = mx.image.CastAug() #self.color_aug = mx.image.ColorJitterAug(0.4, 0.4, 0.4) self.ctx_num = ctx_num self.per_batch_size = int(self.batch_size/self.ctx_num) self.images_per_identity = images_per_identity if self.images_per_identity>0: self.identities = int(self.per_batch_size/self.images_per_identity) self.per_identities = self.identities self.repeat = 3000000.0/(self.images_per_identity*len(self.id2range)) self.repeat = int(self.repeat) print(self.images_per_identity, self.identities, self.repeat) self.mx_model = mx_model self.triplet_params = triplet_params self.triplet_mode = False #self.provide_label = None self.provide_label = [(label_name, (batch_size,))] if self.triplet_params is not None: assert self.images_per_identity>0 assert self.mx_model is not None self.triplet_bag_size = self.triplet_params[0] self.triplet_alpha = self.triplet_params[1] self.triplet_max_ap = self.triplet_params[2] assert self.triplet_bag_size>0 assert self.triplet_alpha>=0.0 assert self.triplet_alpha<=1.0 self.triplet_mode = True self.triplet_cur = 0 self.triplet_seq = [] self.triplet_reset() self.seq_min_size = self.batch_size*2 self.cur = 0 self.nbatch = 0 self.is_init = False self.times = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] #self.reset() def pairwise_dists(self, embeddings): nd_embedding_list = [] for i in range(self.ctx_num): nd_embedding = mx.nd.array(embeddings, mx.gpu(i)) nd_embedding_list.append(nd_embedding) nd_pdists = [] pdists = [] for idx in range(embeddings.shape[0]): emb_idx = idx%self.ctx_num nd_embedding = nd_embedding_list[emb_idx] a_embedding = nd_embedding[idx] body = mx.nd.broadcast_sub(a_embedding, nd_embedding) body = body*body body = mx.nd.sum_axis(body, axis=1) nd_pdists.append(body) if len(nd_pdists)==self.ctx_num or idx==embeddings.shape[0]-1: for x in nd_pdists: pdists.append(x.asnumpy()) nd_pdists = [] return pdists def pick_triplets(self, embeddings, nrof_images_per_class): emb_start_idx = 0 triplets = [] people_per_batch = len(nrof_images_per_class) #self.time_reset() pdists = self.pairwise_dists(embeddings) #self.times[3] += self.time_elapsed() for i in range(people_per_batch): nrof_images = int(nrof_images_per_class[i]) for j in range(1,nrof_images): #self.time_reset() a_idx = emb_start_idx + j - 1 #neg_dists_sqr = np.sum(np.square(embeddings[a_idx] - embeddings), 1) neg_dists_sqr = pdists[a_idx] #self.times[3] += self.time_elapsed() for pair in range(j, nrof_images): # For every possible positive pair. p_idx = emb_start_idx + pair #self.time_reset() pos_dist_sqr = np.sum(np.square(embeddings[a_idx]-embeddings[p_idx])) #self.times[4] += self.time_elapsed() #self.time_reset() neg_dists_sqr[emb_start_idx:emb_start_idx+nrof_images] = np.NaN if self.triplet_max_ap>0.0: if pos_dist_sqr>self.triplet_max_ap: continue all_neg = np.where(np.logical_and(neg_dists_sqr-pos_dist_sqr<self.triplet_alpha, pos_dist_sqr<neg_dists_sqr))[0] # FaceNet selection #self.times[5] += self.time_elapsed() #self.time_reset() #all_neg = np.where(neg_dists_sqr-pos_dist_sqr<alpha)[0] # VGG Face selecction nrof_random_negs = all_neg.shape[0] if nrof_random_negs>0: rnd_idx = np.random.randint(nrof_random_negs) n_idx = all_neg[rnd_idx] triplets.append( (a_idx, p_idx, n_idx) ) emb_start_idx += nrof_images np.random.shuffle(triplets) return triplets def triplet_reset(self): #reset self.oseq by identities seq self.triplet_cur = 0 ids = [] for k in self.id2range: ids.append(k) random.shuffle(ids) self.triplet_seq = [] for _id in ids: v = self.id2range[_id] _list = range(*v) random.shuffle(_list) if len(_list)>self.images_per_identity: _list = _list[0:self.images_per_identity] self.triplet_seq += _list print('triplet_seq', len(self.triplet_seq)) assert len(self.triplet_seq)>=self.triplet_bag_size def time_reset(self): self.time_now = datetime.datetime.now() def time_elapsed(self): time_now = datetime.datetime.now() diff = time_now - self.time_now return diff.total_seconds() def select_triplets(self): self.seq = [] while len(self.seq)<self.seq_min_size: self.time_reset() embeddings = None bag_size = self.triplet_bag_size batch_size = self.batch_size #data = np.zeros( (bag_size,)+self.data_shape ) #label = np.zeros( (bag_size,) ) tag = [] #idx = np.zeros( (bag_size,) ) print('eval %d images..'%bag_size, self.triplet_cur) print('triplet time stat', self.times) if self.triplet_cur+bag_size>len(self.triplet_seq): self.triplet_reset() #bag_size = min(bag_size, len(self.triplet_seq)) print('eval %d images..'%bag_size, self.triplet_cur) self.times[0] += self.time_elapsed() self.time_reset() #print(data.shape) data = nd.zeros( self.provide_data[0][1] ) label = None if self.provide_label is not None: label = nd.zeros( self.provide_label[0][1] ) ba = 0 while True: bb = min(ba+batch_size, bag_size) if ba>=bb: break _count = bb-ba #data = nd.zeros( (_count,)+self.data_shape ) #_batch = self.data_iter.next() #_data = _batch.data[0].asnumpy() #print(_data.shape) #_label = _batch.label[0].asnumpy() #data[ba:bb,:,:,:] = _data #label[ba:bb] = _label for i in range(ba, bb): #print(ba, bb, self.triplet_cur, i, len(self.triplet_seq)) _idx = self.triplet_seq[i+self.triplet_cur] s = self.imgrec.read_idx(_idx) header, img = recordio.unpack(s) img = self.imdecode(img) data[i-ba][:] = self.postprocess_data(img) _label = header.label if not isinstance(_label, numbers.Number): _label = _label[0] if label is not None: label[i-ba][:] = _label tag.append( ( int(_label), _idx) ) #idx[i] = _idx db = mx.io.DataBatch(data=(data,)) self.mx_model.forward(db, is_train=False) net_out = self.mx_model.get_outputs() #print('eval for selecting triplets',ba,bb) #print(net_out) #print(len(net_out)) #print(net_out[0].asnumpy()) net_out = net_out[0].asnumpy() #print(net_out) #print('net_out', net_out.shape) if embeddings is None: embeddings = np.zeros( (bag_size, net_out.shape[1])) embeddings[ba:bb,:] = net_out ba = bb assert len(tag)==bag_size self.triplet_cur+=bag_size embeddings = sklearn.preprocessing.normalize(embeddings) self.times[1] += self.time_elapsed() self.time_reset() nrof_images_per_class = [1] for i in range(1, bag_size): if tag[i][0]==tag[i-1][0]: nrof_images_per_class[-1]+=1 else: nrof_images_per_class.append(1) triplets = self.pick_triplets(embeddings, nrof_images_per_class) # shape=(T,3) print('found triplets', len(triplets)) ba = 0 while True: bb = ba+self.per_batch_size//3 if bb>len(triplets): break _triplets = triplets[ba:bb] for i in range(3): for triplet in _triplets: _pos = triplet[i] _idx = tag[_pos][1] self.seq.append(_idx) ba = bb self.times[2] += self.time_elapsed() def hard_mining_reset(self): #import faiss from annoy import AnnoyIndex data = nd.zeros( self.provide_data[0][1] ) label = nd.zeros( self.provide_label[0][1] ) #label = np.zeros( self.provide_label[0][1] ) X = None ba = 0 batch_num = 0 while ba<len(self.oseq): batch_num+=1 if batch_num%10==0: print('loading batch',batch_num, ba) bb = min(ba+self.batch_size, len(self.oseq)) _count = bb-ba for i in range(_count): idx = self.oseq[i+ba] s = self.imgrec.read_idx(idx) header, img = recordio.unpack(s) img = self.imdecode(img) data[i][:] = self.postprocess_data(img) label[i][:] = header.label db = mx.io.DataBatch(data=(data,self.data_extra), label=(label,)) self.mx_model.forward(db, is_train=False) net_out = self.mx_model.get_outputs() embedding = net_out[0].asnumpy() nembedding = sklearn.preprocessing.normalize(embedding) if _count<self.batch_size: nembedding = nembedding[0:_count,:] if X is None: X = np.zeros( (len(self.id2range), nembedding.shape[1]), dtype=np.float32 ) nplabel = label.asnumpy() for i in range(_count): ilabel = int(nplabel[i]) #print(ilabel, ilabel.__class__) X[ilabel] += nembedding[i] ba = bb X = sklearn.preprocessing.normalize(X) d = X.shape[1] t = AnnoyIndex(d, metric='euclidean') for i in range(X.shape[0]): t.add_item(i, X[i]) print('start to build index') t.build(20) print(X.shape) k = self.per_identities self.seq = [] for i in range(X.shape[0]): nnlist = t.get_nns_by_item(i, k) assert nnlist[0]==i for _label in nnlist: assert _label<len(self.id2range) _id = self.header0[0]+_label v = self.id2range[_id] _list = range(*v) if len(_list)<self.images_per_identity: random.shuffle(_list) else: _list = np.random.choice(_list, self.images_per_identity, replace=False) for i in range(self.images_per_identity): _idx = _list[i%len(_list)] self.seq.append(_idx) #faiss_params = [20,5] #quantizer = faiss.IndexFlatL2(d) # the other index #index = faiss.IndexIVFFlat(quantizer, d, faiss_params[0], faiss.METRIC_L2) #assert not index.is_trained #index.train(X) #index.add(X) #assert index.is_trained #print('trained') #index.nprobe = faiss_params[1] #D, I = index.search(X, k) # actual search #print(I.shape) #self.seq = [] #for i in range(I.shape[0]): # #assert I[i][0]==i # for j in range(k): # _label = I[i][j] # assert _label<len(self.id2range) # _id = self.header0[0]+_label # v = self.id2range[_id] # _list = range(*v) # if len(_list)<self.images_per_identity: # random.shuffle(_list) # else: # _list = np.random.choice(_list, self.images_per_identity, replace=False) # for i in range(self.images_per_identity): # _idx = _list[i%len(_list)] # self.seq.append(_idx) def reset(self): """Resets the iterator to the beginning of the data.""" print('call reset()') self.cur = 0 if self.images_per_identity>0: if self.triplet_mode: self.select_triplets() elif not self.hard_mining: self.seq = [] idlist = [] for _id in self.id2range: v = self.id2range[_id] idlist.append((_id,range(*v))) for r in range(self.repeat): if r%10==0: print('repeat', r) if self.shuffle: random.shuffle(idlist) for item in idlist: _id = item[0] _list = item[1] #random.shuffle(_list) if len(_list)<self.images_per_identity: random.shuffle(_list) else: _list = np.random.choice(_list, self.images_per_identity, replace=False) for i in range(self.images_per_identity): _idx = _list[i%len(_list)] self.seq.append(_idx) else: self.hard_mining_reset() print('seq len', len(self.seq)) else: if self.shuffle: random.shuffle(self.seq) if self.seq is None and self.imgrec is not None: self.imgrec.reset() def num_samples(self): return len(self.seq) def next_sample(self): while True: 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) label = header.label if not isinstance(label, numbers.Number): label = label[0] return label, img, None, None def brightness_aug(self, src, x): alpha = 1.0 + random.uniform(-x, x) src *= alpha return src def contrast_aug(self, src, x): alpha = 1.0 + random.uniform(-x, x) coef = np.array([[[0.299, 0.587, 0.114]]]) gray = src * coef gray = (3.0 * (1.0 - alpha) / gray.size) * np.sum(gray) src *= alpha src += gray return src def saturation_aug(self, src, x): alpha = 1.0 + random.uniform(-x, x) coef = np.array([[[0.299, 0.587, 0.114]]]) gray = src * coef gray = np.sum(gray, axis=2, keepdims=True) gray *= (1.0 - alpha) src *= alpha src += gray return src def color_aug(self, img, x): augs = [self.brightness_aug, self.contrast_aug, self.saturation_aug] random.shuffle(augs) for aug in augs: #print(img.shape) img = aug(img, x) #print(img.shape) return img def mirror_aug(self, img): _rd = random.randint(0,1) if _rd==1: for c in range(img.shape[2]): img[:,:,c] = np.fliplr(img[:,:,c]) return img def next(self): if not self.is_init: self.reset() self.is_init = True """Returns the next batch of data.""" #print('in next', self.cur, self.labelcur) self.nbatch+=1 batch_size = self.batch_size c, h, w = self.data_shape batch_data = nd.empty((batch_size, c, h, w)) if self.provide_label is not None: batch_label = nd.empty(self.provide_label[0][1]) i = 0 try: while i < batch_size: label, s, bbox, landmark = self.next_sample() _data = self.imdecode(s) if self.rand_mirror: _rd = random.randint(0,1) if _rd==1: _data = mx.ndarray.flip(data=_data, axis=1) if self.cutoff>0: centerh = random.randint(0, _data.shape[0]-1) centerw = random.randint(0, _data.shape[1]-1) half = self.cutoff//2 starth = max(0, centerh-half) endh = min(_data.shape[0], centerh+half) startw = max(0, centerw-half) endw = min(_data.shape[1], centerw+half) _data = _data.astype('float32') #print(starth, endh, startw, endw, _data.shape) _data[starth:endh, startw:endw, :] = 127.5 #_npdata = _data.asnumpy() #if landmark is not None: # _npdata = face_preprocess.preprocess(_npdata, bbox = bbox, landmark=landmark, image_size=self.image_size) #if self.rand_mirror: # _npdata = self.mirror_aug(_npdata) #if self.mean is not None: # _npdata = _npdata.astype(np.float32) # _npdata -= self.mean # _npdata *= 0.0078125 #nimg = np.zeros(_npdata.shape, dtype=np.float32) #nimg[self.patch[1]:self.patch[3],self.patch[0]:self.patch[2],:] = _npdata[self.patch[1]:self.patch[3], self.patch[0]:self.patch[2], :] #_data = mx.nd.array(nimg) data = [_data] try: self.check_valid_image(data) except RuntimeError as e: logging.debug('Invalid image, skipping: %s', str(e)) continue #print('aa',data[0].shape) #data = self.augmentation_transform(data) #print('bb',data[0].shape) for datum in data: assert i < batch_size, 'Batch size must be multiples of augmenter output length' #print(datum.shape) batch_data[i][:] = self.postprocess_data(datum) if self.provide_label is not None: batch_label[i][:] = label i += 1 except StopIteration: if i<batch_size: raise StopIteration #print('next end', batch_size, i) _label = None if self.provide_label is not None: _label = [batch_label] return io.DataBatch([batch_data], _label, batch_size - i) def check_data_shape(self, data_shape): """Checks if the input data shape is valid""" if not len(data_shape) == 3: raise ValueError('data_shape should have length 3, with dimensions CxHxW') if not data_shape[0] == 3: raise ValueError('This iterator expects inputs to have 3 channels.') def check_valid_image(self, data): """Checks if the input data is valid""" if len(data[0].shape) == 0: raise RuntimeError('Data shape is wrong') def imdecode(self, s): """Decodes a string or byte string to an NDArray. See mx.img.imdecode for more details.""" img = mx.image.imdecode(s) #mx.ndarray return img def read_image(self, fname): """Reads an input image `fname` and returns the decoded raw bytes. Example usage: ---------- >>> dataIter.read_image('Face.jpg') # returns decoded raw bytes. """ with open(os.path.join(self.path_root, fname), 'rb') as fin: img = fin.read() return img def augmentation_transform(self, data): """Transforms input data with specified augmentation.""" for aug in self.auglist: data = [ret for src in data for ret in aug(src)] return data def postprocess_data(self, datum): """Final postprocessing step before image is loaded into the batch.""" return nd.transpose(datum, axes=(2, 0, 1)) class FaceImageIterList(io.DataIter): def __init__(self, iter_list): assert len(iter_list)>0 self.provide_data = iter_list[0].provide_data self.provide_label = iter_list[0].provide_label self.iter_list = iter_list self.cur_iter = None def reset(self): self.cur_iter.reset() def next(self): self.cur_iter = random.choice(self.iter_list) while True: try: ret = self.cur_iter.next() except StopIteration: self.cur_iter.reset() continue return ret