def get_mod(gpuid):
if not use_torch:
import mxnet as mx
ctx = mx.gpu(gpuid)
batch_size = args.batch_size
model_prefix, epoch = args.model.split(',')
epoch = int(epoch)
print('loading %s %d' % (model_prefix, epoch))
sym, arg_params, aux_params = mx.model.load_checkpoint(
model_prefix, epoch)
sym = sym.get_internals()['fc1_output']
mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
mod.bind(for_training=False,
data_shapes=[('data', (batch_size, 3, 112, 112))])
mod.set_params(arg_params, aux_params)
return mod
else:
sys.path.insert(0, '/home/zl/prj/InsightFace_Pytorch/')
from config import conf
conf.ipabn = False
conf.need_log = False
from Learner import face_learner, FaceInfer
learner = FaceInfer(conf, gpuid)
learner.load_state(
resume_path=('work_space/emore.r100.bs.ft.tri.dop/save/'), )
return learner
def get_mod(gpuid):
if not args.use_torch:
import mxnet as mx
from mxnet import ndarray as nd
ctx = mx.gpu(gpuid)
batch_size = args.batch_size
model_prefix, epoch = args.model.split(',')
epoch = int(epoch)
print('loading %s %d' % (model_prefix, epoch))
sym, arg_params, aux_params = mx.model.load_checkpoint(
model_prefix, epoch)
sym = sym.get_internals()['fc1_output']
mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
mod.bind(for_training=False,
data_shapes=[('data', (batch_size, 3, 112, 112))])
mod.set_params(arg_params, aux_params)
return mod
else:
sys.path.insert(0, args.code)
from config import conf
conf.need_log = False
from Learner import face_learner, FaceInfer
learner = FaceInfer(conf, (gpuid, ))
learner.load_state(resume_path=args.model.split(',')[0])
learner.model.eval()
return learner
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_num)
batch_size = args.batch_size
print('#####', args.model, args.output_root)
gpuid = 0
import torch, torch.utils.data
model_prefix, epoch = args.model.split(',')
sys.path.insert(0, args.code)
from config import conf
conf.need_log = False
from Learner import face_learner, FaceInfer
learner = FaceInfer(conf, (gpuid, ))
learner.load_state(resume_path=model_prefix, )
learner.model.eval()
loader = torch.utils.data.DataLoader(TestDataset(args),
batch_size=batch_size,
num_workers=12,
shuffle=False,
pin_memory=True,
drop_last=False)
bin_filename = os.path.join(
args.images_list.split('/')[-2],
args.images_list.split('/')[-1].split('.')[0] + '.bin')
if args.use_torch:
model_name = model_prefix.strip('/').split('/')[-2]
else:
model_name = os.path.basename(model_prefix)
if args.model_name is not None:
model_name = args.model_name
dump_path = os.path.join(args.output_root, model_name, bin_filename)
print('###### features will be dumped to:%s' % dump_path)
dirname = os.path.dirname(dump_path)
if not os.path.exists(dirname):
os.makedirs(dirname)
dump = open(dump_path, 'wb')
for batch in loader:
import lz, torch
dev = torch.device(f'cuda:{gpuid}')
batch = batch - 127.5
batch = batch / 127.5
with torch.no_grad():
embs = learner.model(lz.to_torch(batch).to(dev)).cpu().numpy()
if flip_test:
embs += learner.model(
lz.to_torch(
batch[..., ::-1].copy()).to(dev)).cpu().numpy()
# from IPython import embed; embed()
embs = sklearn.preprocessing.normalize(embs)
for k in range(embs.shape[0]):
dump.write(embs[k].astype(np.float32))
dump.flush()
dump.flush()
dump.close()
use_mxnet = True
if not use_mxnet:
from config import conf
conf.need_log = False
conf.batch_size *= 4
conf.fp16 = False
conf.ipabn = False
conf.cvt_ipabn = True
# conf.upgrade_irse = False
# conf.net_mode = 'ir'
conf.net_depth = 50
from Learner import l2_norm, FaceInfer, get_rec, unpack_auto
learner = FaceInfer(conf, )
learner.load_state(
resume_path='../work_space/asia.emore.r50.ada/models/',
latest=True,
)
# learner.load_model_only('work_space/backbone_ir50_ms1m_epoch120.pth')
# learner.load_model_only('work_space/backbone_ir50_asia.pth')
learner.model.eval()
else:
from config import conf
from recognition.embedding import Embedding
from Learner import l2_norm, FaceInfer, get_rec, unpack_auto
learner = Embedding(
prefix=
'/home/xinglu/prj/insightface/logs/MS1MV2-ResNet100-Arcface/model',
)
else:
mtcnn = MTCNN()
logging.info('mtcnn ok')
if 'th' in conf_str:
from config import conf
conf.need_log = False
conf.ipabn = False
conf.cvt_ipabn = True
conf.net_depth = 50
from Learner import face_learner, FaceInfer
from models.model import l2_norm
learner = FaceInfer(conf, )
learner.load_state(
conf,
resume_path=root_path + 'work_space/asia.emore.r50.5/models'
# resume_path= root_path+'work_space/emore.r152.cont/models'
)
learner.model.eval()
logging.info('learner loaded')
else:
from recognition.embedding import Embedding
import os
model_path = '/home/xinglu/prj/insightface/logs/model-r100-arcface-ms1m-refine-v2/model'
assert os.path.exists(
os.path.dirname(model_path)), os.path.dirname(model_path)
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_num)
batch_size = args.batch_size
print('#####', args.model, args.output_root)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
for_test = np.array(['1231', '213213'], dtype='str')
test_ims = tf.placeholder(for_test.dtype, [None])
def input_parser2(filename):
image_string = tf.read_file(filename)
image_decoded = tf.image.decode_jpeg(image_string,
dct_method="INTEGER_ACCURATE")
image = tf.cast(image_decoded, tf.float32)
image = tf.transpose(image, perm=[2, 0, 1])
if preprocess_img:
image = image - 127.5
image = image / 128.
return image
test_data = tf.data.Dataset.from_tensor_slices((test_ims))
test_data = test_data.map(input_parser2, num_parallel_calls=48)
test_data = test_data.prefetch(batch_size * 100)
test_data = test_data.batch(batch_size)
iterator2 = test_data.make_initializable_iterator()
next_element2 = iterator2.get_next()
sess.run(iterator2.initializer, feed_dict={test_ims: for_test})
if not args.use_torch:
gpuid = 0
ctx = mx.gpu(gpuid)
net = edict()
net.ctx = ctx
model_prefix, epoch = args.model.split(',')
epoch = int(epoch)
print('loading %s %d' % (model_prefix, epoch))
net.sym, net.arg_params, net.aux_params = mx.model.load_checkpoint(
model_prefix, epoch)
net.sym = net.sym.get_internals()['fc1_output']
net.model = mx.mod.Module(symbol=net.sym,
context=net.ctx,
label_names=None)
net.model.bind(for_training=False,
data_shapes=[('data', (batch_size, 3, 112, 112))])
net.model.set_params(net.arg_params, net.aux_params)
else:
model_prefix, epoch = args.model.split(',')
sys.path.insert(0, os.environ['HOME'] + '/prj/InsightFace_Pytorch/')
from config import conf
gpuid = 0
conf.ipabn = False
conf.need_log = False
from Learner import face_learner, FaceInfer
learner = FaceInfer(conf, gpuid)
learner.load_state(resume_path=model_prefix, )
# data = mx.nd.array(np.random.normal(size=(batch_size,3,112,112)))
# db = mx.io.DataBatch(data=(data,))
# net.model.forward(db, is_train=False)
spisok = open(args.images_list).read().split('\n')[:-1]
for i in range(len(spisok)):
spisok[i] = args.prefix + spisok[i].split(' ')[0]
for_test = np.array(spisok, dtype='str')
sess.run(iterator2.initializer, feed_dict={test_ims: for_test})
bin_filename = os.path.join(
args.images_list.split('/')[-2],
args.images_list.split('/')[-1].split('.')[0] + '.bin')
# embed()
if args.use_torch:
model_name = model_prefix.strip('/').split('/')[-2]
else:
model_name = os.path.basename(model_prefix)
dump_path = os.path.join(args.output_root, model_name, bin_filename)
print('###### features will be dumped to:%s' % dump_path)
dirname = os.path.dirname(dump_path)
if not os.path.exists(dirname):
os.makedirs(dirname)
dump = open(dump_path, 'wb')
for i in tqdm(range(int(np.ceil(len(spisok) / float(batch_size))))):
batch = sess.run(next_element2)
data = mx.nd.array(batch)
# im = np.transpose(batch[0],(1,2,0))
# cv2.imshow('x',im.astype(np.uint8))
# cv2.waitKey(0)
db = mx.io.DataBatch(data=(data, ))
if not args.use_torch:
net.model.forward(db, is_train=False)
embs = net.model.get_outputs()[0].asnumpy()
if flip_test:
batch = batch[:, :, :, ::-1]
data = mx.nd.array(batch)
db = mx.io.DataBatch(data=(data, ))
net.model.forward(db, is_train=False)
embs += net.model.get_outputs()[0].asnumpy()
else:
import lz, torch
dev = torch.device(f'cuda:{gpuid}')
# batch = test_transform(batch)
# from IPython import embed; embed()
batch = batch - 127.5
batch = batch / 127.5
with torch.no_grad():
embs = learner.model(lz.to_torch(batch).to(dev)).cpu().numpy()
if flip_test:
embs += learner.model(
lz.to_torch(
batch[..., ::-1].copy()).to(dev)).cpu().numpy()
embs = sklearn.preprocessing.normalize(embs)
for k in range(embs.shape[0]):
dump.write(embs[k].astype(np.float32))
dump.close()
def main(args):
print(args)
gpuid = args.gpu
ctx = mx.gpu(gpuid)
nets = []
image_shape = [int(x) for x in args.image_size.split(',')]
for model in args.model.split('|'):
if use_mxnet:
vec = 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)
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', (1, 3, image_shape[1],
image_shape[2]))])
net.model.set_params(net.arg_params, net.aux_params)
nets.append(net)
else:
from config import conf
conf.need_log = False
conf.batch_size *= 2
from Learner import l2_norm, FaceInfer
learner = FaceInfer(conf, )
learner.load_state(
resume_path=args.model,
latest=False,
)
learner.model.eval()
nets.append(learner)
facescrub_out = os.path.join(args.output, 'facescrub')
megaface_out = os.path.join(args.output, 'megaface')
i = 0
succ = 0
buffer = []
imgfns = []
imgfns += json_load(
'/data/share/megaface/devkit/templatelists/facescrub_features_list_10000.2.json'
)['path']
imgfns += open(args.facescrub_lst, 'r').readlines()
imgfns = ['/'.join(imgfn.split('/')[-2:]).strip() for imgfn in imgfns]
imgfns = np.unique(imgfns).tolist()
# if True:
# imgfns = []
for line in imgfns:
if i % 1000 == 0:
print("writing fs", i, succ)
i += 1
image_path = line.strip()
_path = image_path.split('/')
a, b = _path[-2], _path[-1]
out_dir = os.path.join(facescrub_out, a)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
image_path = os.path.join(args.facescrub_root, image_path)
img = read_img(image_path)
assert img is not None, image_path
assert img.shape == (112, 112, 3), image_path
# if img is None:
# print('read error:', image_path)
# continue
out_path = os.path.join(out_dir, b + "_%s.bin" % (args.algo))
item = (img, out_path)
buffer.append(item)
if len(buffer) == args.batch_size:
get_and_write(buffer, nets)
buffer = []
succ += 1
if len(buffer) > 0:
get_and_write(buffer, nets)
buffer = []
print('fs stat', i, succ)
# if True:
# return
i = 0
succ = 0
buffer = []
imgfns = []
for mega_lst in [
'/data/share/megaface/devkit/templatelists/megaface_features_list.json_1000000_1',
'/data/share/megaface/devkit/templatelists/megaface_features_list.json_100000_1',
'/data/share/megaface/devkit/templatelists/megaface_features_list.json_100_1'
]:
imgfns += json_load(mega_lst)['path']
imgfns = np.unique(imgfns).tolist()
for line in imgfns:
# for line in open(args.megaface_lst, 'r'):
if i % 1000 == 0:
print("writing mf", i, succ)
i += 1
image_path = line.strip()
_path = image_path.split('/')
a1, a2, b = _path[-3], _path[-2], _path[-1]
out_dir = os.path.join(megaface_out, a1, a2)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# continue
# print(landmark)
image_path = os.path.join(args.megaface_root, image_path)
img = read_img(image_path)
if img is None:
print('read error:', image_path)
continue
out_path = os.path.join(out_dir, b + "_%s.bin" % (args.algo))
item = (img, out_path)
buffer.append(item)
if len(buffer) == args.batch_size:
get_and_write(buffer, nets)
buffer = []
succ += 1
if len(buffer) > 0:
get_and_write(buffer, nets)
buffer = []
print('mf stat', i, succ)
'prj/insightface/logs/r50-arcface-retina/model',
epoch=16,
ctx_id=0)
else:
from config import conf
conf.need_log = False
bs = conf.batch_size * 2
# conf.ipabn = False
# conf.cvt_ipabn = False
# conf.arch_ft = False
# conf.use_act = 'prelu'
# conf.net_depth = 100
# conf.net_mode = 'ir_se'
# conf.embedding_size = 512
# conf.input_size = 128
# conf.ds = False
# conf.use_bl = False
# conf.mid_type = '' # 'gpool'
from Learner import FaceInfer, face_learner
# todo
gpuid = list(map(int, os.environ['CUDA_VISIBLE_DEVICES'].split(',')))
print(gpuid)
learner = FaceInfer(conf, gpuid)
learner.load_state(
resume_path=f'work_space/{args.modelp}/models/',
latest=True,
)
learner.model.eval()
test_ijbc3(conf, learner)
glob.glob(src_folder + '/**/*.jpg', recursive=True),
glob.glob(src_folder + '/**/*.JPEG', recursive=True))
loader = torch.utils.data.DataLoader(TestData(imgfn_iter), batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True,
drop_last=False
)
conf.need_log = False
conf.fp16 = False
conf.ipabn = False
conf.cvt_ipabn = True
conf.net_depth = 50
conf.use_chkpnt = False
learner = FaceInfer(conf, )
learner.load_state(
# resume_path='work_space/emore.r152.cont/save/',
resume_path='work_space/asia.emore.r50.5/save/',
latest=True,
)
learner.model.eval()
from sklearn.preprocessing import normalize
for ind, data in enumerate(loader):
if (data['finish'] == 1).all().item():
logging.info('finish')
break
# if ind % 10 == 0:
print(f'proc batch {ind}')
img = data['img'].cuda()
def main(args):
global image_shape, net, env, glargs
print(args)
glargs = args
env = lmdb.open(
args.input + '/imgs_lmdb',
readonly=True,
# max_readers=1, lock=False,
# readahead=False, meminit=False
)
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(',')]
if use_mxnet:
net = edict()
vec = args.model.split(',')
assert len(vec) > 1
prefix = vec[0]
epoch = int(vec[1])
print('loading', prefix, epoch)
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)
else:
# sys.path.insert(0, lz.home_path + 'prj/InsightFace_Pytorch/')
from config import conf
from Learner import FaceInfer
conf.need_log = False
conf.batch_size = args.batch_size
conf.fp16 = True
conf.ipabn = False
conf.cvt_ipabn = False
conf.use_chkpnt = False
net = FaceInfer(conf, gpuid=range(conf.num_devs))
net.load_state(
resume_path=args.model,
latest=False,
)
net.model.eval()
features_all = None
i = 0
fstart = 0
buffer = []
for line in open(os.path.join(args.input, 'filelist.txt'), 'r'):
if i % 1000 == 0:
print("processing ", i, data_size, 1. * i / data_size)
i += 1
line = line.strip()
image_path = os.path.join(args.input, line)
buffer.append(image_path)
if len(buffer) == args.batch_size:
embedding = get_feature(buffer)
buffer = []
fend = fstart + embedding.shape[0]
if features_all is None:
features_all = np.zeros((data_size, emb_size),
dtype=np.float32)
# 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)
def main(args):
global image_shape
global net
global ctx_num, env, glargs
print(args)
glargs = args
env = lmdb.open(
args.input + '/imgs_lmdb',
readonly=True,
# max_readers=1, lock=False,
# readahead=False, meminit=False
)
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))
ctx_num = len(ctx)
image_shape = [int(x) for x in args.image_size.split(',')]
if use_mxnet:
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)
else:
# sys.path.insert(0, lz.home_path + 'prj/InsightFace_Pytorch/')
from config import conf
conf.need_log = False
conf.fp16 = True # maybe faster ?
conf.ipabn = False
conf.cvt_ipabn = False
conf.use_chkpnt = False
# conf.net_mode = 'ir_se'
# conf.net_depth = 100
from Learner import FaceInfer
net = FaceInfer(conf, gpuid=range(conf.num_devs))
net.load_state(
resume_path=args.model,
latest=False,
)
net.model.eval()
features_all = None
i = 0
filelist = os.path.join(args.input, 'filelist.txt')
lines = open(filelist, 'r').readlines()
buffer_images = []
buffer_embedding = np.zeros((0, 0), dtype=np.float32)
aggr_nums = []
row_idx = 0
for line in lines:
# if i < 203000:
# i += 1
# continue
if i % 1000 == 0:
print(
"processing ",
i,
len(lines),
i / len(lines),
)
i += 1
# print('stat', i, len(buffer_images), buffer_embedding.shape, aggr_nums, row_idx)
videoname = line.strip().split()[0]
images = glob.glob("%s/%s/*.jpg" % (args.input, videoname))
# images = np.sort(images).tolist()
assert len(images) > 0
image_features = []
for image_path in images:
buffer_images.append(image_path)
aggr_nums.append(len(images))
while len(buffer_images) >= args.batch_size:
embedding = get_feature(buffer_images[0:args.batch_size])
buffer_images = buffer_images[args.batch_size:]
if buffer_embedding.shape[1] == 0:
buffer_embedding = embedding.copy().astype('float32')
else:
buffer_embedding = np.concatenate(
(buffer_embedding, embedding), axis=0)
buffer_idx = 0
acount = 0
for anum in aggr_nums:
if buffer_embedding.shape[0] >= anum + buffer_idx:
image_features = buffer_embedding[buffer_idx:buffer_idx + anum]
video_feature = np.sum(image_features, axis=0, keepdims=True)
video_feature = sklearn.preprocessing.normalize(video_feature)
if features_all is None:
features_all = np.zeros(
(data_size, video_feature.shape[1]), dtype=np.float32)
# print('write to', row_idx, anum, buffer_embedding.shape)
features_all[row_idx] = video_feature.flatten()
row_idx += 1
buffer_idx += anum
acount += 1
else:
break
aggr_nums = aggr_nums[acount:]
buffer_embedding = buffer_embedding[buffer_idx:]
if len(buffer_images) > 0:
embedding = get_feature(buffer_images)
buffer_images = buffer_images[args.batch_size:]
buffer_embedding = np.concatenate((buffer_embedding, embedding),
axis=0)
buffer_idx = 0
acount = 0
for anum in aggr_nums:
assert buffer_embedding.shape[0] >= anum + buffer_idx
image_features = buffer_embedding[buffer_idx:buffer_idx + anum]
video_feature = np.sum(image_features, axis=0, keepdims=True)
video_feature = sklearn.preprocessing.normalize(video_feature)
# print('last write to', row_idx, anum, buffer_embedding.shape)
features_all[row_idx] = video_feature.flatten()
row_idx += 1
buffer_idx += anum
acount += 1
aggr_nums = aggr_nums[acount:]
buffer_embedding = buffer_embedding[buffer_idx:]
# embed()
assert len(aggr_nums) == 0
assert buffer_embedding.shape[0] == 0
write_bin(args.output, features_all)
print(row_idx, features_all.shape)
def main_allimg(args):
global image_shape
global net
global ctx_num, env, glargs
print(args)
glargs = args
env = lmdb.open(
args.input + '/imgs_lmdb',
readonly=True,
# max_readers=1, lock=False,
# readahead=False, meminit=False
)
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))
ctx_num = len(ctx)
image_shape = [int(x) for x in args.image_size.split(',')]
if use_mxnet:
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)
else:
# sys.path.insert(0, lz.home_path + 'prj/InsightFace_Pytorch/')
from config import conf
lz.init_dev(use_devs)
conf.need_log = False
conf.fp16 = True # maybe faster ?
conf.ipabn = False
conf.cvt_ipabn = False
conf.use_chkpnt = False
conf.net_mode = 'ir_se'
conf.net_depth = 100
conf.input_size = 128
conf.embedding_size = 512
from Learner import FaceInfer
net = FaceInfer(
conf,
gpuid=range(len(use_devs)),
)
net.load_state(
resume_path=args.model,
latest=True,
)
net.model.eval()
features_all = None
filelist = os.path.join(args.input, 'filelist.txt')
lines = open(filelist, 'r').readlines()
buffer_images = []
buffer_embedding = np.zeros((0, emb_size), dtype=np.float16)
row_idx = 0
import h5py
f = h5py.File(args.output, 'w')
chunksize = 80 * 10**3
dst = f.create_dataset("feas", (chunksize, 512),
maxshape=(None, emb_size),
dtype='f2')
ind_dst = 0
vdonm2imgs = lz.msgpack_load(args.input + '/../vdonm2imgs.pk')
for line in lines:
if row_idx % 1000 == 0:
logging.info(
f"processing {(row_idx, len(lines), row_idx / len(lines),)}")
row_idx += 1
# if row_idx<203000:continue
# print('stat', i, len(buffer_images), buffer_embedding.shape, aggr_nums, row_idx)
videoname = line.strip().split()[0]
# images2 = glob.glob("%s/%s/*.jpg" % (args.input, videoname))
# images2 = np.sort(images2).tolist()
images = vdonm2imgs[videoname]
assert len(images) > 0
for image_path in images:
buffer_images.append(image_path)
while len(buffer_images) >= args.batch_size:
embedding = get_feature(buffer_images[0:args.batch_size])
buffer_images = buffer_images[args.batch_size:]
if ind_dst + args.batch_size > dst.shape[0]:
dst.resize((dst.shape[0] + chunksize, emb_size), )
dst[ind_dst:ind_dst +
args.batch_size, :] = embedding.astype('float16')
ind_dst += args.batch_size
# buffer_embedding = np.concatenate((buffer_embedding, embedding), axis=0).astype('float16')
if len(buffer_images) != 0:
embedding = get_feature(buffer_images)
if ind_dst + args.batch_size > dst.shape[0]:
dst.resize((dst.shape[0] + chunksize, emb_size), )
dst[ind_dst:ind_dst + args.batch_size, :] = embedding.astype('float16')
# lz.save_mat(args.output, buffer_embedding)
f.flush()
f.close()