优化代码

.gitignore

@@ -6,4 +6,5 @@ AFAD/
 AgeDB/
 megaage_asian/
 temp/
-test.py
+test.py
+result.jpg

README.md

@@ -10,4 +10,57 @@
 
 **测试数据8254张，准确率：**
  - 性别准确率：0.972868
- - 年龄准确率：0.761628 (+-3)
+ - 年龄准确率：0.761628 (+-3)
+
+
+```shell
+100%|██████████| 1032/1032 [00:06<00:00, 153.75it/s]
+性别准确率：0.972868
+年龄准确率：0.761628
+```
+
+![年龄分布](https://img-blog.csdnimg.cn/20210407162913663.png)
+
+![识别结果](https://img-blog.csdnimg.cn/20210407165918268.jpg)
+
+```shell
+第1张人脸，位置(160, 32, 204, 84), 性别：男, 年龄：30
+第2张人脸，位置(545, 162, 579, 206), 性别：女, 年龄：31
+第3张人脸，位置(632, 118, 666, 158), 性别：男, 年龄：28
+第4张人脸，位置(91, 159, 151, 237), 性别：男, 年龄：38
+第5张人脸，位置(723, 123, 760, 169), 性别：男, 年龄：26
+第6张人脸，位置(263, 120, 317, 191), 性别：男, 年龄：27
+第7张人脸，位置(438, 134, 481, 190), 性别：男, 年龄：46
+第8张人脸，位置(908, 160, 963, 224), 性别：男, 年龄：35
+第9张人脸，位置(39, 51, 81, 102), 性别：女, 年龄：31
+第10张人脸，位置(807, 148, 847, 196), 性别：女, 年龄：26
+第11张人脸，位置(449, 40, 485, 84), 性别：男, 年龄：29
+第12张人脸，位置(378, 46, 412, 86), 性别：女, 年龄：33
+第13张人脸，位置(534, 46, 567, 83), 性别：男, 年龄：30
+第14张人脸，位置(272, 20, 311, 67), 性别：男, 年龄：28
+第15张人脸，位置(358, 216, 375, 237), 性别：男, 年龄：27
+```
+
+```shell
+gpu num: 1
+num_layers 50
+data_shape [3, 112, 112]
+Called with argument: Namespace(batch_size=128, color=0, ctx_num=1, cutoff=0, data_dir='dataset', data_shape='3,112,112', end_epoch=200, gpu_ids='0', image_channel=3, image_h=112, image_w=112, lr=0.1, lr_steps='10,30,80,150,200', network='m50', num_layers=50, prefix='temp/model', pretrained='', rand_mirror=1, rescale_threshold=0, version_input=1, version_output='GAP')
+1 GAP 32
+INFO:root:loading recordio dataset\train.rec...
+INFO:root:dataset\train.rec 数据大小：303018
+INFO:root:是否随机翻转图片：1
+INFO:root:loading recordio dataset\val.rec...
+INFO:root:dataset\val.rec 数据大小：1032
+INFO:root:是否随机翻转图片：False
+call reset()
+开始训练...
+INFO:root:Epoch[0] Batch [0-20]	Speed: 520.85 samples/sec	acc=0.572545	MAE=10.734747	CUM_5=0.240699
+INFO:root:Epoch[0] Batch [20-40]	Speed: 518.95 samples/sec	acc=0.589844	MAE=9.351172	CUM_5=0.289844
+INFO:root:Epoch[0] Batch [40-60]	Speed: 516.86 samples/sec	acc=0.603125	MAE=9.184766	CUM_5=0.303906
+INFO:root:Epoch[0] Batch [60-80]	Speed: 508.44 samples/sec	acc=0.609766	MAE=8.759375	CUM_5=0.336719
+INFO:root:Epoch[0] Batch [80-100]	Speed: 461.26 samples/sec	acc=0.656250	MAE=8.224609	CUM_5=0.361328
+INFO:root:Epoch[0] Batch [100-120]	Speed: 518.43 samples/sec	acc=0.696875	MAE=7.611328	CUM_5=0.400391
+INFO:root:Epoch[0] Batch [120-140]	Speed: 514.88 samples/sec	acc=0.715234	MAE=7.224609	CUM_5=0.426172
+INFO:root:Epoch[0] Batch [140-160]	Speed: 517.80 samples/sec	acc=0.722266	MAE=6.976172	CUM_5=0.437500
+```

create_dataset.py

@@ -3,28 +3,23 @@
 import mxnet as mx
 import numpy as np
 from tqdm import tqdm
-from mtcnn_detector import MtcnnDetector
+from utils.mtcnn_detector import MtcnnDetector
 import random
-import face_preprocess
+from utils import face_preprocess
 
 
 class FaceModel:
-    def __init__(self, det):
-        self.det = det
+    def __init__(self):
         ctx = mx.gpu(0)
         self.det_threshold = [0.6, 0.7, 0.8]
         mtcnn_path = 'mtcnn-model'
-        if det == 0:
-            detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark=True,
-                                     threshold=self.det_threshold)
-        else:
-            detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark=True,
-                                     threshold=[0.0, 0.0, 0.2])
+        detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark=True,
+                                 threshold=self.det_threshold)
         print("加载模型：%s" % mtcnn_path)
         self.detector = detector
 
     def get_face(self, face_img, random_bg=False):
-        ret = self.detector.detect_face(face_img, det_type=self.det)
+        ret = self.detector.detect_face(face_img)
         if ret is None:
             return None
         bbox, points = ret
@@ -84,7 +79,7 @@ def get_ga(self, nimg):
 def create_face(image_dir, random_bg=False):
     if not os.path.exists("dataset"):
         os.mkdir("dataset")
-    faceModel = FaceModel(0)
+    faceModel = FaceModel()
     images = []
     for root, dirs, files in os.walk(image_dir):
         for image in files:

eval.py

@@ -0,0 +1,64 @@
+import argparse
+from io import BytesIO
+
+import mxnet as mx
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+from utils.data import FaceImageIter
+from infer import FaceAgeGenderModel
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--image_size',  default='112,112',         help='models input size.')
+parser.add_argument('--model',       default='model/model,200', help='path to load model.')
+parser.add_argument('--mtcnn_model', default='mtcnn-model',     help='path to load model.')
+parser.add_argument('--gpu',         default=0,   type=int,     help='gpu id')
+args = parser.parse_args()
+
+
+def eval():
+    model = FaceAgeGenderModel(args)
+
+    val_dataiter = FaceImageIter(batch_size=1,
+                                 data_shape=(3, 112, 112),
+                                 path_imgrec="dataset/val.rec",
+                                 shuffle=False,
+                                 rand_mirror=False,
+                                 mean=0)
+    gender_correct = 0
+    age_correct = 0
+    pbar = tqdm(total=len(val_dataiter.seq))
+    while True:
+        try:
+            label, s, _, _ = val_dataiter.next_sample()
+            gender = int(label[0])
+            age = int(label[1])
+            img = val_dataiter.imdecode(s)
+            buf = BytesIO()
+            img = Image.fromarray(img.asnumpy(), 'RGB')
+            img.save(buf, format='JPEG', quality=100)
+            buf = buf.getvalue()
+            img = Image.open(BytesIO(buf))
+
+            img = np.array(img).astype(np.float32)
+            img = np.transpose(img, (2, 0, 1))
+            img = np.expand_dims(img, axis=0)
+            img = mx.nd.array(img)
+            img = mx.io.DataBatch(data=(img,))
+
+            gender1, age1 = model.get_ga(img)
+            if gender == gender1:
+                gender_correct += 1
+            if age == age1 or age == age1 + 1 or age == age1 + 2 or age == age1 + 3 or \
+                    age == age1 - 1 or age == age1 - 2 or age == age1 - 3:
+                age_correct += 1
+            pbar.update()
+        except StopIteration:
+            break
+    print("性别准确率：%f" % (gender_correct / len(val_dataiter.seq)))
+    print("年龄准确率：%f" % (age_correct / len(val_dataiter.seq)))
+
+
+if __name__ == '__main__':
+    eval()

infer.py

@@ -1,21 +1,18 @@
 import argparse
+
 import cv2
 import mxnet as mx
 import numpy as np
-import face_preprocess
-from data import FaceImageIter
-from io import BytesIO
-from PIL import Image
-from tqdm import tqdm
-from mtcnn_detector import MtcnnDetector
 
-parser = argparse.ArgumentParser(description='face model test')
-parser.add_argument('--image-size', default='112,112', help='')
-parser.add_argument('--image', default='temp/8DC8CBE1.png', help='')
-parser.add_argument('--model', default='model/model,0', help='path to load model.')
-parser.add_argument('--mtcnn-model', default='mtcnn-model', help='path to load model.')
-parser.add_argument('--gpu', default=0, type=int, help='gpu id')
-parser.add_argument('--det', default=0, type=int, help='mtcnn option, 1 means using R+O, 0 means detect from begining')
+from utils import face_preprocess
+from utils.mtcnn_detector import MtcnnDetector
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--image_size',   default='112,112',           help='models input size.')
+parser.add_argument('--image',        default='test.jpg',          help='infer image path.')
+parser.add_argument('--model',        default='model/model,200',   help='path to load model.')
+parser.add_argument('--mtcnn_model',  default='mtcnn-model',       help='path to load model.')
+parser.add_argument('--gpu',          default=0, type=int,         help='gpu id')
 args = parser.parse_args()
 
 
@@ -36,15 +33,12 @@ def __init__(self, args):
         self.det_minsize = 50
         self.det_threshold = [0.6, 0.7, 0.8]
         self.image_size = image_size
-        if args.det == 0:
-            detector = MtcnnDetector(model_folder=args.mtcnn_model, ctx=ctx, num_worker=1, accurate_landmark=True,
-                                     threshold=self.det_threshold)
-        else:
-            detector = MtcnnDetector(model_folder=args.mtcnn_model, ctx=ctx, num_worker=1, accurate_landmark=True,
-                                     threshold=[0.0, 0.0, 0.2])
+        detector = MtcnnDetector(model_folder=args.mtcnn_model, ctx=ctx, num_worker=1, accurate_landmark=True,
+                                 threshold=self.det_threshold)
         print("加载模型：%s" % args.mtcnn_model)
         self.detector = detector
 
+    # 加载模型
     def get_model(self, ctx, image_size, model_str, layer):
         _vec = model_str.split(',')
         assert len(_vec) == 2
@@ -59,23 +53,32 @@ def get_model(self, ctx, image_size, model_str, layer):
         model.set_params(arg_params, aux_params)
         return model
 
-    def get_face(self, face_img):
-        ret = self.detector.detect_face(face_img, det_type=self.args.det)
+    # 识别人脸
+    def get_faces(self, face_img):
+        ret = self.detector.detect_face(face_img)
         if ret is None:
             return None
         bbox, points = ret
         if bbox.shape[0] == 0:
-            return None
-        bbox = bbox[0, 0:4]
-        points = points[0, :].reshape((2, 5)).T
-        nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
-        nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
-        aligned = np.transpose(nimg, (2, 0, 1))
-        input_blob = np.expand_dims(aligned, axis=0)
-        data = mx.nd.array(input_blob)
-        db = mx.io.DataBatch(data=(data,))
-        return db
-
+            return [], [], []
+        bboxes = []
+        pointses = []
+        faces = []
+        for i in range(len(bbox)):
+            b = bbox[i, 0:4]
+            bboxes.append(b)
+            p = points[i, :].reshape((2, 5)).T
+            pointses.append(p)
+            nimg = face_preprocess.preprocess(face_img, b, p, image_size='112,112')
+            nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
+            aligned = np.transpose(nimg, (2, 0, 1))
+            input_blob = np.expand_dims(aligned, axis=0)
+            data = mx.nd.array(input_blob)
+            db = mx.io.DataBatch(data=(data,))
+            faces.append(db)
+        return faces, bboxes, pointses
+
+    # 性别年龄识别
     def get_ga(self, data):
         self.model.forward(data, is_train=False)
         ret = self.model.get_outputs()[0].asnumpy()
@@ -87,60 +90,47 @@ def get_ga(self, data):
         return gender, age
 
 
-def main():
-    model = FaceAgeGenderModel(args)
-    img = cv2.imread(args.image)
-    img = model.get_face(img)
-    print(img)
-    if img is not None:
-        gender, age = model.get_ga(img)
-        print('gender is', gender)
-        print('age is', age)
+# 画出人脸框和关键点
+def draw_face(image_path, bboxes, pointses, genderes, ages):
+    img = cv2.imread(image_path)
+    for i in range(len(bboxes)):
+        bbox = bboxes[i]
+        corpbbox = [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])]
+        # 画人脸框
+        cv2.rectangle(img, (corpbbox[0], corpbbox[1]),
+                      (corpbbox[2], corpbbox[3]), (255, 0, 0), 1)
+        # 判别为人脸的置信度
+        cv2.putText(img, '{},{}'.format('m' if genderes[i] == 0 else 'f', ages[i]),
+                    (corpbbox[0], corpbbox[1] - 2),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+    # 画关键点
+    for i in range(len(pointses)):
+        points = pointses[i]
+        for p in points:
+            cv2.circle(img, (int(p[0]), int(p[1])), 2, (0, 0, 255))
+    cv2.imwrite('result.jpg', img)
+    cv2.imshow('result', img)
+    cv2.waitKey(0)
 
 
-def eval_val():
+def main():
+    # 加载模型
     model = FaceAgeGenderModel(args)
-
-    val_dataiter = FaceImageIter(batch_size=1,
-                                 data_shape=(3, 112, 112),
-                                 path_imgrec="dataset/val.rec",
-                                 shuffle=False,
-                                 rand_mirror=False,
-                                 mean=0)
-    gender_correct = 0
-    age_correct = 0
-    pbar = tqdm(total=len(val_dataiter.seq))
-    while True:
-        try:
-            label, s, _, _ = val_dataiter.next_sample()
-            gender = int(label[0])
-            age = int(label[1])
-            img = val_dataiter.imdecode(s)
-            buf = BytesIO()
-            img = Image.fromarray(img.asnumpy(), 'RGB')
-            img.save(buf, format='JPEG', quality=100)
-            buf = buf.getvalue()
-            img = Image.open(BytesIO(buf))
-
-            img = np.array(img).astype(np.float32)
-            img = np.transpose(img, (2, 0, 1))
-            img = np.expand_dims(img, axis=0)
-            img = mx.nd.array(img)
-            img = mx.io.DataBatch(data=(img,))
-
-            gender1, age1 = model.get_ga(img)
-            if gender == gender1:
-                gender_correct += 1
-            if age == age1 or age == age1 + 1 or age == age1 + 2 or age == age1 + 3 or\
-                    age == age1 - 1 or age == age1 - 2 or age == age1 - 3:
-                age_correct += 1
-            pbar.update()
-        except StopIteration:
-            break
-    print("性别准确率：%f" % (gender_correct / len(val_dataiter.seq)))
-    print("年龄准确率：%f" % (age_correct / len(val_dataiter.seq)))
+    # 读取图片
+    img = cv2.imread(args.image)
+    # 检测人脸
+    faces, bboxes, pointses = model.get_faces(img)
+    genderes = []
+    ages = []
+    for i in range(len(faces)):
+        # 识别性别年龄
+        gender, age = model.get_ga(faces[i])
+        print('第%d张人脸，位置(%d, %d, %d, %d), 性别：%s, 年龄：%d' % (i + 1, bboxes[i][0], bboxes[i][1], bboxes[i][2],
+                                                           bboxes[i][3], '女' if gender == 0 else '男', age))
+        genderes.append(gender)
+        ages.append(age)
+    draw_face(args.image, bboxes, pointses, genderes, ages)
 
 
 if __name__ == '__main__':
     main()
-    # eval_val()