From be4b0eac36791d1a73bf73c7c155043bd2ffcbfd Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Mon, 28 Jan 2019 14:01:55 +0800 Subject: [PATCH] move config.py to sample_config; add more test code --- alignment/metric.py | 99 +++++++++++++++++++++++ alignment/{config.py => sample_config.py} | 14 +++- alignment/symbol/sym_heatmap.py | 7 +- alignment/test.py | 10 ++- alignment/test_rec_nme.py | 65 +++++++++++++++ alignment/train.py | 99 +---------------------- 6 files changed, 189 insertions(+), 105 deletions(-) create mode 100644 alignment/metric.py rename alignment/{config.py => sample_config.py} (87%) create mode 100644 alignment/test_rec_nme.py diff --git a/alignment/metric.py b/alignment/metric.py new file mode 100644 index 0000000..3a4933d --- /dev/null +++ b/alignment/metric.py @@ -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 diff --git a/alignment/config.py b/alignment/sample_config.py similarity index 87% rename from alignment/config.py rename to alignment/sample_config.py index a575b64..0870638 100644 --- a/alignment/config.py +++ b/alignment/sample_config.py @@ -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 diff --git a/alignment/symbol/sym_heatmap.py b/alignment/symbol/sym_heatmap.py index 9870128..2be3759 100644 --- a/alignment/symbol/sym_heatmap.py +++ b/alignment/symbol/sym_heatmap.py @@ -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 diff --git a/alignment/test.py b/alignment/test.py index 84f8a01..8b66c2d 100644 --- a/alignment/test.py +++ b/alignment/test.py @@ -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]): diff --git a/alignment/test_rec_nme.py b/alignment/test_rec_nme.py new file mode 100644 index 0000000..af11971 --- /dev/null +++ b/alignment/test_rec_nme.py @@ -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)) + diff --git a/alignment/train.py b/alignment/train.py index ba1420b..ee2608b 100644 --- a/alignment/train.py +++ b/alignment/train.py @@ -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