def main(args): # sys.path.append("/home/gaomingda/insightface/recognition") from image_iter import FaceImageIter global image_shape global net print(args) ctx = [] cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip() if len(cvd)>0: for i in xrange(len(cvd.split(','))): ctx.append(mx.gpu(i)) if len(ctx)==0: ctx = [mx.cpu()] print('use cpu') else: print('gpu num:', len(ctx)) image_shape = [int(x) for x in args.image_size.split(',')] vec = args.model.split(',') assert len(vec)>1 prefix = vec[0] epoch = int(vec[1]) print('loading',prefix, epoch) net = edict() net.ctx = ctx net.sym, net.arg_params, net.aux_params = mx.model.load_checkpoint(prefix, epoch) #net.arg_params, net.aux_params = ch_dev(net.arg_params, net.aux_params, net.ctx) all_layers = net.sym.get_internals() net.sym = all_layers['fc1_output'] net.model = mx.mod.Module(symbol=net.sym, context=net.ctx, label_names = None) net.model.bind(data_shapes=[('data', (args.batch_size, 3, image_shape[1], image_shape[2]))]) net.model.set_params(net.arg_params, net.aux_params) train_dataiter = FaceImageIter( batch_size=4, data_shape=(3, 112, 112), path_imgrec=args.input_data, shuffle=True, rand_mirror=False, mean=None, cutoff=False, color_jittering=0, images_filter=0, ) data_size = train_dataiter.num_samples() i = 0 fstart = 0 features_all = np.zeros((data_size, 512), dtype=np.float32) features_all_flip = np.zeros((data_size, 512), dtype=np.float32) # features_all = np.zeros((102, 512), dtype=np.float32) # features_all_flip = np.zeros((102, 512), dtype=np.float32) data_buff = nd.empty((args.batch_size, 3, 112, 112)) count = 0 for i in range(train_dataiter.num_samples()): if i%1000==0: print("processing ",i) label, s, box, landmark = train_dataiter.next_sample() img = train_dataiter.imdecode(s) img = nd.transpose(img, axes=(2, 0, 1)) data_buff[count] = img count += 1 if count==args.batch_size: embedding = get_feature(data_buff, args.batch_size) count = 0 fend = fstart+embedding.shape[0] #print('writing', fstart, fend) features_all[fstart:fend,:] = embedding # flipped image data_buff_flip = mx.ndarray.flip(data=data_buff, axis=3) embedding_fliped = get_feature(data_buff_flip, args.batch_size) features_all_flip[fstart:fend, :] = embedding_fliped fstart = fend # if i==102: # break if count>0: embedding = get_feature(data_buff, args.batch_size) fend = fstart+count print('writing', fstart, fend) features_all[fstart:fend,:] = embedding[:count, :] # flipped image data_buff_flip = mx.ndarray.flip(data=data_buff, axis=3) embedding_fliped = get_feature(data_buff_flip, args.batch_size) features_all_flip[fstart:fend, :] = embedding_fliped[:count, :] # write_bin(args.output, features_all) #os.system("bypy upload %s"%args.output) print("save features ...") features_all.tofile('train_features_oct200') print("save train_features_flip ...") features_all_flip.tofile('train_features_flip_oct200')
path_imgrec= "/home/gaomingda/insightface/datasets/ms1m-retinaface-t1/train.rec", shuffle=True, rand_mirror=False, mean=None, cutoff=False, color_jittering=0, images_filter=0, ) data_nums = train_dataiter.num_samples() train_dataiter.reset() train_dataiter.is_init = True f = open(os.path.join(save_root, "train.txt"), 'w') f.truncate() for i in range(data_nums): label, s, _, _ = train_dataiter.next_sample() img_ = mx.image.imdecode(s) #mx.ndarray img = np.array(img_.asnumpy(), dtype=np.uint8) img = img[:, :, ::-1] #print("img shape", img.shape) cv2.imwrite(save_root + '/traindata/{}.jpg'.format(i), img) f.writelines("{}.jpg".format(i) + " " + str(int(label)) + "\n") if i % 1000 == 0: print(i) f.close() print("done")