ZF/insightface
1
0
Fork 0
mirror of https://github.com/deepinsight/insightface.git synced 2026-05-20 00:10:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

move config.py to sample_config; add more test code

Browse Source
This commit is contained in:
Jia Guo
2019-01-28 14:01:55 +08:00
parent ebc51de83a
commit be4b0eac36
6 changed files with 189 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
alignment/metric.py Normal file
View File

@@ -0,0 +1,99 @@
import mxnet as mx
import numpy as np
import math
import cv2
from config import config
class LossValueMetric(mx.metric.EvalMetric):
def __init__(self):
self.axis = 1
super(LossValueMetric, self).__init__(
'lossvalue', axis=self.axis,
output_names=None, label_names=None)
self.losses = []
def update(self, labels, preds):
loss = preds[0].asnumpy()[0]
self.sum_metric += loss
self.num_inst += 1.0
class NMEMetric(mx.metric.EvalMetric):
def __init__(self):
self.axis = 1
super(NMEMetric, self).__init__(
'NME', axis=self.axis,
output_names=None, label_names=None)
#self.losses = []
self.count = 0
def cal_nme(self, label, pred_label):
nme = []
for b in xrange(pred_label.shape[0]):
record = [None]*6
item = []
if label.ndim==4:
_heatmap = label[b][36]
if np.count_nonzero(_heatmap)==0:
continue
else:#ndim==3
#print(label[b])
if np.count_nonzero(label[b])==0:
continue
for p in xrange(pred_label.shape[1]):
if label.ndim==4:
heatmap_gt = label[b][p]
ind_gt = np.unravel_index(np.argmax(heatmap_gt, axis=None), heatmap_gt.shape)
ind_gt = np.array(ind_gt)
else:
ind_gt = label[b][p]
#ind_gt = ind_gt.astype(np.int)
#print(ind_gt)
heatmap_pred = pred_label[b][p]
heatmap_pred = cv2.resize(heatmap_pred, (config.input_img_size, config.input_img_size))
ind_pred = np.unravel_index(np.argmax(heatmap_pred, axis=None), heatmap_pred.shape)
ind_pred = np.array(ind_pred)
#print(ind_gt.shape)
#print(ind_pred)
if p==36:
#print('b', b, p, ind_gt, np.count_nonzero(heatmap_gt))
record[0] = ind_gt
elif p==39:
record[1] = ind_gt
elif p==42:
record[2] = ind_gt
elif p==45:
record[3] = ind_gt
if record[4] is None or record[5] is None:
record[4] = ind_gt
record[5] = ind_gt
else:
record[4] = np.minimum(record[4], ind_gt)
record[5] = np.maximum(record[5], ind_gt)
#print(ind_gt.shape, ind_pred.shape)
value = np.sqrt(np.sum(np.square(ind_gt - ind_pred)))
item.append(value)
_nme = np.mean(item)
if config.landmark_type=='2d':
left_eye = (record[0]+record[1])/2
right_eye = (record[2]+record[3])/2
_dist = np.sqrt(np.sum(np.square(left_eye - right_eye)))
#print('eye dist', _dist, left_eye, right_eye)
_nme /= _dist
else:
#_dist = np.sqrt(float(label.shape[2]*label.shape[3]))
_dist = np.sqrt(np.sum(np.square(record[5] - record[4])))
#print(_dist)
_nme /= _dist
nme.append(_nme)
return np.mean(nme)
def update(self, labels, preds):
self.count+=1
label = labels[0].asnumpy()
pred_label = preds[-1].asnumpy()
nme = self.cal_nme(label, pred_label)
#print('nme', nme)
#nme = np.mean(nme)
self.sum_metric += np.mean(nme)
self.num_inst += 1.0

14
alignment/config.py → alignment/sample_config.py
View File

@@ -3,11 +3,19 @@ from easydict import EasyDict as edict
config = edict()
#default training/dataset config
config.num_classes = 68
config.record_img_size = 384
config.base_scale = 256
config.input_img_size = 128
config.output_label_size = 64
config.label_xfirst = False
config.losstype = 'heatmap'
config.net_coherent = False
config.multiplier = 1.0
config.gaussian = 0
# network settings
network = edict()
@@ -23,7 +31,7 @@ network.hourglass.losstype = 'heatmap'
network.sdu = edict()
network.sdu.net_coherent = False
network.sdu.net_sta = 2
network.sdu.net_sta = 1
network.sdu.net_n = 3
network.sdu.net_dcn = 3
network.sdu.net_stacks = 2

7
alignment/symbol/sym_heatmap.py
View File

@@ -465,9 +465,10 @@ def ce_loss(x, y):
return loss
def get_symbol(num_classes):
sFilters = 64
mFilters = 128
nFilters = 256
m = config.multiplier
sFilters = max(int(64*m), 32)
mFilters = max(int(128*m), 32)
nFilters = int(256*m)
nModules = 1
nStacks = config.net_stacks

10
alignment/test.py
View File

@@ -13,7 +13,10 @@ from mtcnn_detector import MtcnnDetector
class Handler:
def __init__(self, prefix, epoch, ctx_id=0):
print('loading',prefix, epoch)
ctx = mx.gpu(ctx_id)
if ctx_id>=0:
ctx = mx.gpu(ctx_id)
else:
ctx = mx.cpu()
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
all_layers = sym.get_internals()
sym = all_layers['heatmap_output']
@@ -49,7 +52,7 @@ class Handler:
self.model.forward(db, is_train=False)
alabel = self.model.get_outputs()[-1].asnumpy()[0]
tb = datetime.datetime.now()
print('time cost', (tb-ta).total_seconds())
print('module time cost', (tb-ta).total_seconds())
ret = np.zeros( (alabel.shape[0], 2), dtype=np.float32)
for i in xrange(alabel.shape[0]):
a = cv2.resize(alabel[i], (self.image_size[1], self.image_size[0]))
@@ -65,7 +68,10 @@ img = cv2.imread(img_path)
handler = Handler('./model/HG', 1, ctx_id)
for _ in range(10):
ta = datetime.datetime.now()
landmark, M = handler.get(img)
tb = datetime.datetime.now()
print('get time cost', (tb-ta).total_seconds())
#visualize landmark
IM = cv2.invertAffineTransform(M)
for i in range(landmark.shape[0]):

65
alignment/test_rec_nme.py Normal file
View File

@@ -0,0 +1,65 @@
import argparse
import cv2
import sys
import numpy as np
import os
import mxnet as mx
import datetime
import img_helper
from config import config
from data import FaceSegIter
from metric import LossValueMetric, NMEMetric
parser = argparse.ArgumentParser(description='test nme on rec data')
# general
parser.add_argument('--rec', default='./data_2d/ibug.rec', help='rec data path')
parser.add_argument('--prefix', default='', help='model prefix')
parser.add_argument('--epoch', type=int, default=1, help='model epoch')
parser.add_argument('--gpu', type=int, default=0, help='')
parser.add_argument('--landmark-type', default='2d', help='')
parser.add_argument('--image-size', type=int, default=128, help='')
args = parser.parse_args()
rec_path = args.rec
ctx_id = args.gpu
prefix = args.prefix
epoch = args.epoch
image_size = (args.image_size, args.image_size)
config.landmark_type = args.landmark_type
config.input_img_size = image_size[0]
if ctx_id>=0:
ctx = mx.gpu(ctx_id)
else:
ctx = mx.cpu()
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
all_layers = sym.get_internals()
sym = all_layers['heatmap_output']
#model = mx.mod.Module(symbol=sym, context=ctx, data_names=['data'], label_names=['softmax_label'])
model = mx.mod.Module(symbol=sym, context=ctx, data_names=['data'], label_names=None)
#model = mx.mod.Module(symbol=sym, context=ctx)
model.bind(for_training=False, data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
model.set_params(arg_params, aux_params)
val_iter = FaceSegIter(path_imgrec = rec_path,
batch_size = 1,
aug_level = 0,
)
_metric = NMEMetric()
#val_metric = mx.metric.create(_metric)
#val_metric.reset()
#val_iter.reset()
nme = []
for i, eval_batch in enumerate(val_iter):
if i%10==0:
print('processing', i)
#print(eval_batch.data[0].shape, eval_batch.label[0].shape)
batch_data = mx.io.DataBatch(eval_batch.data)
model.forward(batch_data, is_train=False)
#model.update_metric(val_metric, eval_batch.label, True)
pred_label = model.get_outputs()[-1].asnumpy()
label = eval_batch.label[0].asnumpy()
_nme = _metric.cal_nme(label, pred_label)
nme.append(_nme)
print(np.mean(nme))

99
alignment/train.py
View File

@@ -14,113 +14,18 @@ import math
import random
import cv2
from config import config, default, generate_config
from optimizer import ONadam
from metric import LossValueMetric, NMEMetric
sys.path.append(os.path.join(os.path.dirname(__file__), 'symbol'))
import sym_heatmap
#import sym_fc
#from symbol import fc
from optimizer import ONadam
args = None
logger = logging.getLogger()
logger.setLevel(logging.INFO)
class LossValueMetric(mx.metric.EvalMetric):
def __init__(self):
self.axis = 1
super(LossValueMetric, self).__init__(
'lossvalue', axis=self.axis,
output_names=None, label_names=None)
self.losses = []
def update(self, labels, preds):
loss = preds[0].asnumpy()[0]
self.sum_metric += loss
self.num_inst += 1.0
class NMEMetric(mx.metric.EvalMetric):
def __init__(self):
self.axis = 1
super(NMEMetric, self).__init__(
'NME', axis=self.axis,
output_names=None, label_names=None)
#self.losses = []
self.count = 0
def update(self, labels, preds):
self.count+=1
preds = [preds[-1]]
for label, pred_label in zip(labels, preds):
label = label.asnumpy()
pred_label = pred_label.asnumpy()
#print('label', np.count_nonzero(label[0][36]))
#print('acc',label.shape, pred_label.shape)
#print(label.ndim)
nme = []
for b in xrange(pred_label.shape[0]):
record = [None]*6
item = []
if label.ndim==4:
_heatmap = label[b][36]
if np.count_nonzero(_heatmap)==0:
continue
else:#ndim==3
#print(label[b])
if np.count_nonzero(label[b])==0:
continue
for p in xrange(pred_label.shape[1]):
if label.ndim==4:
heatmap_gt = label[b][p]
ind_gt = np.unravel_index(np.argmax(heatmap_gt, axis=None), heatmap_gt.shape)
ind_gt = np.array(ind_gt)
else:
ind_gt = label[b][p]
#ind_gt = ind_gt.astype(np.int)
#print(ind_gt)
heatmap_pred = pred_label[b][p]
heatmap_pred = cv2.resize(heatmap_pred, (config.input_img_size, config.input_img_size))
ind_pred = np.unravel_index(np.argmax(heatmap_pred, axis=None), heatmap_pred.shape)
ind_pred = np.array(ind_pred)
#print(ind_gt.shape)
#print(ind_pred)
if p==36:
#print('b', b, p, ind_gt, np.count_nonzero(heatmap_gt))
record[0] = ind_gt
elif p==39:
record[1] = ind_gt
elif p==42:
record[2] = ind_gt
elif p==45:
record[3] = ind_gt
if record[4] is None or record[5] is None:
record[4] = ind_gt
record[5] = ind_gt
else:
record[4] = np.minimum(record[4], ind_gt)
record[5] = np.maximum(record[5], ind_gt)
#print(ind_gt.shape, ind_pred.shape)
value = np.sqrt(np.sum(np.square(ind_gt - ind_pred)))
item.append(value)
_nme = np.mean(item)
if config.landmark_type=='2d':
left_eye = (record[0]+record[1])/2
right_eye = (record[2]+record[3])/2
_dist = np.sqrt(np.sum(np.square(left_eye - right_eye)))
#print('eye dist', _dist, left_eye, right_eye)
_nme /= _dist
else:
#_dist = np.sqrt(float(label.shape[2]*label.shape[3]))
_dist = np.sqrt(np.sum(np.square(record[5] - record[4])))
#print(_dist)
_nme /= _dist
nme.append(_nme)
#print('nme', nme)
#nme = np.mean(nme)
if len(nme)>0:
self.sum_metric += np.mean(nme)
self.num_inst += 1.0
def main(args):
_seed = 727