mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-15 12:52:47 +00:00
tiny
This commit is contained in:
@@ -1,239 +0,0 @@
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
#os.environ["CUDA_VISIBLE_DEVICES"] = "7"
|
||||
|
||||
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
|
||||
from sklearn.preprocessing import normalize
|
||||
#import facenet
|
||||
#import lfw
|
||||
import mxnet as mx
|
||||
from mxnet import ndarray as nd
|
||||
#from caffe.proto import caffe_pb2
|
||||
|
||||
facescrub_aligned_root = '/raid5data/dplearn/megaface/facescrub_aligned_mtcnn'
|
||||
facescrub_aligned_root = None
|
||||
megaface_out = '/raid5data/dplearn/megaface/MegaFace_Features'
|
||||
#facescrub_out = '/raid5data/dplearn/megaface/FaceScrubSubset_Features'
|
||||
facescrub_out = '/raid5data/dplearn/megaface/FaceScrub_Features'
|
||||
|
||||
|
||||
def do_flip(data):
|
||||
for idx in xrange(data.shape[0]):
|
||||
data[idx,:,:] = np.fliplr(data[idx,:,:])
|
||||
|
||||
def ch_dev(arg_params, aux_params, ctx):
|
||||
new_args = dict()
|
||||
new_auxs = dict()
|
||||
for k, v in arg_params.items():
|
||||
new_args[k] = v.as_in_context(ctx)
|
||||
for k, v in aux_params.items():
|
||||
new_auxs[k] = v.as_in_context(ctx)
|
||||
return new_args, new_auxs
|
||||
|
||||
def get_feature(image_path, bbox, landmark, nets, image_shape, use_align, aligned, use_mean):
|
||||
img = face_preprocess.read_image(image_path, mode='rgb')
|
||||
#print(img.shape)
|
||||
if img is None:
|
||||
print('parse image',image_path,'error')
|
||||
return None
|
||||
if not aligned:
|
||||
_landmark = landmark
|
||||
if not use_align:
|
||||
_landmark = None
|
||||
#cv2.imwrite("./align/origin_%s"%image_path.split('/')[-1], img)
|
||||
img = face_preprocess.preprocess(img, bbox=bbox, landmark=_landmark, image_size='%d,%d'%(image_shape[1], image_shape[2]))
|
||||
else:
|
||||
assert img.shape==(image_shape[1],image_shape[2],image_shape[0])
|
||||
#print('already aligned', image_path, img.shape)
|
||||
#img = cv2.resize(img, (image_shape[2], image_shape[1]))
|
||||
#cv2.imwrite("./align/%s"%image_path.split('/')[-1], img)
|
||||
if use_mean>0:
|
||||
v_mean = np.array([127.5,127.5,127.5], dtype=np.float32).reshape( (1,1,3) )
|
||||
img = img.astype(np.float32) - v_mean
|
||||
img *= 0.0078125
|
||||
img = np.transpose( img, (2,0,1) )
|
||||
F = None
|
||||
for net in nets:
|
||||
embedding = None
|
||||
#ppatch = net.patch
|
||||
for flipid in [0,1]:
|
||||
_img = np.copy(img)
|
||||
if flipid==1:
|
||||
do_flip(_img)
|
||||
#nimg = np.zeros(_img.shape, dtype=np.float32)
|
||||
#nimg[:,ppatch[1]:ppatch[3],ppatch[0]:ppatch[2]] = _img[:, ppatch[1]:ppatch[3], ppatch[0]:ppatch[2]]
|
||||
#_img = nimg
|
||||
input_blob = np.expand_dims(_img, axis=0)
|
||||
net.arg_params["data"] = mx.nd.array(input_blob, net.ctx)
|
||||
net.arg_params["softmax_label"] = mx.nd.empty((1,), net.ctx)
|
||||
exe = net.sym.bind(net.ctx, net.arg_params ,args_grad=None, grad_req="null", aux_states=net.aux_params)
|
||||
exe.forward(is_train=False)
|
||||
_embedding = exe.outputs[0].asnumpy().flatten()
|
||||
#print(_embedding.shape)
|
||||
if embedding is None:
|
||||
embedding = _embedding
|
||||
else:
|
||||
embedding += _embedding
|
||||
_norm=np.linalg.norm(embedding)
|
||||
embedding /= _norm
|
||||
if F is None:
|
||||
F = embedding
|
||||
else:
|
||||
#F += embedding
|
||||
F = np.concatenate((F,embedding), axis=0)
|
||||
_norm=np.linalg.norm(F)
|
||||
F /= _norm
|
||||
#print(F.shape)
|
||||
return F
|
||||
|
||||
|
||||
def write_bin(path, feature):
|
||||
feature = list(feature)
|
||||
with open(path, 'wb') as f:
|
||||
f.write(struct.pack('4i', len(feature),1,4,5))
|
||||
f.write(struct.pack("%df"%len(feature), *feature))
|
||||
|
||||
def main(args):
|
||||
|
||||
print(args)
|
||||
gpuid = args.gpu
|
||||
ctx = mx.gpu(gpuid)
|
||||
nets = []
|
||||
for model in args.model.split('|'):
|
||||
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)
|
||||
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']
|
||||
#_pp = prefix.rfind('p')+1
|
||||
#_pp = prefix[_pp:]
|
||||
#net.patch = [int(x) for x in _pp.split('_')]
|
||||
#assert len(net.patch)==5
|
||||
#print('patch', net.patch)
|
||||
nets.append(net)
|
||||
image_shape = [int(x) for x in args.image_size.split(',')]
|
||||
|
||||
#megaface_lst = "/raid5data/dplearn/faceinsight_align_megaface.lst"
|
||||
megaface_lst = "/raid5data/dplearn/megaface/megaface_mtcnn_112x112/lst"
|
||||
facescrub_lst = "/raid5data/dplearn/faceinsight_align_facescrub.lst"
|
||||
|
||||
if args.skip==0:
|
||||
i = 0
|
||||
succ = 0
|
||||
for line in open(facescrub_lst, 'r'):
|
||||
if i%1000==0:
|
||||
print("writing fs",i, succ)
|
||||
i+=1
|
||||
image_path, label, bbox, landmark, aligned = face_preprocess.parse_lst_line(line)
|
||||
_path = image_path.split('/')
|
||||
a,b = _path[-2], _path[-1]
|
||||
#a = a.replace(' ', '_')
|
||||
#b = b.replace(' ', '_')
|
||||
file, ext = os.path.splitext(b)
|
||||
image_id = int(file.split('_')[-1])
|
||||
if image_id==40499 or image_id==10788 or image_id==2367:
|
||||
b = file
|
||||
#b = file+".png"
|
||||
out_dir = os.path.join(facescrub_out, a)
|
||||
#print(out_dir)
|
||||
if not os.path.exists(out_dir):
|
||||
os.makedirs(out_dir)
|
||||
#continue
|
||||
if len(ext)==0:
|
||||
print(image_path)
|
||||
image_path = image_path+".jpg"
|
||||
if facescrub_aligned_root is not None:
|
||||
_vec = image_path.split('/')
|
||||
_image_path = os.path.join(facescrub_aligned_root, _vec[-2], _vec[-1])
|
||||
_base, _ext = os.path.splitext(_image_path)
|
||||
if _ext=='.gif':
|
||||
_image_path = _base+".jpg"
|
||||
print('changing', _image_path)
|
||||
if os.path.exists(_image_path):
|
||||
image_path = _image_path
|
||||
bbox = None
|
||||
landmark = None
|
||||
aligned = True
|
||||
else:
|
||||
print("not aligned:",_image_path)
|
||||
feature = get_feature(image_path, bbox, landmark, nets, image_shape, True, aligned, args.mean)
|
||||
if feature is None:
|
||||
continue
|
||||
#print(np.linalg.norm(feature))
|
||||
out_path = os.path.join(out_dir, b+"_%s_%dx%d.bin"%(args.algo, image_shape[1], image_shape[2]))
|
||||
write_bin(out_path, feature)
|
||||
succ+=1
|
||||
print('fs stat',i, succ)
|
||||
#return
|
||||
|
||||
i = 0
|
||||
succ = 0
|
||||
for line in open(megaface_lst, 'r'):
|
||||
if i%1000==0:
|
||||
print("writing mf",i, succ)
|
||||
i+=1
|
||||
if i<=args.skip:
|
||||
continue
|
||||
image_path, label, bbox, landmark, aligned = face_preprocess.parse_lst_line(line)
|
||||
assert aligned==True
|
||||
_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)
|
||||
feature = get_feature(image_path, bbox, landmark, nets, image_shape, True, aligned, args.mean)
|
||||
if feature is None:
|
||||
continue
|
||||
out_path = os.path.join(out_dir, b+"_%s_%dx%d.bin"%(args.algo, image_shape[1], image_shape[2]))
|
||||
#print(out_path)
|
||||
write_bin(out_path, feature)
|
||||
succ+=1
|
||||
print('mf stat',i, succ)
|
||||
|
||||
|
||||
def parse_arguments(argv):
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--batch_size', type=int, help='', default=100)
|
||||
parser.add_argument('--image_size', type=str, help='', default='3,112,112')
|
||||
parser.add_argument('--gpu', type=int, help='', default=0)
|
||||
parser.add_argument('--mean', type=int, help='', default=0)
|
||||
parser.add_argument('--seed', type=int, help='', default=727)
|
||||
parser.add_argument('--skip', type=int, help='', default=0)
|
||||
parser.add_argument('--algo', type=str, help='', default='mxsphereface20c')
|
||||
#parser.add_argument('--model', type=str, help='', default='../model/sphereface-20-p0_0_96_112_0,22|../model/sphereface-20-p0_0_96_95_0,21|../model/sphereface-20-p0_0_80_95_0,21')
|
||||
#parser.add_argument('--model', type=str, help='', default='../model/sphereface-s60-p0_0_96_112_0,31|../model/sphereface-s60-p0_0_96_95_0,21|../model/sphereface2-s60-p0_0_96_112_0,21|../model/sphereface3-s60-p0_0_96_95_0,23')
|
||||
#parser.add_argument('--model', type=str, help='', default='../model/sphereface-s60-p0_0_96_112_0,31|../model/sphereface-s60-p0_0_96_95_0,21|../model/sphereface2-s60-p0_0_96_112_0,21|../model/sphereface3-s60-p0_0_96_95_0,23|../model/sphereface-20-p0_0_96_112_0,22|../model/sphereface-20-p0_0_96_95_0,21|../model/sphereface-20-p0_0_80_95_0,21')
|
||||
#parser.add_argument('--model', type=str, help='', default='../model/spherefacei-s60-p0_0_96_112_0,135')
|
||||
#parser.add_argument('--model', type=str, help='', default='../model/spherefacei-s60-p0_0_96_95_0,95')
|
||||
parser.add_argument('--model', type=str, help='', default='../model/spherefacei-s60-p0_15_96_112_0,95')
|
||||
return parser.parse_args(argv)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(parse_arguments(sys.argv[1:]))
|
||||
|
||||
Reference in New Issue
Block a user