mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-15 04:37:50 +00:00
149 lines
4.2 KiB
Python
149 lines
4.2 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import os
|
|
|
|
from datetime import datetime
|
|
import os.path
|
|
from easydict import EasyDict as edict
|
|
import time
|
|
import json
|
|
import sys
|
|
import numpy as np
|
|
import importlib
|
|
import itertools
|
|
import argparse
|
|
import struct
|
|
import cv2
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'common'))
|
|
import face_preprocess
|
|
import sklearn
|
|
from sklearn.preprocessing import normalize
|
|
import mxnet as mx
|
|
from mxnet import ndarray as nd
|
|
|
|
image_shape = None
|
|
net = None
|
|
data_size = 1862120
|
|
emb_size = 512
|
|
use_flip = True
|
|
|
|
|
|
def do_flip(data):
|
|
for idx in xrange(data.shape[0]):
|
|
data[idx,:,:] = np.fliplr(data[idx,:,:])
|
|
|
|
def get_feature(buffer):
|
|
embedding = np.zeros( (len(buffer), emb_size), dtype=np.float32 )
|
|
if use_flip:
|
|
input_blob = np.zeros( (len(buffer)*2, 3, image_shape[1], image_shape[2] ) )
|
|
else:
|
|
input_blob = np.zeros( (len(buffer), 3, image_shape[1], image_shape[2] ) )
|
|
idx = 0
|
|
for item in buffer:
|
|
img = face_preprocess.read_image(item[0], mode='rgb')
|
|
img = face_preprocess.preprocess(img, bbox=None, landmark=item[1], image_size='%d,%d'%(image_shape[1], image_shape[2]))
|
|
img = np.transpose( img, (2,0,1) )
|
|
attempts = [0,1] if use_flip else [0]
|
|
for flipid in attempts:
|
|
_img = np.copy(img)
|
|
if flipid==1:
|
|
do_flip(_img)
|
|
input_blob[idx] = _img
|
|
idx+=1
|
|
data = mx.nd.array(input_blob)
|
|
db = mx.io.DataBatch(data=(data,))
|
|
net.model.forward(db, is_train=False)
|
|
_embedding = net.model.get_outputs()[0].asnumpy()
|
|
if use_flip:
|
|
embedding1 = _embedding[0::2]
|
|
embedding2 = _embedding[1::2]
|
|
embedding = embedding1+embedding2
|
|
else:
|
|
embedding = _embedding
|
|
embedding = sklearn.preprocessing.normalize(embedding)
|
|
return embedding
|
|
|
|
|
|
|
|
def write_bin(path, m):
|
|
rows, cols = m.shape
|
|
with open(path, 'wb') as f:
|
|
f.write(struct.pack('4i', rows,cols,cols*4,5))
|
|
f.write(m.data)
|
|
|
|
def main(args):
|
|
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)
|
|
|
|
|
|
i = 0
|
|
features_all = np.zeros( (data_size, emb_size), dtype=np.float32 )
|
|
fstart = 0
|
|
buffer = []
|
|
for line in open(args.input, 'r'):
|
|
if i%1000==0:
|
|
print("processing ",i)
|
|
i+=1
|
|
image_path, label, bbox, landmark, aligned = face_preprocess.parse_lst_line(line)
|
|
buffer.append( (image_path, landmark) )
|
|
if len(buffer)==args.batch_size:
|
|
embedding = get_feature(buffer)
|
|
buffer = []
|
|
fend = fstart+embedding.shape[0]
|
|
#print('writing', fstart, fend)
|
|
features_all[fstart:fend,:] = embedding
|
|
fstart = fend
|
|
if len(buffer)>0:
|
|
embedding = get_feature(buffer)
|
|
fend = fstart+embedding.shape[0]
|
|
print('writing', fstart, fend)
|
|
features_all[fstart:fend,:] = embedding
|
|
write_bin(args.output, features_all)
|
|
os.system("bypy upload %s"%args.output)
|
|
|
|
|
|
|
|
def parse_arguments(argv):
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument('--batch_size', type=int, help='', default=32)
|
|
parser.add_argument('--image_size', type=str, help='', default='3,112,112')
|
|
parser.add_argument('--input', type=str, help='', default='')
|
|
parser.add_argument('--output', type=str, help='', default='')
|
|
parser.add_argument('--model', type=str, help='', default='')
|
|
return parser.parse_args(argv)
|
|
|
|
if __name__ == '__main__':
|
|
main(parse_arguments(sys.argv[1:]))
|
|
|
|
|