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")
