From 668d6c748ab00a6b53d2b4c241c03c18d227585b Mon Sep 17 00:00:00 2001
From: "Eric.Lee2021" <305141918@qq.com>
Date: Tue, 27 Apr 2021 23:38:16 +0800
Subject: [PATCH] add face au

---
 applications/face_bioassay_local_app.py     |   6 +-
 components/face_au/face_au_c.py             |  68 +++++
 components/face_au/models/mobilenetv2.py    | 143 +++++++++++
 components/face_au/models/resnet.py         | 262 ++++++++++++++++++++
 components/face_au/utils/common_utils.py    | 132 ++++++++++
 components/face_au/utils/model_utils.py     |  61 +++++
 lib/face_bioassay_lib/cores/face_fuction.py |  40 ++-
 7 files changed, 706 insertions(+), 6 deletions(-)
 create mode 100644 components/face_au/face_au_c.py
 create mode 100644 components/face_au/models/mobilenetv2.py
 create mode 100644 components/face_au/models/resnet.py
 create mode 100644 components/face_au/utils/common_utils.py
 create mode 100644 components/face_au/utils/model_utils.py

diff --git a/applications/face_bioassay_local_app.py b/applications/face_bioassay_local_app.py
index 0d8eccf..846448d 100644
--- a/applications/face_bioassay_local_app.py
+++ b/applications/face_bioassay_local_app.py
@@ -28,6 +28,7 @@ import shutil
 from face_detect.yolo_v3_face import yolo_v3_face_model
 from face_multi_task.face_multi_task_component import FaceMuitiTask_Model
 from face_euler_angle.face_euler_angle_component import FaceAngle_Model
+from face_au.face_au_c import FaceAu_Model
 # 加载工具库
 import sys
 sys.path.append("./lib/face_bioassay_lib/")
@@ -47,6 +48,8 @@ def main_face_bioassay(video_path,cfg_file):
 
     face_euler_model = FaceAngle_Model(model_path = config["face_euler_model_path"])
 
+    face_au_model = FaceAu_Model(model_path = config["face_au_model_path"])
+
     cap = cv2.VideoCapture(video_path)
     frame_idx = 0
     while cap.isOpened():
@@ -58,9 +61,10 @@ def main_face_bioassay(video_path,cfg_file):
             faces_bboxes =face_detect_model.predict(img,vis = False) # 检测手，获取手的边界框
 
             faces_message = get_faces_batch_attribute(
-                face_multitask_model,face_euler_model,faces_bboxes,img,use_cuda = True,vis = True)
+                face_multitask_model,face_euler_model,face_au_model,faces_bboxes,img,use_cuda = True,vis = True)
             if faces_message is not None:
                 print("faces_message : {} \n".format(faces_message))
+                pass
 
             cv2.namedWindow("DriverFatigueMonitor",0)
             cv2.imshow("DriverFatigueMonitor",img)
diff --git a/components/face_au/face_au_c.py b/components/face_au/face_au_c.py
new file mode 100644
index 0000000..4c036b9
--- /dev/null
+++ b/components/face_au/face_au_c.py
@@ -0,0 +1,68 @@
+#-*-coding:utf-8-*-
+# date:2021-03-09
+# Author: Eric.Lee
+# function: handpose_x 21 keypoints 2D
+
+import os
+import torch
+import cv2
+import numpy as np
+import json
+
+import torch
+import torch.nn as nn
+
+import time
+import math
+from datetime import datetime
+
+from face_au.models.resnet import resnet18, resnet34, resnet50, resnet101
+from face_au.models.mobilenetv2 import MobileNetV2
+
+#
+class FaceAu_Model(object):
+    def __init__(self,
+        model_path = './components/face_au/weights/face_au-resnet50-size256-20210427.pth',
+        img_size= 256,
+        num_classes = 24,
+        model_arch = "resnet_50",
+        ):
+        # print("face au loading : ",model_path)
+        self.use_cuda = torch.cuda.is_available()
+        self.device = torch.device("cuda:0" if self.use_cuda else "cpu") # 可选的设备类型及序号
+        self.img_size = img_size
+        #-----------------------------------------------------------------------
+
+        if model_arch == 'resnet_50':
+            model_ = resnet50(num_classes = num_classes,img_size = self.img_size)
+        elif model_arch == 'resnet_18':
+            model_ = resnet18(num_classes = num_classes,img_size = self.img_size)
+        elif model_arch == 'resnet_34':
+            model_ = resnet34(num_classes = num_classes,img_size = self.img_size)
+        elif model_arch == 'resnet_101':
+            model_ = resnet101(num_classes = num_classes,img_size = self.img_size)
+
+        #-----------------------------------------------------------------------
+        model_ = model_.to(self.device)
+        # 加载测试模型
+        if os.access(model_path,os.F_OK):# checkpoint
+            chkpt = torch.load(model_path, map_location=self.device)
+            model_.load_state_dict(chkpt)
+            print('face au model loading : {}'.format(model_path))
+
+        model_.eval() # 设置为前向推断模式
+        self.model_au = model_
+
+    def predict(self, img, vis = False):
+        with torch.no_grad():
+            # img_ = img_ - [123.67, 116.28, 103.53]
+            img_ = torch.from_numpy(img)
+            # img_ = img_.unsqueeze_(0)
+            if self.use_cuda:
+                img_ = img_.cuda()  # (bs, 3, h, w)
+
+            output_ = self.model_au(img_.float())
+            # print(pre_.size())
+            output_ = output_.cpu().detach().numpy()
+
+        return output_
diff --git a/components/face_au/models/mobilenetv2.py b/components/face_au/models/mobilenetv2.py
new file mode 100644
index 0000000..b3f5492
--- /dev/null
+++ b/components/face_au/models/mobilenetv2.py
@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+import math
+
+# A PyTorch implementation of MobileNet V2 architecture and pretrained model.
+# from https://github.com/tonylins/pytorch-mobilenet-v2
+
+def conv_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, n_class=1000, input_size=224,dropout_factor=0.5, width_mult=1.):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        self.dropout_factor = dropout_factor
+        self.dropout = nn.Dropout(self.dropout_factor)
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        assert input_size % 32 == 0
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
+        self.features = [conv_bn(3, input_channel, 2)]
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
+                input_channel = output_channel
+        # building last several layers
+        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
+        # make it nn.Sequential
+        self.features = nn.Sequential(*self.features)
+
+        # building classifier
+        self.classifier = nn.Sequential(
+            # nn.Dropout(0.2),
+            nn.Linear(self.last_channel, n_class),
+        )
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        # print("1 x.size()", x.size())  # bs,1280,7,7
+        x = x.mean(3).mean(2)
+        # print("2 x.size()", x.size())  # bs,1280
+        x = self.dropout(x)
+        x = self.classifier(x)
+        # print("3 x.size()", x.size())  # bs,num_classes
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                n = m.weight.size(1)
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+
+if __name__ == "__main__":
+    img = torch.rand((9, 3, 160, 160))
+    net = MobileNetV2(n_class=120)
+    print(net)
+    print("***"*50)
+    pred = net(img)
+    print(pred.size())
diff --git a/components/face_au/models/resnet.py b/components/face_au/models/resnet.py
new file mode 100644
index 0000000..433b22f
--- /dev/null
+++ b/components/face_au/models/resnet.py
@@ -0,0 +1,262 @@
+import torch
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, img_size=224,dropout_factor = 1.):
+        self.inplanes = 64
+        self.dropout_factor = dropout_factor
+        super(ResNet, self).__init__()
+    
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        # see this issue: https://github.com/xxradon/PytorchToCaffe/issues/16
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        assert img_size % 32 == 0
+        pool_kernel = int(img_size / 32)
+        self.avgpool = nn.AvgPool2d(pool_kernel, stride=1, ceil_mode=True)
+
+        self.dropout = nn.Dropout(self.dropout_factor)
+
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+
+        x = self.dropout(x)
+
+        x = self.fc(x)
+
+        return x
+
+
+def load_model(model, pretrained_state_dict):
+    model_dict = model.state_dict()
+    pretrained_dict = {k: v for k, v in pretrained_state_dict.items() if
+                       k in model_dict and model_dict[k].size() == pretrained_state_dict[k].size()}
+    model.load_state_dict(pretrained_dict, strict=False)
+    if len(pretrained_dict) == 0:
+        print("[INFO] No params were loaded ...")
+    else:
+        for k, v in pretrained_state_dict.items():
+            if k in pretrained_dict:
+                print("==>> Load {} {}".format(k, v.size()))
+            else:
+                print("[INFO] Skip {} {}".format(k, v.size()))
+    return model
+
+
+def resnet18(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        # model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+        print("Load pretrained model from {}".format(model_urls['resnet18']))
+        pretrained_state_dict = model_zoo.load_url(model_urls['resnet18'])
+        model = load_model(model, pretrained_state_dict)
+    return model
+
+
+def resnet34(pretrained=False, **kwargs):
+    """Constructs a ResNet-34 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        # model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+        print("Load pretrained model from {}".format(model_urls['resnet34']))
+        pretrained_state_dict = model_zoo.load_url(model_urls['resnet34'])
+        model = load_model(model, pretrained_state_dict)
+    return model
+
+
+def resnet50(pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        # model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+        print("Load pretrained model from {}".format(model_urls['resnet50']))
+        pretrained_state_dict = model_zoo.load_url(model_urls['resnet50'])
+        model = load_model(model, pretrained_state_dict)
+    return model
+
+
+def resnet101(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        # model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+        print("Load pretrained model from {}".format(model_urls['resnet101']))
+        pretrained_state_dict = model_zoo.load_url(model_urls['resnet101'])
+        model = load_model(model, pretrained_state_dict)
+    return model
+
+
+def resnet152(pretrained=False, **kwargs):
+    """Constructs a ResNet-152 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        # model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+        print("Load pretrained model from {}".format(model_urls['resnet152']))
+        pretrained_state_dict = model_zoo.load_url(model_urls['resnet152'])
+        model = load_model(model, pretrained_state_dict)
+    return model
+
+if __name__ == "__main__":
+    input = torch.randn([32, 3, 256,256])
+    model = resnet34(False, num_classes=2, img_size=256)
+    output = model(input)
+    print(output.size())
diff --git a/components/face_au/utils/common_utils.py b/components/face_au/utils/common_utils.py
new file mode 100644
index 0000000..1e5b172
--- /dev/null
+++ b/components/face_au/utils/common_utils.py
@@ -0,0 +1,132 @@
+#-*-coding:utf-8-*-
+# date:2020-04-11
+# Author: Eric.Lee
+# function: common utils
+
+import os
+import shutil
+import cv2
+import numpy as np
+import json
+
+def mkdir_(path, flag_rm=False):
+    if os.path.exists(path):
+        if flag_rm == True:
+            shutil.rmtree(path)
+            os.mkdir(path)
+            print('remove {} done ~ '.format(path))
+    else:
+        os.mkdir(path)
+
+def plot_box(bbox, img, color=None, label=None, line_thickness=None):
+    tl = line_thickness or round(0.002 * max(img.shape[0:2])) + 1
+    color = color or [random.randint(0, 255) for _ in range(3)]
+    c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
+    cv2.rectangle(img, c1, c2, color, thickness=tl)# 目标的bbox
+    if label:
+        tf = max(tl - 2, 1)
+        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] # label size
+        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 # 字体的bbox
+        cv2.rectangle(img, c1, c2, color, -1)  # label 矩形填充
+        # 文本绘制
+        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 4, [225, 255, 255],thickness=tf, lineType=cv2.LINE_AA)
+
+class JSON_Encoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return obj.tolist()
+        else:
+            return super(JSON_Encoder, self).default(obj)
+
+def draw_landmarks(img,output,draw_circle):
+    img_width = img.shape[1]
+    img_height = img.shape[0]
+    dict_landmarks = {}
+    for i in range(int(output.shape[0]/2)):
+        x = output[i*2+0]*float(img_width)
+        y = output[i*2+1]*float(img_height)
+        if 41>= i >=33:
+            if 'left_eyebrow' not in dict_landmarks.keys():
+                dict_landmarks['left_eyebrow'] = []
+            dict_landmarks['left_eyebrow'].append([int(x),int(y),(0,255,0)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (0,255,0),-1)
+        elif 50>= i >=42:
+            if 'right_eyebrow' not in dict_landmarks.keys():
+                dict_landmarks['right_eyebrow'] = []
+            dict_landmarks['right_eyebrow'].append([int(x),int(y),(0,255,0)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (0,255,0),-1)
+        elif 67>= i >=60:
+            if 'left_eye' not in dict_landmarks.keys():
+                dict_landmarks['left_eye'] = []
+            dict_landmarks['left_eye'].append([int(x),int(y),(255,0,255)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
+        elif 75>= i >=68:
+            if 'right_eye' not in dict_landmarks.keys():
+                dict_landmarks['right_eye'] = []
+            dict_landmarks['right_eye'].append([int(x),int(y),(255,0,255)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
+        elif 97>= i >=96:
+            cv2.circle(img, (int(x),int(y)), 2, (0,0,255),-1)
+        elif 54>= i >=51:
+            if 'bridge_nose' not in dict_landmarks.keys():
+                dict_landmarks['bridge_nose'] = []
+            dict_landmarks['bridge_nose'].append([int(x),int(y),(0,170,255)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (0,170,255),-1)
+        elif 32>= i >=0:
+            if 'basin' not in dict_landmarks.keys():
+                dict_landmarks['basin'] = []
+            dict_landmarks['basin'].append([int(x),int(y),(255,30,30)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (255,30,30),-1)
+        elif 59>= i >=55:
+            if 'wing_nose' not in dict_landmarks.keys():
+                dict_landmarks['wing_nose'] = []
+            dict_landmarks['wing_nose'].append([int(x),int(y),(0,255,255)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (0,255,255),-1)
+        elif 87>= i >=76:
+            if 'out_lip' not in dict_landmarks.keys():
+                dict_landmarks['out_lip'] = []
+            dict_landmarks['out_lip'].append([int(x),int(y),(255,255,0)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (255,255,0),-1)
+        elif 95>= i >=88:
+            if 'in_lip' not in dict_landmarks.keys():
+                dict_landmarks['in_lip'] = []
+            dict_landmarks['in_lip'].append([int(x),int(y),(50,220,255)])
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (50,220,255),-1)
+        else:
+            if draw_circle:
+                cv2.circle(img, (int(x),int(y)), 2, (255,0,255),-1)
+
+    return dict_landmarks
+
+def draw_contour(image,dict):
+    for key in dict.keys():
+        # print(key)
+        _,_,color = dict[key][0]
+
+        if 'basin' == key or 'wing_nose' == key:
+            pts = np.array([[dict[key][i][0],dict[key][i][1]] for i in range(len(dict[key]))],np.int32)
+            # print(pts)
+            cv2.polylines(image,[pts],False,color)
+
+        else:
+            points_array = np.zeros((1,len(dict[key]),2),dtype = np.int32)
+            for i in range(len(dict[key])):
+                x,y,_ = dict[key][i]
+                points_array[0,i,0] = x
+                points_array[0,i,1] = y
+
+            # cv2.fillPoly(image, points_array, color)
+            cv2.drawContours(image,points_array,-1,color,thickness=1)
diff --git a/components/face_au/utils/model_utils.py b/components/face_au/utils/model_utils.py
new file mode 100644
index 0000000..48cc3d3
--- /dev/null
+++ b/components/face_au/utils/model_utils.py
@@ -0,0 +1,61 @@
+#-*-coding:utf-8-*-
+# date:2020-04-11
+# Author: Eric.Lee
+# function: model utils
+
+import os
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+import random
+
+def get_acc(output, label):
+    total = output.shape[0]
+    _, pred_label = output.max(1)
+    num_correct = (pred_label == label).sum().item()
+    return num_correct / float(total)
+
+def set_learning_rate(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+def set_seed(seed = 666):
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        cudnn.deterministic = True
+
+def split_trainval_datasets(ops):
+    print(' --------------->>> split_trainval_datasets ')
+    train_split_datasets = []
+    train_split_datasets_label = []
+
+    val_split_datasets = []
+    val_split_datasets_label = []
+    for idx,doc in enumerate(sorted(os.listdir(ops.train_path), key=lambda x:int(x.split('.')[0]), reverse=False)):
+        # print(' %s label is %s \n'%(doc,idx))
+
+        data_list = os.listdir(ops.train_path+doc)
+        random.shuffle(data_list)
+
+        cal_split_num = int(len(data_list)*ops.val_factor)
+
+        for i,file in enumerate(data_list):
+            if '.jpg' in file:
+                if i < cal_split_num:
+                    val_split_datasets.append(ops.train_path+doc + '/' + file)
+                    val_split_datasets_label.append(idx)
+                else:
+                    train_split_datasets.append(ops.train_path+doc + '/' + file)
+                    train_split_datasets_label.append(idx)
+
+                print(ops.train_path+doc + '/' + file,idx)
+
+    print('\n')
+    print('train_split_datasets len {}'.format(len(train_split_datasets)))
+    print('val_split_datasets len {}'.format(len(val_split_datasets)))
+
+    return train_split_datasets,train_split_datasets_label,val_split_datasets,val_split_datasets_label
diff --git a/lib/face_bioassay_lib/cores/face_fuction.py b/lib/face_bioassay_lib/cores/face_fuction.py
index 19fc61b..d6a497e 100644
--- a/lib/face_bioassay_lib/cores/face_fuction.py
+++ b/lib/face_bioassay_lib/cores/face_fuction.py
@@ -163,7 +163,6 @@ def draw_contour(image,dict,vis = False):
             if vis:
                 cv2.drawContours(image,points_array,-1,color,thickness=2)
 
-
 def plot_box(x, img, color=None, label=None, line_thickness=None):
     # Plots one bounding box on image img
     tl = line_thickness or round(0.002 * max(img.shape[0:2])) + 1  # line thickness
@@ -259,7 +258,7 @@ def refine_face_bbox(bbox,img_shape):
 
     return (x1,y1,x2,y2)
 
-def get_faces_batch_attribute(face_multitask_model,face_euler_model,dets,img_raw,use_cuda,face_size = 256,vis = False):
+def get_faces_batch_attribute(face_multitask_model,face_euler_model,face_au_model,dets,img_raw,use_cuda,face_size = 256,vis = False):
     if len(dets) == 0:
         return None
     img_align = img_raw.copy()
@@ -291,6 +290,7 @@ def get_faces_batch_attribute(face_multitask_model,face_euler_model,dets,img_raw
     euler_angles = face_euler_model.predict(image_batch)
 
     faces_message = None
+    faces_align_batch = None
     for i in range(len(dets)):
         x0,y0 = r_bboxes[i][0],r_bboxes[i][1]
         face_w = r_bboxes[i][2]-r_bboxes[i][0]
@@ -318,10 +318,22 @@ def get_faces_batch_attribute(face_multitask_model,face_euler_model,dets,img_raw
 
         if abs(yaw)<45.:
             face_align_output = face_alignment(img_align,eyes_center[0],eyes_center[1],
-                desiredLeftEye=(0.365, 0.38),desiredFaceWidth=112, desiredFaceHeight=None)
+                desiredLeftEye=(0.365, 0.38),desiredFaceWidth=256, desiredFaceHeight=None)
         else:
             face_align_output = face_alignment(img_align,eyes_center[0],eyes_center[1],
-                desiredLeftEye=(0.38, 0.40),desiredFaceWidth=112, desiredFaceHeight=None)
+                desiredLeftEye=(0.38, 0.40),desiredFaceWidth=256, desiredFaceHeight=None)
+        # au 输入 预处理
+        face_align_output_ = cv2.cvtColor(face_align_output, cv2.COLOR_BGR2RGB).astype(np.float32)
+        face_align_output_ = face_align_output_ - [123.67, 116.28, 103.53]
+        face_align_output_ = face_align_output_.transpose(2, 0, 1)
+        face_align_output_ = np.expand_dims(face_align_output_,0)
+        if faces_align_batch is None:
+            faces_align_batch = face_align_output_
+        else:
+            faces_align_batch = np.concatenate((faces_align_batch,face_align_output_),axis=0)
+
+
+        #
         if faces_message is None:
             faces_message = []
         faces_message.append(
@@ -329,12 +341,30 @@ def get_faces_batch_attribute(face_multitask_model,face_euler_model,dets,img_raw
                 "xyxy":r_bboxes[i],
                 "age":age_pre[i][0],
                 "gender":gender_str,
-                "head_3d_angle":{"yaw":yaw,"pitch":pitch,"roll":roll}
+                "head_3d_angle":{"yaw":yaw,"pitch":pitch,"roll":roll},
+                "open_mouth":None,
+                "eye_close":None
             }
             )
         # plot_box(r_bboxes[i][0:4], img_raw,label="{}, age: {:.1f}, unblur:{}".format(gender_str,age_pre[i][0],int(blur_)), color=(255,90,90), line_thickness=2)
         if vis:
             plot_box(r_bboxes[i][0:4], img_raw,label="{}, age: {:.1f}".format(gender_str,age_pre[i][0]), color=(255,90,90), line_thickness=2)
+    # print("faces_align_batch shape : ",faces_align_batch.shape)
+    au_features = face_au_model.predict(faces_align_batch)
+
 
+    for i in range(len(dets)):
+        # print("au_features : ",au_features[i])
+        open_mouth = np.clip(au_features[i][8],0.,1.)
+        close_eye_r = np.clip(au_features[i][0],0.,1.)
+        close_eye_l = np.clip(au_features[i][1],0.,1.)
+        print("open_mouth : {:.2f},eye_close: {:.3f} , {:.3f}".format(open_mouth,close_eye_r,close_eye_l))
+        faces_message[i]["open_mouth"] = open_mouth
+        faces_message[i]["eye_close"] = (close_eye_r,close_eye_l)
+        cv2.putText(img_raw, "{:.2f}".format(open_mouth),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+15)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,185,25), 3)
+        cv2.putText(img_raw, "{:.2f}".format(open_mouth),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+15)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,185,255), 1)
+        
+        cv2.putText(img_raw, "{:.2f},{:.2f}".format(close_eye_r,close_eye_l),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+35)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,33,255), 3)
+        cv2.putText(img_raw, "{:.2f},{:.2f}".format(close_eye_r,close_eye_l),(int(r_bboxes[i][0]-1),int(r_bboxes[i][3]+35)),cv2.FONT_HERSHEY_DUPLEX, 0.65, (20,185,255), 1)
 
     return faces_message
-- 
GitLab