初始化：换发型/换发色/训练发型服务

包含：
- hair_service_sd: 主服务（换发型/换发色/生发，端口8801）
- photo_service: LoRA调度+训练（端口32678）
- hair_grow_service: 调试测试页（端口8888，含4个测试页）
- 批量训练脚本（batch_train_hairstyles.py）
- 发际线mask自动识别（hairline_mask.py，4种方案）
- 手绘mask换发型（hair_swap_manual.py）
- 文档：README.md + LARGE_FILES.md + docs/

大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md
OSS/COS密钥已脱敏为环境变量，原文件备份在本地

project/hair_service_sd/core/MMCVFaceRecognitionServer.py

@@ -0,0 +1,196 @@
+import torch
+import torch.nn as nn
+import math
+import os
+import cv2
+import numpy as np
+from core.utils import landmark_processor
+import glob
+# from models.BigResNetStable import MomocvFaceAlignment
+
+
+def op_name(op_name, m):
+    m.op_name = op_name
+    return m
+
+class Flatten(nn.Module):
+    def __init__(self, axis=1):
+        super(Flatten, self).__init__()
+        self.axis = axis
+
+    def forward(self, x):
+        assert self.axis == 1
+        x = x.reshape(x.shape[0], -1)
+        return x
+
+def flatten(name, axis=1):
+    return op_name(name, Flatten(axis))
+
+def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
+    return nn.Sequential(
+        op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
+        op_name(name.replace('conv', 'relu'), nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
+    )
+
+def conv(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
+    return op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True))
+
+def bn_conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.9, track_running_stats=True):
+    return nn.Sequential(
+        op_name(name + '_bn1', nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats)),
+        op_name(name + '_conv1', nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
+        op_name(name + '_relu', nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
+    )
+
+def bn(name, in_channels, momentum = 0.9, track_running_stats=True):
+    return op_name(name, nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats))
+
+class BasicBlock(nn.Module):
+    def __init__(self, stage, unit, inplanes, outplanes, stride):
+        super(BasicBlock, self).__init__()
+
+        self.bn = bn('stage{}_unit{}_bn1'.format(stage, unit), inplanes)
+        self.conv1 = conv_relu('stage{}_unit{}_conv1'.format(stage, unit), inplanes, outplanes, kernel_size=3, stride=1, padding=1)
+        self.conv2 = conv('stage{}_unit{}_conv2'.format(stage, unit), outplanes, outplanes, kernel_size=3, stride=stride, padding=1)
+
+        if inplanes != outplanes or stride != 1:
+            self.conv_sc = conv('stage{}_unit{}_conv1sc'.format(stage, unit), inplanes, outplanes, kernel_size=1, stride=stride, padding=0)
+
+        self.inplanes = inplanes
+        self.outplanes = outplanes
+        self.stride = stride
+
+    def forward(self, x):
+        residule = x
+        out = self.bn(x)
+        out = self.conv1(out)
+        out = self.conv2(out)
+        if self.inplanes != self.outplanes or self.stride != 1:
+            residule = self.conv_sc(residule)
+        ret = out + residule
+        return ret
+
+class ResnetBlock(nn.Module):
+    def __init__(self, stage, inplanes, outplanes, stride=2, n_blocks=1):
+        super(ResnetBlock, self).__init__()
+
+        self.conv = []
+        for m in range(n_blocks):
+            if m == 0:
+                self.conv.append(BasicBlock(stage, m + 1, inplanes, outplanes, stride))
+            else:
+                self.conv.append(BasicBlock(stage, m + 1, outplanes, outplanes, 1))
+        self.conv = nn.Sequential(*self.conv)
+
+    def forward(self, x):
+        ret = self.conv(x)
+        return ret
+
+class FaceRecognitionServer(nn.Module):
+    def __init__(self):
+        super(FaceRecognitionServer, self).__init__()
+        ch_num = [64, 64, 128, 256, 512]
+        # ch_num = [64, 64]
+        strides = [2, 2, 2, 2]
+        n_blocks = [3, 4, 14, 3]
+        op_list = [conv_relu('conv0', 3, 64, 3, 1, 1)]
+        op_list += [ResnetBlock(i + 1, ch_num[i], ch_num[i + 1], strides[i], n_blocks[i]) for i in range(len(ch_num) - 1)]
+        op_list += [bn('bn1', 512)]
+        op_list += [flatten('flatten', 1)]
+        op_list += [op_name('pre_fc1', nn.Linear(25088, 512))]
+        self.features = nn.Sequential(*op_list)
+
+        model_path = os.path.dirname(os.path.split(os.path.realpath(__file__))[0])
+        weights = torch.load('weights/MMCVFaceRecognitionServer.pth',
+                             map_location=lambda storage, loc: storage)
+        self.load_state_dict(weights)
+        self.eval()
+
+    def forward(self, x):
+        features = self.features(x)
+        return features
+
+class MomocvFaceRecognitionServer(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        self.face_recognition_net = FaceRecognitionServer()
+        # self.mmcv = MomocvFaceAlignment()
+        # self.model_path, _ = os.path.split(os.path.realpath(__file__))
+        # weights = torch.load(os.path.join(self.model_path, 'MMCVFaceRecognitionServer.pth'),
+        #                      map_location=lambda storage, loc: storage)
+        # self.face_recognition_net.load_state_dict(weights)
+        self.face_recognition_net.to(self.device)
+        # self.face_recognition_net.eval()
+
+    def forward(self, images, landmarks):
+        dst_size = 112
+
+        with torch.no_grad():
+            input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
+            assert len(images) == len(landmarks)
+            for ix, img in enumerate(images):
+                landmark = landmarks[ix]
+
+
+                # get angel
+                # M = landmark_processor.get_transform_mat_full_face(landmark, 576, scale=1, offset=(0, 0.3))
+                # pt1k_crop = landmark_processor.transform_points(landmark, M)
+                # crop_face = cv2.warpAffine(img, M, (576, 768),
+                #                            flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_REFLECT)
+                # t0 = time.time()
+                # landmarks87, poselayer, tracking_probe, occlusion_probe = self.mmcv.detect(crop_face, [pt1k_crop])
+                # # print('time:', time.time() - t0)
+                # if tracking_probe[0] < 0.5:
+                #     continue
+
+                mat = landmark_processor.get_transform_mat_for_face_recognition(landmark, dst_size)[0:2]
+
+                tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
+
+                input_numpy[ix, :, :, :] = (tmp.transpose((2, 0, 1)).astype(np.float32) - 127.5) / 128
+
+                # cv2.imshow('tmp', tmp)
+                # cv2.waitKey()
+
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            features = self.face_recognition_net(in_tensor)
+            features = features.cpu().numpy()
+            norm_factor = np.sqrt(np.sum(features ** 2, axis=1))
+            norm_factor = norm_factor.reshape(-1, 1)
+            features /= norm_factor
+            return features
+
+    def cos_sim(self, a, b):
+        a_norm = np.linalg.norm(a)
+        b_norm = np.linalg.norm(b)
+        cos = np.dot(a, b) / (a_norm * b_norm)
+        return cos
+
+if __name__ == '__main__':
+    all_jpegs = glob.glob(r'D:\data\deepface_example\data\02b59e75ce91bda300ff827a85a74687d633cdfb5d830e7d3855e31b75201fa8\*.jpg')
+    for s_filename_path in all_jpegs:
+        img = cv2.imread(s_filename_path)
+
+        # cv2.imshow('img', img)
+        # cv2.waitKey()
+
+        # dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True)
+        # if dflpng is None:
+        #     print('ERROR')
+        #
+        # landmarks = dflpng.get_landmarks_mmcv_137()
+        #
+        # mmcv = MomocvFaceRecognitionServer()
+        # features = mmcv.forward([img, img], [landmarks, landmarks])
+        # print(features)
+        # print('conansherry')
+        # fullyconnected1 = fullyconnected1[0]
+        # fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32)
+        # occlusion_probe = occlusion_probe[0]
+        # for ix, pt in enumerate(fullyconnected1):
+        #     cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2)
+        #     # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
+        # cv2.imshow('img', img)
+        # cv2.waitKey()

project/hair_service_sd/core/bodyseg/backbone/backbone.py

@@ -0,0 +1,75 @@
+import numpy as np
+import torch
+from torch import nn
+
+def get_norm(norm, out_channels=None):
+    """
+    Args:
+        norm (str or callable):
+
+    Returns:
+        nn.Module or None: the normalization layer
+    """
+    if isinstance(norm, str):
+        if len(norm) == 0:
+            return None
+        norm = {
+            "BN": nn.BatchNorm2d,
+            "IN": nn.InstanceNorm2d,
+            "GN": lambda channels: nn.GroupNorm(32, channels),
+            "nnSyncBN": nn.SyncBatchNorm,  # keep for debugging
+        }[norm]
+    if out_channels is not None:
+        return norm(out_channels)
+    else:
+        return norm
+
+class Conv2d(torch.nn.Conv2d):
+    """
+    A wrapper around :class:`torch.nn.Conv2d` to support zero-size tensor and more features.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
+
+        Args:
+            norm (nn.Module, optional): a normalization layer
+            activation (callable(Tensor) -> Tensor): a callable activation function
+
+        It assumes that norm layer is used before activation.
+        """
+        norm = kwargs.pop("norm", None)
+        activation = kwargs.pop("activation", None)
+        super().__init__(*args, **kwargs)
+
+        self.norm = norm
+        self.activation = activation
+
+    def forward(self, x):
+        x = super().forward(x)
+        if self.norm is not None:
+            x = self.norm(x)
+        if self.activation is not None:
+            x = self.activation(x)
+        return x
+
+class Backbone(nn.Module):
+    """
+    Abstract base class for network backbones.
+    """
+
+    def __init__(self):
+        """
+        The `__init__` method of any subclass can specify its own set of arguments.
+        """
+        super().__init__()
+
+    def forward(self):
+        """
+        Subclasses must override this method, but adhere to the same return type.
+
+        Returns:
+            dict[str: Tensor]: mapping from feature name (e.g., "res2") to tensor
+        """
+        pass

project/hair_service_sd/core/bodyseg/backbone/fpn.py

@@ -0,0 +1,268 @@
+import math
+import core.utils.weight_init as weight_init
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
+from core.bodyseg.backbone.resnet import build_resnet_backbone
+
+class FPN(Backbone):
+    """
+    This module implements Feature Pyramid Network.
+    It creates pyramid features built on top of some input feature maps.
+    """
+
+    def __init__(
+        self, bottom_up, in_features, out_channels, norm="", top_block=None, fuse_type="sum"
+    ):
+        """
+        Args:
+            bottom_up (Backbone): module representing the bottom up subnetwork.
+                Must be a subclass of :class:`Backbone`. The multi-scale feature
+                maps generated by the bottom up network, and listed in `in_features`,
+                are used to generate FPN levels.
+            in_features (list[str]): names of the input feature maps coming
+                from the backbone to which FPN is attached. For example, if the
+                backbone produces ["res2", "res3", "res4"], any *contiguous* sublist
+                of these may be used; order must be from high to low resolution.
+            out_channels (int): number of channels in the output feature maps.
+            norm (str): the normalization to use.
+            top_block (nn.Module or None): if provided, an extra operation will
+                be performed on the output of the last (smallest resolution)
+                FPN output, and the result will extend the result list. The top_block
+                further downsamples the feature map. It must have an attribute
+                "num_levels", meaning the number of extra FPN levels added by
+                this block, and "in_feature", which is a string representing
+                its input feature (e.g., p5).
+            fuse_type (str): types for fusing the top down features and the lateral
+                ones. It can be "sum" (default), which sums up element-wise; or "avg",
+                which takes the element-wise mean of the two.
+        """
+        super(FPN, self).__init__()
+        assert isinstance(bottom_up, Backbone)
+
+        # Feature map strides and channels from the bottom up network (e.g. ResNet)
+        in_strides = [bottom_up._out_feature_strides[f] for f in in_features]
+        in_channels = [bottom_up._out_feature_channels[f] for f in in_features]
+
+        _assert_strides_are_log2_contiguous(in_strides)
+        lateral_convs = []
+        output_convs = []
+
+        use_bias = norm == ""
+        for idx, in_channels in enumerate(in_channels):
+            lateral_norm = get_norm(norm, out_channels)
+            output_norm = get_norm(norm, out_channels)
+
+            lateral_conv = Conv2d(
+                in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
+            )
+            output_conv = Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=use_bias,
+                norm=output_norm,
+            )
+            weight_init.c2_xavier_fill(lateral_conv)
+            weight_init.c2_xavier_fill(output_conv)
+            stage = int(math.log2(in_strides[idx]))
+
+            lateral_convs.append(lateral_conv)
+            output_convs.append(output_conv)
+        # Place convs into top-down order (from low to high resolution)
+        # to make the top-down computation in forward clearer.
+        self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
+        self.output_convs = nn.ModuleList(output_convs[::-1])
+        self.top_block = top_block
+        self.in_features = in_features
+        self.bottom_up = bottom_up
+        # Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
+        self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in in_strides}
+        # top block output feature maps.
+        if self.top_block is not None:
+            for s in range(stage, stage + self.top_block.num_levels):
+                self._out_feature_strides["p{}".format(s + 1)] = 2 ** (s + 1)
+
+        self._out_features = list(self._out_feature_strides.keys())
+        self._out_feature_channels = {k: out_channels for k in self._out_features}
+        assert fuse_type in {"avg", "sum"}
+        self._fuse_type = fuse_type
+
+    def forward(self, x):
+        """
+        Args:
+            input (dict[str: Tensor]): mapping feature map name (e.g., "res5") to
+                feature map tensor for each feature level in high to low resolution order.
+
+        Returns:
+            dict[str: Tensor]:
+                mapping from feature map name to FPN feature map tensor
+                in high to low resolution order. Returned feature names follow the FPN
+                paper convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
+                ["p2", "p3", ..., "p6"].
+        """
+        # Reverse feature maps into top-down order (from low to high resolution)
+        bottom_up_features = self.bottom_up(x)
+        x = [bottom_up_features[f] for f in self.in_features[::-1]]
+        results = []
+        prev_features = self.lateral_convs[0](x[0])
+        results.append(self.output_convs[0](prev_features))
+        for features, lateral_conv, output_conv in zip(
+            x[1:], self.lateral_convs[1:], self.output_convs[1:]
+        ):
+            top_down_features = F.interpolate(prev_features, scale_factor=2, mode="nearest")
+            lateral_features = lateral_conv(features)
+            prev_features = lateral_features + top_down_features
+            if self._fuse_type == "avg":
+                prev_features /= 2
+            results.insert(0, output_conv(prev_features))
+
+        if self.top_block is not None:
+            top_block_in_feature = bottom_up_features.get(self.top_block.in_feature, None)
+            if top_block_in_feature is None:
+                top_block_in_feature = results[self._out_features.index(self.top_block.in_feature)]
+            results.extend(self.top_block(top_block_in_feature))
+        assert len(self._out_features) == len(results)
+        return dict(zip(self._out_features, results))
+
+def _assert_strides_are_log2_contiguous(strides):
+    """
+    Assert that each stride is 2x times its preceding stride, i.e. "contiguous in log2".
+    """
+    for i, stride in enumerate(strides[1:], 1):
+        assert stride == 2 * strides[i - 1], "Strides {} {} are not log2 contiguous".format(
+            stride, strides[i - 1]
+        )
+
+
+class LastLevelMaxPool(nn.Module):
+    """
+    This module is used in the original FPN to generate a downsampled
+    P6 feature from P5.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.num_levels = 1
+        self.in_feature = "p5"
+
+    def forward(self, x):
+        return [F.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
+
+
+class LastLevelP6P7(nn.Module):
+    """
+    This module is used in RetinaNet to generate extra layers, P6 and P7 from
+    C5 feature.
+    """
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.num_levels = 2
+        self.in_feature = "res5"
+        self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
+        self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
+        for module in [self.p6, self.p7]:
+            weight_init.c2_xavier_fill(module)
+
+    def forward(self, c5):
+        p6 = self.p6(c5)
+        p7 = self.p7(F.relu(p6))
+        return [p6, p7]
+
+
+def build_resnet_fpn_backbone(in_channels=3):
+    """
+    Args:
+        cfg: a detectron2 CfgNode
+
+    Returns:
+        backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
+    """
+    bottom_up = build_resnet_backbone(in_channels)
+    in_features = ["res2", "res3", "res4"]
+    out_channels = 256
+    backbone = FPN(
+        bottom_up=bottom_up,
+        in_features=in_features,
+        out_channels=out_channels,
+        norm="BN",
+        # top_block=LastLevelMaxPool(),
+        top_block=None,
+        fuse_type="sum",
+    )
+    return backbone
+
+def build_retinanet_resnet_fpn_backbone(cfg, in_channels=3):
+    """
+    Args:
+        cfg: a detectron2 CfgNode
+
+    Returns:
+        backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
+    """
+    bottom_up = build_resnet_backbone(cfg, in_channels)
+    in_features = cfg.MODEL.FPN.IN_FEATURES
+    out_channels = cfg.MODEL.FPN.OUT_CHANNELS
+    in_channels_p6p7 = bottom_up._out_feature_channels["res5"]
+    backbone = FPN(
+        bottom_up=bottom_up,
+        in_features=in_features,
+        out_channels=out_channels,
+        norm=cfg.MODEL.FPN.NORM,
+        top_block=LastLevelP6P7(in_channels_p6p7, out_channels),
+        fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
+    )
+    return backbone
+
+if __name__ == "__main__":
+    import argparse
+    from config.default import get_cfg
+
+    def setup(args):
+        """
+        Create configs and perform basic setups.
+        """
+        cfg = get_cfg()
+        cfg.merge_from_file(args.cfg)
+        cfg.merge_from_list(args.opts)
+        cfg.freeze()
+        return cfg
+
+    parser = argparse.ArgumentParser(description='Train ImageNet network')
+    # general
+    parser.add_argument('--cfg',
+                        help='experiment configure file name',
+                        required=True,
+                        type=str)
+
+    parser.add_argument('opts',
+                        help="Modify config options using the command-line",
+                        default=None,
+                        nargs=argparse.REMAINDER)
+
+    args = parser.parse_args()
+    cfg = setup(args)
+    print(cfg)
+
+    model = build_resnet_fpn_backbone(cfg, 3)
+    # model = build_retinanet_resnet_fpn_backbone(cfg, 3)
+    print(model)
+    # model = torch.nn.DataParallel(model, list(range(2))).cuda()
+    dummy_input = torch.randn(4, 3, 512, 512)
+
+    out = model(dummy_input)
+
+    for k, v in out.items():
+        print(k, v.shape)
+
+    # torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True,
+    #                   input_names=['input'],
+    #                   output_names=['output'])
+
+    pass
+

project/hair_service_sd/core/bodyseg/backbone/resnet.py

@@ -0,0 +1,298 @@
+import numpy as np
+import core.utils.weight_init as weight_init
+import torch
+import torch.nn.functional as F
+from torch import nn
+from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
+
+class BasicStem(nn.Module):
+    def __init__(self, in_channels=3, out_channels=64, norm="BN"):
+        """
+        Args:
+            norm (str or callable): a callable that takes the number of
+                channels and return a `nn.Module`, or a pre-defined string
+                (one of {"FrozenBN", "BN", "GN"}).
+        """
+        super().__init__()
+        self.conv1 = Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False,
+            norm=get_norm(norm, out_channels),
+        )
+        weight_init.c2_msra_fill(self.conv1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu_(x)
+        x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
+        return x
+
+    @property
+    def out_channels(self):
+        return self.conv1.out_channels
+
+    @property
+    def stride(self):
+        return 4  # = stride 2 conv -> stride 2 max pool
+
+class ResNetBlockBase(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        """
+        The `__init__` method of any subclass should also contain these arguments.
+
+        Args:
+            in_channels (int):
+            out_channels (int):
+            stride (int):
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.stride = stride
+
+class BottleneckBlock(ResNetBlockBase):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        *,
+        bottleneck_channels,
+        stride=1,
+        num_groups=1,
+        norm="BN",
+        stride_in_1x1=False,
+        dilation=1,
+    ):
+        """
+        Args:
+            norm (str or callable): a callable that takes the number of
+                channels and return a `nn.Module`, or a pre-defined string
+                (one of {"FrozenBN", "BN", "GN"}).
+            stride_in_1x1 (bool): when stride==2, whether to put stride in the
+                first 1x1 convolution or the bottleneck 3x3 convolution.
+        """
+        super().__init__(in_channels, out_channels, stride)
+
+        if in_channels != out_channels:
+            self.shortcut = Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=1,
+                stride=stride,
+                bias=False,
+                norm=get_norm(norm, out_channels),
+            )
+        else:
+            self.shortcut = None
+
+        # The original MSRA ResNet models have stride in the first 1x1 conv
+        # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
+        # stride in the 3x3 conv
+        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
+
+        self.conv1 = Conv2d(
+            in_channels,
+            bottleneck_channels,
+            kernel_size=1,
+            stride=stride_1x1,
+            bias=False,
+            norm=get_norm(norm, bottleneck_channels),
+        )
+
+        self.conv2 = Conv2d(
+            bottleneck_channels,
+            bottleneck_channels,
+            kernel_size=3,
+            stride=stride_3x3,
+            padding=1 * dilation,
+            bias=False,
+            groups=num_groups,
+            dilation=dilation,
+            norm=get_norm(norm, bottleneck_channels),
+        )
+
+        self.conv3 = Conv2d(
+            bottleneck_channels,
+            out_channels,
+            kernel_size=1,
+            bias=False,
+            norm=get_norm(norm, out_channels),
+        )
+
+        for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
+            if layer is not None:  # shortcut can be None
+                weight_init.c2_msra_fill(layer)
+
+        # Zero-initialize the last normalization in each residual branch,
+        # so that at the beginning, the residual branch starts with zeros,
+        # and each residual block behaves like an identity.
+        # See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
+        # "For BN layers, the learnable scaling coefficient γ is initialized
+        # to be 1, except for each residual block's last BN
+        # where γ is initialized to be 0."
+
+        # nn.init.constant_(self.conv3.norm.weight, 0)
+        # TODO this somehow hurts performance when training GN models from scratch.
+        # Add it as an option when we need to use this code to train a backbone.
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = F.relu_(out)
+
+        out = self.conv2(out)
+        out = F.relu_(out)
+
+        out = self.conv3(out)
+
+        if self.shortcut is not None:
+            shortcut = self.shortcut(x)
+        else:
+            shortcut = x
+
+        out += shortcut
+        out = F.relu_(out)
+        return out
+
+def make_stage(block_class, num_blocks, first_stride, **kwargs):
+    """
+    Create a resnet stage by creating many blocks.
+    Args:
+        block_class (class): a subclass of ResNetBlockBase
+        num_blocks (int):
+        first_stride (int): the stride of the first block. The other blocks will have stride=1.
+            A `stride` argument will be passed to the block constructor.
+        kwargs: other arguments passed to the block constructor.
+
+    Returns:
+        list[nn.Module]: a list of block module.
+    """
+    blocks = []
+    for i in range(num_blocks):
+        blocks.append(block_class(stride=first_stride if i == 0 else 1, **kwargs))
+        kwargs["in_channels"] = kwargs["out_channels"]
+    return blocks
+
+class ResNet(Backbone):
+    def __init__(self, stem, stages, num_classes=None, out_features=None):
+        """
+        Args:
+            stem (nn.Module): a stem module
+            stages (list[list[ResNetBlock]]): several (typically 4) stages,
+                each contains multiple :class:`ResNetBlockBase`.
+            num_classes (None or int): if None, will not perform classification.
+            out_features (list[str]): name of the layers whose outputs should
+                be returned in forward. Can be anything in "stem", "linear", or "res2" ...
+                If None, will return the output of the last layer.
+        """
+        super(ResNet, self).__init__()
+        self.stem = stem
+        self.num_classes = num_classes
+
+        current_stride = self.stem.stride
+        self._out_feature_strides = {"stem": current_stride}
+        self._out_feature_channels = {"stem": self.stem.out_channels}
+
+        self.stages = []
+        self.names = []
+        for i, blocks in enumerate(stages):
+            for block in blocks:
+                assert isinstance(block, ResNetBlockBase), block
+                curr_channels = block.out_channels
+            stage = nn.Sequential(*blocks)
+            name = "res" + str(i + 2)
+            self.stages.append(stage)
+            self.names.append(name)
+            self._out_feature_strides[name] = current_stride = int(
+                current_stride * np.prod([k.stride for k in blocks])
+            )
+            self._out_feature_channels[name] = blocks[-1].out_channels
+        self.stages = nn.ModuleList(self.stages)
+
+        if num_classes is not None:
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            self.linear = nn.Linear(curr_channels, num_classes)
+
+            # Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
+            # "The 1000-way fully-connected layer is initialized by
+            # drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
+            nn.init.normal_(self.linear.weight, stddev=0.01)
+            name = "linear"
+
+        if out_features is None:
+            out_features = [name]
+        self._out_features = out_features
+        assert len(self._out_features)
+        for out_feature in self._out_features:
+            assert out_feature in self.names, "Available children: {}".format(", ".join(self.names))
+
+    def forward(self, x):
+        outputs = {}
+        x = self.stem(x)
+        if "stem" in self._out_features:
+            outputs["stem"] = x
+        for ix, stage in enumerate(self.stages):
+            name = self.names[ix]
+            x = stage(x)
+            if name in self._out_features:
+                outputs[name] = x
+        if self.num_classes is not None:
+            x = self.avgpool(x)
+            x = self.linear(x)
+            if "linear" in self._out_features:
+                outputs["linear"] = x
+        return outputs
+
+def build_resnet_backbone(in_channels=3):
+    norm = "BN"
+    stem = BasicStem(
+        in_channels=in_channels,
+        out_channels=64,
+        norm=norm,
+    )
+
+    # fmt: off
+    out_features = ["res2", "res3", "res4"]
+    depth = 101
+    num_groups = 1
+    bottleneck_channels = 64
+    in_channels = 64
+    out_channels = 256
+    stride_in_1x1 = True
+    res5_dilation = 1
+    # fmt: on
+    assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
+
+    num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
+
+    stages = []
+
+    # Avoid creating variables without gradients
+    # It consumes extra memory and may cause allreduce to fail
+    out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
+    max_stage_idx = max(out_stage_idx)
+    for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
+        dilation = res5_dilation if stage_idx == 5 else 1
+        first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
+        stage_kargs = dict()
+        stage_kargs.update({
+            "num_blocks": num_blocks_per_stage[idx],
+            "first_stride": first_stride,
+            "in_channels": in_channels,
+            "bottleneck_channels": bottleneck_channels,
+            "out_channels": out_channels,
+            "num_groups": num_groups,
+            "norm": norm,
+            "stride_in_1x1": stride_in_1x1,
+            "dilation": dilation,
+        })
+        stage_kargs["block_class"] = BottleneckBlock
+        blocks = make_stage(**stage_kargs)
+        in_channels = out_channels
+        out_channels *= 2
+        bottleneck_channels *= 2
+        stages.append(blocks)
+    return ResNet(stem, stages, out_features=out_features)

project/hair_service_sd/core/bodyseg/backbone/xception.py

@@ -0,0 +1,244 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+from core.bodyseg.backbone.backbone import Backbone
+
+def fixed_padding(inputs, kernel_size, dilation):
+    kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+    pad_total = kernel_size_effective - 1
+    pad_beg = pad_total // 2
+    pad_end = pad_total - pad_beg
+    padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
+    return padded_inputs
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
+        super(SeparableConv2d, self).__init__()
+
+        self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
+                               groups=inplanes, bias=bias)
+        self.bn = BatchNorm(inplanes)
+        self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
+
+    def forward(self, x):
+        x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
+        x = self.conv1(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+        return x
+
+class Block(nn.Module):
+    def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
+                 start_with_relu=True, grow_first=True, is_last=False):
+        super(Block, self).__init__()
+
+        if planes != inplanes or stride != 1:
+            self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
+            self.skipbn = BatchNorm(planes)
+        else:
+            self.skip = None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep = []
+
+        filters = inplanes
+        if grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+            filters = planes
+
+        for i in range(reps - 1):
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride != 1:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride == 1 and is_last:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if not start_with_relu:
+            rep = rep[1:]
+
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x = x + skip
+
+        return x
+
+class AlignedXception(Backbone):
+    """
+    Modified Alighed Xception
+    """
+    def __init__(self, output_stride, BatchNorm):
+        super(AlignedXception, self).__init__()
+
+        if output_stride == 16:
+            entry_block3_stride = 2
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 2)
+        elif output_stride == 8:
+            entry_block3_stride = 1
+            middle_block_dilation = 2
+            exit_block_dilations = (2, 4)
+        else:
+            raise NotImplementedError
+
+
+        # Entry flow
+        self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
+        self.bn1 = BatchNorm(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
+        self.bn2 = BatchNorm(64)
+
+        self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
+        self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
+                            grow_first=True)
+        self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
+                            start_with_relu=True, grow_first=True, is_last=True)
+
+        # Middle flow
+        self.block4  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block5  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block6  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block7  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block8  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block9  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+
+        # Exit flow
+        self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
+
+        self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn3 = BatchNorm(1536)
+
+        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn4 = BatchNorm(1536)
+
+        self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn5 = BatchNorm(2048)
+
+        # Init weights
+        self._init_weight()
+
+    def forward(self, x):
+        # Entry flow
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.block1(x)
+        # add relu here
+        x = self.relu(x)
+        low_level_feat = x
+        x = self.block2(x)
+        x = self.block3(x)
+
+        # Middle flow
+        x = self.block4(x)
+        x = self.block5(x)
+        x = self.block6(x)
+        x = self.block7(x)
+        x = self.block8(x)
+        x = self.block9(x)
+        x = self.block10(x)
+        x = self.block11(x)
+        x = self.block12(x)
+        x = self.block13(x)
+        x = self.block14(x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = self.block17(x)
+        x = self.block18(x)
+        x = self.block19(x)
+
+        # Exit flow
+        x = self.block20(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+
+        return x, low_level_feat
+
+    def _init_weight(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))
+            elif isinstance(m, nn.SyncBatchNorm):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+if __name__ == "__main__":
+    import torch
+    model = AlignedXception(BatchNorm=nn.BatchNorm2d, output_stride=16)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

project/hair_service_sd/core/bodyseg/msc_distilling.py

@@ -0,0 +1,330 @@
+import sys
+# sys.path.append("/Users/momo/human_seg_train")
+# print(sys.path)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from core.bodyseg.backbone.backbone import get_norm
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
+        padding = (kernel_size - 1) // 2
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            norm_layer(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            norm_layer(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+class UpSampleBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, expand_ratio=6):
+        super(UpSampleBlock, self).__init__()
+        self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
+
+    def forward(self, x0, x1):
+        x = torch.cat([x0, x1], dim=1)
+        x = self.refine(x)
+        return x
+
+class BodySegNet_32_thin_sigmod(nn.Module):
+    def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
+        super(BodySegNet_32_thin_sigmod, self).__init__()
+        self.class_nums = class_nums
+        self.output_onnx = output_onnx
+
+        self.stage_1 = nn.Sequential(
+            ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
+        )
+        self.stage_2 = nn.Sequential(
+            ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
+            ConvBNReLU(16, 16, kernel_size=1, stride=1),
+        )
+        self.stage_3 = nn.Sequential(
+            InvertedResidual(16, 24, stride=2, expand_ratio=6),
+            InvertedResidual(24, 24, stride=1, expand_ratio=6),
+            InvertedResidual(24, 24, stride=1, expand_ratio=6),
+        )
+        self.stage_4 = nn.Sequential(
+            InvertedResidual(24, 32, stride=2, expand_ratio=6),
+            InvertedResidual(32, 32, stride=1, expand_ratio=6),
+            InvertedResidual(32, 32, stride=1, expand_ratio=6),
+            InvertedResidual(32, 32, stride=1, expand_ratio=6),
+        )
+        self.stage_5 = nn.Sequential(
+            InvertedResidual(32, 48, stride=2, expand_ratio=6),
+            InvertedResidual(48, 48, stride=1, expand_ratio=6),
+            InvertedResidual(48, 48, stride=1, expand_ratio=6),
+            InvertedResidual(48, 48, stride=1, expand_ratio=6)
+        )
+        self.up_to_4 = UpSampleBlock(48 + 32, 16)
+        self.up_to_3 = UpSampleBlock(16 + 24, 16)
+        self.up_to_2 = UpSampleBlock(16 + 16, 16)
+        self.up_to_1 = UpSampleBlock(16 + 16, 16)
+        self.last_layer = nn.Sequential(
+            ConvBNReLU(16, 16, kernel_size=1, stride=1),
+            nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
+        )
+        self._initialize_weights()
+
+    def forward(self, x):
+        feature_S = []
+        x1 = self.stage_1(x)
+        x2 = self.stage_2(x1)
+        x3 = self.stage_3(x2)
+        x4 = self.stage_4(x3)
+        feature = self.stage_5(x4)
+
+        feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
+        feature = self.up_to_4(x4, feature)
+        feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
+        feature = self.up_to_3(x3, feature)
+        feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
+        feature = self.up_to_2(x2, feature)
+        feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
+        feature = self.up_to_1(x1, feature)
+        feature_S.append(feature)
+        feature = self.last_layer(feature)
+        feature_S.append(feature)
+        output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
+        # output = torch.sigmoid(output)
+        if self.output_onnx:
+            output = torch.argmax(output, dim=1).to(torch.float32)
+
+        return output, feature_S
+
+    def _initialize_weights(self):
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Conv2d):
+                if 'first' in name:
+                    nn.init.normal_(m.weight, 0, 0.01)
+                else:
+                    nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+class _ASPPModule(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
+        super(_ASPPModule, self).__init__()
+        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
+                                            stride=1, padding=padding, dilation=dilation, bias=False)
+        self.bn = BatchNorm(planes)
+        self.relu = nn.ReLU()
+
+        self._init_weight()
+
+    def forward(self, x):
+        x = self.atrous_conv(x)
+        x = self.bn(x)
+
+        return self.relu(x)
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+
+class ASPP(nn.Module):
+    def __init__(self, backbone, output_stride, BatchNorm):
+        super(ASPP, self).__init__()
+        if backbone == 'drn':
+            inplanes = 512
+        elif backbone == 'mobilenet':
+            inplanes = 320
+        elif backbone == 'resnet_fpn':
+            inplanes = 256
+        else:
+            inplanes = 2048
+        if output_stride == 16:
+            dilations = [1, 6, 12, 18]
+        elif output_stride == 8:
+            dilations = [1, 12, 24, 36]
+        else:
+            raise NotImplementedError
+
+        self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
+        self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
+        self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
+        self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
+
+        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                             nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
+                                             BatchNorm(256),
+                                             nn.ReLU())
+        self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
+        self.bn1 = BatchNorm(256)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.5)
+        self._init_weight()
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x2 = self.aspp2(x)
+        x3 = self.aspp3(x)
+        x4 = self.aspp4(x)
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
+        # x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
+        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        return self.dropout(x)
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+
+class Decoder(nn.Module):
+    def __init__(self,num_classes, backbone, BatchNorm):
+        super(Decoder, self).__init__()
+        if backbone == 'resnet_fpn' or backbone == 'drn':
+            low_level_inplanes = 256
+        elif backbone == 'xception':
+            low_level_inplanes = 128
+        elif backbone == 'mobilenet':
+            low_level_inplanes = 24
+        else:
+            raise NotImplementedError
+
+        self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
+        self.bn1 = BatchNorm(16)
+        self.relu = nn.ReLU()
+        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.5),
+                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.1),
+                                       nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
+
+        self.up_to_1 = UpSampleBlock(16 + 16, 16)
+        self.last_layer = nn.Sequential(
+            ConvBNReLU(16, 16, kernel_size=1, stride=1),
+            nn.Conv2d(16, 1, kernel_size=1, stride=1)
+        )
+        self._init_weight()
+
+    #  x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
+    def forward(self, x, low_level_feat):
+        feature_T = []
+        # deeplab part
+        #(1,256,32,24) -> (1,16,64,48)
+        low_level_feat = self.conv1(low_level_feat)
+        low_level_feat = self.bn1(low_level_feat)
+        low_level_feat = self.relu(low_level_feat)
+
+        # x(1, 256, 8, 6)-> (1,16,32,24)
+        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
+        low_level_feat = F.interpolate(low_level_feat,
+                                       size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
+                                       mode='bilinear', align_corners=True)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        # bodyseg part
+        feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
+        # 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
+        feature = self.up_to_1(low_level_feat, feature)
+        feature_T.append(feature)
+        # (1,1,64,48)
+        feature = self.last_layer(feature)
+        feature_T.append(feature)
+        # (1,1,128,96)
+        output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
+        # output = torch.sigmoid(output)
+
+        return output, feature_T
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+
+def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
+    from core.bodyseg.backbone.fpn import build_resnet_fpn_backbone
+    from core.bodyseg.backbone.xception import AlignedXception
+
+    return build_resnet_fpn_backbone(input_channel)
+
+def build_aspp(backbone, output_stride, BatchNorm):
+    return ASPP(backbone, output_stride, BatchNorm)
+
+def build_decoder(num_classes, backbone, BatchNorm):
+    return Decoder(num_classes, backbone, BatchNorm)
+
+class DeepLab(nn.Module):
+    def __init__(self, input_channel=3, class_num=1):
+        super(DeepLab, self).__init__()
+
+        BatchNorm = get_norm("BN")
+
+        self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
+        self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
+        self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
+
+    def forward(self, input):
+        # input(1,3,128,96)
+        #output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
+        output = self.backbone(input)
+        # "p4"(1,256,8,6)  "p2"(1,256,32,24)
+        x, low_level_feat = output['p4'], output['p2']
+        # x(1,256,8,6)-》(1,256,8,6)
+        x = self.aspp(x)
+        # x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
+        x, feature_T = self.decoder(x, low_level_feat)
+        # x(1,1,128,96)
+        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
+        # x = F.interpolate(x, size=input.size()[2:], mode='nearest')
+        return x, feature_T

project/hair_service_sd/core/cos_module.py

@@ -0,0 +1,59 @@
+from qcloud_cos import CosConfig
+from qcloud_cos import CosS3Client
+import sys
+import os
+import os.path as osp
+import time
+import logging
+from common.logger import config as confccc
+
+
+class COS_object():
+    def __init__(self):
+        # 正常情况日志级别使用 INFO，需要定位时可以修改为 DEBUG，此时 SDK 会打印和服务端的通信信息
+        logging.basicConfig(level=logging.INFO, stream=sys.stdout)
+
+        # 1. 设置用户属性, 包括 secret_id, secret_key, region等。从环境变量读取，避免硬编码密钥。
+        secret_id = os.getenv('COS_SECRET_ID', '<your-cos-secret-id>')
+        secret_key = os.getenv('COS_SECRET_KEY', '<your-cos-secret-key>')
+        self.region = os.getenv('COS_REGION', '<your-cos-region>')      # 例: ap-beijing
+        token = None               # 如果使用永久密钥不需要填入 token，如果使用临时密钥需要填入，临时密钥生成和使用指引参见 https://cloud.tencent.com/document/product/436/14048
+        self.scheme = 'https'           # 指定使用 http/https 协议来访问 COS，默认为 https，可不填
+        self.BucketName= os.getenv('COS_BUCKET', '<your-cos-bucket>')
+        for param in (secret_id, secret_key, self.region, self.BucketName):
+            assert '<' not in param, '请设置环境变量 COS_SECRET_ID / COS_SECRET_KEY / COS_REGION / COS_BUCKET'
+        config = CosConfig(Region=self.region, SecretId=secret_id, SecretKey=secret_key, Token=token, Scheme=self.scheme)
+        self.client = CosS3Client(config)
+
+    def upload_file(self, file, target_name):
+        t0 = time.time()
+        with open(file, 'rb') as fp:
+            response = self.client.put_object(
+                Bucket=self.BucketName,  # Bucket 由 BucketName-APPID 组成
+                Body=fp,
+                Key=target_name,
+                StorageClass='STANDARD',
+                ContentType='image/jpeg;image/jpg;image/png;image/gif'
+            )
+        ret_url = self.scheme + '://' + self.BucketName + '.cos.' + self.region + '.myqcloud.com/' + target_name
+        print('time costs: {}'.format(time.time() - t0), ret_url)
+        return ret_url
+
+    def download_img(self, img_url, tmp_dir):
+        # user_img_save_dir = confccc.get('default', 'userDir')
+        # user_img_tmp_dir = confccc.get('default', 'tmp_dir')
+        fileName = img_url.split('.com/')[-1]
+        print(fileName)
+        response = self.client.get_object(
+            Bucket=self.BucketName,
+            Key=fileName,
+        )
+        response['Body'].get_stream_to_file(tmp_dir)
+
+
+if __name__=='__main__':
+    t0 = time.time()
+    cos = COS_object()
+    cos.upload_file('/home/chinatszrn/Pictures/vaffflue12.png', 'hair_mz/images/vaffflue12.png')
+    # cos.download_img('https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/fb7d7a58-3228-40f2-84a6-337088ee31e2/2023031623425988.jpg', 'a0')
+    print(time.time() - t0)

project/hair_service_sd/core/face3d/__init__.py

@@ -0,0 +1,2 @@
+from . import mesh
+from . import morphable_model

project/hair_service_sd/core/face3d/mesh/__init__.py

@@ -0,0 +1,7 @@
+# from .cython import mesh_core_cython
+from . import io
+from . import vis
+from . import transform
+from . import light
+from . import render
+

project/hair_service_sd/core/face3d/mesh/cython/.gitignore

@@ -0,0 +1 @@
+*.cpython-36m-x86_64-linux-gnu.so

project/hair_service_sd/core/face3d/mesh/cython/mesh_core.cpp

@@ -0,0 +1,375 @@
+/*
+functions that can not be optimazed by vertorization in python.
+1. rasterization.(need process each triangle)
+2. normal of each vertex.(use one-ring, need process each vertex)
+3. write obj(seems that it can be verctorized? anyway, writing it in c++ is simple, so also add function here. --> however, why writting in c++ is still slow?)
+
+Author: Yao Feng 
+Mail: yaofeng1995@gmail.com
+*/
+
+#include "mesh_core.h"
+
+
+/* Judge whether the point is in the triangle
+Method:
+    http://blackpawn.com/texts/pointinpoly/
+Args:
+    point: [x, y] 
+    tri_points: three vertices(2d points) of a triangle. 2 coords x 3 vertices
+Returns:
+    bool: true for in triangle
+*/
+bool isPointInTri(point p, point p0, point p1, point p2)
+{   
+    // vectors
+    point v0, v1, v2;
+    v0 = p2 - p0;
+    v1 = p1 - p0;
+    v2 = p - p0;
+
+    // dot products
+    float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
+    float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
+    float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
+    float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
+    float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
+
+    // barycentric coordinates
+    float inverDeno;
+    if(dot00*dot11 - dot01*dot01 == 0)
+        inverDeno = 0;
+    else
+        inverDeno = 1/(dot00*dot11 - dot01*dot01);
+
+    float u = (dot11*dot02 - dot01*dot12)*inverDeno;
+    float v = (dot00*dot12 - dot01*dot02)*inverDeno;
+
+    // check if point in triangle
+    return (u >= 0) && (v >= 0) && (u + v < 1);
+}
+
+
+void get_point_weight(float* weight, point p, point p0, point p1, point p2)
+{   
+    // vectors
+    point v0, v1, v2;
+    v0 = p2 - p0; 
+    v1 = p1 - p0; 
+    v2 = p - p0; 
+
+    // dot products
+    float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
+    float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
+    float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
+    float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
+    float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
+
+    // barycentric coordinates
+    float inverDeno;
+    if(dot00*dot11 - dot01*dot01 == 0)
+        inverDeno = 0;
+    else
+        inverDeno = 1/(dot00*dot11 - dot01*dot01);
+
+    float u = (dot11*dot02 - dot01*dot12)*inverDeno;
+    float v = (dot00*dot12 - dot01*dot02)*inverDeno;
+
+    // weight
+    weight[0] = 1 - u - v;
+    weight[1] = v;
+    weight[2] = u;
+}
+
+
+void _get_normal_core(
+    float* normal, float* tri_normal, int* triangles,
+    int ntri)
+{
+    int i, j;
+    int tri_p0_ind, tri_p1_ind, tri_p2_ind;
+
+    for(i = 0; i < ntri; i++)
+    {
+        tri_p0_ind = triangles[3*i];
+        tri_p1_ind = triangles[3*i + 1];
+        tri_p2_ind = triangles[3*i + 2];
+
+        for(j = 0; j < 3; j++)
+        {
+            normal[3*tri_p0_ind + j] = normal[3*tri_p0_ind + j] + tri_normal[3*i + j];
+            normal[3*tri_p1_ind + j] = normal[3*tri_p1_ind + j] + tri_normal[3*i + j];
+            normal[3*tri_p2_ind + j] = normal[3*tri_p2_ind + j] + tri_normal[3*i + j];
+        }
+    }
+}
+
+
+void _rasterize_triangles_core(
+    float* vertices, int* triangles, 
+    float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
+    int nver, int ntri,
+    int h, int w)
+{
+    int i;
+    int x, y, k;
+    int tri_p0_ind, tri_p1_ind, tri_p2_ind;
+    point p0, p1, p2, p;
+    int x_min, x_max, y_min, y_max;
+    float p_depth, p0_depth, p1_depth, p2_depth;
+    float weight[3];
+
+    for(i = 0; i < ntri; i++)
+    {
+        tri_p0_ind = triangles[3*i];
+        tri_p1_ind = triangles[3*i + 1];
+        tri_p2_ind = triangles[3*i + 2];
+
+        p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
+        p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
+        p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
+        
+        x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
+        x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
+      
+        y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
+        y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
+
+        if(x_max < x_min || y_max < y_min)
+        {
+            continue;
+        }
+
+        for(y = y_min; y <= y_max; y++) //h
+        {
+            for(x = x_min; x <= x_max; x++) //w
+            {
+                p.x = x; p.y = y;
+                if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
+                {
+                    get_point_weight(weight, p, p0, p1, p2);
+                    p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
+
+                    if((p_depth > depth_buffer[y*w + x]))
+                    {
+                        depth_buffer[y*w + x] = p_depth;
+                        triangle_buffer[y*w + x] = i;
+                        for(k = 0; k < 3; k++)
+                        {
+                            barycentric_weight[y*w*3 + x*3 + k] = weight[k];
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+
+void _render_colors_core(
+    float* image, float* vertices, int* triangles, 
+    float* colors, 
+    float* depth_buffer,
+    int nver, int ntri,
+    int h, int w, int c)
+{
+    int i;
+    int x, y, k;
+    int tri_p0_ind, tri_p1_ind, tri_p2_ind;
+    point p0, p1, p2, p;
+    int x_min, x_max, y_min, y_max;
+    float p_depth, p0_depth, p1_depth, p2_depth;
+    float p_color, p0_color, p1_color, p2_color;
+    float weight[3];
+
+    for(i = 0; i < ntri; i++)
+    {
+        tri_p0_ind = triangles[3*i];
+        tri_p1_ind = triangles[3*i + 1];
+        tri_p2_ind = triangles[3*i + 2];
+
+        p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
+        p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
+        p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
+        
+        x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
+        x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
+      
+        y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
+        y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
+
+        if(x_max < x_min || y_max < y_min)
+        {
+            continue;
+        }
+
+        for(y = y_min; y <= y_max; y++) //h
+        {
+            for(x = x_min; x <= x_max; x++) //w
+            {
+                p.x = x; p.y = y;
+                if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
+                {
+                    get_point_weight(weight, p, p0, p1, p2);
+                    p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
+
+                    if((p_depth > depth_buffer[y*w + x]))
+                    {
+                        for(k = 0; k < c; k++) // c
+                        {   
+                            p0_color = colors[c*tri_p0_ind + k];
+                            p1_color = colors[c*tri_p1_ind + k];
+                            p2_color = colors[c*tri_p2_ind + k]; 
+
+                            p_color = weight[0]*p0_color + weight[1]*p1_color + weight[2]*p2_color;
+                            image[y*w*c + x*c + k] = p_color;
+                        }
+
+                        depth_buffer[y*w + x] = p_depth;
+                    }
+                }
+            }
+        }
+    }
+}
+
+
+void _render_texture_core(
+    float* image, float* vertices, int* triangles, 
+    float* texture, float* tex_coords, int* tex_triangles, 
+    float* depth_buffer,
+    int nver, int tex_nver, int ntri, 
+    int h, int w, int c, 
+    int tex_h, int tex_w, int tex_c, 
+    int mapping_type)
+{
+    int i;
+    int x, y, k;
+    int tri_p0_ind, tri_p1_ind, tri_p2_ind;
+    int tex_tri_p0_ind, tex_tri_p1_ind, tex_tri_p2_ind;
+    point p0, p1, p2, p;
+    point tex_p0, tex_p1, tex_p2, tex_p;
+    int x_min, x_max, y_min, y_max;
+    float weight[3];
+    float p_depth, p0_depth, p1_depth, p2_depth;
+    float xd, yd;
+    float ul, ur, dl, dr;
+    for(i = 0; i < ntri; i++)
+    {
+        // mesh
+        tri_p0_ind = triangles[3*i];
+        tri_p1_ind = triangles[3*i + 1];
+        tri_p2_ind = triangles[3*i + 2];
+
+        p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
+        p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
+        p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
+       
+        // texture
+        tex_tri_p0_ind = tex_triangles[3*i];
+        tex_tri_p1_ind = tex_triangles[3*i + 1];
+        tex_tri_p2_ind = tex_triangles[3*i + 2];
+
+        tex_p0.x = tex_coords[3*tex_tri_p0_ind]; tex_p0.y = tex_coords[3*tri_p0_ind + 1];
+        tex_p1.x = tex_coords[3*tex_tri_p1_ind]; tex_p1.y = tex_coords[3*tri_p1_ind + 1];
+        tex_p2.x = tex_coords[3*tex_tri_p2_ind]; tex_p2.y = tex_coords[3*tri_p2_ind + 1];
+
+
+        x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
+        x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
+      
+        y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
+        y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
+
+
+        if(x_max < x_min || y_max < y_min)
+        {
+            continue;
+        }
+
+        for(y = y_min; y <= y_max; y++) //h
+        {
+            for(x = x_min; x <= x_max; x++) //w
+            {
+                p.x = x; p.y = y;
+                if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
+                {
+                    get_point_weight(weight, p, p0, p1, p2);
+                    p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
+                    
+                    if((p_depth > depth_buffer[y*w + x]))
+                    {
+                        // -- color from texture
+                        // cal weight in mesh tri
+                        get_point_weight(weight, p, p0, p1, p2);
+                        // cal coord in texture
+                        tex_p = tex_p0*weight[0] + tex_p1*weight[1] + tex_p2*weight[2];
+                        tex_p.x = max(min(tex_p.x, float(tex_w - 1)), float(0)); 
+                        tex_p.y = max(min(tex_p.y, float(tex_h - 1)), float(0)); 
+
+                        yd = tex_p.y - floor(tex_p.y);
+                        xd = tex_p.x - floor(tex_p.x);
+                        for(k = 0; k < c; k++)
+                        {
+                            if(mapping_type==0)// nearest
+                            {   
+                                image[y*w*c + x*c + k] = texture[int(round(tex_p.y))*tex_w*tex_c + int(round(tex_p.x))*tex_c + k];
+                            }
+                            else//bilinear interp
+                            { 
+                                ul = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
+                                ur = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
+                                dl = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
+                                dr = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
+
+                                image[y*w*c + x*c + k] = ul*(1-xd)*(1-yd) + ur*xd*(1-yd) + dl*(1-xd)*yd + dr*xd*yd;
+                            }
+
+                        }
+
+                        depth_buffer[y*w + x] = p_depth;
+                    } 
+                }
+            }
+        }
+    }
+}
+
+
+
+// ------------------------------------------------- write
+// obj write
+// Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/core/Mesh.hpp
+void _write_obj_with_colors_texture(string filename, string mtl_name, 
+    float* vertices, int* triangles, float* colors, float* uv_coords,
+    int nver, int ntri, int ntexver)
+{
+    int i;
+
+    ofstream obj_file(filename.c_str());
+
+    // first line of the obj file: the mtl name
+    obj_file << "mtllib " << mtl_name << endl;
+    
+    // write vertices 
+    for (i = 0; i < nver; ++i) 
+    {
+        obj_file << "v " << vertices[3*i] << " " << vertices[3*i + 1] << " " << vertices[3*i + 2] << colors[3*i] << " " << colors[3*i + 1] << " " << colors[3*i + 2] <<  endl;
+    }
+
+    // write uv coordinates
+    for (i = 0; i < ntexver; ++i) 
+    {
+        //obj_file << "vt " << uv_coords[2*i] << " " << (1 - uv_coords[2*i + 1]) << endl;
+        obj_file << "vt " << uv_coords[2*i] << " " << uv_coords[2*i + 1] << endl;
+    }
+
+    obj_file << "usemtl FaceTexture" << endl;
+    // write triangles
+    for (i = 0; i < ntri; ++i) 
+    {
+        // obj_file << "f " << triangles[3*i] << "/" << triangles[3*i] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << endl;
+        obj_file << "f " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i] << "/" << triangles[3*i] << endl;
+    }
+
+}

project/hair_service_sd/core/face3d/mesh/cython/mesh_core.h

@@ -0,0 +1,83 @@
+#ifndef MESH_CORE_HPP_
+#define MESH_CORE_HPP_
+
+#include <stdio.h>
+#include <cmath>
+#include <algorithm>  
+#include <string>
+#include <iostream>
+#include <fstream>
+
+using namespace std;
+
+class point
+{
+ public:
+    float x;
+    float y;
+
+    float dot(point p)
+    {
+        return this->x * p.x + this->y * p.y;
+    }
+
+    point operator-(const point& p)
+    {
+        point np;
+        np.x = this->x - p.x;
+        np.y = this->y - p.y;
+        return np;
+    }
+
+    point operator+(const point& p)
+    {
+        point np;
+        np.x = this->x + p.x;
+        np.y = this->y + p.y;
+        return np;
+    }
+
+    point operator*(float s)
+    {
+        point np;
+        np.x = s * this->x;
+        np.y = s * this->y;
+        return np;
+    }
+}; 
+
+
+bool isPointInTri(point p, point p0, point p1, point p2, int h, int w);
+void get_point_weight(float* weight, point p, point p0, point p1, point p2);
+
+void _get_normal_core(
+    float* normal, float* tri_normal, int* triangles,
+    int ntri);
+
+void _rasterize_triangles_core(
+    float* vertices, int* triangles, 
+    float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
+    int nver, int ntri,
+    int h, int w);
+
+void _render_colors_core(
+    float* image, float* vertices, int* triangles, 
+    float* colors, 
+    float* depth_buffer,
+    int nver, int ntri,
+    int h, int w, int c);
+
+void _render_texture_core(
+    float* image, float* vertices, int* triangles, 
+    float* texture, float* tex_coords, int* tex_triangles, 
+    float* depth_buffer,
+    int nver, int tex_nver, int ntri, 
+    int h, int w, int c, 
+    int tex_h, int tex_w, int tex_c, 
+    int mapping_type);
+
+void _write_obj_with_colors_texture(string filename, string mtl_name, 
+    float* vertices, int* triangles, float* colors, float* uv_coords,
+    int nver, int ntri, int ntexver);
+
+#endif

project/hair_service_sd/core/face3d/mesh/cython/setup.py

@@ -0,0 +1,22 @@
+'''
+python setup.py build_ext -i
+to compile
+'''
+
+# setup.py
+from distutils.core import setup
+from setuptools import Extension
+
+from Cython.Build import cythonize
+from Cython.Distutils import build_ext
+import numpy
+
+setup(
+	name = 'mesh_core_cython',
+    cmdclass={'build_ext': build_ext},
+    ext_modules=[Extension("mesh_core_cython",
+                 sources=["mesh_core_cython.pyx", "mesh_core.cpp"],
+                 language='c++',
+                 include_dirs=[numpy.get_include()])],
+)
+

project/hair_service_sd/core/face3d/mesh/io.py

@@ -0,0 +1,142 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import os
+# from skimage import io
+from time import time
+
+# from .cython import mesh_core_cython
+
+## TODO
+## TODO: c++ version
+def read_obj(obj_name):
+	''' read mesh
+	'''
+	return 0
+
+# ------------------------- write
+def write_asc(path, vertices):
+    '''
+    Args:
+        vertices: shape = (nver, 3)
+    '''
+    if path.split('.')[-1] == 'asc':
+        np.savetxt(path, vertices)
+    else:
+        np.savetxt(path + '.asc', vertices)
+
+def write_obj_with_colors(obj_name, vertices, triangles, colors):
+    ''' Save 3D face model with texture represented by colors.
+    Args:
+        obj_name: str
+        vertices: shape = (nver, 3)
+        triangles: shape = (ntri, 3)
+        colors: shape = (nver, 3)
+    '''
+    triangles = triangles.copy()
+    triangles += 1 # meshlab start with 1
+    
+    if obj_name.split('.')[-1] != 'obj':
+        obj_name = obj_name + '.obj'
+        
+    # write obj
+    with open(obj_name, 'w') as f:
+        
+        # write vertices & colors
+        for i in range(vertices.shape[0]):
+            # s = 'v {} {} {} \n'.format(vertices[0,i], vertices[1,i], vertices[2,i])
+            s = 'v {} {} {} {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2], colors[i, 0], colors[i, 1], colors[i, 2])
+            f.write(s)
+
+        # write f: ver ind/ uv ind
+        [k, ntri] = triangles.shape
+        for i in range(triangles.shape[0]):
+            # s = 'f {} {} {}\n'.format(triangles[i, 0], triangles[i, 1], triangles[i, 2])
+            s = 'f {} {} {}\n'.format(triangles[i, 2], triangles[i, 1], triangles[i, 0])
+            f.write(s)
+
+## TODO: c++ version
+def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords):
+    ''' Save 3D face model with texture represented by texture map.
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
+    Args:
+        obj_name: str
+        vertices: shape = (nver, 3)
+        triangles: shape = (ntri, 3)
+        texture: shape = (256,256,3)
+        uv_coords: shape = (nver, 3) max value<=1
+    '''
+    if obj_name.split('.')[-1] != 'obj':
+        obj_name = obj_name + '.obj'
+    mtl_name = obj_name.replace('.obj', '.mtl')
+    texture_name = obj_name.replace('.obj', '_texture.png')
+    
+    triangles = triangles.copy()
+    triangles += 1 # mesh lab start with 1
+    
+    # write obj
+    with open(obj_name, 'w') as f:
+        # first line: write mtlib(material library)
+        s = "mtllib {}\n".format(os.path.abspath(mtl_name))
+        f.write(s)
+
+        # write vertices
+        for i in range(vertices.shape[0]):
+            s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2])
+            f.write(s)
+        
+        # write uv coords
+        for i in range(uv_coords.shape[0]):
+            s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1])
+            f.write(s)
+
+        f.write("usemtl FaceTexture\n")
+
+        # write f: ver ind/ uv ind
+        for i in range(triangles.shape[0]):
+            s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0])
+            f.write(s)
+
+    # write mtl
+    with open(mtl_name, 'w') as f:
+        f.write("newmtl FaceTexture\n")
+        s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
+        f.write(s)
+
+    # write texture as png
+    imsave(texture_name, texture)
+
+# c++ version
+def write_obj_with_colors_texture(obj_name, vertices, triangles, colors, texture, uv_coords):
+    ''' Save 3D face model with texture. 
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
+    Args:
+        obj_name: str
+        vertices: shape = (nver, 3)
+        triangles: shape = (ntri, 3)
+        colors: shape = (nver, 3)
+        texture: shape = (256,256,3)
+        uv_coords: shape = (nver, 3) max value<=1
+    '''
+    if obj_name.split('.')[-1] != 'obj':
+        obj_name = obj_name + '.obj'
+    mtl_name = obj_name.replace('.obj', '.mtl')
+    texture_name = obj_name.replace('.obj', '_texture.png')
+    
+    triangles = triangles.copy()
+    triangles += 1 # mesh lab start with 1
+    
+    # write obj
+    vertices, colors, uv_coords = vertices.astype(np.float32).copy(), colors.astype(np.float32).copy(), uv_coords.astype(np.float32).copy()
+    mesh_core_cython.write_obj_with_colors_texture_core(str.encode(obj_name), str.encode(os.path.abspath(mtl_name)), vertices, triangles, colors, uv_coords, vertices.shape[0], triangles.shape[0], uv_coords.shape[0])
+   
+    # write mtl
+    with open(mtl_name, 'w') as f:
+        f.write("newmtl FaceTexture\n")
+        s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
+        f.write(s)
+
+    # write texture as png
+    # io.imsave(texture_name, texture)

project/hair_service_sd/core/face3d/mesh/light.py

@@ -0,0 +1,213 @@
+'''
+Functions about lighting mesh(changing colors/texture of mesh).
+1. add light to colors/texture (shade each vertex)
+2. fit light according to colors/texture & image.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+# from .cython import mesh_core_cython
+
+def get_normal(vertices, triangles):
+    ''' calculate normal direction in each vertex
+    Args:
+        vertices: [nver, 3]
+        triangles: [ntri, 3]
+    Returns:
+        normal: [nver, 3]
+    '''
+    pt0 = vertices[triangles[:, 0], :] # [ntri, 3]
+    pt1 = vertices[triangles[:, 1], :] # [ntri, 3]
+    pt2 = vertices[triangles[:, 2], :] # [ntri, 3]
+    tri_normal = np.cross(pt0 - pt1, pt0 - pt2) # [ntri, 3]. normal of each triangle
+
+    normal = np.zeros_like(vertices, dtype = np.float32).copy() # [nver, 3]
+    # for i in range(triangles.shape[0]):
+    #     normal[triangles[i, 0], :] = normal[triangles[i, 0], :] + tri_normal[i, :]
+    #     normal[triangles[i, 1], :] = normal[triangles[i, 1], :] + tri_normal[i, :]
+    #     normal[triangles[i, 2], :] = normal[triangles[i, 2], :] + tri_normal[i, :]
+    mesh_core_cython.get_normal_core(normal, tri_normal.astype(np.float32).copy(), triangles.copy(), triangles.shape[0])
+
+    # normalize to unit length
+    mag = np.sum(normal**2, 1) # [nver]
+    zero_ind = (mag == 0)
+    mag[zero_ind] = 1
+    normal[zero_ind, 0] = np.ones((np.sum(zero_ind)))
+
+    normal = normal/np.sqrt(mag[:,np.newaxis])
+
+    return normal
+
+# TODO: test
+def add_light_sh(vertices, triangles, colors, sh_coeff):
+    ''' 
+    In 3d face, usually assume:
+    1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
+    2. Lighting can be an arbitrary combination of point sources
+    --> can be expressed in terms of spherical harmonics(omit the lighting coefficients)
+    I = albedo * (sh(n) x sh_coeff)
+    
+    albedo: n x 1
+    sh_coeff: 9 x 1
+    Y(n) = (1, n_x, n_y, n_z, n_xn_y, n_xn_z, n_yn_z, n_x^2 - n_y^2, 3n_z^2 - 1)': n x 9 
+    # Y(n) = (1, n_x, n_y, n_z)': n x 4
+
+    Args:
+        vertices: [nver, 3]
+        triangles: [ntri, 3]
+        colors: [nver, 3] albedo
+        sh_coeff: [9, 1] spherical harmonics coefficients
+
+    Returns:
+        lit_colors: [nver, 3]
+    '''
+    assert vertices.shape[0] == colors.shape[0]
+    nver = vertices.shape[0]
+    normal = get_normal(vertices, triangles) # [nver, 3]
+    sh = np.array((np.ones(nver), n[:,0], n[:,1], n[:,2], n[:,0]*n[:,1], n[:,0]*n[:,2], n[:,1]*n[:,2], n[:,0]**2 - n[:,1]**2, 3*(n[:,2]**2) - 1)) # [nver, 9]
+    ref = sh.dot(sh_coeff) #[nver, 1]
+    lit_colors = colors*ref
+    return lit_colors
+
+
+def add_light(vertices, triangles, colors, light_positions = 0, light_intensities = 0):
+    ''' Gouraud shading. add point lights.
+    In 3d face, usually assume:
+    1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
+    2. Lighting can be an arbitrary combination of point sources
+    3. No specular (unless skin is oil, 23333)
+
+    Ref: https://cs184.eecs.berkeley.edu/lecture/pipeline    
+    Args:
+        vertices: [nver, 3]
+        triangles: [ntri, 3]
+        light_positions: [nlight, 3] 
+        light_intensities: [nlight, 3]
+    Returns:
+        lit_colors: [nver, 3]
+    '''
+    nver = vertices.shape[0]
+    normals = get_normal(vertices, triangles) # [nver, 3]
+
+    # ambient
+    # La = ka*Ia
+
+    # diffuse
+    # Ld = kd*(I/r^2)max(0, nxl)
+    direction_to_lights = vertices[np.newaxis, :, :] - light_positions[:, np.newaxis, :] # [nlight, nver, 3]
+    direction_to_lights_n = np.sqrt(np.sum(direction_to_lights**2, axis = 2)) # [nlight, nver]
+    direction_to_lights = direction_to_lights/direction_to_lights_n[:, :, np.newaxis]
+    normals_dot_lights = normals[np.newaxis, :, :]*direction_to_lights # [nlight, nver, 3]
+    normals_dot_lights = np.sum(normals_dot_lights, axis = 2) # [nlight, nver]
+    diffuse_output = colors[np.newaxis, :, :]*normals_dot_lights[:, :, np.newaxis]*light_intensities[:, np.newaxis, :]
+    diffuse_output = np.sum(diffuse_output, axis = 0) # [nver, 3]
+    
+    # specular
+    # h = (v + l)/(|v + l|) bisector
+    # Ls = ks*(I/r^2)max(0, nxh)^p
+    # increasing p narrows the reflectionlob
+
+    lit_colors = diffuse_output # only diffuse part here.
+    lit_colors = np.minimum(np.maximum(lit_colors, 0), 1)
+    return lit_colors
+
+
+
+## TODO. estimate light(sh coeff)
+## -------------------------------- estimate. can not use now. 
+def fit_light(image, vertices, colors, triangles, vis_ind, lamb = 10, max_iter = 3):
+    [h, w, c] = image.shape
+
+    # surface normal
+    norm = get_normal(vertices, triangles)
+    
+    nver = vertices.shape[1]
+
+    # vertices --> corresponding image pixel
+    pt2d = vertices[:2, :]
+
+    pt2d[0,:] = np.minimum(np.maximum(pt2d[0,:], 0), w - 1)
+    pt2d[1,:] = np.minimum(np.maximum(pt2d[1,:], 0), h - 1)
+    pt2d = np.round(pt2d).astype(np.int32) # 2 x nver
+
+    image_pixel = image[pt2d[1,:], pt2d[0,:], :] # nver x 3
+    image_pixel = image_pixel.T # 3 x nver
+
+    # vertices --> corresponding mean texture pixel with illumination
+    # Spherical Harmonic Basis
+    harmonic_dim = 9
+    nx = norm[0,:];
+    ny = norm[1,:];
+    nz = norm[2,:];
+    harmonic = np.zeros((nver, harmonic_dim))
+
+    pi = np.pi
+    harmonic[:,0] = np.sqrt(1/(4*pi)) * np.ones((nver,));
+    harmonic[:,1] = np.sqrt(3/(4*pi)) * nx;
+    harmonic[:,2] = np.sqrt(3/(4*pi)) * ny;
+    harmonic[:,3] = np.sqrt(3/(4*pi)) * nz;
+    harmonic[:,4] = 1/2. * np.sqrt(3/(4*pi)) * (2*nz**2 - nx**2 - ny**2);
+    harmonic[:,5] = 3 * np.sqrt(5/(12*pi)) * (ny*nz);
+    harmonic[:,6] = 3 * np.sqrt(5/(12*pi)) * (nx*nz);
+    harmonic[:,7] = 3 * np.sqrt(5/(12*pi)) * (nx*ny);
+    harmonic[:,8] = 3/2. * np.sqrt(5/(12*pi)) * (nx*nx - ny*ny);
+    
+    '''
+    I' = sum(albedo * lj * hj) j = 0:9 (albedo = tex)
+    set A = albedo*h (n x 9)
+        alpha = lj (9 x 1)
+        Y = I (n x 1)
+        Y' = A.dot(alpha)
+
+    opt function:
+        ||Y - A*alpha|| + lambda*(alpha'*alpha)
+    result:
+        A'*(Y - A*alpha) + lambda*alpha = 0
+        ==>
+        (A'*A*alpha - lambda)*alpha = A'*Y
+        left: 9 x 9
+        right: 9 x 1
+    '''
+    n_vis_ind = len(vis_ind)
+    n = n_vis_ind*c
+
+    Y = np.zeros((n, 1))
+    A = np.zeros((n, 9))
+    light = np.zeros((3, 1))
+
+    for k in range(c):
+        Y[k*n_vis_ind:(k+1)*n_vis_ind, :] = image_pixel[k, vis_ind][:, np.newaxis]
+        A[k*n_vis_ind:(k+1)*n_vis_ind, :] = texture[k, vis_ind][:, np.newaxis] * harmonic[vis_ind, :]
+        Ac = texture[k, vis_ind][:, np.newaxis]
+        Yc = image_pixel[k, vis_ind][:, np.newaxis]
+        light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
+
+    for i in range(max_iter):
+
+        Yc = Y.copy()
+        for k in range(c):
+            Yc[k*n_vis_ind:(k+1)*n_vis_ind, :]  /= light[k]
+
+        # update alpha
+        equation_left = np.dot(A.T, A) + lamb*np.eye(harmonic_dim); # why + ?
+        equation_right = np.dot(A.T, Yc) 
+        alpha = np.dot(np.linalg.inv(equation_left), equation_right)
+
+        # update light
+        for k in range(c):
+            Ac = A[k*n_vis_ind:(k+1)*n_vis_ind, :].dot(alpha)
+            Yc = Y[k*n_vis_ind:(k+1)*n_vis_ind, :]
+            light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
+
+    appearance = np.zeros_like(texture)
+    for k in range(c):
+        tmp = np.dot(harmonic*texture[k, :][:, np.newaxis], alpha*light[k])
+        appearance[k,:] = tmp.T
+
+    appearance = np.minimum(np.maximum(appearance, 0), 1)
+
+    return appearance
+

project/hair_service_sd/core/face3d/mesh/render.py

@@ -0,0 +1,135 @@
+'''
+functions about rendering mesh(from 3d obj to 2d image).
+only use rasterization render here.
+Note that:
+1. Generally, render func includes camera, light, raterize. Here no camera and light(I write these in other files)
+2. Generally, the input vertices are normalized to [-1,1] and cetered on [0, 0]. (in world space)
+   Here, the vertices are using image coords, which centers on [w/2, h/2] with the y-axis pointing to oppisite direction.
+ Means: render here only conducts interpolation.(I just want to make the input flexible)
+
+Author: Yao Feng 
+Mail: yaofeng1995@gmail.com
+'''
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from time import time
+
+# from .cython import mesh_core_cython
+
+def rasterize_triangles(vertices, triangles, h, w):
+    ''' 
+    Args:
+        vertices: [nver, 3]
+        triangles: [ntri, 3]
+        h: height
+        w: width
+    Returns:
+        depth_buffer: [h, w] saves the depth, here, the bigger the z, the fronter the point.
+        triangle_buffer: [h, w] saves the tri id(-1 for no triangle). 
+        barycentric_weight: [h, w, 3] saves corresponding barycentric weight.
+
+    # Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z.
+    # h, w is the size of rendering
+    '''
+
+    # initial 
+    depth_buffer = np.zeros([h, w]) - 999999. #set the initial z to the farest position
+    triangle_buffer = np.zeros([h, w], dtype = np.int32) - 1  # if tri id = -1, the pixel has no triangle correspondance
+    barycentric_weight = np.zeros([h, w, 3], dtype = np.float32)  # 
+    
+    vertices = vertices.astype(np.float32).copy()
+    triangles = triangles.astype(np.int32).copy()
+
+    mesh_core_cython.rasterize_triangles_core(
+                vertices, triangles,
+                depth_buffer, triangle_buffer, barycentric_weight, 
+                vertices.shape[0], triangles.shape[0], 
+                h, w)
+
+def render_colors(vertices, triangles, colors, h, w, c = 3, BG = None):
+    ''' render mesh with colors
+    Args:
+        vertices: [nver, 3]
+        triangles: [ntri, 3] 
+        colors: [nver, 3]
+        h: height
+        w: width  
+        c: channel
+        BG: background image
+    Returns:
+        image: [h, w, c]. rendered image./rendering.
+    '''
+
+    # initial 
+    if BG is None:
+        image = np.zeros((h, w, c), dtype = np.float32)
+    else:
+        assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
+        image = BG
+    depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
+
+    # change orders. --> C-contiguous order(column major)
+    vertices = vertices.astype(np.float32).copy()
+    triangles = triangles.astype(np.int32).copy()
+    colors = colors.astype(np.float32).copy()
+    ###
+    st = time()
+    mesh_core_cython.render_colors_core(
+                image, vertices, triangles,
+                colors,
+                depth_buffer,
+                vertices.shape[0], triangles.shape[0], 
+                h, w, c)
+    return image
+
+
+def render_texture(vertices, triangles, texture, tex_coords, tex_triangles, h, w, c = 3, mapping_type = 'nearest', BG = None):
+    ''' render mesh with texture map
+    Args:
+        vertices: [3, nver]
+        triangles: [3, ntri]
+        texture: [tex_h, tex_w, 3]
+        tex_coords: [ntexcoords, 3]
+        tex_triangles: [ntri, 3]
+        h: height of rendering
+        w: width of rendering
+        c: channel
+        mapping_type: 'bilinear' or 'nearest'
+    '''
+    # initial 
+    if BG is None:
+        image = np.zeros((h, w, c), dtype = np.float32)
+    else:
+        assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
+        image = BG.astype(np.float32)
+
+    depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
+    
+    tex_h, tex_w, tex_c = texture.shape
+    if mapping_type == 'nearest':
+        mt = int(0)
+    elif mapping_type == 'bilinear':
+        mt = int(1)
+    else:
+        mt = int(0)
+    
+    # -> C order
+    vertices = vertices.astype(np.float32).copy()
+    triangles = triangles.astype(np.int32).copy()
+    texture = texture.astype(np.float32).copy()
+    tex_coords = tex_coords.astype(np.float32).copy()
+    tex_triangles = tex_triangles.astype(np.int32).copy()
+
+    mesh_core_cython.render_texture_core(
+                image, vertices, triangles,
+                texture, tex_coords, tex_triangles,
+                depth_buffer,
+                vertices.shape[0], tex_coords.shape[0], triangles.shape[0], 
+                h, w, c,
+                tex_h, tex_w, tex_c,
+                mt)
+    return image
+

project/hair_service_sd/core/face3d/mesh/transform.py

@@ -0,0 +1,383 @@
+'''
+Functions about transforming mesh(changing the position: modify vertices).
+1. forward: transform(transform, camera, project).
+2. backward: estimate transform matrix from correspondences.
+
+Author: Yao Feng 
+Mail: yaofeng1995@gmail.com
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import math
+from math import cos, sin
+
+def angle2matrix(angles):
+    ''' get rotation matrix from three rotation angles(degree). right-handed.
+    Args:
+        angles: [3,]. x, y, z angles
+        x: pitch. positive for looking down.
+        y: yaw. positive for looking left. 
+        z: roll. positive for tilting head right. 
+    Returns:
+        R: [3, 3]. rotation matrix.
+    '''
+    x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
+    # x
+    Rx=np.array([[1,      0,       0],
+                 [0, cos(x),  -sin(x)],
+                 [0, sin(x),   cos(x)]])
+    # y
+    Ry=np.array([[ cos(y), 0, sin(y)],
+                 [      0, 1,      0],
+                 [-sin(y), 0, cos(y)]])
+    # z
+    Rz=np.array([[cos(z), -sin(z), 0],
+                 [sin(z),  cos(z), 0],
+                 [     0,       0, 1]])
+    
+    R=Rz.dot(Ry.dot(Rx))
+    return R.astype(np.float32)
+
+def angle2matrix_3ddfa(angles):
+    ''' get rotation matrix from three rotation angles(radian). The same as in 3DDFA.
+    Args:
+        angles: [3,]. x, y, z angles
+        x: pitch.
+        y: yaw. 
+        z: roll. 
+    Returns:
+        R: 3x3. rotation matrix.
+    '''
+    # x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
+    x, y, z = angles[0], angles[1], angles[2]
+    
+    # x
+    Rx=np.array([[1,      0,       0],
+                 [0, cos(x),  sin(x)],
+                 [0, -sin(x),   cos(x)]])
+    # y
+    Ry=np.array([[ cos(y), 0, -sin(y)],
+                 [      0, 1,      0],
+                 [sin(y), 0, cos(y)]])
+    # z
+    Rz=np.array([[cos(z), sin(z), 0],
+                 [-sin(z),  cos(z), 0],
+                 [     0,       0, 1]])
+    R = Rx.dot(Ry).dot(Rz)
+    return R.astype(np.float32)
+
+
+## ------------------------------------------ 1. transform(transform, project, camera).
+## ---------- 3d-3d transform. Transform obj in world space
+def rotate(vertices, angles):
+    ''' rotate vertices. 
+    X_new = R.dot(X). X: 3 x 1   
+    Args:
+        vertices: [nver, 3]. 
+        rx, ry, rz: degree angles
+        rx: pitch. positive for looking down 
+        ry: yaw. positive for looking left
+        rz: roll. positive for tilting head right
+    Returns:
+        rotated vertices: [nver, 3]
+    '''
+    R = angle2matrix(angles)
+    rotated_vertices = vertices.dot(R.T)
+
+    return rotated_vertices
+
+def similarity_transform(vertices, s, R, t3d):
+    ''' similarity transform. dof = 7.
+    3D: s*R.dot(X) + t
+    Homo: M = [[sR, t],[0^T, 1]].  M.dot(X)
+    Args:(float32)
+        vertices: [nver, 3]. 
+        s: [1,]. scale factor.
+        R: [3,3]. rotation matrix.
+        t3d: [3,]. 3d translation vector.
+    Returns:
+        transformed vertices: [nver, 3]
+    '''
+    t3d = np.squeeze(np.array(t3d, dtype = np.float32))
+    transformed_vertices = s * vertices.dot(R.T) + t3d[np.newaxis, :]
+
+    return transformed_vertices
+
+
+## -------------- Camera. from world space to camera space
+# Ref: https://cs184.eecs.berkeley.edu/lecture/transforms-2
+def normalize(x):
+    epsilon = 1e-12
+    norm = np.sqrt(np.sum(x**2, axis = 0))
+    norm = np.maximum(norm, epsilon)
+    return x/norm
+
+def lookat_camera(vertices, eye, at = None, up = None):
+    """ 'look at' transformation: from world space to camera space
+    standard camera space: 
+        camera located at the origin. 
+        looking down negative z-axis. 
+        vertical vector is y-axis.
+    Xcam = R(X - C)
+    Homo: [[R, -RC], [0, 1]]
+    Args:
+      vertices: [nver, 3] 
+      eye: [3,] the XYZ world space position of the camera.
+      at: [3,] a position along the center of the camera's gaze.
+      up: [3,] up direction 
+    Returns:
+      transformed_vertices: [nver, 3]
+    """
+    if at is None:
+      at = np.array([0, 0, 0], np.float32)
+    if up is None:
+      up = np.array([0, 1, 0], np.float32)
+
+    eye = np.array(eye).astype(np.float32)
+    at = np.array(at).astype(np.float32)
+    z_aixs = -normalize(at - eye) # look forward
+    x_aixs = normalize(np.cross(up, z_aixs)) # look right
+    y_axis = np.cross(z_aixs, x_aixs) # look up
+
+    R = np.stack((x_aixs, y_axis, z_aixs))#, axis = 0) # 3 x 3
+    transformed_vertices = vertices - eye # translation
+    transformed_vertices = transformed_vertices.dot(R.T) # rotation
+    return transformed_vertices
+
+## --------- 3d-2d project. from camera space to image plane
+# generally, image plane only keeps x,y channels, here reserve z channel for calculating z-buffer.
+def orthographic_project(vertices):
+    ''' scaled orthographic projection(just delete z)
+        assumes: variations in depth over the object is small relative to the mean distance from camera to object
+        x -> x*f/z, y -> x*f/z, z -> f.
+        for point i,j. zi~=zj. so just delete z
+        ** often used in face
+        Homo: P = [[1,0,0,0], [0,1,0,0], [0,0,1,0]]
+    Args:
+        vertices: [nver, 3]
+    Returns:
+        projected_vertices: [nver, 3] if isKeepZ=True. [nver, 2] if isKeepZ=False.
+    '''
+    return vertices.copy()
+
+def perspective_project(vertices, fovy, aspect_ratio = 1., near = 0.1, far = 1000.):
+    ''' perspective projection.
+    Args:
+        vertices: [nver, 3]
+        fovy: vertical angular field of view. degree.
+        aspect_ratio : width / height of field of view
+        near : depth of near clipping plane
+        far : depth of far clipping plane
+    Returns:
+        projected_vertices: [nver, 3] 
+    '''
+    fovy = np.deg2rad(fovy)
+    top = near*np.tan(fovy)
+    bottom = -top 
+    right = top*aspect_ratio
+    left = -right
+
+    #-- homo
+    P = np.array([[near/right, 0, 0, 0],
+                 [0, near/top, 0, 0],
+                 [0, 0, -(far+near)/(far-near), -2*far*near/(far-near)],
+                 [0, 0, -1, 0]])
+    vertices_homo = np.hstack((vertices, np.ones((vertices.shape[0], 1)))) # [nver, 4]
+    projected_vertices = vertices_homo.dot(P.T)
+    projected_vertices = projected_vertices/projected_vertices[:,3:]
+    projected_vertices = projected_vertices[:,:3]
+    projected_vertices[:,2] = -projected_vertices[:,2]
+
+    #-- non homo. only fovy
+    # projected_vertices = vertices.copy()
+    # projected_vertices[:,0] = -(near/right)*vertices[:,0]/vertices[:,2]
+    # projected_vertices[:,1] = -(near/top)*vertices[:,1]/vertices[:,2]
+    return projected_vertices
+
+
+def to_image(vertices, h, w, is_perspective = False):
+    ''' change vertices to image coord system
+    3d system: XYZ, center(0, 0, 0)
+    2d image: x(u), y(v). center(w/2, h/2), flip y-axis. 
+    Args:
+        vertices: [nver, 3]
+        h: height of the rendering
+        w : width of the rendering
+    Returns:
+        projected_vertices: [nver, 3]  
+    '''
+    image_vertices = vertices.copy()
+    if is_perspective:
+        # if perspective, the projected vertices are normalized to [-1, 1]. so change it to image size first.
+        image_vertices[:,0] = image_vertices[:,0]*w/2
+        image_vertices[:,1] = image_vertices[:,1]*h/2
+    # move to center of image
+    image_vertices[:,0] = image_vertices[:,0] + w/2
+    image_vertices[:,1] = image_vertices[:,1] + h/2
+    # flip vertices along y-axis.
+    image_vertices[:,1] = h - image_vertices[:,1] - 1
+    return image_vertices
+
+
+#### -------------------------------------------2. estimate transform matrix from correspondences.
+def estimate_affine_matrix_3d23d(X, Y):
+    ''' Using least-squares solution 
+    Args:
+        X: [n, 3]. 3d points(fixed)
+        Y: [n, 3]. corresponding 3d points(moving). Y = PX
+    Returns:
+        P_Affine: (3, 4). Affine camera matrix (the third row is [0, 0, 0, 1]).
+    '''
+    X_homo = np.hstack((X, np.ones([X.shape[1],1]))) #n x 4
+    P = np.linalg.lstsq(X_homo, Y)[0].T # Affine matrix. 3 x 4
+    return P
+    
+def estimate_affine_matrix_3d22d(X, x):
+    ''' Using Golden Standard Algorithm for estimating an affine camera
+        matrix P from world to image correspondences.
+        See Alg.7.2. in MVGCV 
+        Code Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/fitting/affine_camera_estimation.hpp
+        x_homo = X_homo.dot(P_Affine)
+    Args:
+        X: [n, 3]. corresponding 3d points(fixed)
+        x: [n, 2]. n>=4. 2d points(moving). x = PX
+    Returns:
+        P_Affine: [3, 4]. Affine camera matrix
+    '''
+    X = X.T; x = x.T
+    assert(x.shape[1] == X.shape[1])
+    n = x.shape[1]
+    assert(n >= 4)
+
+    #--- 1. normalization
+    # 2d points
+    mean = np.mean(x, 1) # (2,)
+    x = x - np.tile(mean[:, np.newaxis], [1, n])
+    average_norm = np.mean(np.sqrt(np.sum(x**2, 0)))
+    scale = np.sqrt(2) / average_norm
+    x = scale * x
+
+    T = np.zeros((3,3), dtype = np.float32)
+    T[0, 0] = T[1, 1] = scale
+    T[:2, 2] = -mean*scale
+    T[2, 2] = 1
+
+    # 3d points
+    X_homo = np.vstack((X, np.ones((1, n))))
+    mean = np.mean(X, 1) # (3,)
+    X = X - np.tile(mean[:, np.newaxis], [1, n])
+    m = X_homo[:3,:] - X
+    average_norm = np.mean(np.sqrt(np.sum(X**2, 0)))
+    scale = np.sqrt(3) / average_norm
+    X = scale * X
+
+    U = np.zeros((4,4), dtype = np.float32)
+    U[0, 0] = U[1, 1] = U[2, 2] = scale
+    U[:3, 3] = -mean*scale
+    U[3, 3] = 1
+
+    # --- 2. equations
+    A = np.zeros((n*2, 8), dtype = np.float32);
+    X_homo = np.vstack((X, np.ones((1, n)))).T
+    A[:n, :4] = X_homo
+    A[n:, 4:] = X_homo
+    b = np.reshape(x, [-1, 1])
+ 
+    # --- 3. solution
+    p_8 = np.linalg.pinv(A).dot(b)
+    P = np.zeros((3, 4), dtype = np.float32)
+    P[0, :] = p_8[:4, 0]
+    P[1, :] = p_8[4:, 0]
+    P[-1, -1] = 1
+
+    # --- 4. denormalization
+    P_Affine = np.linalg.inv(T).dot(P.dot(U))
+    return P_Affine
+
+def P2sRt(P):
+    ''' decompositing camera matrix P
+    Args: 
+        P: (3, 4). Affine Camera Matrix.
+    Returns:
+        s: scale factor.
+        R: (3, 3). rotation matrix.
+        t: (3,). translation. 
+    '''
+    t = P[:, 3]
+    R1 = P[0:1, :3]
+    R2 = P[1:2, :3]
+    s = (np.linalg.norm(R1) + np.linalg.norm(R2))/2.0
+    r1 = R1/np.linalg.norm(R1)
+    r2 = R2/np.linalg.norm(R2)
+    r3 = np.cross(r1, r2)
+
+    R = np.concatenate((r1, r2, r3), 0)
+    return s, R, t
+
+#Ref: https://www.learnopencv.com/rotation-matrix-to-euler-angles/
+def isRotationMatrix(R):
+    ''' checks if a matrix is a valid rotation matrix(whether orthogonal or not)
+    '''
+    Rt = np.transpose(R)
+    shouldBeIdentity = np.dot(Rt, R)
+    I = np.identity(3, dtype = R.dtype)
+    n = np.linalg.norm(I - shouldBeIdentity)
+    return n < 1e-6
+
+def matrix2angle(R):
+    ''' get three Euler angles from Rotation Matrix
+    Args:
+        R: (3,3). rotation matrix
+    Returns:
+        x: pitch
+        y: yaw
+        z: roll
+    '''
+    assert(isRotationMatrix)
+    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
+     
+    singular = sy < 1e-6
+ 
+    if  not singular :
+        x = math.atan2(R[2,1] , R[2,2])
+        y = math.atan2(-R[2,0], sy)
+        z = math.atan2(R[1,0], R[0,0])
+    else :
+        x = math.atan2(-R[1,2], R[1,1])
+        y = math.atan2(-R[2,0], sy)
+        z = 0
+
+    # rx, ry, rz = np.rad2deg(x), np.rad2deg(y), np.rad2deg(z)
+    rx, ry, rz = x*180/np.pi, y*180/np.pi, z*180/np.pi
+    return rx, ry, rz
+
+# def matrix2angle(R):
+#     ''' compute three Euler angles from a Rotation Matrix. Ref: http://www.gregslabaugh.net/publications/euler.pdf
+#     Args:
+#         R: (3,3). rotation matrix
+#     Returns:
+#         x: yaw
+#         y: pitch
+#         z: roll
+#     '''
+#     # assert(isRotationMatrix(R))
+
+#     if R[2,0] !=1 or R[2,0] != -1:
+#         x = math.asin(R[2,0])
+#         y = math.atan2(R[2,1]/cos(x), R[2,2]/cos(x))
+#         z = math.atan2(R[1,0]/cos(x), R[0,0]/cos(x))
+        
+#     else:# Gimbal lock
+#         z = 0 #can be anything
+#         if R[2,0] == -1:
+#             x = np.pi/2
+#             y = z + math.atan2(R[0,1], R[0,2])
+#         else:
+#             x = -np.pi/2
+#             y = -z + math.atan2(-R[0,1], -R[0,2])
+
+#     return x, y, z

project/hair_service_sd/core/face3d/mesh/vis.py

@@ -0,0 +1,24 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import matplotlib.pyplot as plt
+# from skimage import measure
+from mpl_toolkits.mplot3d import Axes3D
+
+def plot_mesh(vertices, triangles, subplot = [1,1,1], title = 'mesh', el = 90, az = -90, lwdt=.1, dist = 6, color = "grey"):
+	'''
+	plot the mesh 
+	Args:
+		vertices: [nver, 3]
+		triangles: [ntri, 3]
+	'''
+	ax = plt.subplot(subplot[0], subplot[1], subplot[2], projection = '3d')
+	ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles = triangles, lw = lwdt, color = color, alpha = 1)
+	ax.axis("off")
+	ax.view_init(elev = el, azim = az)
+	ax.dist = dist
+	plt.title(title)
+
+### -------------- Todo: use vtk to visualize mesh? or visvis? or VisPy?

project/hair_service_sd/core/face3d/morphable_model/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .. import mesh
+from .morphabel_model import MorphabelModel
+from . import load

project/hair_service_sd/core/face3d/morphable_model/fit.py

@@ -0,0 +1,270 @@
+'''
+Estimating parameters about vertices: shape para, exp para, pose para(s, R, t)
+'''
+import numpy as np
+
+''' TODO: a clear document. 
+Given: image_points, 3D Model, Camera Matrix(s, R, t2d)
+Estimate: shape parameters, expression parameters
+
+Inference: 
+
+    projected_vertices = s*P*R(mu + shape + exp) + t2d  --> image_points
+    s*P*R*shape + s*P*R(mu + exp) + t2d --> image_poitns
+
+    # Define:
+    X = vertices
+    x_hat = projected_vertices
+    x = image_points
+    A = s*P*R
+    b = s*P*R(mu + exp) + t2d
+    ==>
+    x_hat = A*shape + b  (2 x n)
+
+    A*shape (2 x n)
+    shape = reshape(shapePC * sp) (3 x n)
+    shapePC*sp : (3n x 1)
+
+    * flatten:
+    x_hat_flatten = A*shape + b_flatten  (2n x 1)
+    A*shape (2n x 1)
+    --> A*shapePC (2n x 199)  sp: 199 x 1
+    
+    # Define:
+    pc_2d = A* reshape(shapePC)
+    pc_2d_flatten = flatten(pc_2d) (2n x 199)
+
+    =====>
+    x_hat_flatten = pc_2d_flatten * sp + b_flatten ---> x_flatten (2n x 1)
+
+    Goals:
+    (ignore flatten, pc_2d-->pc)
+    min E = || x_hat - x || + lambda*sum(sp/sigma)^2
+          = || pc * sp + b - x || + lambda*sum(sp/sigma)^2
+
+    Solve:
+    d(E)/d(sp) = 0
+    2 * pc' * (pc * sp + b - x) + 2 * lambda * sp / (sigma' * sigma) = 0
+
+    Get:
+    (pc' * pc + lambda / (sigma'* sigma)) * sp  = pc' * (x - b)
+
+'''
+
+def estimate_shape(x, shapeMU, shapePC, shapeEV, expression, s, R, t2d, lamb = 3000):
+    '''
+    Args:
+        x: (2, n). image points (to be fitted)
+        shapeMU: (3n, 1)
+        shapePC: (3n, n_sp)
+        shapeEV: (n_sp, 1)
+        expression: (3, n)
+        s: scale
+        R: (3, 3). rotation matrix
+        t2d: (2,). 2d translation
+        lambda: regulation coefficient
+
+    Returns:
+        shape_para: (n_sp, 1) shape parameters(coefficients)
+    '''
+    x = x.copy()
+    assert(shapeMU.shape[0] == shapePC.shape[0])
+    assert(shapeMU.shape[0] == x.shape[1]*3)
+
+    dof = shapePC.shape[1]
+
+    n = x.shape[1]
+    sigma = shapeEV
+    t2d = np.array(t2d)
+    P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
+    A = s*P.dot(R)
+
+    # --- calc pc
+    pc_3d = np.resize(shapePC.T, [dof, n, 3]) # 199 x n x 3
+    pc_3d = np.reshape(pc_3d, [dof*n, 3]) 
+    pc_2d = pc_3d.dot(A.T.copy()) # 199 x n x 2
+    
+    pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 199
+
+    # --- calc b
+    # shapeMU
+    mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
+    # expression
+    exp_3d = expression
+    # 
+    b = A.dot(mu_3d + exp_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
+    b = np.reshape(b.T, [-1, 1]) # 2n x 1
+
+    # --- solve
+    equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
+    x = np.reshape(x.T, [-1, 1])
+    equation_right = np.dot(pc.T, x - b)
+
+    shape_para = np.dot(np.linalg.inv(equation_left), equation_right)
+
+    return shape_para
+
+def estimate_expression(x, shapeMU, expPC, expEV, shape, s, R, t2d, lamb = 2000):
+    '''
+    Args:
+        x: (2, n). image points (to be fitted)
+        shapeMU: (3n, 1)
+        expPC: (3n, n_ep)
+        expEV: (n_ep, 1)
+        shape: (3, n)
+        s: scale
+        R: (3, 3). rotation matrix
+        t2d: (2,). 2d translation
+        lambda: regulation coefficient
+
+    Returns:
+        exp_para: (n_ep, 1) shape parameters(coefficients)
+    '''
+    x = x.copy()
+    assert(shapeMU.shape[0] == expPC.shape[0])
+    assert(shapeMU.shape[0] == x.shape[1]*3)
+
+    dof = expPC.shape[1]
+
+    n = x.shape[1]
+    sigma = expEV
+    t2d = np.array(t2d)
+    P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
+    A = s*P.dot(R)
+
+    # --- calc pc
+    pc_3d = np.resize(expPC.T, [dof, n, 3]) 
+    pc_3d = np.reshape(pc_3d, [dof*n, 3]) 
+    pc_2d = pc_3d.dot(A.T) 
+    pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 29
+
+    # --- calc b
+    # shapeMU
+    mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
+    # expression
+    shape_3d = shape
+    # 
+    b = A.dot(mu_3d + shape_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
+    b = np.reshape(b.T, [-1, 1]) # 2n x 1
+
+    # --- solve
+    equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
+    x = np.reshape(x.T, [-1, 1])
+    equation_right = np.dot(pc.T, x - b)
+
+    exp_para = np.dot(np.linalg.inv(equation_left), equation_right)
+    
+    return exp_para
+
+
+# ---------------- fit 
+def fit_points(x, X_ind, model, n_sp, n_ep, max_iter = 4):
+    '''
+    Args:
+        x: (n, 2) image points
+        X_ind: (n,) corresponding Model vertex indices
+        model: 3DMM
+        max_iter: iteration
+    Returns:
+        sp: (n_sp, 1). shape parameters
+        ep: (n_ep, 1). exp parameters
+        s, R, t
+    '''
+    x = x.copy().T
+
+    #-- init
+    sp = np.zeros((n_sp, 1), dtype = np.float32)
+    ep = np.zeros((n_ep, 1), dtype = np.float32)
+
+    #-------------------- estimate
+    X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
+    X_ind_all[1, :] += 1
+    X_ind_all[2, :] += 2
+    valid_ind = X_ind_all.flatten('F')
+
+    shapeMU = model['shapeMU'][valid_ind, :]
+    shapePC = model['shapePC'][valid_ind, :n_sp]
+    expPC = model['expPC'][valid_ind, :n_ep]
+
+    for i in range(max_iter):
+        X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
+        X = np.reshape(X, [int(len(X)/3), 3]).T
+        
+        #----- estimate pose
+        P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
+        s, R, t = Face.face3d.mesh.transform.P2sRt(P)
+        rx, ry, rz = Face.face3d.mesh.transform.matrix2angle(R)
+        # print('Iter:{}; estimated pose: s {}, rx {}, ry {}, rz {}, t1 {}, t2 {}'.format(i, s, rx, ry, rz, t[0], t[1]))
+
+        #----- estimate shape
+        # expression
+        shape = shapePC.dot(sp)
+        shape = np.reshape(shape, [int(len(shape)/3), 3]).T
+        ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
+
+        # shape
+        expression = expPC.dot(ep)
+        expression = np.reshape(expression, [int(len(expression)/3), 3]).T
+        sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
+
+    return sp, ep, s, R, t
+
+
+# ---------------- fitting process
+def fit_points_for_show(x, X_ind, model, n_sp, n_ep, max_iter = 4):
+    '''
+    Args:
+        x: (n, 2) image points
+        X_ind: (n,) corresponding Model vertex indices
+        model: 3DMM
+        max_iter: iteration
+    Returns:
+        sp: (n_sp, 1). shape parameters
+        ep: (n_ep, 1). exp parameters
+        s, R, t
+    '''
+    x = x.copy().T
+
+    #-- init
+    sp = np.zeros((n_sp, 1), dtype = np.float32)
+    ep = np.zeros((n_ep, 1), dtype = np.float32)
+
+    #-------------------- estimate
+    X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
+    X_ind_all[1, :] += 1
+    X_ind_all[2, :] += 2
+    valid_ind = X_ind_all.flatten('F')
+
+    shapeMU = model['shapeMU'][valid_ind, :]
+    shapePC = model['shapePC'][valid_ind, :n_sp]
+    expPC = model['expPC'][valid_ind, :n_ep]
+
+    s = 4e-04
+    R = Face.face3d.mesh.transform.angle2matrix([0, 0, 0])
+    t = [0, 0, 0]
+    lsp = []; lep = []; ls = []; lR = []; lt = []
+    for i in range(max_iter):
+        X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
+        X = np.reshape(X, [int(len(X)/3), 3]).T
+        lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
+        
+        #----- estimate pose
+        P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
+        s, R, t = Face.face3d.mesh.transform.P2sRt(P)
+        lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
+
+        #----- estimate shape
+        # expression
+        shape = shapePC.dot(sp)
+        shape = np.reshape(shape, [int(len(shape)/3), 3]).T
+        ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
+        lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
+
+        # shape
+        expression = expPC.dot(ep)
+        expression = np.reshape(expression, [int(len(expression)/3), 3]).T
+        sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
+
+    # print('ls', ls)
+    # print('lR', lR)
+    return np.array(lsp), np.array(lep), np.array(ls), np.array(lR), np.array(lt)

project/hair_service_sd/core/face3d/morphable_model/load.py

@@ -0,0 +1,110 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import scipy.io as sio
+
+### ---------------------------------  load BFM data
+def load_BFM(model_path):
+    ''' load BFM 3DMM model
+    Args:
+        model_path: path to BFM model. 
+    Returns:
+        model: (nver = 53215, ntri = 105840). nver: number of vertices. ntri: number of triangles.
+            'shapeMU': [3*nver, 1]
+            'shapePC': [3*nver, 199]
+            'shapeEV': [199, 1]
+            'expMU': [3*nver, 1]
+            'expPC': [3*nver, 29]
+            'expEV': [29, 1]
+            'texMU': [3*nver, 1]
+            'texPC': [3*nver, 199]
+            'texEV': [199, 1]
+            'tri': [ntri, 3] (start from 1, should sub 1 in python and c++)
+            'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles)
+            'kpt_ind': [68,] (start from 1)
+    PS:
+        You can change codes according to your own saved data.
+        Just make sure the model has corresponding attributes.
+    '''
+    C = sio.loadmat(model_path)
+    model = C['model']
+    model = model[0,0]
+
+    # change dtype from double(np.float64) to np.float32, 
+    # since big matrix process(espetially matrix dot) is too slow in python.
+    model['shapeMU'] = (model['shapeMU'] + model['expMU']).astype(np.float32)
+    model['shapePC'] = model['shapePC'].astype(np.float32)
+    model['shapeEV'] = model['shapeEV'].astype(np.float32)
+    model['expEV'] = model['expEV'].astype(np.float32)
+    model['expPC'] = model['expPC'].astype(np.float32)
+
+    # matlab start with 1. change to 0 in python.
+    model['tri'] = model['tri'].T.copy(order = 'C').astype(np.int32) - 1
+    model['tri_mouth'] = model['tri_mouth'].T.copy(order = 'C').astype(np.int32) - 1
+    
+    # kpt ind
+    model['kpt_ind'] = (np.squeeze(model['kpt_ind']) - 1).astype(np.int32)
+
+    return model
+
+def load_BFM_info(path = 'BFM_info.mat'):
+    ''' load 3DMM model extra information
+    Args:
+        path: path to BFM info. 
+    Returns:  
+        model_info:
+            'symlist': 2 x 26720
+            'symlist_tri': 2 x 52937
+            'segbin': 4 x n (0: nose, 1: eye, 2: mouth, 3: cheek)
+            'segbin_tri': 4 x ntri 
+            'face_contour': 1 x 28
+            'face_contour_line': 1 x 512
+            'face_contour_front': 1 x 28
+            'face_contour_front_line': 1 x 512
+            'nose_hole': 1 x 142
+            'nose_hole_right': 1 x 71
+            'nose_hole_left': 1 x 71
+            'parallel': 17 x 1 cell
+            'parallel_face_contour': 28 x 1 cell
+            'uv_coords': n x 2
+    '''
+    C = sio.loadmat(path)
+    model_info = C['model_info']
+    model_info = model_info[0,0]
+    return model_info
+
+def load_uv_coords(path = 'BFM_UV.mat'):
+    ''' load uv coords of BFM
+    Args:
+        path: path to data.
+    Returns:  
+        uv_coords: [nver, 2]. range: 0-1
+    '''
+    C = sio.loadmat(path)
+    uv_coords = C['UV'].copy(order = 'C')
+    return uv_coords
+
+def load_pncc_code(path = 'pncc_code.mat'):
+    ''' load pncc code of BFM
+    PNCC code: Defined in 'Face Alignment Across Large Poses: A 3D Solution Xiangyu'
+    download at http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm.
+    Args:
+        path: path to data.
+    Returns:  
+        pncc_code: [nver, 3]
+    '''
+    C = sio.loadmat(path)
+    pncc_code = C['vertex_code'].T
+    return pncc_code
+
+## 
+def get_organ_ind(model_info):
+    ''' get nose, eye, mouth index
+    '''
+    valid_bin = model_info['segbin'].astype(bool)
+    organ_ind = np.nonzero(valid_bin[0,:])[0]
+    for i in range(1, valid_bin.shape[0] - 1):
+        organ_ind = np.union1d(organ_ind, np.nonzero(valid_bin[i,:])[0])
+    return organ_ind.astype(np.int32)

project/hair_service_sd/core/face3d/morphable_model/morphabel_model.py

@@ -0,0 +1,141 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from . import fit
+from . import load
+
+class  MorphabelModel(object):
+    """docstring for  MorphabelModel
+    model: nver: number of vertices. ntri: number of triangles. *: must have. ~: can generate ones array for place holder.
+            'shapeMU': [3*nver, 1]. *
+            'shapePC': [3*nver, n_shape_para]. *
+            'shapeEV': [n_shape_para, 1]. ~
+            'expMU': [3*nver, 1]. ~ 
+            'expPC': [3*nver, n_exp_para]. ~
+            'expEV': [n_exp_para, 1]. ~
+            'texMU': [3*nver, 1]. ~
+            'texPC': [3*nver, n_tex_para]. ~
+            'texEV': [n_tex_para, 1]. ~
+            'tri': [ntri, 3] (start from 1, should sub 1 in python and c++). *
+            'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles). ~
+            'kpt_ind': [68,] (start from 1). ~
+    """
+    def __init__(self, model_path, model_type = 'BFM'):
+        super( MorphabelModel, self).__init__()
+        if model_type=='BFM':
+            self.model = load.load_BFM(model_path)
+        else:
+            print('sorry, not support other 3DMM model now')
+            exit()
+            
+        # fixed attributes
+        self.nver = self.model['shapePC'].shape[0]/3
+        self.ntri = self.model['tri'].shape[0]
+        self.n_shape_para = self.model['shapePC'].shape[1]
+        self.n_exp_para = self.model['expPC'].shape[1]
+        self.n_tex_para = self.model['texMU'].shape[1]
+        
+        self.kpt_ind = self.model['kpt_ind']
+        self.triangles = self.model['tri']
+        self.full_triangles = np.vstack((self.model['tri'], self.model['tri_mouth']))
+
+    # ------------------------------------- shape: represented with mesh(vertices & triangles(fixed))
+    def get_shape_para(self, type = 'random'):
+        if type == 'zero':
+            sp = np.random.zeros((self.n_shape_para, 1))
+        elif type == 'random':
+            sp = np.random.rand(self.n_shape_para, 1)*1e04
+        return sp
+
+    def get_exp_para(self, type = 'random'):
+        if type == 'zero':
+            ep = np.zeros((self.n_exp_para, 1))
+        elif type == 'random':
+            ep = -1.5 + 3*np.random.random([self.n_exp_para, 1])
+            ep[6:, 0] = 0
+
+        return ep 
+
+    def generate_vertices(self, shape_para, exp_para):
+        '''
+        Args:
+            shape_para: (n_shape_para, 1)
+            exp_para: (n_exp_para, 1) 
+        Returns:
+            vertices: (nver, 3)
+        '''
+        vertices = self.model['shapeMU'] + self.model['shapePC'].dot(shape_para) + self.model['expPC'].dot(exp_para)
+        vertices = np.reshape(vertices, [int(3), int(len(vertices)/3)], 'F').T
+
+        return vertices
+
+    # -------------------------------------- texture: here represented with rgb value(colors) in vertices.
+    def get_tex_para(self, type = 'random'):
+        if type == 'zero':
+            tp = np.zeros((self.n_tex_para, 1))
+        elif type == 'random':
+            tp = np.random.rand(self.n_tex_para, 1)
+        return tp
+
+    def generate_colors(self, tex_para):
+        '''
+        Args:
+            tex_para: (n_tex_para, 1)
+        Returns:
+            colors: (nver, 3)
+        '''
+        colors = self.model['texMU'] + self.model['texPC'].dot(tex_para*self.model['texEV'])
+        colors = np.reshape(colors, [int(3), int(len(colors)/3)], 'F').T/255.  
+        
+        return colors
+
+
+    # ------------------------------------------- transformation
+    # -------------  transform
+    def rotate(self, vertices, angles):
+        ''' rotate face
+        Args:
+            vertices: [nver, 3]
+            angles: [3] x, y, z rotation angle(degree)
+            x: pitch. positive for looking down 
+            y: yaw. positive for looking left
+            z: roll. positive for tilting head right
+        Returns:
+            vertices: rotated vertices
+        '''
+        return Face.face3d.mesh.transform.rotate(vertices, angles)
+
+    def transform(self, vertices, s, angles, t3d):
+        R = Face.face3d.mesh.transform.angle2matrix(angles)
+        return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
+
+    def transform_3ddfa(self, vertices, s, angles, t3d): # only used for processing 300W_LP data
+        R = Face.face3d.mesh.transform.angle2matrix_3ddfa(angles)
+        return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
+
+    # --------------------------------------------------- fitting
+    def fit(self, x, X_ind, max_iter = 4, isShow = False):
+        ''' fit 3dmm & pose parameters
+        Args:
+            x: (n, 2) image points
+            X_ind: (n,) corresponding Model vertex indices
+            max_iter: iteration
+            isShow: whether to reserve middle results for show
+        Returns:
+            fitted_sp: (n_sp, 1). shape parameters
+            fitted_ep: (n_ep, 1). exp parameters
+            s, angles, t
+        '''
+        if isShow:
+            fitted_sp, fitted_ep, s, R, t = fit.fit_points_for_show(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
+            angles = np.zeros((R.shape[0], 3))
+            for i in range(R.shape[0]):
+                angles[i] = Face.face3d.mesh.transform.matrix2angle(R[i])
+        else:
+            fitted_sp, fitted_ep, s, R, t = fit.fit_points(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
+            angles = Face.face3d.mesh.transform.matrix2angle(R)
+        return fitted_sp, fitted_ep, s, angles, t
+
+

project/hair_service_sd/core/face_enhance/face_enhancement.py

@@ -0,0 +1,96 @@
+import os
+import cv2
+import glob
+import numpy as np
+from core.utils import landmark_processor
+
+from core.face_enhance.face_gan_pt import FaceGAN
+from time import time
+
+class FaceEnhancement(object):
+    def __init__(self,  size=512, gpu_id=None):
+        self.facegan = FaceGAN(size, gpu_id)
+        self.size = size
+        self.threshold = 0.9
+
+        # the mask for pasting restored faces back
+        self.mask = np.zeros((512, 512), np.float32)
+        cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
+        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
+        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
+
+        self.kernel = np.array((
+                [0.0625, 0.125, 0.0625],
+                [0.125, 0.25, 0.125],
+                [0.0625, 0.125, 0.0625]), dtype="float32")
+
+    def process(self, img, landmarks1k):
+
+        assert len(landmarks1k) == 1000
+
+        image_to_face_mat = landmark_processor.get_transform_mat_face_restore(landmarks1k, self.size)
+        tfm_inv = cv2.invertAffineTransform(image_to_face_mat)
+
+        height, width = img.shape[:2]
+        full_mask = np.zeros((height, width), dtype=np.float32)
+        full_img = np.zeros(img.shape, dtype=np.uint8)
+
+        fh, fw = (landmarks1k[0][1]-landmarks1k[154][1]), (landmarks1k[95][0]-landmarks1k[215][0])
+
+        of = cv2.warpAffine(img, image_to_face_mat, (self.size, self.size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[0, 0, 0])
+
+        # enhance the face
+        t0 = time()
+
+        ef = self.facegan.process(of)
+        # print('facegan process costs:', time() - t0)
+        tmp_mask = self.mask
+        tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
+        tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)
+
+        if min(fh, fw)<100: # gaussian filter for small faces
+            ef = cv2.filter2D(ef, -1, self.kernel)
+
+        # tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
+        tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), dst=img.copy(), borderMode=cv2.BORDER_TRANSPARENT)
+
+        # cv2.imshow("tmp_img: ", tmp_img)
+
+        mask = tmp_mask - full_mask
+        full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]
+        full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]
+
+        full_mask = full_mask[:, :, np.newaxis]
+        img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)
+
+        return img
+
+if __name__=='__main__':
+    
+    indir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res'
+    outdir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res_outs2'
+    os.makedirs(outdir, exist_ok=True)
+
+    faceenhancer = FaceEnhancement(base_dir="/", size=512, model="GPEN-512", channel_multiplier=2)
+
+    files = sorted(glob.glob(os.path.join(indir, '*.*g')))
+    for n, file in enumerate(files[:]):
+        filename = os.path.basename(file)
+        txtname = file.replace(".jpg", "_landmark1k.txt")
+
+        im = cv2.imread(file, cv2.IMREAD_COLOR) # BGR
+        print(txtname)
+        landmark = np.loadtxt(txtname)
+        if not isinstance(im, np.ndarray): print(filename, 'error'); continue
+
+        start = time()
+
+        img = faceenhancer.process(im, landmark)
+
+        end = time()
+
+        print("Time cost: {:.4f}".format(end - start))
+
+        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_2.jpg'), img)
+
+

project/hair_service_sd/core/face_enhance/face_gan_pt.py

@@ -0,0 +1,64 @@
+'''
+@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
+@author: yangxy (yangtao9009@gmail.com)
+'''
+import torch
+import os
+import cv2
+import numpy as np
+class FaceGAN(object):
+    def __init__(self, size=512, gpu_id=0):
+        # self.mfile = os.path.join(base_dir, model+'.pth')
+        self.n_mlp = 8
+        self.resolution = size
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.load_model()
+
+    def load_model(self):
+        self.face_gan_model = os.path.join('weights', "face_enhance_0630.pt")
+        self.model = torch.jit.load(self.face_gan_model, torch.device('cpu')).to(self.device)
+        # self.model = torch.jit.load(self.face_gan_model).to(self.device)
+        # self.model = torch.load(self.face_gan_model,map_location=self.device)
+
+        self.model.eval()
+
+    def process_o(self, img):
+        img = cv2.resize(img, (self.resolution, self.resolution))
+        img_t = self.img2tensor(img)
+
+        with torch.no_grad():
+            out, __ = self.model(img_t)
+
+        out = self.tensor2img(out)
+
+        return out
+
+    def process(self, img):
+        img = cv2.resize(img, (self.resolution, self.resolution))
+        img_t = self.img2tensor(img)
+
+        with torch.no_grad():
+            out = self.forward(img_t)
+
+        out = self.tensor2img(out)
+
+        return out
+
+    def forward(self, img_t):
+        with torch.no_grad():
+            out = self.model(img_t)
+
+        return out
+
+    def img2tensor(self, img):
+        img_t = (torch.from_numpy(img).to(self.device)/255. - 0.5) / 0.5
+        img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB
+        return img_t
+
+    def tensor2img(self, image_tensor, pmax=255.0, imtype=np.uint8):
+        image_tensor = image_tensor * 0.5 + 0.5
+        image_tensor = image_tensor.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
+        image_numpy = np.clip(image_tensor.float().cpu().numpy(), 0, 1) * pmax
+
+        return image_numpy.astype(imtype)

project/hair_service_sd/core/face_enhance/setup.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# @File    : setup.py
+# @Time    : 2020/1/15
+# @Author  : yangchaojie (yangchaojie@immomo.com)
+
+import os
+import sys
+import shutil
+import numpy
+import tempfile
+
+from setuptools import setup
+from setuptools.extension import Extension
+
+from Cython.Build import cythonize
+from Cython.Distutils import build_ext
+
+import platform
+
+
+def get_root_path(root):
+    if os.path.dirname(root) in ['', '.']:
+        return os.path.basename(root)
+    else:
+        return get_root_path(os.path.dirname(root))
+
+
+def copy_file(src, dest):
+    if os.path.exists(dest):
+        return
+
+    if not os.path.exists(os.path.dirname(dest)):
+        os.makedirs(os.path.dirname(dest))
+    if os.path.isdir(src):
+        shutil.copytree(src, dest)
+    else:
+        shutil.copyfile(src, dest)
+
+
+def touch_init_file():
+    init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
+    with open(init_file_name, 'w'):
+        pass
+    return init_file_name
+
+
+
+
+def compose_extensions(root='.'):
+    for file_ in os.listdir(root):
+        abs_file = os.path.join(root, file_)
+
+        if os.path.isfile(abs_file):
+            if abs_file.endswith('.py'):
+                extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
+            elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
+                continue
+            else:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            if abs_file.endswith('__init__.py'):
+                copy_file(init_file, os.path.join(build_root_dir, abs_file))
+
+        else:
+            if os.path.basename(abs_file) in ignore_folders :
+                continue
+            if os.path.basename(abs_file) in conf_folders:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            compose_extensions(abs_file)
+
+
+
+# if __name__ == '__main__':
+build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
+    sys.version_info.major) + '.' + str(sys.version_info.minor)
+
+print(build_root_dir)
+
+extensions = []
+ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
+conf_folders = ['conf']
+
+
+init_file = touch_init_file()
+print(init_file)
+
+
+compose_extensions()
+os.remove(init_file)
+
+setup(
+    name='moxie_hairstyle',
+    version='1.0',
+    ext_modules=cythonize(
+        extensions,
+        nthreads=16,
+        compiler_directives=dict(always_allow_keywords=True),
+        include_path=[numpy.get_include()]),
+    cmdclass=dict(build_ext=build_ext))
+
+    # python setup.py build_ext

project/hair_service_sd/core/face_quality.py

@@ -0,0 +1,186 @@
+import cv2
+
+
+class SingleFaceQualityInfo():
+    def __int__(self, gpu_id):
+        print('Model 3DDFA init success')
+
+    def getSpotProportion(self, grayImage):
+        Proportion = 0
+        hist = cv2.calcHist([grayImage], [0], None, [256], [0, 256])
+        for i in range(253, 256):
+            Proportion = Proportion + hist[i, 0]
+        return Proportion
+
+    def GetSpotRatio(self, wrap_img, wrap_face_rect):
+        scale = 0.9
+
+        x = max(0, int(wrap_face_rect[0]))
+        y = max(0, int(wrap_face_rect[1]))
+        w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
+        h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
+
+        height, width, _ = wrap_img.shape
+        if (x + w) > width:
+            w = width - x
+        if (y + h) > height:
+            h = height - y
+        p_cen = (x + w/2, y+h/2)
+        new_w = w*scale
+        new_h = h*scale
+        na_box = [int(p_cen[0]-new_w/2), int(p_cen[1]) - int(new_h/2), int(new_w), int(new_h)]
+        imgCropGray = cv2.cvtColor(wrap_img[na_box[1]:na_box[1] + na_box[3], na_box[0]:na_box[0] + na_box[2]],cv2.COLOR_RGB2GRAY)
+
+        Spot_Ratio = self.getSpotProportion(imgCropGray) / (imgCropGray.shape[0] * imgCropGray.shape[1])
+        return Spot_Ratio
+
+    def BoxClarityAndBrightValue(self, wrap_img, wrap_face_rect):
+        x = max(0, int(wrap_face_rect[0]))
+        y = max(0, int(wrap_face_rect[1]))
+        w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
+        h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
+        # cv2.imshow('face', wrap_img[y:y+h,x:x+w])
+        # cv2.waitKey()
+        height, width, _ = wrap_img.shape
+        if (x + w) > width:
+            w = width - x
+        if (y + h) > height:
+            h = height - y
+        box_crop = [int(x), int(y), int(w), int(h)]
+        imgCropGray = cv2.cvtColor(wrap_img[box_crop[1]:box_crop[1] + box_crop[3], box_crop[0]:box_crop[0] + box_crop[2]],cv2.COLOR_RGB2GRAY)
+        left_ = [0,0,int(w/2.0),int(h)]
+        right_ = [int(w/2.0),0,int(w/2.0),int(h)]
+        imgR = imgCropGray[right_[1]:right_[1] + right_[3], right_[0]:right_[0] + right_[2]]
+        imgL = imgCropGray[left_[1]:left_[1] + left_[3], left_[0]:left_[0] + left_[2]]
+        mean_R, var_R = cv2.meanStdDev(imgR)
+        mean_L, var_L = cv2.meanStdDev(imgL)
+
+        bV_R = mean_R[0]
+        bV_L = mean_L[0]
+
+        imageCropSobel = cv2.Laplacian(imgCropGray, cv2.CV_64F, 3)
+        mean, var = cv2.meanStdDev(imageCropSobel)
+        clarityValue = var[0]
+        result = [bV_R, bV_L, clarityValue]
+        return result
+
+    def qualityTest(self, wrap_img, wrap_face_rect, single_eulers):
+        single_face_quality_info = {
+            'quality_score': -1,
+            'face_quality_flag': -1,
+            'SpotRatio': 0,
+            'Brightness': 0,
+            'Clarity': 0,
+            'SpotScore': 0,
+            'BrightnessScore': 0,
+            'ClarityScore': 0
+        }
+
+        if wrap_img is None:
+            single_face_quality_info['face_quality_flag'] = 0
+        qp_img_size = wrap_img.shape[1]
+        Spot_Ratio = self.GetSpotRatio(wrap_img, wrap_face_rect)
+        result = self.BoxClarityAndBrightValue(wrap_img, wrap_face_rect)
+        bv_R, bv_L, clarityValue = result
+        single_face_quality_info['SpotRatio'] = Spot_Ratio
+        single_face_quality_info['Brightness'] = min(bv_L, bv_R)
+        single_face_quality_info['Clarity'] = clarityValue
+        if qp_img_size == 100:
+            clarityValue_thresholdDown = 35
+            clarityValue_thresholdUp = 300
+        else:
+            clarityValue_thresholdDown = 24
+            clarityValue_thresholdUp = 300
+        Bright_area_thre = 0.06
+        brightnessValue_thresholddown = 40
+        brightnessValue_thresholdUp = 230
+
+        whitespot = (Spot_Ratio > Bright_area_thre)
+        brightness =(bv_R > brightnessValue_thresholddown and bv_R < brightnessValue_thresholdUp) and (bv_L > brightnessValue_thresholddown and bv_L < brightnessValue_thresholdUp)
+        clarity = clarityValue < clarityValue_thresholdUp and clarityValue > clarityValue_thresholdDown
+        if whitespot:
+            single_face_quality_info['face_qality_flag'] = 2
+        if brightness:
+            single_face_quality_info['face_qality_flag'] = 4
+        if not clarity:
+            single_face_quality_info['face_qality_flag'] = 3
+        if brightness and clarity and not whitespot:
+            single_face_quality_info['face_qality_flag'] = 1
+
+
+        if single_face_quality_info['face_qality_flag'] == 1:
+            PerfectSpotRatio = 0.02
+            if qp_img_size == 100:
+                PerfectClarity = 130
+            else:
+                PerfectClarity = 60
+            PerfectBrightness_down = 90
+            PerfectBrightness_up = 200
+            PerfectBrightness_standard = PerfectBrightness_down - brightnessValue_thresholddown
+            PerfectClarity_standard = PerfectClarity - clarityValue_thresholdDown
+            PerfectSpot_standard = Bright_area_thre - PerfectSpotRatio
+            if single_face_quality_info['SpotRatio'] <= PerfectSpotRatio:
+                single_face_quality_info['SpotScore'] = 1
+            else:
+                single_face_quality_info['SpotRatio'] = 1 - (single_face_quality_info['SpotRatio'] - PerfectSpotRatio) / PerfectSpot_standard
+            if single_face_quality_info['Clarity'] >= PerfectClarity:
+                single_face_quality_info['ClarityScore'] = 1
+            else:
+                single_face_quality_info['ClarityScore'] = 1 - (PerfectClarity - single_face_quality_info['Clarity']) / PerfectClarity_standard
+            if single_face_quality_info['Brightness'] >= PerfectBrightness_down and single_face_quality_info['Brightness']<= PerfectBrightness_up:
+                single_face_quality_info['BrightnessScore'] = 1
+            elif single_face_quality_info['Brightness'] < PerfectBrightness_down:
+                single_face_quality_info['BrightnessScore'] = 1 - (PerfectBrightness_down - single_face_quality_info['Brightness']) / PerfectBrightness_standard
+            elif single_face_quality_info['Brightness'] > PerfectBrightness_up:
+                single_face_quality_info['BrightnessScore'] = 1 - (single_face_quality_info['Brightness'] - PerfectBrightness_up)/PerfectBrightness_standard
+            else:
+                print("Wrong Brightness :%f",single_face_quality_info['Brightness'])
+        else:
+            single_face_quality_info['ClarityScore'] = 0
+            single_face_quality_info['BrightnessScore'] = 0
+            single_face_quality_info['SpotScore'] = 0
+
+        if len(single_eulers) > 2:
+            euler_pitch_perfect = 10
+            euler_yaw_perfect = 10
+
+            euler_pitch_max = 25
+            euler_yaw_max = 25
+            euler_roll_max = 30
+
+            image_ClarityScore = float(single_face_quality_info['ClarityScore'])
+            image_BrightnessScore = single_face_quality_info['BrightnessScore']
+            image_SpotScore = single_face_quality_info['SpotScore']
+
+            euler_pitch, euler_yaw, euler_roll = single_eulers
+
+            if abs(euler_pitch) > euler_pitch_perfect and abs(euler_pitch) < euler_pitch_max and single_face_quality_info['face_qality_flag'] == 1:
+                image_PitchScore = 1 - (abs(euler_pitch) - euler_pitch_perfect) / (euler_pitch_max - euler_pitch_perfect)
+            elif abs(euler_pitch) < euler_pitch_perfect and single_face_quality_info['face_qality_flag'] == 1:
+                image_PitchScore = 1.0
+            else:
+                image_PitchScore = 0
+            if abs(euler_yaw) > euler_yaw_perfect and abs(euler_yaw) < euler_yaw_max and single_face_quality_info['face_qality_flag']==1:
+                image_YawScore = 1 - (abs(euler_yaw) - euler_yaw_perfect) / (euler_yaw_max - euler_yaw_perfect)
+            elif abs(euler_yaw) < euler_yaw_perfect and single_face_quality_info['face_qality_flag']==1:
+                image_YawScore = 1.0
+            else:
+                image_YawScore = 0
+
+            if image_PitchScore < 0.3 and image_YawScore < 0.3:
+                image_spot_weight = 0
+                image_clarity_weight = 0.1
+                image_brightness_weight = 0.1
+                euler_pitch_weight = 0.4
+                euler_yaw_weight = 0.4
+            else:
+                image_spot_weight = 0.05
+                image_clarity_weight = 0.35
+                image_brightness_weight = 0.1
+                euler_pitch_weight = 0.25
+                euler_yaw_weight = 0.25
+            if abs(euler_pitch) < euler_pitch_max and abs(euler_yaw) < euler_yaw_max and abs(euler_roll) < euler_roll_max and single_face_quality_info['face_qality_flag'] == 1:
+                single_face_quality_info['quality_score'] = (image_spot_weight*image_SpotScore + image_clarity_weight * image_ClarityScore + image_brightness_weight * image_BrightnessScore + euler_pitch_weight * image_PitchScore + euler_yaw_weight* image_YawScore)
+            else:
+                single_face_quality_info['quality_score'] = 0
+        return single_face_quality_info

project/hair_service_sd/core/faceseg/face_seg.py

@@ -0,0 +1,75 @@
+import os
+import pickle
+
+import torch
+from core.faceseg.u2net import U2NET
+from core.utils import landmark_processor
+import cv2
+import numpy as np
+
+
+class FaceSeg:
+    def __init__(self, gpu_id = 0):
+        model = U2NET(in_ch=4, out_ch=1)
+        weights = torch.load('weights/faceseg_20210927_01.pth', map_location='cpu')
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        for ix, (k, v) in enumerate(model_dict.items()):
+            if k in weights and weights[k].data.shape == v.data.shape:
+                pretrained_dict[k] = weights[k]
+            else:
+                print('ignore {}'.format(k))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print('update success')
+        model.cuda(gpu_id)
+        model.eval()
+        self.model = model
+        self.last_mask = None
+        self.output_img_size = 320
+        self.gpu_id = gpu_id
+
+    def inference(self, frame, pt1k, video_mode = False):
+        image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
+        face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
+        if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
+        if video_mode and self.last_mask is not None:
+            last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
+                                             (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
+            input_img = np.concatenate([face_image, last_small_mask], axis=2)
+        else:
+            zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
+            input_img = np.concatenate([face_image, zero_mask], axis=2)
+
+        face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
+        face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)
+        with torch.no_grad():
+            mask = self.model.test(face_image_tensor)
+        mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))
+
+        origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
+                                     flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
+        if video_mode: self.last_mask = origin_mask.copy()
+
+        return origin_mask
+
+
+if __name__ == '__main__':
+    face_segmentor = FaceSeg(gpu_id=0)
+    testdata_dir = "/home/yangchaojie/Desktop/faceswap_hd/datasets/origin/8171"
+    for picname in os.listdir(testdata_dir):
+        img_path = os.path.join(testdata_dir, picname)
+        pkl_path = img_path[:-4]+".pkl"
+        if not picname.endswith(".jpg"):
+            continue
+        if not os.path.exists(pkl_path):
+            continue
+        img = cv2.imread(img_path)
+        with open(pkl_path, "rb") as fp:
+            info = pickle.load(fp)
+        pt1k = info["pt1k"]
+
+        face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
+        # cv2.imshow("face_seg_mask", face_seg_mask)
+        # cv2.imshow("img", img)
+        # cv2.waitKey()

project/hair_service_sd/core/faceseg/stm.py

@@ -0,0 +1,299 @@
+from __future__ import division
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import torch.utils.model_zoo as model_zoo
+from torchvision import models
+
+# general libs
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image
+import numpy as np
+import math
+import time
+import tqdm
+import os
+import argparse
+import copy
+import sys
+
+from utils.helpers import *
+
+
+class ResBlock(nn.Module):
+    def __init__(self, backbone, indim, outdim=None, stride=1):
+        super(ResBlock, self).__init__()
+        self.backbone = backbone
+        if outdim == None:
+            outdim = indim
+        if indim == outdim and stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
+
+        self.conv1 = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(outdim, outdim, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        if self.backbone == 'resnest101':
+            r = self.conv1(F.relu(x, inplace=True))
+            r = self.conv2(F.relu(r, inplace=True))
+        else:
+            r = self.conv1(F.relu(x))
+            r = self.conv2(F.relu(r))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return x + r
+
+
+class Encoder_M(nn.Module):
+    def __init__(self, backbone):
+        super(Encoder_M, self).__init__()
+        if backbone == 'resnest101':
+            self.conv1_m = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
+            self.conv1_o = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
+        else:
+            self.conv1_m = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
+            self.conv1_o = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
+
+        if backbone == 'resnet50':
+            resnet = models.resnet50(pretrained=True)
+        elif backbone == 'resnet18':
+            resnet = models.resnet18(pretrained=True)
+
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.relu = resnet.relu  # 1/2, 64
+        self.maxpool = resnet.maxpool
+
+        self.res2 = resnet.layer1  # 1/4, 256
+        self.res3 = resnet.layer2  # 1/8, 512
+        self.res4 = resnet.layer3  # 1/8, 1024
+
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def forward(self, in_f, in_m, in_o):
+        f = (in_f - self.mean) / self.std
+        m = torch.unsqueeze(in_m, dim=1).float()  # add channel dim
+        o = torch.unsqueeze(in_o, dim=1).float()  # add channel dim
+
+        x = self.conv1(f) + self.conv1_m(m) + self.conv1_o(o)
+        x = self.bn1(x)
+        c1 = self.relu(x)  # 1/2, 64
+        x = self.maxpool(c1)  # 1/4, 64
+        r2 = self.res2(x)  # 1/4, 256
+        r3 = self.res3(r2)  # 1/8, 512
+        r4 = self.res4(r3)  # 1/8, 1024
+        return r4, r3, r2, c1, f
+
+
+class Encoder_Q(nn.Module):
+    def __init__(self, backbone):
+        super(Encoder_Q, self).__init__()
+
+        if backbone == 'resnet50':
+            resnet = models.resnet50(pretrained=True)
+        elif backbone == 'resnet18':
+            resnet = models.resnet18(pretrained=True)
+
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.relu = resnet.relu  # 1/2, 64
+        self.maxpool = resnet.maxpool
+
+        self.res2 = resnet.layer1  # 1/4, 256
+        self.res3 = resnet.layer2  # 1/8, 512
+        self.res4 = resnet.layer3  # 1/8, 1024
+
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def forward(self, in_f):
+        f = (in_f - self.mean) / self.std
+
+        x = self.conv1(f)
+        x = self.bn1(x)
+        c1 = self.relu(x)  # 1/2, 64
+        x = self.maxpool(c1)  # 1/4, 64
+        r2 = self.res2(x)  # 1/4, 256
+        r3 = self.res3(r2)  # 1/8, 512
+        r4 = self.res4(r3)  # 1/8, 1024
+        return r4, r3, r2, c1, f
+
+
+class Refine(nn.Module):
+    def __init__(self, backbone, inplanes, planes, scale_factor=2):
+        super(Refine, self).__init__()
+        self.convFS = nn.Conv2d(inplanes, planes, kernel_size=(3, 3), padding=(1, 1), stride=1)
+        self.ResFS = ResBlock(backbone, planes, planes)
+        self.ResMM = ResBlock(backbone, planes, planes)
+        self.scale_factor = scale_factor
+
+    def forward(self, f, pm):
+        s = self.ResFS(self.convFS(f))
+        m = s + F.interpolate(pm, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
+        m = self.ResMM(m)
+        return m
+
+
+class Decoder(nn.Module):
+    def __init__(self, mdim, scale_rate, backbone):
+        super(Decoder, self).__init__()
+        self.backbone = backbone
+        if backbone == 'resnest101':
+            self.convFM = nn.Conv2d(256, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
+        else:
+            self.convFM = nn.Conv2d(1024 // scale_rate, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
+        self.ResMM = ResBlock(backbone, mdim, mdim)
+        self.RF3 = Refine(backbone, 512 // scale_rate, mdim)  # 1/8 -> 1/4
+        self.RF2 = Refine(backbone, 256 // scale_rate, mdim)  # 1/4 -> 1
+
+        self.pred2 = nn.Conv2d(mdim, 2, kernel_size=(3, 3), padding=(1, 1), stride=1)
+
+    def forward(self, r4, r3, r2):
+        m4 = self.ResMM(self.convFM(r4))
+        m3 = self.RF3(r3, m4)  # out: 1/8, 256
+        m2 = self.RF2(r2, m3)  # out: 1/4, 256
+
+        if self.backbone == 'resnest101':
+            p2 = self.pred2(F.relu(m2, inplace=True))
+        else:
+            p2 = self.pred2(F.relu(m2))
+
+        p = F.interpolate(p2, scale_factor=4, mode='bilinear', align_corners=False)
+        return p  # , p2, p3, p4
+
+
+class Memory(nn.Module):
+    def __init__(self):
+        super(Memory, self).__init__()
+
+    def forward(self, m_in, m_out, q_in, q_out):  # m_in: o,c,t,h,w
+        B, D_e, T, H, W = m_in.size()
+        _, D_o, _, _, _ = m_out.size()
+
+        mi = m_in.view(B, D_e, T * H * W)
+        mi = torch.transpose(mi, 1, 2)  # b, THW, emb
+
+        qi = q_in.view(B, D_e, H * W)  # b, emb, HW
+
+        p = torch.bmm(mi, qi)  # b, THW, HW
+        p = p / math.sqrt(D_e)
+        p = F.softmax(p, dim=1)  # b, THW, HW
+
+        mo = m_out.view(B, D_o, T * H * W)
+        mem = torch.bmm(mo, p)  # Weighted-sum B, D_o, HW
+        mem = mem.view(B, D_o, H, W)
+
+        mem_out = torch.cat([mem, q_out], dim=1)
+
+        return mem_out, p
+
+
+class KeyValue(nn.Module):
+    # Not using location
+    def __init__(self, indim, keydim, valdim):
+        super(KeyValue, self).__init__()
+        self.Key = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
+        self.Value = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
+
+    def forward(self, x):
+        return self.Key(x), self.Value(x)
+
+
+class STM(nn.Module):
+    def __init__(self, backbone='resnet50'):
+        super(STM, self).__init__()
+        self.backbone = backbone
+        assert backbone == 'resnet50' or backbone == 'resnet18' or backbone == 'resnest101'
+        scale_rate = (1 if (backbone == 'resnet50' or backbone == 'resnest101') else 4)
+
+        self.Encoder_M = Encoder_M(backbone)
+        self.Encoder_Q = Encoder_Q(backbone)
+
+        self.KV_M_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
+        self.KV_Q_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
+
+        self.Memory = Memory()
+        self.Decoder = Decoder(256, scale_rate, backbone)
+
+    def Pad_memory(self, mems, num_objects, K):
+        pad_mems = []
+        for mem in mems:
+            pad_mem = ToCuda(torch.zeros(1, K, mem.size()[1], 1, mem.size()[2], mem.size()[3]))
+            pad_mem[0, 1:num_objects + 1, :, 0] = mem
+            pad_mems.append(pad_mem)
+        return pad_mems
+
+    def memorize(self, frame, masks, num_objects):
+        # memorize a frame
+        num_objects = num_objects[0].item()
+        _, K, H, W = masks.shape  # B = 1
+
+        (frame, masks), pad = pad_divide_by([frame, masks], 16, (frame.size()[2], frame.size()[3]))
+
+        # make batch arg list
+        B_list = {'f': [], 'm': [], 'o': []}
+        for o in range(1, num_objects + 1):  # 1 - no
+            B_list['f'].append(frame)
+            B_list['m'].append(masks[:, o])
+            B_list['o'].append((torch.sum(masks[:, 1:o], dim=1) + \
+                                torch.sum(masks[:, o + 1:num_objects + 1], dim=1)).clamp(0, 1))
+
+        # make Batch
+        B_ = {}
+        for arg in B_list.keys():
+            B_[arg] = torch.cat(B_list[arg], dim=0)
+
+        r4, _, _, _, _ = self.Encoder_M(B_['f'], B_['m'], B_['o'])
+        k4, v4 = self.KV_M_r4(r4)  # num_objects, 128 and 512, H/16, W/16
+        k4, v4 = self.Pad_memory([k4, v4], num_objects=num_objects, K=K)
+        return k4, v4
+
+    def Soft_aggregation(self, ps, K):
+        num_objects, H, W = ps.shape
+        em = ToCuda(torch.zeros(1, K, H, W))
+        em[0, 0] = torch.prod(1 - ps, dim=0)  # bg prob
+        em[0, 1:num_objects + 1] = ps  # obj prob
+        em = torch.clamp(em, 1e-7, 1 - 1e-7)
+        logit = torch.log((em / (1 - em)))
+        return logit
+
+    def segment(self, frame, keys, values, num_objects):
+        num_objects = num_objects[0].item()
+        _, K, keydim, T, H, W = keys.shape  # B = 1
+        # pad
+        [frame], pad = pad_divide_by([frame], 16, (frame.size()[2], frame.size()[3]))
+
+        r4, r3, r2, _, _ = self.Encoder_Q(frame)
+        k4, v4 = self.KV_Q_r4(r4)  # 1, dim, H/16, W/16
+
+        # expand to ---  no, c, h, w
+        k4e, v4e = k4.expand(num_objects, -1, -1, -1), v4.expand(num_objects, -1, -1, -1)
+        r3e, r2e = r3.expand(num_objects, -1, -1, -1), r2.expand(num_objects, -1, -1, -1)
+
+        # memory select kv:(1, K, C, T, H, W)
+        m4, viz = self.Memory(keys[0, 1:num_objects + 1], values[0, 1:num_objects + 1], k4e, v4e)
+        logits = self.Decoder(m4, r3e, r2e)
+        ps = F.softmax(logits, dim=1)[:, 1]  # no, h, w
+        # ps = indipendant possibility to belong to each object
+
+        logit = self.Soft_aggregation(ps, K)  # 1, K, H, W
+
+        if pad[2] + pad[3] > 0:
+            logit = logit[:, :, pad[2]:-pad[3], :]
+        if pad[0] + pad[1] > 0:
+            logit = logit[:, :, :, pad[0]:-pad[1]]
+
+        return logit
+
+    def forward(self, *args, **kwargs):
+        if args[1].dim() > 4:  # keys
+            return self.segment(*args, **kwargs)
+        else:
+            return self.memorize(*args, **kwargs)

project/hair_service_sd/core/faceseg/tma.py

@@ -0,0 +1,185 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+import numpy as np
+
+class SequenceConv(nn.ModuleList):
+    """Sequence conv module.
+
+    Args:
+        in_channels (int): input tensor channel.
+        out_channels (int): output tensor channel.
+        kernel_size (int): convolution kernel size.
+        sequence_num (int): sequence length.
+        conv_cfg (dict): convolution config dictionary.
+        norm_cfg (dict): normalization config dictionary.
+        act_cfg (dict): activation config dictionary.
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, sequence_num):
+        super(SequenceConv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.sequence_num = sequence_num
+        for _ in range(sequence_num):
+            self.append(
+                nn.Sequential(
+                    nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, 1, self.kernel_size // 2, bias=False),
+                    nn.BatchNorm2d(self.out_channels),
+                    nn.ReLU()
+                )
+            )
+
+    def forward(self, sequence_imgs):
+        """
+
+        Args:
+            sequence_imgs (Tensor): TxBxCxHxW
+
+        Returns:
+            sequence conv output: TxBxCxHxW
+        """
+        sequence_outs = []
+        assert sequence_imgs.shape[0] == self.sequence_num
+        for i, sequence_conv in enumerate(self):
+            sequence_out = sequence_conv(sequence_imgs[i, ...])
+            sequence_out = sequence_out.unsqueeze(0)
+            sequence_outs.append(sequence_out)
+
+        sequence_outs = torch.cat(sequence_outs, dim=0)  # TxBxCxHxW
+        return sequence_outs
+
+class MemoryModule(nn.Module):
+    """Memory read module.
+    Args:
+
+    """
+
+    def __init__(self,
+                 matmul_norm=False):
+        super(MemoryModule, self).__init__()
+        self.matmul_norm = matmul_norm
+
+    def forward(self, memory_keys, memory_values, query_key, query_value):
+        """
+        Memory Module forward.
+        Args:
+            memory_keys (Tensor): memory keys tensor, shape: TxBxCxHxW
+            memory_values (Tensor): memory values tensor, shape: TxBxCxHxW
+            query_key (Tensor): query keys tensor, shape: BxCxHxW
+            query_value (Tensor): query values tensor, shape: BxCxHxW
+
+        Returns:
+            Concat query and memory tensor.
+        """
+        sequence_num, batch_size, key_channels, height, width = memory_keys.shape
+        _, _, value_channels, _, _ = memory_values.shape
+        assert query_key.shape[1] == key_channels and query_value.shape[1] == value_channels
+        memory_keys = memory_keys.permute(1, 2, 0, 3, 4).contiguous()  # BxCxTxHxW
+        memory_keys = memory_keys.view(batch_size, key_channels, sequence_num * height * width)  # BxCxT*H*W
+
+        query_key = query_key.view(batch_size, key_channels, height * width).permute(0, 2, 1).contiguous()  # BxH*WxCk
+        key_attention = torch.bmm(query_key, memory_keys)  # BxH*WxT*H*W
+        if self.matmul_norm:
+            key_attention = (key_channels ** -.5) * key_attention
+        key_attention = F.softmax(key_attention, dim=-1)  # BxH*WxT*H*W
+
+        memory_values = memory_values.permute(1, 2, 0, 3, 4).contiguous()  # BxCxTxHxW
+        memory_values = memory_values.view(batch_size, value_channels, sequence_num * height * width)
+        memory_values = memory_values.permute(0, 2, 1).contiguous()  # BxT*H*WxC
+        memory = torch.bmm(key_attention, memory_values)  # BxH*WxC
+        memory = memory.permute(0, 2, 1).contiguous()  # BxCxH*W
+        memory = memory.view(batch_size, value_channels, height, width)  # BxCxHxW
+
+        query_memory = torch.cat([query_value, memory], dim=1)
+        return query_memory
+#
+# class TMAHead(nn.Module):
+#     """TMAHead decoder for video semantic segmentation."""
+#
+#     def __init__(self, sequence_num, key_channels, value_channels, num_classes=2, dropout_ratio=0):
+#         super(TMAHead, self).__init__()
+#
+#         self.sequence_num = sequence_num
+#         self.memory_key_conv = nn.Sequential(
+#             SequenceConv(self.in_channels, key_channels, 1, sequence_num),
+#             SequenceConv(key_channels, key_channels, 3, sequence_num)
+#         )
+#         self.memory_value_conv = nn.Sequential(
+#             SequenceConv(self.in_channels, value_channels, 1, sequence_num),
+#             SequenceConv(value_channels, value_channels, 3, sequence_num)
+#         )
+#         self.query_key_conv = nn.Sequential(
+#             nn.Sequential(
+#                 nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
+#                 nn.BatchNorm2d(key_channels),
+#                 nn.ReLU()
+#             ),
+#             nn.Sequential(
+#                 nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
+#                 nn.BatchNorm2d(key_channels),
+#                 nn.ReLU()
+#             ),
+#         )
+#
+#         self.query_value_conv = nn.Sequential(
+#             nn.Sequential(
+#                 nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
+#                 nn.BatchNorm2d(value_channels),
+#                 nn.ReLU()
+#             ),
+#             nn.Sequential(
+#                 nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
+#                 nn.BatchNorm2d(value_channels),
+#                 nn.ReLU()
+#             ),
+#         )
+#         self.memory_module = MemoryModule(matmul_norm=False)
+#         self.bottleneck = nn.Sequential(
+#                 nn.Conv2d(value_channels * 2, self.channels, 3, 1, 1, bias=False),
+#                 nn.BatchNorm2d(value_channels),
+#                 nn.ReLU()
+#             )
+#
+#         self.conv_seg = nn.Conv2d(self.channels, num_classes, kernel_size=1)
+#         if dropout_ratio > 0:
+#             self.dropout = nn.Dropout2d(dropout_ratio)
+#         else:
+#             self.dropout = None
+#
+#     def cls_seg(self, feat):
+#         """Classify each pixel."""
+#         if self.dropout is not None:
+#             feat = self.dropout(feat)
+#         output = self.conv_seg(feat)
+#         return output
+#
+#     def forward(self, inputs, sequence_imgs):
+#         """
+#         Forward fuction.
+#         Args:
+#             inputs (list[Tensor]): backbone multi-level outputs.
+#             sequence_imgs (list[Tensor]): len(sequence_imgs) is equal to batch_size,
+#                 each element is a Tensor with shape of TxCxHxW.
+#
+#         Returns:
+#             decoder logits.
+#         """
+#         x = inputs
+#         sequence_imgs = [y.unsqueeze(0) for y in sequence_imgs]  # T, BxCxHxW
+#         sequence_imgs = torch.cat(sequence_imgs, dim=0)  # TxBxCxHxW
+#         sequence_num, batch_size, channels, height, width = sequence_imgs.shape
+#
+#         assert sequence_num == self.sequence_num
+#         memory_keys = self.memory_key_conv(sequence_imgs)
+#         memory_values = self.memory_value_conv(sequence_imgs)
+#         query_key = self.query_key_conv(x)  # BxCxHxW
+#         query_value = self.query_value_conv(x)  # BxCxHxW
+#
+#         # memory read
+#         output = self.memory_module(memory_keys, memory_values, query_key, query_value)
+#         output = self.bottleneck(output)
+#         output = self.cls_seg(output)
+#
+#         return output

project/hair_service_sd/core/faceseg/u2net.py

@@ -0,0 +1,624 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):#UNet07DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-6 ###
+class RSU6(nn.Module):#UNet06DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):#UNet05DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):#UNet04DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):#UNet04FRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+from core.faceseg.tma import SequenceConv, MemoryModule
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NET, self).__init__()
+
+        self.stage1 = RSU7(in_ch,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+        self.in_channels = 512
+        key_channels = 128
+        value_channels = 512
+        self.sequence_num = sequence_num = 2
+        self.memory_key_conv = nn.Sequential(
+            SequenceConv(self.in_channels, key_channels, 1, sequence_num),
+            SequenceConv(key_channels, key_channels, 3, sequence_num)
+        )
+        self.memory_value_conv = nn.Sequential(
+            SequenceConv(self.in_channels, value_channels, 1, sequence_num),
+            SequenceConv(value_channels, value_channels, 3, sequence_num)
+        )
+        self.query_key_conv = nn.Sequential(
+            nn.Sequential(
+                nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(key_channels),
+                nn.ReLU()
+            ),
+            nn.Sequential(
+                nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
+                nn.BatchNorm2d(key_channels),
+                nn.ReLU()
+            ),
+        )
+        self.query_value_conv = nn.Sequential(
+            nn.Sequential(
+                nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(value_channels),
+                nn.ReLU()
+            ),
+            nn.Sequential(
+                nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
+                nn.BatchNorm2d(value_channels),
+                nn.ReLU()
+            ),
+        )
+        self.memory_module = MemoryModule(matmul_norm=False)
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(value_channels * 2, self.in_channels, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(value_channels),
+            nn.ReLU()
+        )
+
+        self.is_train = True
+
+    def extract_feature(self, x):
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+
+        return hx1, hx2, hx3, hx4, hx5, hx6
+
+    def decoder(self, hx1, hx2, hx3, hx4, hx5, hx6):
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d, hx2d, hx3d, hx4d, hx5d
+
+    def side_output(self, hx1d, hx2d, hx3d, hx4d, hx5d, hx6):
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
+
+    def forward(self, x, memory_sequence=None):
+        if self.is_train:
+            hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
+
+            if memory_sequence is None:
+                memory_hx6 = [hx6 for _ in range(self.sequence_num)]
+            else:
+                memory_hx6 = []
+                for single_memory in memory_sequence:
+                    _, _, _, _, _, hx6 = self.extract_feature(single_memory)
+                    memory_hx6.append(hx6)
+
+            memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6]  # T, BxCxHxW
+            memory_hx6 = torch.cat(memory_hx6, dim=0)
+            memory_keys = self.memory_key_conv(memory_hx6)
+            memory_values = self.memory_value_conv(memory_hx6)
+            query_key = self.query_key_conv(hx6)
+            query_value = self.query_value_conv(hx6)
+            merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
+            merge_hx6 = self.bottleneck(merge_hx6)
+
+            hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
+
+            return self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
+        else:
+            return self.test(x)
+
+    def test(self, x):
+        with torch.no_grad():
+            hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
+
+            memory_hx6 = [hx6 for _ in range(self.sequence_num)]
+
+            memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6]  # T, BxCxHxW
+            memory_hx6 = torch.cat(memory_hx6, dim=0)
+            memory_keys = self.memory_key_conv(memory_hx6)
+            memory_values = self.memory_value_conv(memory_hx6)
+            query_key = self.query_key_conv(hx6)
+            query_value = self.query_value_conv(hx6)
+            merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
+            merge_hx6 = self.bottleneck(merge_hx6)
+
+            hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
+
+            mask, _, _, _, _, _, _ = self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
+
+            return mask
+#
+# ### U^2-Net small ###
+# class U2NETP(nn.Module):
+#
+#     def __init__(self,in_ch=3,out_ch=1):
+#         super(U2NETP,self).__init__()
+#
+#         self.stage1 = RSU7(in_ch,16,64)
+#         self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+#
+#         self.stage2 = RSU6(64,16,64)
+#         self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+#
+#         self.stage3 = RSU5(64,16,64)
+#         self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+#
+#         self.stage4 = RSU4(64,16,64)
+#         self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+#
+#         self.stage5 = RSU4F(64,16,64)
+#         self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+#
+#         self.stage6 = RSU4F(64,16,64)
+#
+#         # decoder
+#         self.stage5d = RSU4F(128,16,64)
+#         self.stage4d = RSU4(128,16,64)
+#         self.stage3d = RSU5(128,16,64)
+#         self.stage2d = RSU6(128,16,64)
+#         self.stage1d = RSU7(128,16,64)
+#
+#         self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+#         self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+#         self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
+#         self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
+#         self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
+#         self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
+#
+#         self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+#
+#     def forward(self,x):
+#
+#         hx = x
+#
+#         #stage 1
+#         hx1 = self.stage1(hx)
+#         hx = self.pool12(hx1)
+#
+#         #stage 2
+#         hx2 = self.stage2(hx)
+#         hx = self.pool23(hx2)
+#
+#         #stage 3
+#         hx3 = self.stage3(hx)
+#         hx = self.pool34(hx3)
+#
+#         #stage 4
+#         hx4 = self.stage4(hx)
+#         hx = self.pool45(hx4)
+#
+#         #stage 5
+#         hx5 = self.stage5(hx)
+#         hx = self.pool56(hx5)
+#
+#         #stage 6
+#         hx6 = self.stage6(hx)
+#         hx6up = _upsample_like(hx6,hx5)
+#
+#         #decoder
+#         hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+#         hx5dup = _upsample_like(hx5d,hx4)
+#
+#         hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+#         hx4dup = _upsample_like(hx4d,hx3)
+#
+#         hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+#         hx3dup = _upsample_like(hx3d,hx2)
+#
+#         hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+#         hx2dup = _upsample_like(hx2d,hx1)
+#
+#         hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+#
+#
+#         #side output
+#         d1 = self.side1(hx1d)
+#
+#         d2 = self.side2(hx2d)
+#         d2 = _upsample_like(d2,d1)
+#
+#         d3 = self.side3(hx3d)
+#         d3 = _upsample_like(d3,d1)
+#
+#         d4 = self.side4(hx4d)
+#         d4 = _upsample_like(d4,d1)
+#
+#         d5 = self.side5(hx5d)
+#         d5 = _upsample_like(d5,d1)
+#
+#         d6 = self.side6(hx6)
+#         d6 = _upsample_like(d6,d1)
+#
+#         d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+#
+#         return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)

project/hair_service_sd/core/hairstyle_local.py

@@ -0,0 +1,210 @@
+import cv2
+import numpy as np
+import torch
+import json
+import os
+import os.path as osp
+from common.logger import LogFactory
+from core.process_modules import Get_Landmark, Process_Data, localtranslationwarpfastwithstrength,\
+    Generator_Hair, chinClass, Generator_Fusion_Res, Change_Hair_Color, GenderClassifyProcessor, BodySeg, \
+    localtranslationwarpfastwithstrength_v2, localtranslationwarpfastwithstrength_v2_soft, updateEndPosition
+from core.face_enhance.face_enhancement import FaceEnhancement
+from datetime import datetime
+import random
+from core.utils import landmark_processor
+from core.cos_module import COS_object as OSS_object
+from core.faceseg.face_seg import FaceSeg
+from common.logger import config
+
+class Prepare_Ref_HairColor_Data(object):
+    def __init__(self, gpu, device_id):
+        if gpu and torch.cuda.is_available():
+            self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
+            self.cuda = True
+        else:
+            self.device = torch.device("cpu")
+            self.cuda = False
+
+        self.process_data = Process_Data(gpu, device_id)
+        self.get_landmark = Get_Landmark(gpu_id=device_id)
+        self.color_output_size = 768
+
+    def get_prepare_ref_haircolor_768_data(self, ref_rgb_8uc3_orisize):
+
+        ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
+
+        ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
+
+        # 光头分割
+        ref_baldseg_8uc3_orisize = self.process_data.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matte_pred_8uc1_orisize,
+                                                                  ref_landmark_1k2_f_orisize)
+
+        color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(ref_landmark_1k2_f_orisize, self.color_output_size, ratio=0.35, h_offset=0.45)
+
+        ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
+
+        # show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_pred_8uc3_orisize, ref_baldseg_8uc3_orisize), axis=1)
+        # show_concat = cv2.resize(show_concat, (0, 0), fx=0.5, fy=0.5)
+        # cv2.imshow("show_concat", show_concat)
+        # cv2.waitKey()
+
+        ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
+        ref_matte_pred_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
+        ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size), flags=cv2.INTER_NEAREST)
+        ref_landmark_1k2_f_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, color_hair_M)
+
+        return ref_rgb_8uc3_768, ref_matte_pred_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_1k2_f_768
+
+class Prepare_Ref_HairStyle_Data(object):
+    def __init__(self, gpu, device_id):
+
+        if gpu and torch.cuda.is_available():
+            self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
+            self.cuda = True
+        else:
+            self.device = torch.device("cpu")
+            self.cuda = False
+
+        self.process_data = Process_Data(gpu, device_id)
+        self.get_landmark = Get_Landmark(gpu_id=device_id)
+        self.gender_classify = GenderClassifyProcessor(gpu_id=device_id)
+
+    def get_prepare_ref_768_color_data(self, ref_rgb_8uc3_orisize):
+        ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize)
+        ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
+            self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
+
+        return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
+
+    def get_prepare_ref_768_color_data_landmark1k(self, ref_rgb_8uc3_orisize,ref_landmark_1k2_f_orisize):
+        # ref_landmark_1k2_f_orisize = self.get_landmark.inference(ref_rgb_8uc3_orisize)
+        ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
+            self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
+
+        return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
+
+    def calculate_hair_ratio_after_align(self, hair_mask, origin_landmark1k, img_size=768):
+        image_to_face_mat = landmark_processor.get_transform_mat_hair_ratio_v1(origin_landmark1k, 768, ratio=0.35, h_offset=0.32)
+        hair_mask_align = cv2.warpAffine(hair_mask, image_to_face_mat, (img_size, img_size))
+
+        hair_rect = cv2.boundingRect(hair_mask_align[:, :, :1])
+        hair_mask_ratio = hair_rect[2] * hair_rect[3] / (img_size * img_size)
+        return hair_mask_ratio
+
+    def check_female_hair_ratio(self, origin_img_8uc3, landmark_1k2_f_orisize):
+        ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img_8uc3, landmark_1k2_f_orisize)
+        ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
+        hairstyle_M = self.process_data.get_hair_M_girl_v1(landmark_1k2_f_orisize)
+        ref_rgb_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, hairstyle_M, (768, 768))
+
+        # cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
+        # cv2.waitKey()
+
+        edge_width = 5
+        if (ref_rgb_8uc3_768[-edge_width:, :, :]).max() > 0 or (ref_rgb_8uc3_768[:, -edge_width:, :]).max() > 0 or (ref_rgb_8uc3_768[:, :edge_width, :]).max() > 0:
+            return False
+        else:
+            return True
+
+    def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize):
+
+        ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
+        # for i in range(1000):
+        #     cv2.circle(ref_rgb_8uc3_orisize, (int(ref_landmark_1k2_f_orisize[i][0]),int(ref_landmark_1k2_f_orisize[i][1])), 1,(255, 0,0), 1)
+        # cv2.imshow('ffff', ref_rgb_8uc3_orisize)
+        # cv2.waitKey()
+        # ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize = self.process_data.generator_matte.matte_inference(
+        #     ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
+        ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize)
+        gender_res = self.gender_classify.forward(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
+        # hair_ratio = self.calculate_hair_ratio_after_align(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
+        # if hair_ratio > 0.3:
+        #     ratio = 2
+        # else:
+        #     if gender_res:
+        #         ratio = 1
+        #     else:
+        #         ratio = 0
+        if not gender_res:
+            ratio = 0
+        else:
+            check_res = self.check_female_hair_ratio(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
+            if check_res:
+                ratio = 1
+            else:
+                ratio = 2
+
+        if gender_res:
+            gender = "girl"
+        else:
+            gender = "boy"
+
+        # show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_fg_8uc3_orisize), axis=1)
+        # resize_ratio = 1024. / max(show_concat.shape[:2])
+        # show_concat = cv2.resize(show_concat, (0, 0), fx=resize_ratio, fy=resize_ratio)
+        # print("gender_res: ", gender_res, " hair_ratio: ", hair_ratio)
+        # cv2.imshow("show_concat", show_concat)
+        # cv2.imshow("ref_matte_pred_8uc1_orisize", ref_matte_pred_8uc1_orisize)
+        # cv2.waitKey()
+
+        ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \
+            self.process_data.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio)
+
+        # cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
+        # cv2.waitKey()
+
+        return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio
+
+    def Generator_reftensor(self, ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768,
+                                 ref_landmark_f1k2_768):
+
+        """
+        input：
+
+        图像尺寸基于： 人脸 512， 图像均为3通道
+        ref_rgb_8uc3_512: 参考图， size 512， uint8 （0-255）
+        ref_matting_8uc3_512：参考图 matting alpha 值， uint8 （0-255）
+        ref_baldseg_8uc3_512: 参考图 光头分割， uint8 （0-255）
+        ref_landmark_f1k2_512： 参考图 关键点 1k*2 float32
+
+        output：
+
+        input_another_pose_hair_image: 参考图 条件图， float32 （0-255）
+
+        """
+
+        # 8 ********************** another pose hair_image **********************
+        another_pose_image = ref_rgb_8uc3_768.copy()
+
+        another_nohair_pose_mask = ref_baldseg_8uc3_768.copy()
+
+        # cv2.imshow("another_nohair_pose_mask", another_nohair_pose_mask)
+        # cv2.waitKey()
+
+        another_nohair_pose_mask[(np.abs(another_nohair_pose_mask - [0, 0, 255]) < 50).all(axis=2)] = [0, 255, 255]
+        another_nohair_pose_mask[(another_nohair_pose_mask == [0, 0, 0]).all(axis=2)] = [255, 255, 255]
+        another_pose_pts137 = landmark_processor.pts_1k_to_137(ref_landmark_f1k2_768).astype(np.int32)
+        cv2.fillPoly(another_nohair_pose_mask,
+                     np.concatenate((another_pose_pts137[47:35:-1], another_pose_pts137[56:64],
+                                     [another_pose_pts137[48], another_pose_pts137[22]]))[
+                     np.newaxis, :, :], (255, 255, 0))
+
+        cv2.fillPoly(another_nohair_pose_mask,
+                     np.concatenate((another_pose_pts137[22:37], another_pose_pts137[56:47:-1]))[np.newaxis, :, :],
+                     (255, 0, 128))
+
+        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[88:104][np.newaxis, :, :], (255, 0, 255))  # eye
+        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[105:121][np.newaxis, :, :], (255, 0, 255))  # eye
+        # Label nose
+        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[64:79][np.newaxis, :, :], (255, 255, 255))
+        # Label eyebrow
+        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[121:129][np.newaxis, :, :], (255, 128, 0))
+        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[129:137][np.newaxis, :, :], (255, 128, 0))
+
+        another_pose_hair_alpha = ref_matting_8uc3_768.copy() / 255.
+        another_pose_hair_image = (another_pose_image * another_pose_hair_alpha + another_nohair_pose_mask * (
+                    1 - another_pose_hair_alpha)).astype(np.uint8)
+
+        input_another_pose_hair_image = another_pose_hair_image.astype(np.float32) / 255
+
+        return input_another_pose_hair_image

project/hair_service_sd/core/matting/gca_matting_hair_single.py

@@ -0,0 +1,208 @@
+import os
+import cv2
+import argparse
+import numpy as np
+
+import torch
+from torch.nn import functional as F
+
+import utils
+from core.matting import networks
+from utils.data_preprocess import *
+from time import time
+
+def single_inference(model, image_dict, device, return_offset=False):
+
+    with torch.no_grad():
+        image, trimap = image_dict['image'], image_dict['trimap']
+        alpha_shape = image_dict['alpha_shape']
+        image = image.to(device)
+        trimap = trimap.to(device)
+
+        matte_start = time()
+        alpha_pred, info_dict = model(image, trimap)
+        matte_end = time()
+        # print("matte time cost : {:.4f}".format(matte_end-matte_start))
+
+        trimap_argmax = trimap.argmax(dim=1, keepdim=True)
+
+        alpha_pred[trimap_argmax == 2] = 1
+        alpha_pred[trimap_argmax == 0] = 0
+
+        h, w = alpha_shape
+        test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
+        test_pred = test_pred.astype(np.uint8)
+        test_pred = test_pred[32:h+32, 32:w+32]
+
+        if return_offset:
+            short_side = h if h < w else w
+            ratio = 512 / short_side
+            offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
+            # write softmax_scale to offset image
+            scale = info_dict['offset_1'][1].cpu()
+            offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
+            text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
+            offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
+
+            offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
+            # write softmax_scale to offset image
+            scale = info_dict['offset_2'][1].cpu()
+            offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
+            text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
+            offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
+
+            return test_pred, (offset_1, offset_2)
+        else:
+            return test_pred, None
+
+def generator_tensor_dict(image, trimap):
+
+    sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
+
+    # reshape
+    h, w = sample["alpha_shape"]
+
+    if h % 32 == 0 and w % 32 == 0:
+        padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
+        padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
+        sample['image'] = padded_image
+        sample['trimap'] = padded_trimap
+    else:
+        target_h = 32 * ((h - 1) // 32 + 1)
+        target_w = 32 * ((w - 1) // 32 + 1)
+        pad_h = target_h - h
+        pad_w = target_w - w
+        padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
+        padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
+        sample['image'] = padded_image
+        sample['trimap'] = padded_trimap
+
+    # ImageNet mean & std
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+    # convert GBR images to RGB
+    image, trimap = sample['image'][:, :, ::-1], sample['trimap']
+    # swap color axis
+    image = image.transpose((2, 0, 1)).astype(np.float32)
+    trimap[trimap < 85] = 0
+    trimap[trimap >= 170] = 2
+    trimap[trimap >= 85] = 1
+    # normalize image
+    image /= 255.
+
+    # to tensor
+    sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
+    sample['image'] = sample['image'].sub_(mean).div_(std)
+
+    # if CONFIG.model.trimap_channel == 3:
+    sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
+    # elif CONFIG.model.trimap_channel == 1:
+    #     sample['trimap'] = sample['trimap'][None, ...].float()
+    # else:
+    #     raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1")
+
+    # add first channel
+    sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
+
+    return sample
+
+if __name__ == '__main__':
+    print('Torch Version: ', torch.__version__)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
+                        help="path of checkpoint")
+    parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
+    parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
+    parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
+
+    # Parse configuration
+    args = parser.parse_args()
+
+    # # Check if toml config file is loaded
+    # if CONFIG.is_default:
+    #     raise ValueError("No .toml config loaded.")
+
+    args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
+    utils.make_dir(args.output)
+
+    # build model
+    model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=1)
+    model.cuda()
+    print("model: ", model)
+
+    # load checkpoint
+    checkpoint = torch.load(args.checkpoint)
+    model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
+
+    # inference
+    model = model.eval()
+
+    for image_name in os.listdir(args.image_dir):
+
+        if not is_image_file(image_name):
+            continue
+        # assume image and trimap have the same file name
+        img_basename, ext = os.path.splitext(image_name)
+
+        image_path = os.path.join(args.image_dir, image_name)
+        trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
+        # trimap_path = os.path.join(args.trimap_dir, image_name)
+        print('Image: ', image_path, ' Tirmap: ', trimap_path)
+
+        # read images
+        img_basename, img_ext = os.path.splitext(image_name)
+        # img_pt_path = image_path.replace(img_ext, "_landmark1k.txt")
+        # img_landmark1k = np.loadtxt(img_pt_path)
+
+        image = cv2.imread(image_path)
+        trimap = cv2.imread(trimap_path, 0)
+
+        ori_h, ori_w, _ = image.shape
+        tri_h, tri_w = trimap.shape
+        if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
+            image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
+            # img_landmark1k = landmark_processor.resize_points(img_landmark1k, ori_w, ori_h,
+            #                                                       tri_w, tri_h)
+
+        # hair_mat = landmark_processor.get_transform_mat_hair(img_landmark1k, 640, ratio=0.3, w_ratio=0.5,
+        #                                                      h_ratio=0.4)
+        # hair_img_landmark = landmark_processor.transform_points(img_landmark1k, hair_mat)
+        #
+        # image = cv2.warpAffine(image, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
+        # trimap = cv2.warpAffine(trimap, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
+
+        if tri_h > 1920 or tri_w > 1920:
+            if tri_h > 1920:
+                new_tri_h = 1920
+                new_tri_w = int(tri_w * 1920 / tri_h)
+            else:
+                new_tri_w = 1920
+                new_tri_h = int(tri_h * 1920 / tri_w)
+
+            image_resize = cv2.resize(image, (new_tri_w, new_tri_h),  interpolation=cv2.INTER_CUBIC)
+            trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h),  interpolation=cv2.INTER_NEAREST)
+        else:
+            image_resize = image
+            trimap_resize = trimap
+
+        # image_dict = generator_tensor_dict(image, trimap)
+        image_dict = generator_tensor_dict(image_resize, trimap_resize)
+        pred, offset = single_inference(model, image_dict)
+
+        # torch.cuda.empty_cache()
+
+        pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
+        # pred = cv2.warpAffine(pred, cv2.invertAffineTransform(hair_mat), (tri_w, tri_h), flags=cv2.INTER_CUBIC)
+
+        # offset[0] = cv2.resize(offset[0], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+        # offset[1] = cv2.resize(offset[1], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+
+        # cv2.imshow("image_resize", image_resize)
+        # cv2.imshow("trimap_resize", trimap_resize)
+        # cv2.imshow("pred", pred)
+        # cv2.waitKey()
+        cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
+        # if offset is not None:
+        #     cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset1.png'), offset[0])
+        #     cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset2.png'), offset[1])

project/hair_service_sd/core/matting/gca_matting_hair_single_fg.py

@@ -0,0 +1,200 @@
+import os
+import cv2
+import argparse
+import numpy as np
+
+import torch
+from torch.nn import functional as F
+
+import utils
+from core.matting import networks
+from utils.data_preprocess import *
+from time import time
+
+def single_inference(model, image_dict, device, return_offset=False):
+
+    with torch.no_grad():
+        image, trimap = image_dict['image'], image_dict['trimap']
+        alpha_shape = image_dict['alpha_shape']
+        image = image.to(device)
+        trimap = trimap.to(device)
+
+        matte_start = time()
+
+        alpha_pred, info_dict = model(image, trimap)
+
+        fg_pred = alpha_pred[:, :-1, :, :]
+        alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1)
+        matte_end = time()
+        # print("matte time cost : {:.4f}".format(matte_end-matte_start))
+
+        trimap_argmax = trimap.argmax(dim=1, keepdim=True)
+
+        alpha_pred[trimap_argmax == 2] = 1
+        alpha_pred[trimap_argmax == 0] = 0
+
+        h, w = alpha_shape
+        test_fg_pred = fg_pred[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255
+        test_fg_pred = test_fg_pred.astype(np.uint8)
+        test_fg_pred = test_fg_pred[32:h+32, 32:w+32]
+
+        test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
+        test_pred = test_pred.astype(np.uint8)
+        test_pred = test_pred[32:h+32, 32:w+32]
+
+        # cv2.imshow('test_fg_pred', test_fg_pred)
+        # cv2.imshow('test_pred', test_pred)
+        # cv2.waitKey()
+
+        if return_offset:
+            short_side = h if h < w else w
+            ratio = 512 / short_side
+            offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
+            # write softmax_scale to offset image
+            scale = info_dict['offset_1'][1].cpu()
+            offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
+            text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
+            offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
+
+            offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
+            # write softmax_scale to offset image
+            scale = info_dict['offset_2'][1].cpu()
+            offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
+            text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
+            offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
+
+            return test_fg_pred, test_pred, (offset_1, offset_2)
+        else:
+            return test_fg_pred, test_pred, None
+
+def generator_tensor_dict(image, trimap):
+
+    sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
+
+    # reshape
+    h, w = sample["alpha_shape"]
+
+    if h % 32 == 0 and w % 32 == 0:
+        padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
+        padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
+        sample['image'] = padded_image
+        sample['trimap'] = padded_trimap
+    else:
+        target_h = 32 * ((h - 1) // 32 + 1)
+        target_w = 32 * ((w - 1) // 32 + 1)
+        pad_h = target_h - h
+        pad_w = target_w - w
+        padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
+        padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
+        sample['image'] = padded_image
+        sample['trimap'] = padded_trimap
+
+    # ImageNet mean & std
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+    # convert GBR images to RGB
+    image, trimap = sample['image'][:, :, ::-1], sample['trimap']
+    # swap color axis
+    image = image.transpose((2, 0, 1)).astype(np.float32)
+    trimap[trimap < 85] = 0
+    trimap[trimap >= 170] = 2
+    trimap[trimap >= 85] = 1
+    # normalize image
+    image /= 255.
+
+    # to tensor
+    sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
+    sample['image'] = sample['image'].sub_(mean).div_(std)
+
+    sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
+
+    # add first channel
+    sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
+
+    return sample
+
+if __name__ == '__main__':
+    print('Torch Version: ', torch.__version__)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
+                        help="path of checkpoint")
+    parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
+    parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
+    parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
+
+    # Parse configuration
+    args = parser.parse_args()
+
+    # # Check if toml config file is loaded
+    # if CONFIG.is_default:
+    #     raise ValueError("No .toml config loaded.")
+
+    args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
+    utils.make_dir(args.output)
+
+    # build model
+    model = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder)
+    model.cuda()
+    print("model: ", model)
+
+    # load checkpoint
+    checkpoint = torch.load(args.checkpoint)
+    model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
+
+    # inference
+    model = model.eval()
+    export_onnx_file = "test.onnx"
+    torch.onnx.export(model, x, export_onnx_file, opset_version=10, do_constant_folding=True, input_names=["image", "trimap"],  # 输入名
+                                                        output_names=["fg_pred", "alpha_pred", "None"])
+
+    for image_name in os.listdir(args.image_dir):
+
+        if not is_image_file(image_name):
+            continue
+        # assume image and trimap have the same file name
+        img_basename, ext = os.path.splitext(image_name)
+
+        image_path = os.path.join(args.image_dir, image_name)
+        trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
+        # trimap_path = os.path.join(args.trimap_dir, image_name)
+        print('Image: ', image_path, ' Tirmap: ', trimap_path)
+
+        # read images
+        img_basename, img_ext = os.path.splitext(image_name)
+
+        image = cv2.imread(image_path)
+        trimap = cv2.imread(trimap_path, 0)
+
+        ori_h, ori_w, _ = image.shape
+        tri_h, tri_w = trimap.shape
+        if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
+            image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
+
+
+        if tri_h > 1920 or tri_w > 1920:
+            if tri_h > 1920:
+                new_tri_h = 1920
+                new_tri_w = int(tri_w * 1920 / tri_h)
+            else:
+                new_tri_w = 1920
+                new_tri_h = int(tri_h * 1920 / tri_w)
+
+            image_resize = cv2.resize(image, (new_tri_w, new_tri_h),  interpolation=cv2.INTER_CUBIC)
+            trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h),  interpolation=cv2.INTER_NEAREST)
+        else:
+            image_resize = image
+            trimap_resize = trimap
+
+        # image_dict = generator_tensor_dict(image, trimap)
+        image_dict = generator_tensor_dict(image_resize, trimap_resize)
+        fg_pred, pred, offset = single_inference(model, image_dict)
+
+        # torch.cuda.empty_cache()
+
+        fg_pred = cv2.resize(fg_pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
+        pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
+
+        cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
+        cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_fg_p.png")), fg_pred)
+

project/hair_service_sd/core/matting/networks/__init__.py

@@ -0,0 +1 @@
+from .generators import *

project/hair_service_sd/core/matting/networks/decoders/__init__.py

@@ -0,0 +1,28 @@
+from .resnet_dec import ResNet_D_Dec, BasicBlock
+from .res_shortcut_dec import ResShortCut_D_Dec
+from .res_gca_dec import ResGuidedCxtAtten_Dec
+
+
+__all__ = ['res_shortcut_decoder_22', 'res_gca_decoder_22']
+
+
+def _res_shortcut_D_dec(block, layers, **kwargs):
+    model = ResShortCut_D_Dec(block, layers, **kwargs)
+    return model
+
+
+def _res_gca_D_dec(block, layers, num_class, **kwargs):
+    model = ResGuidedCxtAtten_Dec(block, layers, num_class, **kwargs)
+    return model
+
+
+def res_shortcut_decoder_22(**kwargs):
+    """Constructs a resnet_encoder_14 model.
+    """
+    return _res_shortcut_D_dec(BasicBlock, [2, 3, 3, 2], **kwargs)
+
+
+def res_gca_decoder_22(num_class=1, **kwargs):
+    """Constructs a resnet_encoder_14 model.
+    """
+    return _res_gca_D_dec(BasicBlock, [2, 3, 3, 2], num_class,  **kwargs)

project/hair_service_sd/core/matting/networks/decoders/res_gca_dec.py

@@ -0,0 +1,28 @@
+from core.matting.networks.ops import GuidedCxtAtten
+from core.matting.networks.decoders.res_shortcut_dec import ResShortCut_D_Dec
+
+
+class ResGuidedCxtAtten_Dec(ResShortCut_D_Dec):
+
+    def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False):
+        super(ResGuidedCxtAtten_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel)
+        self.gca = GuidedCxtAtten(128, 128)
+        self.num_class = num_class
+
+    def forward(self, x, mid_fea):
+        fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
+        im = mid_fea['image_fea']
+        x = self.layer1(x) + fea5  # N x 256 x 32 x 32
+        x = self.layer2(x) + fea4  # N x 128 x 64 x 64
+        x, offset = self.gca(im, x, mid_fea['unknown'])  # contextual attention
+        x = self.layer3(x) + fea3  # N x 64 x 128 x 128
+        x = self.layer4(x) + fea2  # N x 32 x 256 x 256
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.leaky_relu(x) + fea1
+        x = self.conv2(x)
+
+        alpha = (self.tanh(x) + 1.0) / 2.0
+
+        return alpha, {'offset_1': mid_fea['offset_1'], 'offset_2': offset}
+

project/hair_service_sd/core/matting/networks/decoders/res_shortcut_dec.py

@@ -0,0 +1,24 @@
+from core.matting.networks.decoders.resnet_dec import ResNet_D_Dec
+
+
+class ResShortCut_D_Dec(ResNet_D_Dec):
+
+    def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
+        super(ResShortCut_D_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel,
+                                                late_downsample=late_downsample)
+
+    def forward(self, x, mid_fea):
+        fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
+        x = self.layer1(x) + fea5
+        x = self.layer2(x) + fea4
+        x = self.layer3(x) + fea3
+        x = self.layer4(x) + fea2
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.leaky_relu(x) + fea1
+        x = self.conv2(x)
+
+        alpha = (self.tanh(x) + 1.0) / 2.0
+
+        return alpha, None
+

project/hair_service_sd/core/matting/networks/decoders/resnet_dec.py

@@ -0,0 +1,142 @@
+import logging
+import torch.nn as nn
+from core.matting.networks.ops import SpectralNorm
+
+def conv5x5(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """5x5 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
+                     padding=2, groups=groups, bias=False, dilation=dilation)
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, upsample=None, norm_layer=None, large_kernel=False):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.stride = stride
+        conv = conv5x5 if large_kernel else conv3x3
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        if self.stride > 1:
+            self.conv1 = SpectralNorm(nn.ConvTranspose2d(inplanes, inplanes, kernel_size=4, stride=2, padding=1, bias=False))
+        else:
+            self.conv1 = SpectralNorm(conv(inplanes, inplanes))
+        self.bn1 = norm_layer(inplanes)
+        self.activation = nn.LeakyReLU(0.2, inplace=True)
+        self.conv2 = SpectralNorm(conv(inplanes, planes))
+        self.bn2 = norm_layer(planes)
+        self.upsample = upsample
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.activation(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.upsample is not None:
+            identity = self.upsample(x)
+
+        out += identity
+        out = self.activation(out)
+
+        return out
+
+
+class ResNet_D_Dec(nn.Module):
+
+    def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
+        super(ResNet_D_Dec, self).__init__()
+        self.logger = logging.getLogger("Logger")
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+        self.large_kernel = large_kernel
+        self.kernel_size = 5 if self.large_kernel else 3
+
+        self.inplanes = 512 if layers[0] > 0 else 256
+        self.late_downsample = late_downsample
+        self.midplanes = 64 if late_downsample else 32
+
+        self.conv1 = SpectralNorm(nn.ConvTranspose2d(self.midplanes, 32, kernel_size=4, stride=2, padding=1, bias=False))
+        self.bn1 = norm_layer(32)
+        self.leaky_relu = nn.LeakyReLU(0.2, inplace=True)
+        self.conv2 = nn.Conv2d(32, num_class, kernel_size=self.kernel_size, stride=1, padding=self.kernel_size//2)
+        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
+        self.tanh = nn.Tanh()
+        self.layer1 = self._make_layer(block, 256, layers[0], stride=2)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, self.midplanes, layers[3], stride=2)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                if hasattr(m, "weight_bar"):
+                    nn.init.xavier_uniform_(m.weight_bar)
+                else:
+                    nn.init.xavier_uniform_(m.weight)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+        self.logger.debug(self)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        if blocks == 0:
+            return nn.Sequential(nn.Identity())
+        norm_layer = self._norm_layer
+        upsample = None
+        if stride != 1:
+            upsample = nn.Sequential(
+                nn.UpsamplingNearest2d(scale_factor=2),
+                SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
+                norm_layer(planes * block.expansion),
+            )
+        elif self.inplanes != planes * block.expansion:
+            upsample = nn.Sequential(
+                SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = [block(self.inplanes, planes, stride, upsample, norm_layer, self.large_kernel)]
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, norm_layer=norm_layer, large_kernel=self.large_kernel))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x, mid_fea):
+        x = self.layer1(x) # N x 256 x 32 x 32
+        x = self.layer2(x) # N x 128 x 64 x 64
+        x = self.layer3(x) # N x 64 x 128 x 128
+        x = self.layer4(x) # N x 32 x 256 x 256
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.leaky_relu(x)
+        x = self.conv2(x)
+
+        alpha = (self.tanh(x) + 1.0) / 2.0
+
+        return alpha, None

project/hair_service_sd/core/matting/networks/encoders/__init__.py

@@ -0,0 +1,39 @@
+import logging
+from .resnet_enc import ResNet_D, BasicBlock
+from .res_shortcut_enc import ResShortCut_D
+from .res_gca_enc import ResGuidedCxtAtten
+
+
+__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
+
+
+def _res_shortcut_D(block, layers, **kwargs):
+    model = ResShortCut_D(block, layers, **kwargs)
+    return model
+
+
+def _res_gca_D(block, layers, **kwargs):
+    model = ResGuidedCxtAtten(block, layers, **kwargs)
+    return model
+
+
+def resnet_gca_encoder_29(**kwargs):
+    """Constructs a resnet_encoder_29 model.
+    """
+    return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
+
+
+def res_shortcut_encoder_29(**kwargs):
+    """Constructs a resnet_encoder_25 model.
+    """
+    return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
+
+
+if __name__ == "__main__":
+    import torch
+    logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
+                        datefmt='%m-%d %H:%M:%S')
+    resnet_encoder = res_shortcut_encoder_29()
+    x = torch.randn(4,6,512,512)
+    z = resnet_encoder(x)
+    print(z[0].shape)

project/hair_service_sd/core/matting/networks/encoders/res_gca_enc.py

@@ -0,0 +1,97 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+# from utils import CONFIG
+from core.matting.networks.encoders.resnet_enc import ResNet_D
+from core.matting.networks.ops import GuidedCxtAtten, SpectralNorm
+
+class ResGuidedCxtAtten(ResNet_D):
+
+    def __init__(self, block, layers, norm_layer=None, late_downsample=False):
+        super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
+        first_inplane = 3 + 3
+        self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
+        self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
+
+        self.shortcut = nn.ModuleList()
+        for stage, inplane in enumerate(self.shortcut_inplane):
+            self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
+
+        self.guidance_head = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(16),
+            nn.ReflectionPad2d(1),
+            SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(32),
+            nn.ReflectionPad2d(1),
+            SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(128)
+        )
+
+        self.gca = GuidedCxtAtten(128, 128)
+
+        # initialize guidance head
+        for layers in range(len(self.guidance_head)):
+            m = self.guidance_head[layers]
+            if isinstance(m, nn.Conv2d):
+                if hasattr(m, "weight_bar"):
+                    nn.init.xavier_uniform_(m.weight_bar)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_shortcut(self, inplane, planes):
+        return nn.Sequential(
+            SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(planes),
+            SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(planes)
+        )
+
+    def forward(self, x):
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.activation(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        x1 = self.activation(out)  # N x 32 x 256 x 256
+        out = self.conv3(x1)
+        out = self.bn3(out)
+        out = self.activation(out)
+
+        im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
+        # if CONFIG.model.trimap_channel == 3:
+        unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
+        # else:
+        #     unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
+
+        x2 = self.layer1(out)  # N x 64 x 128 x 128
+        x3= self.layer2(x2)  # N x 128 x 64 x 64
+        x3, offset = self.gca(im_fea, x3, unknown)  # contextual attention
+        x4 = self.layer3(x3) # N x 256 x 32 x 32
+        out = self.layer_bottleneck(x4)  # N x 512 x 16 x 16
+
+        fea1 = self.shortcut[0](x)  # input image and trimap
+        fea2 = self.shortcut[1](x1)
+        fea3 = self.shortcut[2](x2)
+        fea4 = self.shortcut[3](x3)
+        fea5 = self.shortcut[4](x4)
+
+        return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
+                     'image_fea': im_fea,
+                     'unknown': unknown,
+                     'offset_1': offset}
+
+
+if __name__ == "__main__":
+    from core.matting.networks.encoders.resnet_enc import BasicBlock
+    m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
+    for m in m.modules():
+        print(m)

project/hair_service_sd/core/matting/networks/encoders/res_shortcut_enc.py

@@ -0,0 +1,51 @@
+import torch.nn as nn
+# from utils import CONFIG
+from core.matting.networks.encoders.resnet_enc import ResNet_D
+from core.matting.networks.ops import SpectralNorm
+
+class ResShortCut_D(ResNet_D):
+
+    def __init__(self, block, layers, norm_layer=None, late_downsample=False):
+        super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
+        first_inplane = 3 + 3
+        self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
+        self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
+
+        self.shortcut = nn.ModuleList()
+        for stage, inplane in enumerate(self.shortcut_inplane):
+            self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
+
+    def _make_shortcut(self, inplane, planes):
+        return nn.Sequential(
+            SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(planes),
+            SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
+            nn.ReLU(inplace=True),
+            self._norm_layer(planes)
+        )
+
+    def forward(self, x):
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.activation(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        x1 = self.activation(out) # N x 32 x 256 x 256
+        out = self.conv3(x1)
+        out = self.bn3(out)
+        out = self.activation(out)
+
+        x2 = self.layer1(out) # N x 64 x 128 x 128
+        x3= self.layer2(x2) # N x 128 x 64 x 64
+        x4 = self.layer3(x3) # N x 256 x 32 x 32
+        out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
+
+        fea1 = self.shortcut[0](x) # input image and trimap
+        fea2 = self.shortcut[1](x1)
+        fea3 = self.shortcut[2](x2)
+        fea4 = self.shortcut[3](x3)
+        fea5 = self.shortcut[4](x4)
+
+        return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}

project/hair_service_sd/core/matting/networks/encoders/resnet_enc.py

@@ -0,0 +1,150 @@
+import logging
+import torch.nn as nn
+from core.matting.networks.ops import SpectralNorm
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
+        self.bn1 = norm_layer(planes)
+        self.activation = nn.ReLU(inplace=True)
+        self.conv2 = SpectralNorm(conv3x3(planes, planes))
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.activation(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.activation(out)
+
+        return out
+
+
+class ResNet_D(nn.Module):
+    """
+    Implement and pre-train on ImageNet with the tricks from
+    https://arxiv.org/abs/1812.01187
+    without the mix-up part.
+    """
+
+    def __init__(self, block, layers, norm_layer=None, late_downsample=False):
+        super(ResNet_D, self).__init__()
+        self.logger = logging.getLogger("Logger")
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.late_downsample = late_downsample
+        self.midplanes = 64 if late_downsample else 32
+        self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
+        self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
+                                            stride=self.start_stride[0], padding=1, bias=False))
+        self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
+                                            bias=False))
+        self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
+                                            padding=1, bias=False))
+        self.bn1 = norm_layer(32)
+        self.bn2 = norm_layer(self.midplanes)
+        self.bn3 = norm_layer(self.inplanes)
+        self.activation = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight_bar)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+        self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
+        self.conv1.module.weight_bar.data[:,3:,:,:] = 0
+
+        self.logger.debug(self)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        if blocks == 0:
+            return nn.Sequential(nn.Identity())
+        norm_layer = self._norm_layer
+        downsample = None
+        if stride != 1:
+            downsample = nn.Sequential(
+                nn.AvgPool2d(2, stride),
+                SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
+                norm_layer(planes * block.expansion),
+            )
+        elif self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.activation(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x1 = self.activation(x) # N x 32 x 256 x 256
+        x = self.conv3(x1)
+        x = self.bn3(x)
+        x2 = self.activation(x) # N x 64 x 128 x 128
+
+        x3 = self.layer1(x2) # N x 64 x 128 x 128
+        x4 = self.layer2(x3) # N x 128 x 64 x 64
+        x5 = self.layer3(x4) # N x 256 x 32 x 32
+        x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
+
+        return x, (x1, x2, x3, x4, x5)
+
+
+if __name__ == "__main__":
+    m = ResNet_D(BasicBlock, [3, 4, 4, 2])
+    for m in m.modules():
+        print(m._get_name())

project/hair_service_sd/core/matting/networks/generators.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+
+# from utils import CONFIG
+from core.matting.networks import decoders, encoders
+
+
+class Generator(nn.Module):
+    def __init__(self, encoder, decoder, num_class=1):
+
+        super(Generator, self).__init__()
+
+        if encoder not in encoders.__all__:
+            raise NotImplementedError("Unknown Encoder {}".format(encoder))
+        self.encoder = encoders.__dict__[encoder]()
+
+        if decoder not in decoders.__all__:
+            raise NotImplementedError("Unknown Decoder {}".format(decoder))
+        self.decoder = decoders.__dict__[decoder](num_class)
+
+    def forward(self, image, trimap):
+        inp = torch.cat((image, trimap), dim=1)
+        embedding, mid_fea = self.encoder(inp)
+        alpha, info_dict = self.decoder(embedding, mid_fea)
+
+        return alpha, info_dict
+
+
+def get_generator(encoder, decoder, num_class=1):
+    generator = Generator(encoder=encoder, decoder=decoder, num_class=num_class)
+    return generator
+
+
+if __name__=="__main__":
+    import time
+    # generator = get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder).cuda().train()
+    batch_size = 12
+    # generator.eval()
+    n_eval = 10
+    # pre run the model
+    # with torch.no_grad():
+    #     for i in range(2):
+    #         x = torch.rand(batch_size, 3, 512, 512, device=device)
+    #         y = torch.rand(batch_size, 3, 512, 512, device=device)
+    #         z = generator(x,y)
+    # test without GPU IO
+
+    # x = torch.zeros(batch_size, 3, 512, 512, device=device)
+    # y = torch.zeros(batch_size, 1, 512, 512, device=device)
+    x = torch.randn(batch_size, 3, 512, 512)
+    y = torch.randn(batch_size, 3, 512, 512)
+    t = time.time()
+    # with torch.no_grad():
+    # for i in range(n_eval):
+    #     a = generator(x.cuda(),y.cuda())
+    # torch.cuda.synchronize()
+    # print(generator.__class__.__name__, 'With IO  \t', f'{(time.time() - t)/n_eval/batch_size:.5f} s')
+    # print(generator.__class__.__name__, 'FPS \t\t', f'{1/((time.time() - t)/n_eval/batch_size):.5f} s')
+    # for n, p in generator.named_parameters():
+    #     print(n)

project/hair_service_sd/core/matting/networks/ops.py

@@ -0,0 +1,256 @@
+import torch
+from   torch import nn
+from   torch.nn import Parameter
+from   torch.autograd import Variable
+from   torch.nn import functional as F
+
+
+def l2normalize(v, eps=1e-12):
+    return v / (v.norm() + eps)
+
+
+class SpectralNorm(nn.Module):
+    """
+    Based on https://github.com/heykeetae/Self-Attention-GAN/blob/master/spectral.py
+    and add _noupdate_u_v() for evaluation
+    """
+    def __init__(self, module, name='weight', power_iterations=1):
+        super(SpectralNorm, self).__init__()
+        self.module = module
+        self.name = name
+        self.power_iterations = power_iterations
+        if not self._made_params():
+            self._make_params()
+
+    def _update_u_v(self):
+        u = getattr(self.module, self.name + "_u")
+        v = getattr(self.module, self.name + "_v")
+        w = getattr(self.module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        for _ in range(self.power_iterations):
+            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
+            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
+
+        sigma = u.dot(w.view(height, -1).mv(v))
+        setattr(self.module, self.name, w / sigma.expand_as(w))
+
+    def _noupdate_u_v(self):
+        u = getattr(self.module, self.name + "_u")
+        v = getattr(self.module, self.name + "_v")
+        w = getattr(self.module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        sigma = u.dot(w.view(height, -1).mv(v))
+        setattr(self.module, self.name, w / sigma.expand_as(w))
+
+    def _made_params(self):
+        try:
+            u = getattr(self.module, self.name + "_u")
+            v = getattr(self.module, self.name + "_v")
+            w = getattr(self.module, self.name + "_bar")
+            return True
+        except AttributeError:
+            return False
+
+    def _make_params(self):
+        w = getattr(self.module, self.name)
+
+        height = w.data.shape[0]
+        width = w.view(height, -1).data.shape[1]
+
+        u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
+        v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
+        u.data = l2normalize(u.data)
+        v.data = l2normalize(v.data)
+        w_bar = Parameter(w.data)
+
+        del self.module._parameters[self.name]
+
+        self.module.register_parameter(self.name + "_u", u)
+        self.module.register_parameter(self.name + "_v", v)
+        self.module.register_parameter(self.name + "_bar", w_bar)
+
+    def forward(self, *args):
+        # if torch.is_grad_enabled() and self.module.training:
+        if self.module.training:
+            self._update_u_v()
+        else:
+            self._noupdate_u_v()
+        return self.module.forward(*args)
+
+
+class GuidedCxtAtten(nn.Module):
+    # based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting/blob/a6fe298fec502bfd9cbc64eb01e39f78a3262a59/models/DeepFill_Models/ops.py#L210
+    def __init__(self, out_channels, guidance_channels, rate=2):
+        super(GuidedCxtAtten, self).__init__()
+        self.rate = rate
+        self.padding = nn.ReflectionPad2d(1)
+        self.up_sample = nn.Upsample(scale_factor=self.rate, mode='nearest')
+
+        self.guidance_conv = nn.Conv2d(in_channels=guidance_channels, out_channels=guidance_channels//2,
+                                       kernel_size=1, stride=1, padding=0)
+
+        self.W = nn.Sequential(
+            nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(out_channels)
+            )
+
+        nn.init.xavier_uniform_(self.guidance_conv.weight)
+        nn.init.constant_(self.guidance_conv.bias, 0)
+        nn.init.xavier_uniform_(self.W[0].weight)
+        nn.init.constant_(self.W[1].weight, 1e-3)
+        nn.init.constant_(self.W[1].bias, 0)
+
+    def forward(self, f, alpha, unknown=None, ksize=3, stride=1, fuse_k=3, softmax_scale=1., training=True):
+
+        f = self.guidance_conv(f)
+        # get shapes
+        raw_int_fs = list(f.size()) # N x 64 x 64 x 64
+        raw_int_alpha = list(alpha.size()) # N x 128 x 64 x 64
+
+        # extract patches from background with stride and rate
+        kernel = 2*self.rate
+        alpha_w = self.extract_patches(alpha, kernel=kernel, stride=self.rate)
+        alpha_w = alpha_w.permute(0, 2, 3, 4, 5, 1)
+        alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], raw_int_alpha[2] // self.rate, raw_int_alpha[3] // self.rate, -1)
+        alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], -1, kernel, kernel, raw_int_alpha[1])
+        alpha_w = alpha_w.permute(0, 1, 4, 2, 3)
+
+        f = F.interpolate(f, scale_factor=1/self.rate, mode='nearest')
+
+        fs = f.size() # B x 64 x 32 x 32
+        f_groups = torch.split(f, 1, dim=0) # Split tensors by batch dimension; tuple is returned
+
+        # from b(B*H*W*C) to w(b*k*k*c*h*w)
+        int_fs = list(fs)
+        w = self.extract_patches(f)
+        w = w.permute(0, 2, 3, 4, 5, 1)
+        w = w.contiguous().view(raw_int_fs[0], raw_int_fs[2] // self.rate, raw_int_fs[3] // self.rate, -1)
+        w = w.contiguous().view(raw_int_fs[0], -1, ksize, ksize, raw_int_fs[1])
+        w = w.permute(0, 1, 4, 2, 3)
+        # process mask
+
+        if unknown is not None:
+            unknown = unknown.clone()
+            unknown = F.interpolate(unknown, scale_factor=1/self.rate, mode='nearest')
+            assert unknown.size(2) == f.size(2), "mask should have same size as f at dim 2,3"
+            unknown_mean = unknown.mean(dim=[2,3])
+            known_mean = 1 - unknown_mean
+            unknown_scale = torch.clamp(torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(alpha)
+            known_scale = torch.clamp(torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(alpha)
+            softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
+        else:
+            unknown = torch.ones([fs[0], 1, fs[2], fs[3]]).to(alpha)
+            softmax_scale = torch.FloatTensor([softmax_scale, softmax_scale]).view(1,2).repeat(fs[0],1).to(alpha)
+
+        m = self.extract_patches(unknown)
+
+        m = m.permute(0, 2, 3, 4, 5, 1)
+        m = m.contiguous().view(raw_int_fs[0], raw_int_fs[2]//self.rate, raw_int_fs[3]//self.rate, -1)
+        m = m.contiguous().view(raw_int_fs[0], -1, ksize, ksize)
+
+        m = self.reduce_mean(m) # smoothing, maybe
+        # mask out the
+        mm = m.gt(0.).float() # (N, 32*32, 1, 1)
+
+        # the correlation with itself should be 0
+        self_mask = F.one_hot(torch.arange(fs[2] * fs[3]).view(fs[2], fs[3]).contiguous().to(alpha).long(),
+                              num_classes=int_fs[2] * int_fs[3])
+        self_mask = self_mask.permute(2, 0, 1).view(1, fs[2] * fs[3], fs[2], fs[3]).float() * (-1e4)
+
+        w_groups = torch.split(w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
+        alpha_w_groups = torch.split(alpha_w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
+        mm_groups = torch.split(mm, 1, dim=0)
+        scale_group = torch.split(softmax_scale, 1, dim=0)
+        y = []
+        offsets = []
+        k = fuse_k
+        y_test = []
+        for xi, wi, alpha_wi, mmi, scale in zip(f_groups, w_groups, alpha_w_groups, mm_groups, scale_group):
+            # conv for compare
+            wi = wi[0]
+            escape_NaN = Variable(torch.FloatTensor([1e-4])).to(alpha)
+            wi_normed = wi / torch.max(self.l2_norm(wi), escape_NaN)
+            xi = F.pad(xi, (1,1,1,1), mode='reflect')
+            yi = F.conv2d(xi, wi_normed, stride=1, padding=0) # yi => (B=1, C=32*32, H=32, W=32)
+            y_test.append(yi)
+            # conv implementation for fuse scores to encourage large patches
+            yi = yi.permute(0, 2, 3, 1)
+            yi = yi.contiguous().view(1, fs[2], fs[3], fs[2] * fs[3])
+            yi = yi.permute(0, 3, 1, 2) # (B=1, C=32*32, H=32, W=32)
+
+            # softmax to match
+            # scale the correlation with predicted scale factor for known and unknown area
+            yi = yi * (scale[0,0] * mmi.gt(0.).float() + scale[0,1] * mmi.le(0.).float()) # mmi => (1, 32*32, 1, 1)
+            # mask itself, self-mask only applied to unknown area
+            yi = yi + self_mask * mmi  # self_mask: (1, 32*32, 32, 32)
+            # for small input inference
+            yi = F.softmax(yi, dim=1)
+
+            _, offset = torch.max(yi, dim=1) # argmax; index
+            offset = torch.stack([offset // fs[3], offset % fs[3]], dim=1)
+
+            wi_center = alpha_wi[0]
+
+            if self.rate == 1:
+                left = (kernel) // 2
+                right = (kernel - 1) // 2
+                yi = F.pad(yi, (left, right, left, right), mode='reflect')
+                wi_center = wi_center.permute(1, 0, 2, 3)
+                yi = F.conv2d(yi, wi_center, padding=0) / 4. # (B=1, C=128, H=64, W=64)
+            else:
+                yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
+            y.append(yi)
+            offsets.append(offset)
+
+        y = torch.cat(y, dim=0) # back to the mini-batch
+        y.contiguous().view(raw_int_alpha)
+        offsets = torch.cat(offsets, dim=0)
+        offsets = offsets.view([int_fs[0]] + [2] + int_fs[2:])
+
+        # # case1: visualize optical flow: minus current position
+        # h_add = Variable(torch.arange(0,float(fs[2]))).to(alpha).view([1, 1, fs[2], 1])
+        # h_add = h_add.expand(fs[0], 1, fs[2], fs[3])
+        # w_add = Variable(torch.arange(0,float(fs[3]))).to(alpha).view([1, 1, 1, fs[3]])
+        # w_add = w_add.expand(fs[0], 1, fs[2], fs[3])
+        #
+        # offsets = offsets - torch.cat([h_add, w_add], dim=1).long()
+
+        # case2: visualize absolute position
+        offsets = offsets - torch.Tensor([fs[2]//2, fs[3]//2]).view(1,2,1,1).to(alpha).long()
+
+        y = self.W(y) + alpha
+
+        return y, (offsets, softmax_scale)
+
+    @staticmethod
+    def extract_patches(x, kernel=3, stride=1):
+        left =(kernel - stride + 1) // 2
+        right =(kernel - stride) // 2
+        x = F.pad(x, (left, right, left, right), mode='reflect')
+        all_patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
+
+        return all_patches
+
+    @staticmethod
+    def reduce_mean(x):
+        for i in range(4):
+            if i <= 1:
+                continue
+            x = torch.mean(x, dim=i, keepdim=True)
+        return x
+
+    @staticmethod
+    def l2_norm(x):
+        def reduce_sum(x):
+            for i in range(4):
+                if i == 0:
+                    continue
+                x = torch.sum(x, dim=i, keepdim=True)
+            return x
+
+        x = x**2
+        x = reduce_sum(x)
+        return torch.sqrt(x)

project/hair_service_sd/core/matting/setup.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# @File    : setup.py
+# @Time    : 2020/1/15
+# @Author  : yangchaojie (yangchaojie@immomo.com)
+
+import os
+import sys
+import shutil
+import numpy
+import tempfile
+
+from setuptools import setup
+from setuptools.extension import Extension
+
+from Cython.Build import cythonize
+from Cython.Distutils import build_ext
+
+import platform
+
+
+def get_root_path(root):
+    if os.path.dirname(root) in ['', '.']:
+        return os.path.basename(root)
+    else:
+        return get_root_path(os.path.dirname(root))
+
+
+def copy_file(src, dest):
+    if os.path.exists(dest):
+        return
+
+    if not os.path.exists(os.path.dirname(dest)):
+        os.makedirs(os.path.dirname(dest))
+    if os.path.isdir(src):
+        shutil.copytree(src, dest)
+    else:
+        shutil.copyfile(src, dest)
+
+
+def touch_init_file():
+    init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
+    with open(init_file_name, 'w'):
+        pass
+    return init_file_name
+
+
+
+
+def compose_extensions(root='.'):
+    for file_ in os.listdir(root):
+        abs_file = os.path.join(root, file_)
+
+        if os.path.isfile(abs_file):
+            if abs_file.endswith('.py'):
+                extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
+            elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
+                continue
+            else:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            if abs_file.endswith('__init__.py'):
+                copy_file(init_file, os.path.join(build_root_dir, abs_file))
+
+        else:
+            if os.path.basename(abs_file) in ignore_folders :
+                continue
+            if os.path.basename(abs_file) in conf_folders:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            compose_extensions(abs_file)
+
+
+
+# if __name__ == '__main__':
+build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
+    sys.version_info.major) + '.' + str(sys.version_info.minor)
+
+print(build_root_dir)
+
+extensions = []
+ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
+conf_folders = ['conf']
+
+
+init_file = touch_init_file()
+print(init_file)
+
+
+compose_extensions()
+os.remove(init_file)
+
+setup(
+    name='moxie_hairstyle',
+    version='1.0',
+    ext_modules=cythonize(
+        extensions,
+        nthreads=16,
+        compiler_directives=dict(always_allow_keywords=True),
+        include_path=[numpy.get_include()]),
+    cmdclass=dict(build_ext=build_ext))
+
+    # python setup.py build_ext

project/hair_service_sd/core/model_3ddfa/model_3ddfa.py

@@ -0,0 +1,319 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import os
+import cv2
+import numpy as np
+from core.utils import utils_3ddfa,params_3ddfa,landmark_processor
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'http://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'http://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'http://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'http://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'http://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)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 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):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        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 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, no_branch=False, no_activate=False, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        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)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        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)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        if no_activate:
+            if no_branch:
+                self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes)])
+            else:
+                self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2)])
+                self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2)])
+        else:
+            if no_branch:
+                self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
+            else:
+                self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
+                self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
+        self.no_branch = no_branch
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ 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)
+
+        if self.no_branch:
+            key = self.avgpool(x)
+            key = key.view(key.size(0), -1)
+            key = self.fc(key)
+            return key
+        else:
+            key, ctrl = torch.chunk(x, 2, dim=1)
+
+            key = self.avgpool(key)
+            key = key.view(key.size(0), -1)
+            key = self.fc_key(key)
+
+            ctrl = self.avgpool(ctrl)
+            ctrl = ctrl.view(ctrl.size(0), -1)
+            ctrl = self.fc_ctrl(ctrl)
+
+            return key, ctrl
+
+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']), strict=False)
+    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']))
+    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']))
+    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']))
+    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']))
+    return model
+
+
+class Model_3DDFA(nn.Module):
+
+
+    def __init__(self, gpu_id=None):
+        super(Model_3DDFA, self).__init__()
+
+        self.init_status = False
+        self.face_alignment_net = resnet18(pretrained=True, num_classes=76, no_branch=True, no_activate=True)
+        self.model_path = 'weights'
+
+        # model_path, _ = os.path.split(os.path.realpath(__file__))
+        weights = torch.load(os.path.join(self.model_path, 'face_3ddfa.pth'), map_location=lambda storage, loc: storage)
+        self.load_state_dict(weights)
+        self.eval()
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        self.to(self.device)
+        self.init_status = True
+
+    def forward(self, imgs):
+        pred_pose_shape_exp = self.face_alignment_net(imgs)
+        return pred_pose_shape_exp
+
+    def running(self):
+        return self.init_status
+
+    def forward_np(self, imgs):
+        pred_pose_shape_exp = self.face_alignment_net(imgs)
+        return pred_pose_shape_exp.detach().cpu().numpy()
+
+    def detect(self, images, landmarks):
+        dst_size = 256
+        with torch.no_grad():
+            input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
+            assert len(images) == len(landmarks)
+            all_mat = []
+            all_res = []
+            all_height = []
+            for ix, img in enumerate(images):
+                landmark = landmarks[ix]
+
+                mat = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
+
+                all_mat.append(mat)
+                all_height.append(img.shape[0])
+
+                tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
+
+                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+
+                # cv2.imshow('3ddfa_', tmp)
+                # cv2.waitKey()
+
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            params = self.face_alignment_net(in_tensor)
+            params = params.cpu().numpy()
+
+            for ix, param in enumerate(params):
+                param[0] = param[0] / params_3ddfa.SCALE_F
+                param[1:4] = param[1:4] / params_3ddfa.SCALE_ROTATE
+                param[4:6] = param[4:6] / params_3ddfa.SCALE_OFFSET
+                param[6:56] = (param[6:56] / params_3ddfa.SCALE_SHAPE)
+                param[56:] = (param[56:] / params_3ddfa.SCALE_EXP)
+                new_param = utils_3ddfa.transform_params(param, cv2.invertAffineTransform(all_mat[ix]), all_height[ix],
+                                                         dst_size)
+                all_res.append(new_param)
+
+            return all_res

project/hair_service_sd/core/models/Generator_BaldSeg.py

@@ -0,0 +1,67 @@
+import os
+import torch
+import torch.nn.parallel
+
+import numpy as np
+from core.utils import landmark_processor
+import cv2
+
+modelRoot = "weights"
+
+label_map = [
+    [0, 0, 0],
+    [255, 0, 0],
+    [0, 0, 255],
+    [0, 255, 0],
+    [0, 255, 255],
+    # [0, 255, 0]
+]
+class Generator_BaldSeg_5c(object):
+    def __init__(self, gpu_flag, gpu_id):
+
+        if not gpu_flag:
+            self.device = torch.device("cpu")
+        else:
+            self.device = torch.device('cuda:{0}'.format(gpu_id))
+
+        # load seg model
+        self.model_dir = modelRoot
+        self.pre_trained_model = os.path.join('weights', "ori_hair_checkpoint_7660_0611.pt")
+        self.net = torch.jit.load(self.pre_trained_model, map_location=self.device).to(self.device)
+        self.net.eval()
+
+        self.output_img_size = 512
+        self.img_ratio = 0.4
+
+    def label_to_mask(self, label_np):
+        label_np = label_np.astype(np.int32)[:, :, np.newaxis]
+        mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
+        for id, color in enumerate(label_map):
+            index = (label_np == id).all(axis=2)
+            mask[index] = color
+        return mask
+    def forward(self, image, alpha, landmarks1k):
+        image_to_face_mat = landmark_processor.get_transform_mat_full_face_ratio_deeplab(landmarks1k, self.output_img_size, self.img_ratio)
+        img_ = (image * (1 - alpha[:, :, np.newaxis].astype(np.float32) / 255)).astype(np.uint8)
+        img_cuted = cv2.warpAffine(img_, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255])
+        inv_image_to_face_mat = cv2.invertAffineTransform(image_to_face_mat)
+
+        img = (img_cuted.astype(np.float32) / 255).transpose((2, 0, 1))
+        img = np.expand_dims(img, axis=0)
+        img = torch.from_numpy(img).to(self.device)  #cuda(self.gpu_id)
+
+        with torch.no_grad():
+            output = self.net(img)
+
+        pred = output.detach().cpu().numpy().squeeze().astype(np.float32) #torch.max(output[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32)
+        mask = self.label_to_mask(pred)
+
+        # cv2.imshow("img_cuted: ", img_cuted)
+        # cv2.imshow("mask: ", mask)
+        # cv2.waitKey()
+
+        black_img = np.zeros(image.shape).astype(np.uint8)
+        cv2.warpAffine(mask, inv_image_to_face_mat, (image.shape[1], image.shape[0]),
+                       dst=black_img, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_TRANSPARENT)
+
+        return black_img

project/hair_service_sd/core/models/MomocvFaceAlignment1K.py

@@ -0,0 +1,441 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import numpy as np
+import os
+from core.utils import landmark_processor
+from core.utils.umeyama import umeyama
+import cv2
+
+modelRoot = "weights"
+
+__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)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 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):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        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 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        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)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        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)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        if is_1k:
+            self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
+        else:
+            self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
+            self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
+        self.is_1k = is_1k
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ 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)
+
+        if self.is_1k:
+            key = self.avgpool(x)
+            key = key.view(key.size(0), -1)
+            key = self.fc(key)
+            return key
+        else:
+            key, ctrl = torch.chunk(x, 2, dim=1)
+
+            key = self.avgpool(key)
+            key = key.view(key.size(0), -1)
+            key = self.fc_key(key)
+
+            ctrl = self.avgpool(ctrl)
+            ctrl = ctrl.view(ctrl.size(0), -1)
+            ctrl = self.fc_ctrl(ctrl)
+
+            return key, ctrl
+
+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']), strict=False)
+    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']))
+    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']))
+    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']))
+    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']))
+    return model
+
+class Model1k(nn.Module):
+    def __init__(self, gpu_id=None):
+        super(Model1k, self).__init__()
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True)
+
+        self.model_path, _ = os.path.split(os.path.realpath(__file__))
+        # weights = torch.load(os.path.join(self.model_path, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
+        weights = torch.load(os.path.join(modelRoot, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
+        self.load_state_dict(weights)
+        self.to(self.device)
+        self.eval()
+
+    def forward(self, imgs):
+        pred_key_pts = self.face_alignment_net(imgs)
+        pred_key_pts = pred_key_pts + 0.5
+        return pred_key_pts
+
+class MomocvFaceAlignment1K(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = Model1k(gpu_id)
+
+        self.trackingFaceRects = []
+
+        print('conansherry MomocvFaceAlignment1K')
+
+    def forward(self, img_tensor):
+        fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy()
+        return fullyconnected1
+
+    def detect(self, img, landmarks):
+        dst_size = 256
+        landmarks_res = []
+        with torch.no_grad():
+            input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
+            all_mat = []
+            for ix, landmark in enumerate(landmarks):
+                M = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
+                all_mat.append(M)
+                tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
+
+                # cv2.imshow('inp', tmp)
+                # cv2.waitKey()
+
+                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            for ix, pts in enumerate(fullyconnected1):
+                orig_pts = (np.reshape(pts, (2, 1000)).transpose((1, 0)) * dst_size)
+                orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
+                landmarks_res.append(orig_pts)
+            return landmarks_res
+    def detect_single_face(self, img):
+        dst_size = 256
+        with torch.no_grad():
+            crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :]
+            tmp = cv2.resize(crop_img, (dst_size, dst_size))
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * crop_img.shape[0])
+            orig_pts[:, 0] += 104
+            orig_pts[:, 1] += 104
+            return orig_pts
+
+    def detect_single_face_old(self, img):
+        dst_size = 256
+        with torch.no_grad():
+            tmp = cv2.resize(img, (dst_size, dst_size))
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * img.shape[0])
+            return orig_pts
+    def detect_according_5pts(self, img, pts5):
+        dst_size = 256
+        with torch.no_grad():
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            eye_dis = 0.34
+            mouth_dis = 0.34
+            g_Average_5point_180 = np.array([
+                eye_dis, 0.3,
+                1 - eye_dis, 0.37,
+                0.5, 0.6,
+                mouth_dis, 0.63,
+                1 - mouth_dis, 0.7
+            ])
+            # print(g_Average_5point_180)
+            # left_eye = np.array([pts5[0], pts5[5]])
+            # right_eye = np.array([pts5[1], pts5[6]])
+            # nose = np.array([pts5[2], pts5[7]])
+            # left_mouth = np.array([pts5[3], pts5[8]])
+            # right_mouth = np.array([pts5[4], pts5[9]])
+            left_eye = np.array([pts5[0], pts5[1]])
+            right_eye = np.array([pts5[2], pts5[3]])
+            nose = np.array([pts5[4], pts5[5]])
+            left_mouth = np.array([pts5[6], pts5[7]])
+            right_mouth = np.array([pts5[8], pts5[9]])
+            pts5_src = np.vstack((left_eye, right_eye,
+                                  nose,
+                                  left_mouth, right_mouth))
+            pts5_src = np.array(pts5_src).astype(np.int32)
+            pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size
+            # print('pts5_src: ', pts5_src)
+            # print('pts5_dst: ', pts5_dst)
+
+            # mat = cv2.estimateAffinePartial2D(pts5_src, pts5_dst, False)[0]
+            mat = umeyama(pts5_src, pts5_dst, True)[0:2]
+            print('mat: ', mat)
+            tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
+
+            # tmp2 = cv2.warpAffine(img, mat2, (dst_size, dst_size))
+            # cv2.imshow("tmp2", tmp2)
+            # cv2.imshow("tmp", tmp)
+            # cv2.waitKey()
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size)
+            orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True)
+        return orig_pts
+
+    def stable_forward(self, image, detected_faces, reset=False):
+        if reset is True:
+            self.trackingFaceRects = []
+
+        if len(self.trackingFaceRects) == 0:
+            for face_rect in detected_faces:
+                    new_tracking_rect = [face_rect, True, [0, 0], 0, None]
+                    self.trackingFaceRects.append(new_tracking_rect)
+
+        with torch.no_grad():
+            landmarks = []
+            for ix, tracking_face_rect in enumerate(self.trackingFaceRects):
+                if tracking_face_rect[1] == True:
+                    d = tracking_face_rect[0]
+                    src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
+                    rotate_degree = tracking_face_rect[3]
+                    scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1])
+                    dst_center = np.array([0.5, 0.5]) * 256
+                    offset = dst_center - src_center
+                    M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
+                    M[:, 2] += offset
+                else:
+                    rotate_degree = 0
+                    M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256)
+                inp = cv2.warpAffine(image, M, (256, 256))
+
+                # cv2.imshow('inp_{}'.format(ix), inp)
+                # cv2.waitKey()
+                orig_inp = inp
+
+                inp = inp.transpose((2, 0, 1)).astype(np.float32)
+                inp = inp[np.newaxis, :, :, :] / 255
+
+                in_tensor = torch.from_numpy(inp)
+                in_tensor = in_tensor.cuda(0)
+                fullyconnected1 = self.forward(in_tensor)
+                fullyconnected1 = fullyconnected1[0]
+                orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256
+
+                t2 = cv2.getTickCount()
+                orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
+
+                # orig_pts = orig_pts.transpose((1, 0))
+                fullyconnected1 = orig_pts
+
+                # update tracking infos
+                tracking_face_rect[1] = False
+                tracking_face_rect[2] = None
+                tracking_face_rect[3] = rotate_degree
+                tracking_face_rect[4] = fullyconnected1
+
+                # fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1)
+
+                # eye_landmark = self.detect_eye(image, fullyconnected1)
+                # fullyconnected1[87:104] = eye_landmark[0]
+                # fullyconnected1[104:121] = eye_landmark[1]
+
+                landmarks.append(fullyconnected1)
+        return landmarks

project/hair_service_sd/core/models/box_utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+import cv2
+
+def nms(boxes, overlap_threshold=0.5, mode='union'):
+    """ Pure Python NMS baseline. """
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    scores = boxes[:, 4]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+
+        if mode is 'min':
+            ovr = inter / np.minimum(areas[i], areas[order[1:]])
+        else:
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= overlap_threshold)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def convert_to_square(bboxes):
+    """
+        Convert bounding boxes to a square form.
+    """
+    square_bboxes = np.zeros_like(bboxes)
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    h = y2 - y1 + 1.0
+    w = x2 - x1 + 1.0
+    max_side = np.maximum(h, w)
+    square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
+    square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
+    square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
+    square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
+    return square_bboxes
+
+
+def calibrate_box(bboxes, offsets):
+    """Transform bounding boxes to be more like true bounding boxes.
+    'offsets' is one of the outputs of the nets.
+    """
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w = x2 - x1 + 1.0
+    h = y2 - y1 + 1.0
+    w = np.expand_dims(w, 1)
+    h = np.expand_dims(h, 1)
+
+    translation = np.hstack([w, h, w, h])*offsets
+    bboxes[:, 0:4] = bboxes[:, 0:4] + translation
+    return bboxes
+
+
+def get_image_boxes(bounding_boxes, img, size=24):
+    """Cut out boxes from the image.
+    """
+    num_boxes = len(bounding_boxes)
+    (height, width, _) = img.shape
+
+    [dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
+    img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
+
+    for i in range(num_boxes):
+        img_box = np.zeros((h[i], w[i], 3), 'uint8')
+
+        img_array = np.asarray(img, 'uint8')
+        img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
+            img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
+
+        img_box = cv2.resize(img_box, (size, size))
+        img_box = np.asarray(img_box, 'float32')
+
+        img_boxes[i, :, :, :] = _preprocess(img_box)
+
+    return img_boxes
+
+
+def correct_bboxes(bboxes, width, height):
+    """Crop boxes that are too big and get coordinates
+    with respect to cutouts.
+    """
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w, h = x2 - x1 + 1.0,  y2 - y1 + 1.0
+    num_boxes = bboxes.shape[0]
+
+    x, y, ex, ey = x1, y1, x2, y2
+    dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
+    edx, edy = w.copy() - 1.0, h.copy() - 1.0
+
+    ind = np.where(ex > width - 1.0)[0]
+    edx[ind] = w[ind] + width - 2.0 - ex[ind]
+    ex[ind] = width - 1.0
+
+    ind = np.where(ey > height - 1.0)[0]
+    edy[ind] = h[ind] + height - 2.0 - ey[ind]
+    ey[ind] = height - 1.0
+
+    ind = np.where(x < 0.0)[0]
+    dx[ind] = 0.0 - x[ind]
+    x[ind] = 0.0
+
+    ind = np.where(y < 0.0)[0]
+    dy[ind] = 0.0 - y[ind]
+    y[ind] = 0.0
+    return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
+    return_list = [i.astype('int32') for i in return_list]
+
+    return return_list
+
+
+def _preprocess(img):
+    """Preprocessing step before feeding the network.
+    """
+    img = img.transpose((2, 0, 1))
+    img = np.expand_dims(img, 0)
+    img = (img - 127.5)*0.0078125
+    return img

project/hair_service_sd/core/models/config.py

@@ -0,0 +1,42 @@
+# config.py
+
+cfg_mnet = {
+    'name': 'mobilenet0.25',
+    'min_sizes': [[16, 32], [64, 128], [256, 512]],
+    'steps': [8, 16, 32],
+    'variance': [0.1, 0.2],
+    'clip': False,
+    'loc_weight': 2.0,
+    'gpu_train': True,
+    'batch_size': 32,
+    'ngpu': 1,
+    'epoch': 250,
+    'decay1': 190,
+    'decay2': 220,
+    'image_size': 640,
+    'pretrain': True,
+    'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
+    'in_channel': 32,
+    'out_channel': 64
+}
+
+cfg_re50 = {
+    'name': 'Resnet50',
+    'min_sizes': [[16, 32], [64, 128], [256, 512]],
+    'steps': [8, 16, 32],
+    'variance': [0.1, 0.2],
+    'clip': False,
+    'loc_weight': 2.0,
+    'gpu_train': True,
+    'batch_size': 24,
+    'ngpu': 4,
+    'epoch': 100,
+    'decay1': 70,
+    'decay2': 90,
+    'image_size': 840,
+    'pretrain': True,
+    'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
+    'in_channel': 256,
+    'out_channel': 256
+}
+

project/hair_service_sd/core/models/detector.py

@@ -0,0 +1,259 @@
+import math
+import numpy as np
+from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
+import torch
+import cv2
+from .nms.py_cpu_nms import py_cpu_nms
+from core.utils import box_utils_Retina
+from .layers.functions.prior_box import PriorBox
+from .config import cfg_re50
+from .retinaface import RetinaFace
+from mtcnn.model import PNet, RNet, ONet
+
+gpu_id = 0
+device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+pnet, rnet, onet= PNet(), RNet(), ONet()
+pnet.to(device)
+rnet.to(device)
+onet.to(device)
+onet.eval()
+
+
+class RetinaFaceDetector(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+        self.cfg = cfg_re50
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        model = RetinaFace(cfg=self.cfg, phase='test')
+        model = self.load_model(model, 'weights/Resnet50_Final.pth', True)
+        model.eval()
+        self.net = model.to(self.device)
+        self.resize = 1
+        self.confidence_threshold = 0.02
+        self.top_k = 5000
+        self.nms_threshold = 0.4
+        self.keep_top_k = 750
+
+    def remove_prefix(self, state_dict, prefix):
+        # print('remove prefix \'{}\''.format(prefix))
+        f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
+        return {f(key): value for key, value in state_dict.items()}
+
+    def check_keys(self, model, pretrained_state_dict):
+        ckpt_keys = set(pretrained_state_dict.keys())
+        model_keys = set(model.state_dict().keys())
+        used_pretrained_keys = model_keys & ckpt_keys
+        # unused_pretrained_keys = ckpt_keys - model_keys
+        # missing_keys = model_keys - ckpt_keys
+        assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
+        return True
+
+    def load_model(self, model, pretrained_path, load_to_cpu):
+        # print('Loading pretrained model from {}'.format(pretrained_path))
+        if load_to_cpu:
+            pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
+        else:
+            device = torch.cuda.current_device()
+            pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
+        if "state_dict" in pretrained_dict.keys():
+            pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
+        else:
+            pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
+        self.check_keys(model, pretrained_dict)
+        model.load_state_dict(pretrained_dict, strict=False)
+        return model
+
+    def forward(self, img_raw, min_face_size=50):
+        img_scale = 640 / max(img_raw.shape[0], img_raw.shape[1])
+        img = cv2.resize(img_raw, (0, 0), fx=img_scale, fy=img_scale)
+        img = np.float32(img)
+
+        im_height, im_width, _ = img.shape
+        scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
+        img -= (104, 117, 123)
+        img = img.transpose(2, 0, 1)
+        img = torch.from_numpy(img).unsqueeze(0)
+        img = img.to(self.device)
+        scale = scale.to(self.device)
+
+        loc, conf, landms = self.net(img)  # forward pass
+
+        priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
+        priors = priorbox.forward()
+        priors = priors.to(self.device)
+        prior_data = priors.data
+        boxes = box_utils_Retina.decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
+
+        boxes = boxes * scale / self.resize
+        boxes = boxes.cpu().numpy()
+        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
+        landms = box_utils_Retina.decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance'])
+        scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
+                               img.shape[3], img.shape[2], img.shape[3], img.shape[2],
+                               img.shape[3], img.shape[2]])
+        scale1 = scale1.to(self.device)
+        landms = landms * scale1 / self.resize
+        landms = landms.cpu().numpy()
+
+        # ignore low scores
+        inds = np.where(scores > self.confidence_threshold)[0]
+        boxes = boxes[inds]
+        landms = landms[inds]
+        scores = scores[inds]
+
+        # keep top-K before NMS
+        order = scores.argsort()[::-1][:self.top_k]
+        boxes = boxes[order]
+        landms = landms[order]
+        scores = scores[order]
+
+        # do NMS
+        dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
+        keep = py_cpu_nms(dets, self.nms_threshold, min_face_size = min_face_size * img_scale)
+        # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
+        dets = dets[keep, :]
+        landms = landms[keep]
+
+        # keep top-K faster NMS
+        dets = dets[:self.keep_top_k, :]
+        landms = landms[:self.keep_top_k, :]
+
+        dets[:, :4] = dets[:, :4] / img_scale
+        landms /= img_scale
+        return dets, landms
+
+    def forward_v2(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        (height, width, _) = image.shape
+        min_length = min(height, width)
+        min_detection_size = 12
+        factor = 0.707  # sqrt(0.5)
+
+        scales = []
+        m = min_detection_size / min_face_size
+        min_length *= m
+
+        factor_count = 0
+        while min_length > min_detection_size:
+            scales.append(m * factor ** factor_count)
+            min_length *= factor
+            factor_count += 1
+
+        # STAGE 1
+        bounding_boxes = []
+        for s in scales:  # run P-Net on different scales
+            boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
+            bounding_boxes.append(boxes)
+        bounding_boxes = [i for i in bounding_boxes if i is not None]
+        if len(bounding_boxes) == 0:
+            return [], []
+        bounding_boxes = np.vstack(bounding_boxes)
+
+        keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 2
+        img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = rnet(img_boxes)
+        offsets = output[0].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[1].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[1])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+
+        keep = nms(bounding_boxes, nms_thresholds[1])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 3
+        img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+        if len(img_boxes) == 0:
+            return [], []
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = onet(img_boxes)
+        landmarks = output[0].to('cpu').data.numpy()  # shape [n_boxes, 10]
+        offsets = output[1].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[2].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[2])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+        landmarks = landmarks[keep]
+
+        # compute landmark points
+        width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+        height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+        xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+        landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
+        landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
+
+        bounding_boxes = calibrate_box(bounding_boxes, offsets)
+        keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+        bounding_boxes = bounding_boxes[keep]
+        landmarks = landmarks[keep]
+
+        return bounding_boxes, landmarks
+
+def run_first_stage(image, net, scale, threshold, gpu_id=0):
+    """ 
+        Run P-Net, generate bounding boxes, and do NMS.
+    """
+    device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+    (height, width, _) = image.shape
+    sw, sh = math.ceil(width * scale), math.ceil(height * scale)
+    img = cv2.resize(image, (sw, sh))
+    # img = image.resize((sw, sh), Image.BILINEAR)
+    img = np.asarray(img, 'float32')
+    img = torch.from_numpy(_preprocess(img))
+    img = img.to(device)
+
+    output = net(img)
+    probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
+    offsets = output[0].to('cpu').data.numpy()
+
+    boxes = _generate_bboxes(probs, offsets, scale, threshold)
+    if len(boxes) == 0:
+        return None
+
+    keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
+    return boxes[keep]
+
+
+def _generate_bboxes(probs, offsets, scale, threshold):
+    """
+       Generate bounding boxes at places where there is probably a face.
+    """
+    stride = 2
+    cell_size = 12
+
+    inds = np.where(probs > threshold)
+
+    if inds[0].size == 0:
+        return np.array([])
+
+    tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
+
+    offsets = np.array([tx1, ty1, tx2, ty2])
+    score = probs[inds[0], inds[1]]
+
+    # P-Net is applied to scaled images, so we need to rescale bounding boxes back
+    bounding_boxes = np.vstack([
+        np.round((stride * inds[1] + 1.0) / scale),
+        np.round((stride * inds[0] + 1.0) / scale),
+        np.round((stride * inds[1] + 1.0 + cell_size) / scale),
+        np.round((stride * inds[0] + 1.0 + cell_size) / scale),
+        score, offsets
+    ])
+
+    return bounding_boxes.T

project/hair_service_sd/core/models/layers/__init__.py

@@ -0,0 +1,2 @@
+from .functions import *
+from .modules import *

project/hair_service_sd/core/models/layers/functions/prior_box.py

@@ -0,0 +1,34 @@
+import torch
+from itertools import product as product
+import numpy as np
+from math import ceil
+
+
+class PriorBox(object):
+    def __init__(self, cfg, image_size=None, phase='train'):
+        super(PriorBox, self).__init__()
+        self.min_sizes = cfg['min_sizes']
+        self.steps = cfg['steps']
+        self.clip = cfg['clip']
+        self.image_size = image_size
+        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
+        self.name = "s"
+
+    def forward(self):
+        anchors = []
+        for k, f in enumerate(self.feature_maps):
+            min_sizes = self.min_sizes[k]
+            for i, j in product(range(f[0]), range(f[1])):
+                for min_size in min_sizes:
+                    s_kx = min_size / self.image_size[1]
+                    s_ky = min_size / self.image_size[0]
+                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
+                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
+                    for cy, cx in product(dense_cy, dense_cx):
+                        anchors += [cx, cy, s_kx, s_ky]
+
+        # back to torch land
+        output = torch.Tensor(anchors).view(-1, 4)
+        if self.clip:
+            output.clamp_(max=1, min=0)
+        return output

project/hair_service_sd/core/models/layers/modules/__init__.py

@@ -0,0 +1,3 @@
+from .multibox_loss import MultiBoxLoss
+
+__all__ = ['MultiBoxLoss']

project/hair_service_sd/core/models/layers/modules/multibox_loss.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from core.utils.box_utils_Retina import match, log_sum_exp
+from core.models.layers.data import cfg_mnet
+GPU = cfg_mnet['gpu_train']
+
+class MultiBoxLoss(nn.Module):
+    """SSD Weighted Loss Function
+    Compute Targets:
+        1) Produce Confidence Target Indices by matching  ground truth boxes
+           with (default) 'priorboxes' that have jaccard index > threshold parameter
+           (default threshold: 0.5).
+        2) Produce localization target by 'encoding' variance into offsets of ground
+           truth boxes and their matched  'priorboxes'.
+        3) Hard negative mining to filter the excessive number of negative examples
+           that comes with using a large number of default bounding boxes.
+           (default negative:positive ratio 3:1)
+    Objective Loss:
+        L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
+        Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss
+        weighted by α which is set to 1 by cross val.
+        Args:
+            c: class confidences,
+            l: predicted boxes,
+            g: ground truth boxes
+            N: number of matched default boxes
+        See: https://arxiv.org/pdf/1512.02325.pdf for more details.
+    """
+
+    def __init__(self, num_classes, overlap_thresh, prior_for_matching, bkg_label, neg_mining, neg_pos, neg_overlap, encode_target):
+        super(MultiBoxLoss, self).__init__()
+        self.num_classes = num_classes
+        self.threshold = overlap_thresh
+        self.background_label = bkg_label
+        self.encode_target = encode_target
+        self.use_prior_for_matching = prior_for_matching
+        self.do_neg_mining = neg_mining
+        self.negpos_ratio = neg_pos
+        self.neg_overlap = neg_overlap
+        self.variance = [0.1, 0.2]
+
+    def forward(self, predictions, priors, targets):
+        """Multibox Loss
+        Args:
+            predictions (tuple): A tuple containing loc preds, conf preds,
+            and prior boxes from SSD net.
+                conf shape: torch.size(batch_size,num_priors,num_classes)
+                loc shape: torch.size(batch_size,num_priors,4)
+                priors shape: torch.size(num_priors,4)
+
+            ground_truth (tensor): Ground truth boxes and labels for a batch,
+                shape: [batch_size,num_objs,5] (last idx is the label).
+        """
+
+        loc_data, conf_data, landm_data = predictions
+        priors = priors
+        num = loc_data.size(0)
+        num_priors = (priors.size(0))
+
+        # match priors (default boxes) and ground truth boxes
+        loc_t = torch.Tensor(num, num_priors, 4)
+        landm_t = torch.Tensor(num, num_priors, 10)
+        conf_t = torch.LongTensor(num, num_priors)
+        for idx in range(num):
+            truths = targets[idx][:, :4].data
+            labels = targets[idx][:, -1].data
+            landms = targets[idx][:, 4:14].data
+            defaults = priors.data
+            match(self.threshold, truths, defaults, self.variance, labels, landms, loc_t, conf_t, landm_t, idx)
+        if GPU:
+            loc_t = loc_t.cuda()
+            conf_t = conf_t.cuda()
+            landm_t = landm_t.cuda()
+
+        zeros = torch.tensor(0).cuda()
+        # landm Loss (Smooth L1)
+        # Shape: [batch,num_priors,10]
+        pos1 = conf_t > zeros
+        num_pos_landm = pos1.long().sum(1, keepdim=True)
+        N1 = max(num_pos_landm.data.sum().float(), 1)
+        pos_idx1 = pos1.unsqueeze(pos1.dim()).expand_as(landm_data)
+        landm_p = landm_data[pos_idx1].view(-1, 10)
+        landm_t = landm_t[pos_idx1].view(-1, 10)
+        loss_landm = F.smooth_l1_loss(landm_p, landm_t, reduction='sum')
+
+
+        pos = conf_t != zeros
+        conf_t[pos] = 1
+
+        # Localization Loss (Smooth L1)
+        # Shape: [batch,num_priors,4]
+        pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
+        loc_p = loc_data[pos_idx].view(-1, 4)
+        loc_t = loc_t[pos_idx].view(-1, 4)
+        loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum')
+
+        # Compute max conf across batch for hard negative mining
+        batch_conf = conf_data.view(-1, self.num_classes)
+        loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
+
+        # Hard Negative Mining
+        loss_c[pos.view(-1, 1)] = 0 # filter out pos boxes for now
+        loss_c = loss_c.view(num, -1)
+        _, loss_idx = loss_c.sort(1, descending=True)
+        _, idx_rank = loss_idx.sort(1)
+        num_pos = pos.long().sum(1, keepdim=True)
+        num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)
+        neg = idx_rank < num_neg.expand_as(idx_rank)
+
+        # Confidence Loss Including Positive and Negative Examples
+        pos_idx = pos.unsqueeze(2).expand_as(conf_data)
+        neg_idx = neg.unsqueeze(2).expand_as(conf_data)
+        conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1,self.num_classes)
+        targets_weighted = conf_t[(pos+neg).gt(0)]
+        loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum')
+
+        # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
+        N = max(num_pos.data.sum().float(), 1)
+        loss_l /= N
+        loss_c /= N
+        loss_landm /= N1
+
+        return loss_l, loss_c, loss_landm

project/hair_service_sd/core/models/net.py

@@ -0,0 +1,137 @@
+import time
+import torch
+import torch.nn as nn
+# import torchvision.models._utils as _utils
+# import torchvision.models as models
+import torch.nn.functional as F
+# from torch.autograd import Variable
+
+def conv_bn(inp, oup, stride = 1, leaky = 0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True)
+    )
+
+def conv_bn_no_relu(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+    )
+
+def conv_bn1X1(inp, oup, stride, leaky=0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True)
+    )
+
+def conv_dw(inp, oup, stride, leaky=0.1):
+    return nn.Sequential(
+        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
+        nn.BatchNorm2d(inp),
+        nn.LeakyReLU(negative_slope= leaky,inplace=True),
+
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope= leaky,inplace=True),
+    )
+
+class SSH(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(SSH, self).__init__()
+        assert out_channel % 4 == 0
+        leaky = 0
+        if (out_channel <= 64):
+            leaky = 0.1
+        self.conv3X3 = conv_bn_no_relu(in_channel, out_channel//2, stride=1)
+
+        self.conv5X5_1 = conv_bn(in_channel, out_channel//4, stride=1, leaky = leaky)
+        self.conv5X5_2 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1)
+
+        self.conv7X7_2 = conv_bn(out_channel//4, out_channel//4, stride=1, leaky = leaky)
+        self.conv7x7_3 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1)
+
+    def forward(self, input):
+        conv3X3 = self.conv3X3(input)
+
+        conv5X5_1 = self.conv5X5_1(input)
+        conv5X5 = self.conv5X5_2(conv5X5_1)
+
+        conv7X7_2 = self.conv7X7_2(conv5X5_1)
+        conv7X7 = self.conv7x7_3(conv7X7_2)
+
+        out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
+        out = F.relu(out)
+        return out
+
+class FPN(nn.Module):
+    def __init__(self,in_channels_list,out_channels):
+        super(FPN,self).__init__()
+        leaky = 0
+        if (out_channels <= 64):
+            leaky = 0.1
+        self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride = 1, leaky = leaky)
+        self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride = 1, leaky = leaky)
+        self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride = 1, leaky = leaky)
+
+        self.merge1 = conv_bn(out_channels, out_channels, leaky = leaky)
+        self.merge2 = conv_bn(out_channels, out_channels, leaky = leaky)
+
+    def forward(self, input):
+        # names = list(input.keys())
+        input = list(input.values())
+
+        output1 = self.output1(input[0])
+        output2 = self.output2(input[1])
+        output3 = self.output3(input[2])
+
+        up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest")
+        output2 = output2 + up3
+        output2 = self.merge2(output2)
+
+        up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest")
+        output1 = output1 + up2
+        output1 = self.merge1(output1)
+
+        out = [output1, output2, output3]
+        return out
+
+
+
+class MobileNetV1(nn.Module):
+    def __init__(self):
+        super(MobileNetV1, self).__init__()
+        self.stage1 = nn.Sequential(
+            conv_bn(3, 8, 2, leaky = 0.1),    # 3
+            conv_dw(8, 16, 1),   # 7
+            conv_dw(16, 32, 2),  # 11
+            conv_dw(32, 32, 1),  # 19
+            conv_dw(32, 64, 2),  # 27
+            conv_dw(64, 64, 1),  # 43
+        )
+        self.stage2 = nn.Sequential(
+            conv_dw(64, 128, 2),  # 43 + 16 = 59
+            conv_dw(128, 128, 1), # 59 + 32 = 91
+            conv_dw(128, 128, 1), # 91 + 32 = 123
+            conv_dw(128, 128, 1), # 123 + 32 = 155
+            conv_dw(128, 128, 1), # 155 + 32 = 187
+            conv_dw(128, 128, 1), # 187 + 32 = 219
+        )
+        self.stage3 = nn.Sequential(
+            conv_dw(128, 256, 2), # 219 +3 2 = 241
+            conv_dw(256, 256, 1), # 241 + 64 = 301
+        )
+        self.avg = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Linear(256, 1000)
+
+    def forward(self, x):
+        x = self.stage1(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.avg(x)
+        # x = self.model(x)
+        x = x.view(-1, 256)
+        x = self.fc(x)
+        return x
+

project/hair_service_sd/core/models/nms/py_cpu_nms.py

@@ -0,0 +1,46 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import numpy as np
+
+def py_cpu_nms(dets, thresh, min_face_size = 50):
+    """Pure Python NMS baseline."""
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+    scores = dets[:, 4]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    #filter_small_faces
+    filter_list = []
+    for idx in keep:
+        w = np.abs(dets[idx, 2] - dets[idx, 0])
+        h = np.abs(dets[idx, 3] - dets[idx, 1])
+        if max(w, h) < min_face_size: continue
+        filter_list.append(idx)
+
+    return filter_list

project/hair_service_sd/core/models/resnet.py

@@ -0,0 +1,342 @@
+import torch
+import torch.nn as nn
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+           'wide_resnet50_2', 'wide_resnet101_2']
+
+
+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',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=2, padding=3, bias=False)   ### ycj
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+        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 = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+
+def _resnet(arch, block, layers, pretrained, progress, n_class, **kwargs):
+    model = ResNet(block, layers, num_classes=n_class, **kwargs)
+    # if pretrained:
+    #     state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+    #     model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(n_class=1000, pretrained=False, progress=True, **kwargs):
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, n_class,
+                   **kwargs)
+
+
+def resnet34(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet50(n_class=1000, pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, n_class, **kwargs)
+
+def resnet101(n_class=1000, pretrained=False, progress=True, **kwargs):
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, n_class, **kwargs)
+
+
+def resnet152(n_class=1000, pretrained=False, progress=True, **kwargs):
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, n_class,
+                   **kwargs)
+
+
+def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+if __name__ == '__main__':
+    net = resnet18(n_class=33)
+    input_size = (1, 3, 224, 224)
+    x = torch.randn(input_size)
+    out = net(x)
+    print(out.shape)

project/hair_service_sd/core/models/retinaface.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn as nn
+# import torchvision.models.detection.backbone_utils as backbone_utils
+import torchvision.models._utils as _utils
+import torch.nn.functional as F
+# from collections import OrderedDict
+
+from core.models.net import MobileNetV1 as MobileNetV1
+from core.models.net import FPN as FPN
+from core.models.net import SSH as SSH
+
+
+
+class ClassHead(nn.Module):
+    def __init__(self,inchannels=512,num_anchors=3):
+        super(ClassHead,self).__init__()
+        self.num_anchors = num_anchors
+        self.conv1x1 = nn.Conv2d(inchannels,self.num_anchors*2,kernel_size=(1,1),stride=1,padding=0)
+
+    def forward(self,x):
+        out = self.conv1x1(x)
+        out = out.permute(0,2,3,1).contiguous()
+        
+        return out.view(out.shape[0], -1, 2)
+
+class BboxHead(nn.Module):
+    def __init__(self,inchannels=512,num_anchors=3):
+        super(BboxHead,self).__init__()
+        self.conv1x1 = nn.Conv2d(inchannels,num_anchors*4,kernel_size=(1,1),stride=1,padding=0)
+
+    def forward(self,x):
+        out = self.conv1x1(x)
+        out = out.permute(0,2,3,1).contiguous()
+
+        return out.view(out.shape[0], -1, 4)
+
+class LandmarkHead(nn.Module):
+    def __init__(self,inchannels=512,num_anchors=3):
+        super(LandmarkHead,self).__init__()
+        self.conv1x1 = nn.Conv2d(inchannels,num_anchors*10,kernel_size=(1,1),stride=1,padding=0)
+
+    def forward(self,x):
+        out = self.conv1x1(x)
+        out = out.permute(0,2,3,1).contiguous()
+
+        return out.view(out.shape[0], -1, 10)
+
+class RetinaFace(nn.Module):
+    def __init__(self, cfg = None, phase = 'train'):
+        """
+        :param cfg:  Network related settings.
+        :param phase: train or test.
+        """
+        super(RetinaFace,self).__init__()
+        self.phase = phase
+        backbone = None
+        if cfg['name'] == 'mobilenet0.25':
+            backbone = MobileNetV1()
+            if cfg['pretrain']:
+                checkpoint = torch.load("./weights/mobilenetV1X0.25_pretrain.tar", map_location=torch.device('cpu'))
+                from collections import OrderedDict
+                new_state_dict = OrderedDict()
+                for k, v in checkpoint['state_dict'].items():
+                    name = k[7:]  # remove module.
+                    new_state_dict[name] = v
+                # load params
+                backbone.load_state_dict(new_state_dict)
+        elif cfg['name'] == 'Resnet50':
+            import torchvision.models as models
+            backbone = models.resnet50(pretrained=cfg['pretrain'])
+
+        self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])
+        in_channels_stage2 = cfg['in_channel']
+        in_channels_list = [
+            in_channels_stage2 * 2,
+            in_channels_stage2 * 4,
+            in_channels_stage2 * 8,
+        ]
+        out_channels = cfg['out_channel']
+        self.fpn = FPN(in_channels_list,out_channels)
+        self.ssh1 = SSH(out_channels, out_channels)
+        self.ssh2 = SSH(out_channels, out_channels)
+        self.ssh3 = SSH(out_channels, out_channels)
+
+        self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
+        self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
+        self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
+
+    def _make_class_head(self,fpn_num=3,inchannels=64,anchor_num=2):
+        classhead = nn.ModuleList()
+        for i in range(fpn_num):
+            classhead.append(ClassHead(inchannels,anchor_num))
+        return classhead
+    
+    def _make_bbox_head(self,fpn_num=3,inchannels=64,anchor_num=2):
+        bboxhead = nn.ModuleList()
+        for i in range(fpn_num):
+            bboxhead.append(BboxHead(inchannels,anchor_num))
+        return bboxhead
+
+    def _make_landmark_head(self,fpn_num=3,inchannels=64,anchor_num=2):
+        landmarkhead = nn.ModuleList()
+        for i in range(fpn_num):
+            landmarkhead.append(LandmarkHead(inchannels,anchor_num))
+        return landmarkhead
+
+    def forward(self,inputs):
+        out = self.body(inputs)
+
+        # FPN
+        fpn = self.fpn(out)
+
+        # SSH
+        feature1 = self.ssh1(fpn[0])
+        feature2 = self.ssh2(fpn[1])
+        feature3 = self.ssh3(fpn[2])
+        features = [feature1, feature2, feature3]
+
+        bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
+        classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)],dim=1)
+        ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1)
+
+        if self.phase == 'train':
+            output = (bbox_regressions, classifications, ldm_regressions)
+        else:
+            output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
+        return output

project/hair_service_sd/core/mtcnn/__init__.py

@@ -0,0 +1 @@
+

project/hair_service_sd/core/mtcnn/box_utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+import cv2
+
+def nms(boxes, overlap_threshold=0.5, mode='union'):
+    """ Pure Python NMS baseline. """
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    scores = boxes[:, 4]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+
+        if mode is 'min':
+            ovr = inter / np.minimum(areas[i], areas[order[1:]])
+        else:
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= overlap_threshold)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def convert_to_square(bboxes):
+    """
+        Convert bounding boxes to a square form.
+    """
+    square_bboxes = np.zeros_like(bboxes)
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    h = y2 - y1 + 1.0
+    w = x2 - x1 + 1.0
+    max_side = np.maximum(h, w)
+    square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
+    square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
+    square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
+    square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
+    return square_bboxes
+
+
+def calibrate_box(bboxes, offsets):
+    """Transform bounding boxes to be more like true bounding boxes.
+    'offsets' is one of the outputs of the nets.
+    """
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w = x2 - x1 + 1.0
+    h = y2 - y1 + 1.0
+    w = np.expand_dims(w, 1)
+    h = np.expand_dims(h, 1)
+
+    translation = np.hstack([w, h, w, h])*offsets
+    bboxes[:, 0:4] = bboxes[:, 0:4] + translation
+    return bboxes
+
+
+def get_image_boxes(bounding_boxes, img, size=24):
+    """Cut out boxes from the image.
+    """
+    num_boxes = len(bounding_boxes)
+    (height, width, _) = img.shape
+
+    [dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
+    img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
+
+    for i in range(num_boxes):
+        img_box = np.zeros((h[i], w[i], 3), 'uint8')
+
+        img_array = np.asarray(img, 'uint8')
+        img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
+            img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
+
+        img_box = cv2.resize(img_box, (size, size))
+        img_box = np.asarray(img_box, 'float32')
+
+        img_boxes[i, :, :, :] = _preprocess(img_box)
+
+    return img_boxes
+
+
+def correct_bboxes(bboxes, width, height):
+    """Crop boxes that are too big and get coordinates
+    with respect to cutouts.
+    """
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w, h = x2 - x1 + 1.0,  y2 - y1 + 1.0
+    num_boxes = bboxes.shape[0]
+
+    x, y, ex, ey = x1, y1, x2, y2
+    dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
+    edx, edy = w.copy() - 1.0, h.copy() - 1.0
+
+    ind = np.where(ex > width - 1.0)[0]
+    edx[ind] = w[ind] + width - 2.0 - ex[ind]
+    ex[ind] = width - 1.0
+
+    ind = np.where(ey > height - 1.0)[0]
+    edy[ind] = h[ind] + height - 2.0 - ey[ind]
+    ey[ind] = height - 1.0
+
+    ind = np.where(x < 0.0)[0]
+    dx[ind] = 0.0 - x[ind]
+    x[ind] = 0.0
+
+    ind = np.where(y < 0.0)[0]
+    dy[ind] = 0.0 - y[ind]
+    y[ind] = 0.0
+    return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
+    return_list = [i.astype('int32') for i in return_list]
+
+    return return_list
+
+
+def _preprocess(img):
+    """Preprocessing step before feeding the network.
+    """
+    img = img.transpose((2, 0, 1))
+    img = np.expand_dims(img, 0)
+    img = (img - 127.5)*0.0078125
+    return img

project/hair_service_sd/core/mtcnn/detector.py

@@ -0,0 +1,242 @@
+import math
+import numpy as np
+import torch
+from .model import PNet, RNet, ONet
+from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
+import torch
+import cv2
+
+def detect_faces(image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8],
+                 nms_thresholds=[0.7, 0.7, 0.7], gpu_id=0):
+    device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+    pnet, rnet, onet= PNet(), RNet(), ONet()
+    pnet.to(device)
+    rnet.to(device)
+    onet.to(device)
+    onet.eval()
+
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    (height, width, _) = image.shape
+    min_length = min(height, width)
+    min_detection_size = 12
+    factor = 0.707  # sqrt(0.5)
+
+    scales = []
+    m = min_detection_size/min_face_size
+    min_length *= m
+
+    factor_count = 0
+    while min_length > min_detection_size:
+        scales.append(m*factor**factor_count)
+        min_length *= factor
+        factor_count += 1
+
+    # STAGE 1
+    bounding_boxes = []
+    for s in scales:    # run P-Net on different scales
+        boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=gpu_id)
+        bounding_boxes.append(boxes)
+    bounding_boxes = [i for i in bounding_boxes if i is not None]
+    bounding_boxes = np.vstack(bounding_boxes)
+
+    keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 2
+    img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+    img_boxes = torch.from_numpy(img_boxes)
+    img_boxes = img_boxes.to(device)
+    output = rnet(img_boxes)
+    offsets = output[0].to('cpu').data.numpy()  # shape [n_boxes, 4]
+    probs = output[1].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[1])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+    offsets = offsets[keep]
+
+    keep = nms(bounding_boxes, nms_thresholds[1])
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 3
+    img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+    if len(img_boxes) == 0: 
+        return [], []
+    img_boxes = torch.from_numpy(img_boxes)
+    img_boxes = img_boxes.to(device)
+    output = onet(img_boxes)
+    landmarks = output[0].to('cpu').data.numpy()  # shape [n_boxes, 10]
+    offsets = output[1].to('cpu').data.numpy()  # shape [n_boxes, 4]
+    probs = output[2].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[2])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+    offsets = offsets[keep]
+    landmarks = landmarks[keep]
+
+    # compute landmark points
+    width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+    height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+    xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+    landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1)*landmarks[:, 0:5]
+    landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1)*landmarks[:, 5:10]
+
+    bounding_boxes = calibrate_box(bounding_boxes, offsets)
+    keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+    bounding_boxes = bounding_boxes[keep]
+    landmarks = landmarks[keep]
+
+    return bounding_boxes, landmarks
+
+class MTCNNFaceDetector(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        self.pnet, self.rnet, self.onet = PNet(), RNet(), ONet()
+        self.pnet.to(self.device)
+        self.rnet.to(self.device)
+        self.onet.to(self.device)
+        self.onet.eval()
+
+    def forward(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        (height, width, _) = image.shape
+        min_length = min(height, width)
+        min_detection_size = 12
+        factor = 0.707  # sqrt(0.5)
+
+        scales = []
+        m = min_detection_size / min_face_size
+        min_length *= m
+
+        factor_count = 0
+        while min_length > min_detection_size:
+            scales.append(m * factor ** factor_count)
+            min_length *= factor
+            factor_count += 1
+
+        # STAGE 1
+        bounding_boxes = []
+        for s in scales:  # run P-Net on different scales
+            boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
+            bounding_boxes.append(boxes)
+        bounding_boxes = [i for i in bounding_boxes if i is not None]
+        if len(bounding_boxes) == 0:
+            return [], []
+        bounding_boxes = np.vstack(bounding_boxes)
+
+        keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 2
+        img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = self.rnet(img_boxes)
+        offsets = output[0].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[1].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[1])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+
+        keep = nms(bounding_boxes, nms_thresholds[1])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 3
+        img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+        if len(img_boxes) == 0:
+            return [], []
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = self.onet(img_boxes)
+        landmarks = output[0].to('cpu').data.numpy()  # shape [n_boxes, 10]
+        offsets = output[1].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[2].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[2])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+        landmarks = landmarks[keep]
+
+        # compute landmark points
+        width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+        height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+        xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+        landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
+        landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
+
+        bounding_boxes = calibrate_box(bounding_boxes, offsets)
+        keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+        bounding_boxes = bounding_boxes[keep]
+        landmarks = landmarks[keep]
+
+        return bounding_boxes, landmarks
+
+def run_first_stage(image, net, scale, threshold, gpu_id=0):
+    """ 
+        Run P-Net, generate bounding boxes, and do NMS.
+    """
+    device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+    (height, width, _) = image.shape
+    sw, sh = math.ceil(width*scale), math.ceil(height*scale)
+    img = cv2.resize(image, (sw, sh))
+    # img = image.resize((sw, sh), Image.BILINEAR)
+    img = np.asarray(img, 'float32')
+    img = torch.from_numpy(_preprocess(img))
+    img = img.to(device)
+
+    output = net(img)
+    probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
+    offsets = output[0].to('cpu').data.numpy()
+
+    boxes = _generate_bboxes(probs, offsets, scale, threshold)
+    if len(boxes) == 0:
+        return None
+
+    keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
+    return boxes[keep]
+
+
+def _generate_bboxes(probs, offsets, scale, threshold):
+    """
+       Generate bounding boxes at places where there is probably a face.
+    """
+    stride = 2
+    cell_size = 12
+
+    inds = np.where(probs > threshold)
+
+    if inds[0].size == 0:
+        return np.array([])
+
+    tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
+
+    offsets = np.array([tx1, ty1, tx2, ty2])
+    score = probs[inds[0], inds[1]]
+
+    # P-Net is applied to scaled images, so we need to rescale bounding boxes back
+    bounding_boxes = np.vstack([
+        np.round((stride*inds[1] + 1.0)/scale),
+        np.round((stride*inds[0] + 1.0)/scale),
+        np.round((stride*inds[1] + 1.0 + cell_size)/scale),
+        np.round((stride*inds[0] + 1.0 + cell_size)/scale),
+        score, offsets
+    ])
+
+    return bounding_boxes.T

project/hair_service_sd/core/mtcnn/detector_ly.py

@@ -0,0 +1,242 @@
+import math
+import numpy as np
+import torch
+from .model import PNet, RNet, ONet
+from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
+import torch
+import cv2
+
+def detect_faces(image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8],
+                 nms_thresholds=[0.7, 0.7, 0.7], gpu_id=0):
+    device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+    pnet, rnet, onet= PNet(), RNet(), ONet()
+    pnet.to(device)
+    rnet.to(device)
+    onet.to(device)
+    onet.eval()
+
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    (height, width, _) = image.shape
+    min_length = min(height, width)
+    min_detection_size = 12
+    factor = 0.707  # sqrt(0.5)
+
+    scales = []
+    m = min_detection_size/min_face_size
+    min_length *= m
+
+    factor_count = 0
+    while min_length > min_detection_size:
+        scales.append(m*factor**factor_count)
+        min_length *= factor
+        factor_count += 1
+
+    # STAGE 1
+    bounding_boxes = []
+    for s in scales:    # run P-Net on different scales
+        boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=gpu_id)
+        bounding_boxes.append(boxes)
+    bounding_boxes = [i for i in bounding_boxes if i is not None]
+    bounding_boxes = np.vstack(bounding_boxes)
+
+    keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 2
+    img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+    img_boxes = torch.from_numpy(img_boxes)
+    img_boxes = img_boxes.to(device)
+    output = rnet(img_boxes)
+    offsets = output[0].to('cpu').data.numpy()  # shape [n_boxes, 4]
+    probs = output[1].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[1])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+    offsets = offsets[keep]
+
+    keep = nms(bounding_boxes, nms_thresholds[1])
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 3
+    img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+    if len(img_boxes) == 0: 
+        return [], []
+    img_boxes = torch.from_numpy(img_boxes)
+    img_boxes = img_boxes.to(device)
+    output = onet(img_boxes)
+    landmarks = output[0].to('cpu').data.numpy()  # shape [n_boxes, 10]
+    offsets = output[1].to('cpu').data.numpy()  # shape [n_boxes, 4]
+    probs = output[2].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[2])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+    offsets = offsets[keep]
+    landmarks = landmarks[keep]
+
+    # compute landmark points
+    width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+    height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+    xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+    landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1)*landmarks[:, 0:5]
+    landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1)*landmarks[:, 5:10]
+
+    bounding_boxes = calibrate_box(bounding_boxes, offsets)
+    keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+    bounding_boxes = bounding_boxes[keep]
+    landmarks = landmarks[keep]
+
+    return bounding_boxes, landmarks
+
+class MTCNNFaceDetector(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        self.pnet, self.rnet, self.onet = PNet(), RNet(), ONet()
+        self.pnet.to(self.device)
+        self.rnet.to(self.device)
+        self.onet.to(self.device)
+        self.onet.eval()
+
+    def forward(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        (height, width, _) = image.shape
+        min_length = min(height, width)
+        min_detection_size = 12
+        factor = 0.707  # sqrt(0.5)
+
+        scales = []
+        m = min_detection_size / min_face_size
+        min_length *= m
+
+        factor_count = 0
+        while min_length > min_detection_size:
+            scales.append(m * factor ** factor_count)
+            min_length *= factor
+            factor_count += 1
+
+        # STAGE 1
+        bounding_boxes = []
+        for s in scales:  # run P-Net on different scales
+            boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
+            bounding_boxes.append(boxes)
+        bounding_boxes = [i for i in bounding_boxes if i is not None]
+        if len(bounding_boxes) == 0:
+            return [], []
+        bounding_boxes = np.vstack(bounding_boxes)
+
+        keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 2
+        img_boxes = get_image_boxes(bounding_boxes, image, size=24)
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = self.rnet(img_boxes)
+        offsets = output[0].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[1].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[1])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+
+        keep = nms(bounding_boxes, nms_thresholds[1])
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+        bounding_boxes = convert_to_square(bounding_boxes)
+        bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+        # STAGE 3
+        img_boxes = get_image_boxes(bounding_boxes, image, size=48)
+        if len(img_boxes) == 0:
+            return [], []
+        img_boxes = torch.from_numpy(img_boxes)
+        img_boxes = img_boxes.to(self.device)
+        output = self.onet(img_boxes)
+        landmarks = output[0].to('cpu').data.numpy()  # shape [n_boxes, 10]
+        offsets = output[1].to('cpu').data.numpy()  # shape [n_boxes, 4]
+        probs = output[2].to('cpu').data.numpy()  # shape [n_boxes, 2]
+
+        keep = np.where(probs[:, 1] > thresholds[2])[0]
+        bounding_boxes = bounding_boxes[keep]
+        bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
+        offsets = offsets[keep]
+        landmarks = landmarks[keep]
+
+        # compute landmark points
+        width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+        height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+        xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+        landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
+        landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
+
+        bounding_boxes = calibrate_box(bounding_boxes, offsets)
+        keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
+        bounding_boxes = bounding_boxes[keep]
+        landmarks = landmarks[keep]
+
+        return bounding_boxes, landmarks
+
+def run_first_stage(image, net, scale, threshold, gpu_id=0):
+    """ 
+        Run P-Net, generate bounding boxes, and do NMS.
+    """
+    device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+    (height, width, _) = image.shape
+    sw, sh = math.ceil(width*scale), math.ceil(height*scale)
+    img = cv2.resize(image, (sw, sh))
+    # img = image.resize((sw, sh), Image.BILINEAR)
+    img = np.asarray(img, 'float32')
+    img = torch.from_numpy(_preprocess(img))
+    img = img.to(device)
+
+    output = net(img)
+    probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
+    offsets = output[0].to('cpu').data.numpy()
+
+    boxes = _generate_bboxes(probs, offsets, scale, threshold)
+    if len(boxes) == 0:
+        return None
+
+    keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
+    return boxes[keep]
+
+
+def _generate_bboxes(probs, offsets, scale, threshold):
+    """
+       Generate bounding boxes at places where there is probably a face.
+    """
+    stride = 2
+    cell_size = 12
+
+    inds = np.where(probs > threshold)
+
+    if inds[0].size == 0:
+        return np.array([])
+
+    tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
+
+    offsets = np.array([tx1, ty1, tx2, ty2])
+    score = probs[inds[0], inds[1]]
+
+    # P-Net is applied to scaled images, so we need to rescale bounding boxes back
+    bounding_boxes = np.vstack([
+        np.round((stride*inds[1] + 1.0)/scale),
+        np.round((stride*inds[0] + 1.0)/scale),
+        np.round((stride*inds[1] + 1.0 + cell_size)/scale),
+        np.round((stride*inds[0] + 1.0 + cell_size)/scale),
+        score, offsets
+    ])
+
+    return bounding_boxes.T

project/hair_service_sd/core/mtcnn/model.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+import numpy as np
+import os
+# from hairstyle_model import modelRoot
+modelRoot = "weights"
+class Flatten(nn.Module):
+    def __init__(self):
+        super(Flatten, self).__init__()
+    def forward(self, x):
+        x = x.transpose(3, 2).contiguous()
+        return x.view(x.size(0), -1)
+
+class PNet(nn.Module):
+    def __init__(self):
+        super(PNet, self).__init__()
+        self.model_path = modelRoot
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 10, 3, 1)),
+            ('prelu1', nn.PReLU(10)),
+            ('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)),
+            ('conv2', nn.Conv2d(10, 16, 3, 1)),
+            ('prelu2', nn.PReLU(16)),
+            ('conv3', nn.Conv2d(16, 32, 3, 1)),
+            ('prelu3', nn.PReLU(32))
+        ]))
+        self.conv4_1 = nn.Conv2d(32, 2, 1, 1)
+        self.conv4_2 = nn.Conv2d(32, 4, 1, 1)
+        weights = np.load(os.path.join(self.model_path, 'pnet.npy'), allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        x = self.features(x)
+        a = self.conv4_1(x)
+        b = self.conv4_2(x)
+        a = F.softmax(a, dim=1)
+        return b, a
+
+class RNet(nn.Module):
+    def __init__(self):
+        super(RNet, self).__init__()
+        self.model_path = modelRoot
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 28, 3, 1)),
+            ('prelu1', nn.PReLU(28)),
+            ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
+            ('conv2', nn.Conv2d(28, 48, 3, 1)),
+            ('prelu2', nn.PReLU(48)),
+            ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
+            ('conv3', nn.Conv2d(48, 64, 2, 1)),
+            ('prelu3', nn.PReLU(64)),
+            ('flatten', Flatten()),
+            ('conv4', nn.Linear(576, 128)),
+            ('prelu4', nn.PReLU(128))
+        ]))
+        self.conv5_1 = nn.Linear(128, 2)
+        self.conv5_2 = nn.Linear(128, 4)
+        weights = np.load(os.path.join(self.model_path, 'rnet.npy'), allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        x = self.features(x)
+        a = self.conv5_1(x)
+        b = self.conv5_2(x)
+        a = F.softmax(a, dim=1)
+        return b, a
+
+class ONet(nn.Module):
+    def __init__(self):
+        super(ONet, self).__init__()
+        self.model_path = modelRoot
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 32, 3, 1)),
+            ('prelu1', nn.PReLU(32)),
+            ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
+            ('conv2', nn.Conv2d(32, 64, 3, 1)),
+            ('prelu2', nn.PReLU(64)),
+            ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
+            ('conv3', nn.Conv2d(64, 64, 3, 1)),
+            ('prelu3', nn.PReLU(64)),
+            ('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)),
+            ('conv4', nn.Conv2d(64, 128, 2, 1)),
+            ('prelu4', nn.PReLU(128)),
+            ('flatten', Flatten()),
+            ('conv5', nn.Linear(1152, 256)),
+            ('drop5', nn.Dropout(0.25)),
+            ('prelu5', nn.PReLU(256)),
+        ]))
+        self.conv6_1 = nn.Linear(256, 2)
+        self.conv6_2 = nn.Linear(256, 4)
+        self.conv6_3 = nn.Linear(256, 10)
+        weights = np.load(os.path.join(self.model_path, 'onet.npy'), allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        x = self.features(x)
+        a = self.conv6_1(x)
+        b = self.conv6_2(x)
+        c = self.conv6_3(x)
+        a = F.softmax(a, dim=1)
+        return c, b, a

project/hair_service_sd/core/oss_module.py

@@ -0,0 +1,33 @@
+import time
+
+import oss2
+import os
+
+class OSS_object():
+    def __init__(self):
+        access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
+        access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
+        bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
+        endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
+        for param in (access_key_id, access_key_secret, bucket_name, endpoint):
+            assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
+
+        self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
+
+    def upload_file(self, file, target_name):
+        t0 = time.time()
+        with open(oss2.to_unicode(file), 'rb') as f:
+            ret = self.bucket.put_object(target_name, f)
+            print(ret.headers['x-oss-request-id'])
+
+        url = "https://oss-aidigitalfield.oss-cn-beijing.aliyuncs.com/{}".format(target_name)
+        print('耗时：{},签名url的地址为：{}'.format(time.time() - t0, url))
+        return url
+
+
+
+if __name__ == '__main__':
+    oss_2 = OSS_object()
+    t0 = time.time()
+    oss_2.upload_file('../data/left.jpg', 'hair_mz/images/test.jpg')
+    print(time.time() - t0)

project/hair_service_sd/core/seg/hairseg_single_model.py

@@ -0,0 +1,86 @@
+import os
+import torch
+
+from core.seg.networks.deeplabv3_plus import get_deeplabv3_plus
+import numpy as np
+import cv2
+from core.utils import landmark_processor
+
+
+def label_to_mask(label_np):
+    label_np = label_np.astype(np.int32)[:, :, np.newaxis]
+    mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
+    for id, color in enumerate(label_map):
+        index = (label_np == id).all(axis=2)
+        mask[index] = color
+    return mask
+
+label_map = [
+    [0, 0, 0],  #
+    [128, 128, 128],
+    [255, 255, 255],
+]
+
+class Evaluator(object):
+    def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None):
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+        # print("gpu_id: ", gpu_id)
+
+        # create network
+        self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass)
+        model_path = os.path.join(seg_model_path)
+        self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
+        # print("seg device: ", self.model.device)
+        self.model.to(self.device)
+        self.model.eval()
+
+        # images = torch.randn((1, 3, 512, 512)).to(self.device)
+        # torch.onnx.export(self.model, images,
+        #                   "deeplabv3_hair512_360_0520_wl.onnx",
+        #                   verbose=True,
+        #                   opset_version=11,
+        #                   input_names=['data'],
+        #                   do_constant_folding=True,
+        #                   output_names=['output'])
+
+        # exit()
+
+        self.output_img_size = output_img_size
+        self.nclass = nclass
+
+
+    def process_data(self, img):
+        img = (img.astype(np.float32) / 255).transpose((2, 0, 1))
+        img = torch.from_numpy(img).unsqueeze(0)
+
+        return img
+
+    def eval(self, img, pts1k):
+        orig_h, orig_w, _ = img.shape
+
+        M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.3)
+        crop_img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
+        crop_img = self.process_data(crop_img)
+        crop_img = crop_img.to(self.device)
+
+        with torch.no_grad():
+            # torch.cuda.synchronize()
+            outputs = self.model(crop_img)
+            pred = torch.argmax(outputs[0], 1)
+
+            pred = pred[0].detach().cpu().numpy()
+            predict = pred.astype(np.float32)
+
+            pred_mask = label_to_mask(predict)
+
+            M1_invert = cv2.invertAffineTransform(M1)
+            img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC)  #flags=cv2.INTER_NEAREST
+            orig_mask = img_pred.copy()
+
+        # show_concat = np.concatenate((img, orig_mask), axis=1)
+        # cv2.imshow("show_concat", show_concat)
+        # cv2.waitKey()
+        return orig_mask
+
+
+

project/hair_service_sd/core/seg/networks/basic.py

@@ -0,0 +1,462 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['_ConvBNReLU', '_DWConvBNReLU', 'InvertedResidual', '_ASPP', '_FCNHead',
+           '_Hswish', '_ConvBNHswish', 'SEModule', 'Bottleneck', 'ShuffleNetUnit',
+           'ShuffleNetV2Unit', 'InvertedIGCV3', 'MBConvBlock']
+
+
+class _ConvBNReLU(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
+                 dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
+        self.bn = norm_layer(out_channels)
+        self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class _FCNHead(nn.Module):
+    def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs):
+        super(_FCNHead, self).__init__()
+        inter_channels = in_channels // 4
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
+            norm_layer(inter_channels, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True),
+            nn.Dropout(0.1),
+            nn.Conv2d(inter_channels, channels, 1)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+# -----------------------------------------------------------------
+#                      For MobileNet
+# -----------------------------------------------------------------
+class _DWConvBNReLU(nn.Module):
+    """Depthwise Separable Convolution in MobileNet.
+    depthwise convolution + pointwise convolution
+    """
+
+    def __init__(self, in_channels, dw_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_DWConvBNReLU, self).__init__()
+        self.conv = nn.Sequential(
+            _ConvBNReLU(in_channels, dw_channels, 3, stride, dilation, dilation, in_channels, norm_layer=norm_layer),
+            _ConvBNReLU(dw_channels, out_channels, 1, norm_layer=norm_layer))
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+# -----------------------------------------------------------------
+#                      For MobileNetV2
+# -----------------------------------------------------------------
+class InvertedResidual(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, expand_ratio,
+                 dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(InvertedResidual, self).__init__()
+        assert stride in [1, 2]
+        self.use_res_connect = stride == 1 and in_channels == out_channels
+
+        layers = list()
+        inter_channels = int(round(in_channels * expand_ratio))
+        if expand_ratio != 1:
+            # pw
+            layers.append(_ConvBNReLU(in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer))
+        layers.extend([
+            # dw
+            _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation,
+                        groups=inter_channels, relu6=True, norm_layer=norm_layer),
+            # pw-linear
+            nn.Conv2d(inter_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels)])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+# -----------------------------------------------------------------
+#                  ASPP: For MobileNetV2
+# -----------------------------------------------------------------
+class _AsppPooling(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer, **kwargs):
+        super(_AsppPooling, self).__init__()
+        self.gap = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+#        size = x.size()[2:]
+        size = (48, 48)
+#        print("size: ", size)
+        pool = self.gap(x)
+#        out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
+        out = F.interpolate(pool, size, mode='nearest')
+        return out
+
+
+class _ASPP(nn.Module):
+    def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_ASPP, self).__init__()
+        out_channels = 256
+        self.b0 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        self.b1 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate1, dilation=rate1, norm_layer=norm_layer)
+        self.b2 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate2, dilation=rate2, norm_layer=norm_layer)
+        self.b3 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate3, dilation=rate3, norm_layer=norm_layer)
+        self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
+
+        self.project = nn.Sequential(
+            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True),
+            nn.Dropout2d(0.5)
+        )
+
+    def forward(self, x):
+        feat1 = self.b0(x)
+        feat2 = self.b1(x)
+        feat3 = self.b2(x)
+        feat4 = self.b3(x)
+        feat5 = self.b4(x)
+        x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        x = self.project(x)
+        return x
+
+
+# -----------------------------------------------------------------
+#                      For MobileNetV3
+# -----------------------------------------------------------------
+class _Hswish(nn.Module):
+    def __init__(self, inplace=True):
+        super(_Hswish, self).__init__()
+        self.relu6 = nn.ReLU6(inplace)
+
+    def forward(self, x):
+        return x * self.relu6(x + 3.) / 6.
+
+
+class _Hsigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(_Hsigmoid, self).__init__()
+        self.relu6 = nn.ReLU6(inplace)
+
+    def forward(self, x):
+        return self.relu6(x + 3.) / 6.
+
+
+class _ConvBNHswish(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
+                 dilation=1, groups=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_ConvBNHswish, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
+        self.bn = norm_layer(out_channels)
+        self.act = _Hswish(True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.act(x)
+        return x
+
+
+class SEModule(nn.Module):
+    def __init__(self, in_channels, reduction=4):
+        super(SEModule, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(in_channels, in_channels // reduction, bias=False),
+            nn.ReLU(True),
+            nn.Linear(in_channels // reduction, in_channels, bias=False),
+            _Hsigmoid(True)
+        )
+
+    def forward(self, x):
+        n, c, _, _ = x.size()
+        out = self.avg_pool(x).view(n, c)
+        out = self.fc(out).view(n, c, 1, 1)
+        return x * out.expand_as(x)
+
+
+class Identity(nn.Module):
+    def __init__(self, in_channels):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+
+class Bottleneck(nn.Module):
+    def __init__(self, in_channels, out_channels, exp_size, kernel_size, stride, dilation=1, se=False, nl='RE',
+                 norm_layer=nn.BatchNorm2d, **kwargs):
+        super(Bottleneck, self).__init__()
+        assert stride in [1, 2]
+        self.use_res_connect = stride == 1 and in_channels == out_channels
+        if nl == 'HS':
+            act = _Hswish
+        else:
+            act = nn.ReLU
+        if se:
+            SELayer = SEModule
+        else:
+            SELayer = Identity
+
+        self.conv = nn.Sequential(
+            # pw
+            nn.Conv2d(in_channels, exp_size, 1, bias=False),
+            norm_layer(exp_size),
+            act(True),
+            # dw
+            nn.Conv2d(exp_size, exp_size, kernel_size, stride, (kernel_size - 1) // 2 * dilation,
+                      dilation, groups=exp_size, bias=False),
+            norm_layer(exp_size),
+            SELayer(exp_size),
+            act(True),
+            # pw-linear
+            nn.Conv2d(exp_size, out_channels, 1, bias=False),
+            norm_layer(out_channels)
+        )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+# -----------------------------------------------------------------
+#                      For ShuffleNet
+# -----------------------------------------------------------------
+def channel_shuffle(x, groups):
+    n, c, h, w = x.size()
+
+    channels_per_group = c // groups
+    x = x.view(n, groups, channels_per_group, h, w)
+    x = torch.transpose(x, 1, 2).contiguous()
+    x = x.view(n, -1, h, w)
+
+    return x
+
+
+class ShuffleNetUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, groups, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(ShuffleNetUnit, self).__init__()
+        self.stride = stride
+        self.groups = groups
+        self.dilation = dilation
+        assert stride in [1, 2, 3]
+
+        inter_channels = out_channels // 4
+
+        if stride > 1:
+            self.shortcut = nn.AvgPool2d(3, stride, 1)
+            out_channels -= in_channels
+        elif dilation > 1:
+            out_channels -= in_channels
+
+        g = 1 if in_channels == 24 else groups
+        self.conv1 = _ConvBNReLU(in_channels, inter_channels, 1, groups=g, norm_layer=norm_layer)
+        self.conv2 = _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation,
+                                 dilation, groups, norm_layer=norm_layer)
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(inter_channels, out_channels, 1, groups=groups, bias=False),
+            norm_layer(out_channels))
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = channel_shuffle(out, self.groups)
+        out = self.conv2(out)
+        out = self.conv3(out)
+        if self.stride > 1:
+            x = self.shortcut(x)
+            out = torch.cat([out, x], dim=1)
+        elif self.dilation > 1:
+            out = torch.cat([out, x], dim=1)
+        else:
+            out = out + x
+        out = F.relu(out)
+
+        return out
+
+
+# -----------------------------------------------------------------
+#                      For ShuffleNetV2
+# -----------------------------------------------------------------
+class _DWConv(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False):
+        super(_DWConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride,
+                              padding, dilation, groups=in_channels, bias=bias)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class ShuffleNetV2Unit(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(ShuffleNetV2Unit, self).__init__()
+        assert stride in [1, 2, 3]
+        self.stride = stride
+        self.dilation = dilation
+
+        inter_channels = out_channels // 2
+
+        if (stride > 1) or (dilation > 1):
+            self.branch1 = nn.Sequential(
+                _DWConv(in_channels, in_channels, 3, stride, dilation, dilation),
+                norm_layer(in_channels),
+                _ConvBNReLU(in_channels, inter_channels, 1, norm_layer=norm_layer))
+        self.branch2 = nn.Sequential(
+            _ConvBNReLU(in_channels if (stride > 1) else inter_channels, inter_channels, 1, norm_layer=norm_layer),
+            _DWConv(inter_channels, inter_channels, 3, stride, dilation, dilation),
+            norm_layer(inter_channels),
+            _ConvBNReLU(inter_channels, inter_channels, 1, norm_layer=norm_layer))
+
+    def forward(self, x):
+        if (self.stride == 1) and (self.dilation == 1):
+            x1, x2 = x.chunk(2, dim=1)
+            out = torch.cat((x1, self.branch2(x2)), dim=1)
+        else:
+            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
+        out = channel_shuffle(out, 2)
+
+        return out
+
+
+# -----------------------------------------------------------------
+#                      For IGCV3
+# -----------------------------------------------------------------
+class PermutationBlock(nn.Module):
+    def __init__(self, groups):
+        super(PermutationBlock, self).__init__()
+        self.groups = groups
+
+    def forward(self, x):
+        n, c, h, w = x.size()
+        x = x.view(n, self.groups, c // self.groups, h, w).permute(0, 2, 1, 3, 4).contiguous().view(n, c, h, w)
+        return x
+
+
+class InvertedIGCV3(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, expand_ratio,
+                 dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(InvertedIGCV3, self).__init__()
+        assert stride in [1, 2]
+        self.use_res_connect = stride == 1 and in_channels == out_channels
+
+        layers = list()
+        inter_channels = int(round(in_channels * expand_ratio))
+        if expand_ratio != 1:
+            # pw
+            layers.append(_ConvBNReLU(in_channels, inter_channels, 1,
+                                      groups=2, relu6=True, norm_layer=norm_layer))
+            # permutation
+            layers.append(PermutationBlock(groups=2))
+        layers.extend([
+            # dw
+            _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation,
+                        groups=inter_channels, relu6=True, norm_layer=norm_layer),
+            # pw-linear
+            nn.Conv2d(inter_channels, out_channels, 1, groups=2, bias=False),
+            norm_layer(out_channels),
+            # permutation
+            PermutationBlock(groups=int(round(out_channels / 2)))
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+# -----------------------------------------------------------------
+#                      For EfficientNet
+# -----------------------------------------------------------------
+class _Swish(nn.Module):
+    def __init__(self):
+        super(_Swish, self).__init__()
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        return x * self.sigmoid(x)
+
+
+class SEModuleV2(nn.Module):
+    def __init__(self, in_channels, se_ratio=0.25):
+        super(SEModuleV2, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        se_channels = max(1, int(in_channels * se_ratio))
+        self.fc = nn.Sequential(
+            nn.Conv2d(in_channels, se_channels, 1, bias=False),
+            _Swish(),
+            nn.Conv2d(se_channels, in_channels, 1, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        n, c, _, _ = x.size()
+        out = self.avg_pool(x)
+        out = self.fc(out)
+        return x * out.expand_as(x)
+
+
+class MBConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio,
+                 dilation=1, se_ratio=0.25, drop_connect_rate=0.2, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(MBConvBlock, self).__init__()
+        assert stride in [1, 2]
+        self.use_res_connect = stride == 1 and in_channels == out_channels
+        self.drop_connect_rate = drop_connect_rate
+        use_se = (se_ratio is not None) and (0 < se_ratio <= 1.)
+        if use_se:
+            SELayer = SEModuleV2
+        else:
+            SELayer = Identity
+
+        layers = list()
+        inter_channels = int(round(in_channels * expand_ratio))
+        if expand_ratio != 1:
+            layers.append(_ConvBNHswish(in_channels, inter_channels, 1, norm_layer=norm_layer))
+        layers.extend([
+            # dw
+            _ConvBNHswish(inter_channels, inter_channels, kernel_size, stride, kernel_size // 2 * dilation, dilation,
+                          groups=inter_channels, norm_layer=norm_layer),  # check act function
+            SELayer(inter_channels, se_ratio),
+            # pw-linear
+            nn.Conv2d(inter_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels)
+        ])
+        self.conv = nn.Sequential(*layers)
+
+        if drop_connect_rate:
+            self.dropout = nn.Dropout2d(drop_connect_rate)
+
+    def forward(self, x):
+        out = self.conv(x)
+        if self.use_res_connect:
+            if self.drop_connect_rate:
+                out = self.dropout(out)
+            out = x + out
+        return out

project/hair_service_sd/core/seg/networks/deeplabv3.py

@@ -0,0 +1,187 @@
+"""Pyramid Scene Parsing Network"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from core.seg.networks.segbase import SegBaseModel
+from core.seg.networks.fcn import _FCNHead
+
+__all__ = ['DeepLabV3', 'get_deeplabv3', 'get_deeplabv3_resnet50_voc', 'get_deeplabv3_resnet101_voc',
+           'get_deeplabv3_resnet152_voc', 'get_deeplabv3_resnet50_ade', 'get_deeplabv3_resnet101_ade',
+           'get_deeplabv3_resnet152_ade']
+
+
+class DeepLabV3(SegBaseModel):
+    r"""DeepLabV3
+
+    Parameters
+    ----------
+    nclass : int
+        Number of categories for the training dataset.
+    backbone : string
+        Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
+        'resnet101' or 'resnet152').
+    norm_layer : object
+        Normalization layer used in backbone network (default: :class:`nn.BatchNorm`;
+        for Synchronized Cross-GPU BachNormalization).
+    aux : bool
+        Auxiliary loss.
+
+    Reference:
+        Chen, Liang-Chieh, et al. "Rethinking atrous convolution for semantic image segmentation."
+        arXiv preprint arXiv:1706.05587 (2017).
+    """
+
+    def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs):
+        super(DeepLabV3, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs)
+        self.head = _DeepLabHead(nclass, **kwargs)
+        if self.aux:
+            self.auxlayer = _FCNHead(1024, nclass, **kwargs)
+
+        self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
+
+    def forward(self, x):
+        size = x.size()[2:]
+        _, _, c3, c4 = self.base_forward(x)
+        outputs = []
+        x = self.head(c4)
+        x = F.interpolate(x, size, mode='bilinear', align_corners=True)
+        outputs.append(x)
+
+        if self.aux:
+            auxout = self.auxlayer(c3)
+            auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
+            outputs.append(auxout)
+        return tuple(outputs)
+
+
+class _DeepLabHead(nn.Module):
+    def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs):
+        super(_DeepLabHead, self).__init__()
+        self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs)
+        self.block = nn.Sequential(
+            nn.Conv2d(256, 256, 3, padding=1, bias=False),
+            norm_layer(256, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True),
+            nn.Dropout(0.1),
+            nn.Conv2d(256, nclass, 1)
+        )
+
+    def forward(self, x):
+        x = self.aspp(x)
+        return self.block(x)
+
+
+class _ASPPConv(nn.Module):
+    def __init__(self, in_channels, out_channels, atrous_rate, norm_layer, norm_kwargs):
+        super(_ASPPConv, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False),
+            norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class _AsppPooling(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer, norm_kwargs, **kwargs):
+        super(_AsppPooling, self).__init__()
+        self.gap = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        size = x.size()[2:]
+#        print("before gap: ", x.size())
+        pool = self.gap(x)
+        out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
+        return out
+
+
+class _ASPP(nn.Module):
+    def __init__(self, in_channels, atrous_rates, norm_layer, norm_kwargs=None, **kwargs):
+        super(_ASPP, self).__init__()
+        out_channels = 256
+        self.b0 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True)
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer, norm_kwargs)
+        self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer, norm_kwargs)
+        self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer, norm_kwargs)
+        self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
+
+        self.project = nn.Sequential(
+            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
+            nn.ReLU(True),
+            nn.Dropout(0.5)
+        )
+
+    def forward(self, x):
+        feat1 = self.b0(x)
+        feat2 = self.b1(x)
+        feat3 = self.b2(x)
+        feat4 = self.b3(x)
+#        print("before b4: ", x.size())
+        feat5 = self.b4(x)
+        x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        x = self.project(x)
+        return x
+
+
+def get_deeplabv3(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.torch/models',
+                  pretrained_base=True, **kwargs):
+    acronyms = {
+        'pascal_voc': 'pascal_voc',
+        'pascal_aug': 'pascal_aug',
+        'ade20k': 'ade',
+        'coco': 'coco',
+        'citys': 'citys',
+    }
+    from ..data.dataloader import datasets
+    model = DeepLabV3(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
+    if pretrained:
+        from .model_store import get_model_file
+        device = torch.device(kwargs['local_rank'])
+        model.load_state_dict(torch.load(get_model_file('deeplabv3_%s_%s' % (backbone, acronyms[dataset]), root=root),
+                              map_location=device))
+    return model
+
+
+def get_deeplabv3_resnet50_voc(**kwargs):
+    return get_deeplabv3('pascal_voc', 'resnet50', **kwargs)
+
+
+def get_deeplabv3_resnet101_voc(**kwargs):
+    return get_deeplabv3('pascal_voc', 'resnet101', **kwargs)
+
+
+def get_deeplabv3_resnet152_voc(**kwargs):
+    return get_deeplabv3('pascal_voc', 'resnet152', **kwargs)
+
+
+def get_deeplabv3_resnet50_ade(**kwargs):
+    return get_deeplabv3('ade20k', 'resnet50', **kwargs)
+
+
+def get_deeplabv3_resnet101_ade(**kwargs):
+    return get_deeplabv3('ade20k', 'resnet101', **kwargs)
+
+
+def get_deeplabv3_resnet152_ade(**kwargs):
+    return get_deeplabv3('ade20k', 'resnet152', **kwargs)
+
+
+if __name__ == '__main__':
+    model = get_deeplabv3_resnet50_voc()
+    img = torch.randn(2, 3, 480, 480)
+    output = model(img)

project/hair_service_sd/core/seg/networks/deeplabv3_plus.py

@@ -0,0 +1,160 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from core.seg.networks.xception import get_xception
+from core.seg.networks.deeplabv3 import _ASPP
+from core.seg.networks.fcn import _FCNHead
+from core.seg.networks.basic import _ConvBNReLU
+
+__all__ = ['DeepLabV3Plus', 'get_deeplabv3_plus', 'get_deeplabv3_plus_xception_voc']
+
+
+class DeepLabV3Plus(nn.Module):
+    r"""DeepLabV3Plus
+    Parameters
+    ----------
+    nclass : int
+        Number of categories for the training dataset.
+    backbone : string
+        Pre-trained dilated backbone network type (default:'xception').
+    norm_layer : object
+        Normalization layer used in backbone network (default: :class:`nn.BatchNorm`;
+        for Synchronized Cross-GPU BachNormalization).
+    aux : bool
+        Auxiliary loss.
+
+    Reference:
+        Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic
+        Image Segmentation."
+    """
+
+    def __init__(self, nclass, backbone='xception', aux=True, pretrained_base=True, dilated=True, **kwargs):
+        super(DeepLabV3Plus, self).__init__()
+        self.aux = aux
+        self.nclass = nclass
+        output_stride = 8 if dilated else 32
+
+        self.pretrained = get_xception(pretrained=pretrained_base, output_stride=output_stride, **kwargs)
+
+        # deeplabv3 plus
+        self.head = _DeepLabHead(nclass, **kwargs)
+        if aux:
+            self.auxlayer = _FCNHead(728, nclass, **kwargs)
+
+    def base_forward(self, x):
+        # Entry flow
+        x = self.pretrained.conv1(x)
+        x = self.pretrained.bn1(x)
+        x = self.pretrained.relu(x)
+
+        x = self.pretrained.conv2(x)
+        x = self.pretrained.bn2(x)
+        x = self.pretrained.relu(x)
+
+        x = self.pretrained.block1(x)
+        # add relu here
+        x = self.pretrained.relu(x)
+        low_level_feat = x
+
+        x = self.pretrained.block2(x)
+        x = self.pretrained.block3(x)
+
+        # Middle flow
+        x = self.pretrained.midflow(x)
+        mid_level_feat = x
+
+        # Exit flow
+        x = self.pretrained.block20(x)
+        x = self.pretrained.relu(x)
+        x = self.pretrained.conv3(x)
+        x = self.pretrained.bn3(x)
+        x = self.pretrained.relu(x)
+
+        x = self.pretrained.conv4(x)
+        x = self.pretrained.bn4(x)
+        x = self.pretrained.relu(x)
+
+        x = self.pretrained.conv5(x)
+        x = self.pretrained.bn5(x)
+        x = self.pretrained.relu(x)
+        return low_level_feat, mid_level_feat, x
+
+    def forward(self, x):
+#        print("x size: ", x.size())
+        size = x.size()[2:]
+        c1, c3, c4 = self.base_forward(x)
+#        print("c1 size: ", c1.size())
+#        print("c4 size: ", c4.size())
+        outputs = list()
+        x = self.head(c4, c1)
+        x = F.interpolate(x, size, mode='bilinear', align_corners=True)
+        outputs.append(x)
+        if self.aux:
+            auxout = self.auxlayer(c3)
+            auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
+            outputs.append(auxout)
+
+        # for save onnx
+        # y = torch.max(x, 1)[1].to(torch.float32)
+        # return y
+
+        return tuple(outputs)
+
+
+class _DeepLabHead(nn.Module):
+    def __init__(self, nclass, c1_channels=128, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_DeepLabHead, self).__init__()
+        self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, **kwargs)
+        self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer)
+        self.block = nn.Sequential(
+            _ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer),
+            nn.Dropout(0.5),
+            _ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer),
+            nn.Dropout(0.1),
+            nn.Conv2d(256, nclass, 1))
+
+    def forward(self, x, c1):
+        size = c1.size()[2:]
+        c1 = self.c1_block(c1)
+#        print("c1", c1.size())
+#        print("before aspp: ", x.size())
+        x = self.aspp(x)
+#        print("after aspp: ", x.size())
+        x = F.interpolate(x, size, mode='bilinear', align_corners=True)
+        return self.block(torch.cat([x, c1], dim=1))
+
+
+def get_deeplabv3_plus(dataset='pascal_voc', backbone='xception', pretrained=False, root='../ckpt',
+                       pretrained_base=False, nclass=3, **kwargs):
+    acronyms = {
+        'pascal_voc': 'pascal_voc',
+        'pascal_aug': 'pascal_aug',
+        'ade20k': 'ade',
+        'coco': 'coco',
+        'citys': 'citys',
+    }
+    #from light.data import datasets
+    
+    model = DeepLabV3Plus(nclass, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
+    if pretrained:
+        pass
+        # if dataset not in acronyms.keys():
+        #     print("root:", root)
+        #     model_path = os.path.join(root, "deeplabv3_plus_28.pth")
+        #     model.load_state_dict(torch.load(model_path), strict=False)
+        # else:
+        #     from .model_store import get_model_file
+        #     device = torch.device(kwargs['local_rank'])
+        #     model.load_state_dict(
+        #     torch.load(get_model_file('deeplabv3_plus_%s_%s' % (backbone, acronyms[dataset]), root=root),
+        #         map_location=device))
+    return model
+
+
+def get_deeplabv3_plus_xception_voc(**kwargs):
+    return get_deeplabv3_plus('pascal_voc', 'xception', **kwargs)
+
+
+if __name__ == '__main__':
+    model = get_deeplabv3_plus_xception_voc()

project/hair_service_sd/core/seg/networks/fcn.py

@@ -0,0 +1,221 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from core.seg.networks.vgg import vgg16
+
+__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s',
+           'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc']
+
+
+class FCN32s(nn.Module):
+    """There are some difference from original fcn"""
+
+    def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True,
+                 norm_layer=nn.BatchNorm2d, **kwargs):
+        super(FCN32s, self).__init__()
+        self.aux = aux
+        if backbone == 'vgg16':
+            self.pretrained = vgg16(pretrained=pretrained_base).features
+        else:
+            raise RuntimeError('unknown backbone: {}'.format(backbone))
+        self.head = _FCNHead(512, nclass, norm_layer)
+        if aux:
+            self.auxlayer = _FCNHead(512, nclass, norm_layer)
+
+        self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
+
+    def forward(self, x):
+        size = x.size()[2:]
+        pool5 = self.pretrained(x)
+
+        outputs = []
+        out = self.head(pool5)
+        out = F.interpolate(out, size, mode='bilinear', align_corners=True)
+        outputs.append(out)
+
+        if self.aux:
+            auxout = self.auxlayer(pool5)
+            auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
+            outputs.append(auxout)
+
+        return tuple(outputs)
+
+
+class FCN16s(nn.Module):
+    def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(FCN16s, self).__init__()
+        self.aux = aux
+        if backbone == 'vgg16':
+            self.pretrained = vgg16(pretrained=pretrained_base).features
+        else:
+            raise RuntimeError('unknown backbone: {}'.format(backbone))
+        self.pool4 = nn.Sequential(*self.pretrained[:24])
+        self.pool5 = nn.Sequential(*self.pretrained[24:])
+        self.head = _FCNHead(512, nclass, norm_layer)
+        self.score_pool4 = nn.Conv2d(512, nclass, 1)
+        if aux:
+            self.auxlayer = _FCNHead(512, nclass, norm_layer)
+
+        self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4'])
+
+    def forward(self, x):
+        pool4 = self.pool4(x)
+        pool5 = self.pool5(pool4)
+
+        outputs = []
+        score_fr = self.head(pool5)
+
+        score_pool4 = self.score_pool4(pool4)
+
+        upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
+        fuse_pool4 = upscore2 + score_pool4
+
+        out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True)
+        outputs.append(out)
+
+        if self.aux:
+            auxout = self.auxlayer(pool5)
+            auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
+            outputs.append(auxout)
+
+        return tuple(outputs)
+
+
+class FCN8s(nn.Module):
+    def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(FCN8s, self).__init__()
+        self.aux = aux
+        if backbone == 'vgg16':
+            self.pretrained = vgg16(pretrained=pretrained_base).features
+        else:
+            raise RuntimeError('unknown backbone: {}'.format(backbone))
+        self.pool3 = nn.Sequential(*self.pretrained[:17])
+        self.pool4 = nn.Sequential(*self.pretrained[17:24])
+        self.pool5 = nn.Sequential(*self.pretrained[24:])
+        self.head = _FCNHead(512, nclass, norm_layer)
+        self.score_pool3 = nn.Conv2d(256, nclass, 1)
+        self.score_pool4 = nn.Conv2d(512, nclass, 1)
+        if aux:
+            self.auxlayer = _FCNHead(512, nclass, norm_layer)
+
+        self.__setattr__('exclusive',
+                         ['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3',
+                                                                                         'score_pool4'])
+
+    def forward(self, x):
+        pool3 = self.pool3(x)
+        pool4 = self.pool4(pool3)
+        pool5 = self.pool5(pool4)
+
+        outputs = []
+        score_fr = self.head(pool5)
+
+        score_pool4 = self.score_pool4(pool4)
+        score_pool3 = self.score_pool3(pool3)
+
+        upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
+        fuse_pool4 = upscore2 + score_pool4
+
+        upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True)
+        fuse_pool3 = upscore_pool4 + score_pool3
+
+        out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True)
+        outputs.append(out)
+
+        if self.aux:
+            auxout = self.auxlayer(pool5)
+            auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
+            outputs.append(auxout)
+
+        return tuple(outputs)
+
+
+class _FCNHead(nn.Module):
+    def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(_FCNHead, self).__init__()
+        inter_channels = in_channels // 4
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
+            norm_layer(inter_channels),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv2d(inter_channels, channels, 1)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
+               pretrained_base=True, **kwargs):
+    acronyms = {
+        'pascal_voc': 'pascal_voc',
+        'pascal_aug': 'pascal_aug',
+        'ade20k': 'ade',
+        'coco': 'coco',
+        'citys': 'citys',
+    }
+    from ..data.dataloader import datasets
+    model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
+    if pretrained:
+        from .model_store import get_model_file
+        device = torch.device(kwargs['local_rank'])
+        model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root),
+                              map_location=device))
+    return model
+
+
+def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
+               pretrained_base=True, **kwargs):
+    acronyms = {
+        'pascal_voc': 'pascal_voc',
+        'pascal_aug': 'pascal_aug',
+        'ade20k': 'ade',
+        'coco': 'coco',
+        'citys': 'citys',
+    }
+    from ..data.dataloader import datasets
+    model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
+    if pretrained:
+        from .model_store import get_model_file
+        device = torch.device(kwargs['local_rank'])
+        model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root),
+                              map_location=device))
+    return model
+
+
+def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
+              pretrained_base=True, **kwargs):
+    acronyms = {
+        'pascal_voc': 'pascal_voc',
+        'pascal_aug': 'pascal_aug',
+        'ade20k': 'ade',
+        'coco': 'coco',
+        'citys': 'citys',
+    }
+    from ..data.dataloader import datasets
+    model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
+    if pretrained:
+        from .model_store import get_model_file
+        device = torch.device(kwargs['local_rank'])
+        model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root),
+                              map_location=device))
+    return model
+
+
+def get_fcn32s_vgg16_voc(**kwargs):
+    return get_fcn32s('pascal_voc', 'vgg16', **kwargs)
+
+
+def get_fcn16s_vgg16_voc(**kwargs):
+    return get_fcn16s('pascal_voc', 'vgg16', **kwargs)
+
+
+def get_fcn8s_vgg16_voc(**kwargs):
+    return get_fcn8s('pascal_voc', 'vgg16', **kwargs)
+
+
+if __name__ == '__main__':
+    model = FCN16s(21)
+    print(model)

project/hair_service_sd/core/seg/networks/jpu.py

@@ -0,0 +1,68 @@
+"""Joint Pyramid Upsampling"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['JPU']
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1,
+                 dilation=1, bias=False, norm_layer=nn.BatchNorm2d):
+        super(SeparableConv2d, self).__init__()
+        self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias)
+        self.bn = norm_layer(inplanes)
+        self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+        return x
+
+
+# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py
+class JPU(nn.Module):
+    def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs):
+        super(JPU, self).__init__()
+
+        self.conv5 = nn.Sequential(
+            nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+
+        self.dilation1 = nn.Sequential(
+            SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+        self.dilation2 = nn.Sequential(
+            SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+        self.dilation3 = nn.Sequential(
+            SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+        self.dilation4 = nn.Sequential(
+            SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False),
+            norm_layer(width),
+            nn.ReLU(True))
+
+    def forward(self, *inputs):
+        feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])]
+        size = feats[-1].size()[2:]
+        feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True)
+        feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True)
+        feat = torch.cat(feats, dim=1)
+        feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)],
+                         dim=1)
+
+        return inputs[0], inputs[1], inputs[2], feat

project/hair_service_sd/core/seg/networks/resnetv1b.py

@@ -0,0 +1,264 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['ResNetV1b', 'resnet18_v1b', 'resnet34_v1b', 'resnet50_v1b',
+           'resnet101_v1b', 'resnet152_v1b', 'resnet152_v1s', 'resnet101_v1s', 'resnet50_v1s']
+
+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',
+}
+
+
+class BasicBlockV1b(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
+                 previous_dilation=1, norm_layer=nn.BatchNorm2d):
+        super(BasicBlockV1b, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, 3, stride,
+                               dilation, dilation, bias=False)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(True)
+        self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation,
+                               dilation=previous_dilation, bias=False)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class BottleneckV1b(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
+                 previous_dilation=1, norm_layer=nn.BatchNorm2d):
+        super(BottleneckV1b, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = nn.Conv2d(planes, planes, 3, stride,
+                               dilation, dilation, bias=False)
+        self.bn2 = norm_layer(planes)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetV1b(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False,
+                 zero_init_residual=False, norm_layer=nn.BatchNorm2d):
+        self.inplanes = 128 if deep_stem else 64
+        super(ResNetV1b, self).__init__()
+        if deep_stem:
+            self.conv1 = nn.Sequential(
+                nn.Conv2d(3, 64, 3, 2, 1, bias=False),
+                norm_layer(64),
+                nn.ReLU(True),
+                nn.Conv2d(64, 64, 3, 1, 1, bias=False),
+                norm_layer(64),
+                nn.ReLU(True),
+                nn.Conv2d(64, 128, 3, 1, 1, bias=False)
+            )
+        else:
+            self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(True)
+        self.maxpool = nn.MaxPool2d(3, 2, 1)
+        self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer)
+        if dilated:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer)
+        else:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, BottleneckV1b):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlockV1b):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        if dilation in (1, 2):
+            layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+        elif dilation == 4:
+            layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+        else:
+            raise RuntimeError("=> unknown dilation size: {}".format(dilation))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+
+        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.fc(x)
+
+        return x
+
+
+def resnet18_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BasicBlockV1b, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        old_dict = model_zoo.load_url(model_urls['resnet18'])
+        model_dict = model.state_dict()
+        old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
+        model_dict.update(old_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet34_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        old_dict = model_zoo.load_url(model_urls['resnet34'])
+        model_dict = model.state_dict()
+        old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
+        model_dict.update(old_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet50_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        old_dict = model_zoo.load_url(model_urls['resnet50'])
+        model_dict = model.state_dict()
+        old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
+        model_dict.update(old_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet101_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        old_dict = model_zoo.load_url(model_urls['resnet101'])
+        model_dict = model.state_dict()
+        old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
+        model_dict.update(old_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet152_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        old_dict = model_zoo.load_url(model_urls['resnet152'])
+        model_dict = model.state_dict()
+        old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
+        model_dict.update(old_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs)
+    if pretrained:
+        from ..model_store import get_resnet_file
+        model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False)
+    return model
+
+
+def resnet101_v1s(pretrained=False, root='~/.torch/models', **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, **kwargs)
+    if pretrained:
+        from ..model_store import get_resnet_file
+        model.load_state_dict(torch.load(get_resnet_file('resnet101', root=root)), strict=False)
+    return model
+
+
+def resnet152_v1s(pretrained=False, root='~/.torch/models', **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, **kwargs)
+    if pretrained:
+        from ..model_store import get_resnet_file
+        model.load_state_dict(torch.load(get_resnet_file('resnet152', root=root)), strict=False)
+    return model
+
+
+if __name__ == '__main__':
+    import torch
+
+    img = torch.randn(4, 3, 224, 224)
+    model = resnet50_v1b(True)
+    output = model(img)

project/hair_service_sd/core/seg/networks/segbase.py

@@ -0,0 +1,60 @@
+"""Base Model for Semantic Segmentation"""
+import torch.nn as nn
+
+from core.seg.networks.jpu import JPU
+from core.seg.networks.resnetv1b import resnet50_v1s, resnet101_v1s, resnet152_v1s
+
+__all__ = ['SegBaseModel']
+
+
+class SegBaseModel(nn.Module):
+    r"""Base Model for Semantic Segmentation
+
+    Parameters
+    ----------
+    backbone : string
+        Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
+        'resnet101' or 'resnet152').
+    """
+
+    def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs):
+        super(SegBaseModel, self).__init__()
+        dilated = False if jpu else True
+        self.aux = aux
+        self.nclass = nclass
+        if backbone == 'resnet50':
+            self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        elif backbone == 'resnet101':
+            self.pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        elif backbone == 'resnet152':
+            self.pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        else:
+            raise RuntimeError('unknown backbone: {}'.format(backbone))
+
+        self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None
+
+    def base_forward(self, x):
+        """forwarding pre-trained network"""
+        x = self.pretrained.conv1(x)
+        x = self.pretrained.bn1(x)
+        x = self.pretrained.relu(x)
+        x = self.pretrained.maxpool(x)
+        c1 = self.pretrained.layer1(x)
+        c2 = self.pretrained.layer2(c1)
+        c3 = self.pretrained.layer3(c2)
+        c4 = self.pretrained.layer4(c3)
+
+        if self.jpu:
+            return self.jpu(c1, c2, c3, c4)
+        else:
+            return c1, c2, c3, c4
+
+    def evaluate(self, x):
+        """evaluating network with inputs and targets"""
+        return self.forward(x)[0]
+
+    def demo(self, x):
+        pred = self.forward(x)
+        if self.aux:
+            pred = pred[0]
+        return pred

project/hair_service_sd/core/seg/networks/vgg.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = [
+    'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
+    'vgg19_bn', 'vgg19',
+]
+
+model_urls = {
+    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
+    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
+    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
+    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
+    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
+    'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
+    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
+    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
+}
+
+
+class VGG(nn.Module):
+    def __init__(self, features, num_classes=1000, init_weights=True):
+        super(VGG, self).__init__()
+        self.features = features
+        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, num_classes)
+        )
+        if init_weights:
+            self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.constant_(m.bias, 0)
+
+
+def make_layers(cfg, batch_norm=False):
+    layers = []
+    in_channels = 3
+    for v in cfg:
+        if v == 'M':
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+        else:
+            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+            if batch_norm:
+                layers += (conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True))
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    return nn.Sequential(*layers)
+
+
+cfg = {
+    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
+    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
+    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
+    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
+}
+
+
+def vgg11(pretrained=False, **kwargs):
+    """VGG 11-layer model (configuration "A")
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['A']), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))
+    return model
+
+
+def vgg11_bn(pretrained=False, **kwargs):
+    """VGG 11-layer model (configuration "A") with batch normalization
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn']))
+    return model
+
+
+def vgg13(pretrained=False, **kwargs):
+    """VGG 13-layer model (configuration "B")
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['B']), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
+    return model
+
+
+def vgg13_bn(pretrained=False, **kwargs):
+    """VGG 13-layer model (configuration "B") with batch normalization
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))
+    return model
+
+
+def vgg16(pretrained=False, **kwargs):
+    """VGG 16-layer model (configuration "D")
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['D']), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
+    return model
+
+
+def vgg16_bn(pretrained=False, **kwargs):
+    """VGG 16-layer model (configuration "D") with batch normalization
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
+    return model
+
+
+def vgg19(pretrained=False, **kwargs):
+    """VGG 19-layer model (configuration "E")
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['E']), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
+    return model
+
+
+def vgg19_bn(pretrained=False, **kwargs):
+    """VGG 19-layer model (configuration 'E') with batch normalization
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    if pretrained:
+        kwargs['init_weights'] = False
+    model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))
+    return model
+
+
+if __name__ == '__main__':
+    img = torch.randn((4, 3, 480, 480))
+    model = vgg16(pretrained=False)
+    out = model(img)

project/hair_service_sd/core/seg/networks/xception.py

@@ -0,0 +1,411 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['Enc', 'FCAttention', 'Xception65', 'Xception71', 'get_xception', 'get_xception_71', 'get_xception_a']
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None):
+        super(SeparableConv2d, self).__init__()
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+
+        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 0, dilation, groups=in_channels,
+                               bias=bias)
+        self.bn = norm_layer(in_channels)
+        self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias)
+
+    def forward(self, x):
+        x = self.fix_padding(x, self.kernel_size, self.dilation)
+        x = self.conv1(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+
+        return x
+
+    def fix_padding(self, x, kernel_size, dilation):
+        kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+        pad_total = kernel_size_effective - 1
+        pad_beg = pad_total // 2
+        pad_end = pad_total - pad_beg
+        padded_inputs = F.pad(x, (pad_beg, pad_end, pad_beg, pad_end))
+        return padded_inputs
+
+
+class Block(nn.Module):
+    def __init__(self, in_channels, out_channels, reps, stride=1, dilation=1, norm_layer=None,
+                 start_with_relu=True, grow_first=True, is_last=False):
+        super(Block, self).__init__()
+        if out_channels != in_channels or stride != 1:
+            self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False)
+            self.skipbn = norm_layer(out_channels)
+        else:
+            self.skip = None
+        self.relu = nn.ReLU(True)
+        rep = list()
+        filters = in_channels
+        if grow_first:
+            if start_with_relu:
+                rep.append(self.relu)
+            rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
+            rep.append(norm_layer(out_channels))
+            filters = out_channels
+        for i in range(reps - 1):
+            if grow_first or start_with_relu:
+                rep.append(self.relu)
+            rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, norm_layer=norm_layer))
+            rep.append(norm_layer(filters))
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
+        if stride != 1:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(out_channels, out_channels, 3, stride, norm_layer=norm_layer))
+            rep.append(norm_layer(out_channels))
+        elif is_last:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(out_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
+            rep.append(norm_layer(out_channels))
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, x):
+        out = self.rep(x)
+        if self.skip is not None:
+            skip = self.skipbn(self.skip(x))
+        else:
+            skip = x
+        out = out + skip
+        return out
+
+
+class Xception65(nn.Module):
+    """Modified Aligned Xception
+    """
+
+    def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d):
+        super(Xception65, self).__init__()
+        if output_stride == 32:
+            entry_block3_stride = 2
+            exit_block20_stride = 2
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 1)
+        elif output_stride == 16:
+            entry_block3_stride = 2
+            exit_block20_stride = 1
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 2)
+        elif output_stride == 8:
+            entry_block3_stride = 1
+            exit_block20_stride = 1
+            middle_block_dilation = 2
+            exit_block_dilations = (2, 4)
+        else:
+            raise NotImplementedError
+        # Entry flow
+        self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False)
+        self.bn1 = norm_layer(32)
+        self.relu = nn.ReLU(True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False)
+        self.bn2 = norm_layer(64)
+
+        self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False)
+        self.block2 = Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True)
+        self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer,
+                            start_with_relu=True, grow_first=True, is_last=True)
+
+        # Middle flow
+        midflow = list()
+        for i in range(4, 20):
+            midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer,
+                                 start_with_relu=True, grow_first=True))
+        self.midflow = nn.Sequential(*midflow)
+
+        # Exit flow
+        self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0],
+                             norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True)
+        self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn3 = norm_layer(1536)
+        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn4 = norm_layer(1536)
+        self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn5 = norm_layer(2048)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(2048, num_classes)
+
+    def forward(self, x):
+        # Entry flow
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.block1(x)
+        x = self.relu(x)
+        # c1 = x
+        x = self.block2(x)
+        # c2 = x
+        x = self.block3(x)
+
+        # Middle flow
+        x = self.midflow(x)
+        # c3 = x
+
+        # Exit flow
+        x = self.block20(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+class Xception71(nn.Module):
+    """Modified Aligned Xception
+    """
+
+    def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d):
+        super(Xception71, self).__init__()
+        if output_stride == 32:
+            entry_block3_stride = 2
+            exit_block20_stride = 2
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 1)
+        elif output_stride == 16:
+            entry_block3_stride = 2
+            exit_block20_stride = 1
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 2)
+        elif output_stride == 8:
+            entry_block3_stride = 1
+            exit_block20_stride = 1
+            middle_block_dilation = 2
+            exit_block_dilations = (2, 4)
+        else:
+            raise NotImplementedError
+        # Entry flow
+        self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False)
+        self.bn1 = norm_layer(32)
+        self.relu = nn.ReLU(True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False)
+        self.bn2 = norm_layer(64)
+
+        self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False)
+        self.block2 = nn.Sequential(
+            Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True),
+            Block(256, 728, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True))
+        self.block3 = Block(728, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer,
+                            start_with_relu=True, grow_first=True, is_last=True)
+
+        # Middle flow
+        midflow = list()
+        for i in range(4, 20):
+            midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer,
+                                 start_with_relu=True, grow_first=True))
+        self.midflow = nn.Sequential(*midflow)
+
+        # Exit flow
+        self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0],
+                             norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True)
+        self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn3 = norm_layer(1536)
+        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn4 = norm_layer(1536)
+        self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
+        self.bn5 = norm_layer(2048)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(2048, num_classes)
+
+    def forward(self, x):
+        # Entry flow
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.block1(x)
+        x = self.relu(x)
+        # c1 = x
+        x = self.block2(x)
+        # c2 = x
+        x = self.block3(x)
+
+        # Middle flow
+        x = self.midflow(x)
+        # c3 = x
+
+        # Exit flow
+        x = self.block20(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+# -------------------------------------------------
+#                   For DFANet
+# -------------------------------------------------
+class BlockA(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm_layer=None, start_with_relu=True):
+        super(BlockA, self).__init__()
+        if out_channels != in_channels or stride != 1:
+            self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False)
+            self.skipbn = norm_layer(out_channels)
+        else:
+            self.skip = None
+        self.relu = nn.ReLU(True)
+        rep = list()
+        inter_channels = out_channels // 4
+
+        if start_with_relu:
+            rep.append(self.relu)
+        rep.append(SeparableConv2d(in_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer))
+        rep.append(norm_layer(inter_channels))
+
+        rep.append(self.relu)
+        rep.append(SeparableConv2d(inter_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer))
+        rep.append(norm_layer(inter_channels))
+
+        if stride != 1:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inter_channels, out_channels, 3, stride, norm_layer=norm_layer))
+            rep.append(norm_layer(out_channels))
+        else:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inter_channels, out_channels, 3, 1, norm_layer=norm_layer))
+            rep.append(norm_layer(out_channels))
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, x):
+        out = self.rep(x)
+        if self.skip is not None:
+            skip = self.skipbn(self.skip(x))
+        else:
+            skip = x
+        out = out + skip
+        return out
+
+
+class Enc(nn.Module):
+    def __init__(self, in_channels, out_channels, blocks, norm_layer=None):
+        super(Enc, self).__init__()
+        block = list()
+        block.append(BlockA(in_channels, out_channels, 2, norm_layer=norm_layer))
+        for i in range(blocks - 1):
+            block.append(BlockA(out_channels, out_channels, 1, norm_layer=norm_layer))
+        self.block = nn.Sequential(*block)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class FCAttention(nn.Module):
+    def __init__(self, in_channels, norm_layer=None):
+        super(FCAttention, self).__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(in_channels, 1000)
+        self.conv = nn.Sequential(
+            nn.Conv2d(1000, in_channels, 1, bias=False),
+            norm_layer(in_channels),
+            nn.ReLU(True))
+
+    def forward(self, x):
+        n, c, _, _ = x.size()
+        att = self.avgpool(x).view(n, c)
+        att = self.fc(att).view(n, 1000, 1, 1)
+        att = self.conv(att)
+        return x * att.expand_as(x)
+
+
+class XceptionA(nn.Module):
+    def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d):
+        super(XceptionA, self).__init__()
+        self.conv1 = nn.Sequential(nn.Conv2d(3, 8, 3, 2, 1, bias=False),
+                                   norm_layer(8),
+                                   nn.ReLU(True))
+
+        self.enc2 = Enc(8, 48, 4, norm_layer=norm_layer)
+        self.enc3 = Enc(48, 96, 6, norm_layer=norm_layer)
+        self.enc4 = Enc(96, 192, 4, norm_layer=norm_layer)
+
+        self.fca = FCAttention(192, norm_layer=norm_layer)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(192, num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        x = self.enc2(x)
+        x = self.enc3(x)
+        x = self.enc4(x)
+        x = self.fca(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+# Constructor
+def get_xception(pretrained=False, root='~/.torch/models', **kwargs):
+    model = Xception65(**kwargs)
+    if pretrained:
+        from ..model_store import get_model_file
+        model.load_state_dict(torch.load(get_model_file('xception', root=root)))
+    return model
+
+
+def get_xception_71(pretrained=False, root='~/.torch/models', **kwargs):
+    model = Xception71(**kwargs)
+    if pretrained:
+        from ..model_store import get_model_file
+        model.load_state_dict(torch.load(get_model_file('xception71', root=root)))
+    return model
+
+
+def get_xception_a(pretrained=False, root='~/.torch/models', **kwargs):
+    model = XceptionA(**kwargs)
+    if pretrained:
+        from ..model_store import get_model_file
+        model.load_state_dict(torch.load(get_model_file('xception_a', root=root)))
+    return model
+
+
+if __name__ == '__main__':
+    model = get_xception_a()

project/hair_service_sd/core/seg/setup.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# @File    : setup.py
+# @Time    : 2020/1/15
+# @Author  : yangchaojie (yangchaojie@immomo.com)
+
+import os
+import sys
+import shutil
+import numpy
+import tempfile
+
+from setuptools import setup
+from setuptools.extension import Extension
+
+from Cython.Build import cythonize
+from Cython.Distutils import build_ext
+
+import platform
+
+
+def get_root_path(root):
+    if os.path.dirname(root) in ['', '.']:
+        return os.path.basename(root)
+    else:
+        return get_root_path(os.path.dirname(root))
+
+
+def copy_file(src, dest):
+    if os.path.exists(dest):
+        return
+
+    if not os.path.exists(os.path.dirname(dest)):
+        os.makedirs(os.path.dirname(dest))
+    if os.path.isdir(src):
+        shutil.copytree(src, dest)
+    else:
+        shutil.copyfile(src, dest)
+
+
+def touch_init_file():
+    init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
+    with open(init_file_name, 'w'):
+        pass
+    return init_file_name
+
+
+
+
+def compose_extensions(root='.'):
+    for file_ in os.listdir(root):
+        abs_file = os.path.join(root, file_)
+
+        if os.path.isfile(abs_file):
+            if abs_file.endswith('.py'):
+                extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
+            elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
+                continue
+            else:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            if abs_file.endswith('__init__.py'):
+                copy_file(init_file, os.path.join(build_root_dir, abs_file))
+
+        else:
+            if os.path.basename(abs_file) in ignore_folders :
+                continue
+            if os.path.basename(abs_file) in conf_folders:
+                copy_file(abs_file, os.path.join(build_root_dir, abs_file))
+            compose_extensions(abs_file)
+
+
+
+# if __name__ == '__main__':
+build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
+    sys.version_info.major) + '.' + str(sys.version_info.minor)
+
+print(build_root_dir)
+
+extensions = []
+ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
+conf_folders = ['conf']
+
+
+init_file = touch_init_file()
+print(init_file)
+
+
+compose_extensions()
+os.remove(init_file)
+
+setup(
+    name='moxie_hairstyle',
+    version='1.0',
+    ext_modules=cythonize(
+        extensions,
+        nthreads=16,
+        compiler_directives=dict(always_allow_keywords=True),
+        include_path=[numpy.get_include()]),
+    cmdclass=dict(build_ext=build_ext))
+
+    # python setup.py build_ext

project/hair_service_sd/core/simsun.ttc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a368113c36d516aac1a825acc4bbbc9906c4d197f7c649e7fe0f6f31e7475dd
+size 10500792

project/hair_service_sd/core/utils/MomocvFaceAlignment1K.py

@@ -0,0 +1,462 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import numpy as np
+import os
+from core.utils import landmark_processor
+from algorithm_conf import ConfFactory
+from core.utils.umeyama import umeyama
+
+import cv2
+
+__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)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 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):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        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 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        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)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        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)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        if is_1k:
+            self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
+        else:
+            self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
+            self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
+        self.is_1k = is_1k
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ 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)
+
+        if self.is_1k:
+            key = self.avgpool(x)
+            key = key.view(key.size(0), -1)
+            key = self.fc(key)
+            return key
+        else:
+            key, ctrl = torch.chunk(x, 2, dim=1)
+
+            key = self.avgpool(key)
+            key = key.view(key.size(0), -1)
+            key = self.fc_key(key)
+
+            ctrl = self.avgpool(ctrl)
+            ctrl = ctrl.view(ctrl.size(0), -1)
+            ctrl = self.fc_ctrl(ctrl)
+
+            return key, ctrl
+
+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']), strict=False)
+    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']))
+    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']))
+    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']))
+    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']))
+    return model
+
+class Model1k(nn.Module):
+    def __init__(self, gpu_id=None):
+        super(Model1k, self).__init__()
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True)
+
+        self.model_dir = ConfFactory.getModelValue("model_dir")
+        weights = torch.load(os.path.join(self.model_dir, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
+
+        self.load_state_dict(weights)
+        self.to(self.device)
+        self.eval()
+
+    def forward(self, imgs):
+        pred_key_pts = self.face_alignment_net(imgs)
+        pred_key_pts = pred_key_pts + 0.5
+        return pred_key_pts
+
+class MomocvFaceAlignment1K(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = Model1k(gpu_id)
+
+        self.trackingFaceRects = []
+
+        print('MomocvFaceAlignment1K success')
+
+    def forward(self, img_tensor):
+        fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy()
+        return fullyconnected1
+
+    def detect(self, img, landmarks):
+        dst_size = 256
+        landmarks_res = []
+        with torch.no_grad():
+            input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
+            all_mat = []
+            for ix, landmark in enumerate(landmarks):
+                M = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
+                all_mat.append(M)
+                tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
+
+                # cv2.imshow('inp', tmp)
+                # cv2.waitKey()
+
+                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            for ix, pts in enumerate(fullyconnected1):
+                orig_pts = (np.reshape(pts, (2, 1000)).transpose((1, 0)) * dst_size)
+                orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
+                landmarks_res.append(orig_pts)
+            return landmarks_res
+
+    def detect_single_face(self, img, crop_M):
+        dst_size = 256
+        with torch.no_grad():
+
+            #  add for change crop for 1024
+            h, w, _ = img.shape
+            if h == 768:
+                crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :]
+                tmp = cv2.resize(crop_img, (dst_size, dst_size))
+
+            else:
+                tmp = cv2.warpAffine(img, crop_M, (dst_size, dst_size), flags=cv2.INTER_CUBIC)
+
+                # cv2.imshow("img_paf_test_crop: ", tmp)
+                # cv2.waitKey()
+                # crop_img = img[220:img.shape[0] - 220, 266:img.shape[1] - 266, :]   # 220 266
+                # tmp = cv2.resize(crop_img, (dst_size, dst_size))
+
+            # cv2.imshow("detect face: h:{:d}".format(h), tmp)
+            # cv2.waitKey()
+
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+
+            orig_pts = np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0))
+
+            if h == 768:
+                orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1]
+                orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0]
+                orig_pts[:, 0] += 104
+                orig_pts[:, 1] += 104
+            else:
+                orig_pts[:, 0] = orig_pts[:, 0] * dst_size
+                orig_pts[:, 1] = orig_pts[:, 1] * dst_size
+                orig_pts = landmark_processor.transform_points(orig_pts, crop_M, invert=True)
+                # orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1]
+                # orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0]
+                # orig_pts[:, 0] += 266
+                # orig_pts[:, 1] += 220
+
+            return orig_pts
+
+    def detect_single_face_old(self, img):
+        dst_size = 256
+        with torch.no_grad():
+            tmp = cv2.resize(img, (dst_size, dst_size))
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * img.shape[0])
+            return orig_pts
+
+    def detect_according_5pts(self, img, pts5):
+        dst_size = 256
+        with torch.no_grad():
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            eye_dis = 0.34
+            mouth_dis = 0.34
+            g_Average_5point_180 = np.array([
+                eye_dis, 0.3,
+                1 - eye_dis, 0.3,
+                0.5, 0.6,
+                mouth_dis, 0.63,
+                1 - mouth_dis, 0.63
+            ])
+            # print(g_Average_5point_180)
+            left_eye = np.array([pts5[0], pts5[5]])
+            right_eye = np.array([pts5[1], pts5[6]])
+            nose = np.array([pts5[2], pts5[7]])
+            left_mouth = np.array([pts5[3], pts5[8]])
+            right_mouth = np.array([pts5[4], pts5[9]])
+
+            pts5_src = np.vstack((left_eye, right_eye,
+                                  nose,
+                                  left_mouth, right_mouth))
+            pts5_src = np.array(pts5_src).astype(np.int32)
+            pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size
+
+            mat = umeyama(pts5_src, pts5_dst, True)[0:2]
+
+            tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
+            # cv2.imshow("tmp", tmp)
+            # cv2.waitKey()
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size)
+            orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True)
+        return orig_pts
+
+
+    def stable_forward(self, image, detected_faces, reset=False):
+        if reset is True:
+            self.trackingFaceRects = []
+
+        if len(self.trackingFaceRects) == 0:
+            for face_rect in detected_faces:
+                    new_tracking_rect = [face_rect, True, [0, 0], 0, None]
+                    self.trackingFaceRects.append(new_tracking_rect)
+
+        with torch.no_grad():
+            landmarks = []
+            for ix, tracking_face_rect in enumerate(self.trackingFaceRects):
+                if tracking_face_rect[1] == True:
+                    d = tracking_face_rect[0]
+                    src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
+                    rotate_degree = tracking_face_rect[3]
+                    scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1])
+                    dst_center = np.array([0.5, 0.5]) * 256
+                    offset = dst_center - src_center
+                    M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
+                    M[:, 2] += offset
+                else:
+                    rotate_degree = 0
+                    M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256)
+                inp = cv2.warpAffine(image, M, (256, 256))
+
+                # cv2.imshow('inp_{}'.format(ix), inp)
+                # cv2.waitKey()
+                orig_inp = inp
+
+                inp = inp.transpose((2, 0, 1)).astype(np.float32)
+                inp = inp[np.newaxis, :, :, :] / 255
+
+                in_tensor = torch.from_numpy(inp)
+                in_tensor = in_tensor.cuda(0)
+                fullyconnected1 = self.forward(in_tensor)
+                fullyconnected1 = fullyconnected1[0]
+                orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256
+
+                t2 = cv2.getTickCount()
+                orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
+
+                # orig_pts = orig_pts.transpose((1, 0))
+                fullyconnected1 = orig_pts
+
+                # update tracking infos
+                tracking_face_rect[1] = False
+                tracking_face_rect[2] = None
+                tracking_face_rect[3] = rotate_degree
+                tracking_face_rect[4] = fullyconnected1
+
+                # fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1)
+
+                # eye_landmark = self.detect_eye(image, fullyconnected1)
+                # fullyconnected1[87:104] = eye_landmark[0]
+                # fullyconnected1[104:121] = eye_landmark[1]
+
+                landmarks.append(fullyconnected1)
+        return landmarks

project/hair_service_sd/core/utils/box_utils_Retina.py

@@ -0,0 +1,330 @@
+import torch
+import numpy as np
+
+
+def point_form(boxes):
+    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
+    representation for comparison to point form ground truth data.
+    Args:
+        boxes: (tensor) center-size default boxes from priorbox layers.
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin
+                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax
+
+
+def center_size(boxes):
+    """ Convert prior_boxes to (cx, cy, w, h)
+    representation for comparison to center-size form ground truth data.
+    Args:
+        boxes: (tensor) point_form boxes
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, 2:] + boxes[:, :2])/2,  # cx, cy
+                     boxes[:, 2:] - boxes[:, :2], 1)  # w, h
+
+
+def intersect(box_a, box_b):
+    """ We resize both tensors to [A,B,2] without new malloc:
+    [A,2] -> [A,1,2] -> [A,B,2]
+    [B,2] -> [1,B,2] -> [A,B,2]
+    Then we compute the area of intersect between box_a and box_b.
+    Args:
+      box_a: (tensor) bounding boxes, Shape: [A,4].
+      box_b: (tensor) bounding boxes, Shape: [B,4].
+    Return:
+      (tensor) intersection area, Shape: [A,B].
+    """
+    A = box_a.size(0)
+    B = box_b.size(0)
+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+    inter = torch.clamp((max_xy - min_xy), min=0)
+    return inter[:, :, 0] * inter[:, :, 1]
+
+
+def jaccard(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.  Here we operate on
+    ground truth boxes and default boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
+    Return:
+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
+    """
+    inter = intersect(box_a, box_b)
+    area_a = ((box_a[:, 2]-box_a[:, 0]) *
+              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+    area_b = ((box_b[:, 2]-box_b[:, 0]) *
+              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+def matrix_iou(a, b):
+    """
+    return iou of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
+    return area_i / (area_a[:, np.newaxis] + area_b - area_i)
+
+
+def matrix_iof(a, b):
+    """
+    return iof of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    return area_i / np.maximum(area_a[:, np.newaxis], 1)
+
+
+def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):
+    """Match each prior box with the ground truth box of the highest jaccard
+    overlap, encode the bounding boxes, then return the matched indices
+    corresponding to both confidence and location preds.
+    Args:
+        threshold: (float) The overlap threshold used when mathing boxes.
+        truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
+        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
+        variances: (tensor) Variances corresponding to each prior coord,
+            Shape: [num_priors, 4].
+        labels: (tensor) All the class labels for the image, Shape: [num_obj].
+        landms: (tensor) Ground truth landms, Shape [num_obj, 10].
+        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
+        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
+        landm_t: (tensor) Tensor to be filled w/ endcoded landm targets.
+        idx: (int) current batch index
+    Return:
+        The matched indices corresponding to 1)location 2)confidence 3)landm preds.
+    """
+    # jaccard index
+    overlaps = jaccard(
+        truths,
+        point_form(priors)
+    )
+    # (Bipartite Matching)
+    # [1,num_objects] best prior for each ground truth
+    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
+
+    # ignore hard gt
+    valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
+    best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
+    if best_prior_idx_filter.shape[0] <= 0:
+        loc_t[idx] = 0
+        conf_t[idx] = 0
+        return
+
+    # [1,num_priors] best ground truth for each prior
+    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
+    best_truth_idx.squeeze_(0)
+    best_truth_overlap.squeeze_(0)
+    best_prior_idx.squeeze_(1)
+    best_prior_idx_filter.squeeze_(1)
+    best_prior_overlap.squeeze_(1)
+    best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior
+    # TODO refactor: index  best_prior_idx with long tensor
+    # ensure every gt matches with its prior of max overlap
+    for j in range(best_prior_idx.size(0)):     # 判别此anchor是预测哪一个boxes
+        best_truth_idx[best_prior_idx[j]] = j
+    matches = truths[best_truth_idx]            # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
+    conf = labels[best_truth_idx]               # Shape: [num_priors]      此处为每一个anchor对应的label取出来
+    conf[best_truth_overlap < threshold] = 0    # label as background   overlap<0.35的全部作为负样本
+    loc = encode(matches, priors, variances)
+
+    matches_landm = landms[best_truth_idx]
+    landm = encode_landm(matches_landm, priors, variances)
+    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
+    conf_t[idx] = conf  # [num_priors] top class label for each prior
+    landm_t[idx] = landm
+
+
+def encode(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = torch.log(g_wh) / variances[1]
+    # return target for smooth_l1_loss
+    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+def encode_landm(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 10].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded landm (tensor), Shape: [num_priors, 10]
+    """
+
+    # dist b/t match center and prior's center
+    matched = torch.reshape(matched, (matched.size(0), 5, 2))
+    priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2)
+    g_cxcy = matched[:, :, :2] - priors[:, :, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, :, 2:])
+    # g_cxcy /= priors[:, :, 2:]
+    g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)
+    # return target for smooth_l1_loss
+    return g_cxcy
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = torch.cat((
+        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+def decode_landm(pre, priors, variances):
+    """Decode landm from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): landm predictions for loc layers,
+            Shape: [num_priors,10]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded landm predictions
+    """
+    landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+                        ), dim=1)
+    return landms
+
+
+def log_sum_exp(x):
+    """Utility function for computing log_sum_exp while determining
+    This will be used to determine unaveraged confidence loss across
+    all examples in a batch.
+    Args:
+        x (Variable(tensor)): conf_preds from conf layers
+    """
+    x_max = x.data.max()
+    return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max
+
+
+# Original author: Francisco Massa:
+# https://github.com/fmassa/object-detection.torch
+# Ported to PyTorch by Max deGroot (02/01/2017)
+def nms(boxes, scores, overlap=0.5, top_k=200):
+    """Apply non-maximum suppression at test time to avoid detecting too many
+    overlapping bounding boxes for a given object.
+    Args:
+        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
+        scores: (tensor) The class predscores for the img, Shape:[num_priors].
+        overlap: (float) The overlap thresh for suppressing unnecessary boxes.
+        top_k: (int) The Maximum number of box preds to consider.
+    Return:
+        The indices of the kept boxes with respect to num_priors.
+    """
+
+    keep = torch.Tensor(scores.size(0)).fill_(0).long()
+    if boxes.numel() == 0:
+        return keep
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    area = torch.mul(x2 - x1, y2 - y1)
+    v, idx = scores.sort(0)  # sort in ascending order
+    # I = I[v >= 0.01]
+    idx = idx[-top_k:]  # indices of the top-k largest vals
+    xx1 = boxes.new()
+    yy1 = boxes.new()
+    xx2 = boxes.new()
+    yy2 = boxes.new()
+    w = boxes.new()
+    h = boxes.new()
+
+    # keep = torch.Tensor()
+    count = 0
+    while idx.numel() > 0:
+        i = idx[-1]  # index of current largest val
+        # keep.append(i)
+        keep[count] = i
+        count += 1
+        if idx.size(0) == 1:
+            break
+        idx = idx[:-1]  # remove kept element from view
+        # load bboxes of next highest vals
+        torch.index_select(x1, 0, idx, out=xx1)
+        torch.index_select(y1, 0, idx, out=yy1)
+        torch.index_select(x2, 0, idx, out=xx2)
+        torch.index_select(y2, 0, idx, out=yy2)
+        # store element-wise max with next highest score
+        xx1 = torch.clamp(xx1, min=x1[i])
+        yy1 = torch.clamp(yy1, min=y1[i])
+        xx2 = torch.clamp(xx2, max=x2[i])
+        yy2 = torch.clamp(yy2, max=y2[i])
+        w.resize_as_(xx2)
+        h.resize_as_(yy2)
+        w = xx2 - xx1
+        h = yy2 - yy1
+        # check sizes of xx1 and xx2.. after each iteration
+        w = torch.clamp(w, min=0.0)
+        h = torch.clamp(h, min=0.0)
+        inter = w*h
+        # IoU = i / (area(a) + area(b) - i)
+        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
+        union = (rem_areas - inter) + area[i]
+        IoU = inter/union  # store result in iou
+        # keep only elements with an IoU <= overlap
+        idx = idx[IoU.le(overlap)]
+    return keep, count
+
+

project/hair_service_sd/core/utils/model_io.py

@@ -0,0 +1,25 @@
+import os
+import torch
+
+def load_model_by_path(model_save_path, model, gpu_id = None):
+    if not os.path.exists(model_save_path): return
+    loc = 'cpu' if gpu_id is None else 'cuda:{}'.format(gpu_id)
+    pretrained_dict = torch.load(model_save_path, map_location=loc)
+    if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict']
+    model_dict = model.state_dict()
+    # pretrained_dict.pop('netG.model.1.weight', '404')
+    pretrained_dict_new = {k: v for k, v in pretrained_dict.items()
+                       if (k in model_dict and model_dict[k].data.shape == v.data.shape)}
+    # if 'netG.model.1.weight' not in pretrained_dict_new and 'netG.model.1.weight' in model_dict:
+    #     pretrained_dict_new['netG.model.1.weight'] = model_dict['netG.model.1.weight']
+    #     pretrained_dict_new['netG.model.1.weight'][:,:12] = pretrained_dict['netG.model.1.weight']
+
+    model_dict.update(pretrained_dict_new)
+    model.load_state_dict(model_dict)
+    print("load model: ", model_save_path)
+
+def save_model_by_path(model_save_path, model):
+    save_dir, _ = os.path.split(model_save_path)
+    if not os.path.isdir(save_dir): os.makedirs(save_dir)
+    model_dic={k.replace('.module', ''):v for k,v in model.state_dict().items()}
+    torch.save(model_dic, model_save_path)

project/hair_service_sd/core/utils/params_3ddfa.py

@@ -0,0 +1,8 @@
+import cv2
+import os
+
+SCALE_F = 1e4
+SCALE_ROTATE = 1e2
+SCALE_OFFSET = 1e-1
+SCALE_SHAPE = 1e-6
+SCALE_EXP = 1

project/hair_service_sd/core/utils/umeyama.py

@@ -0,0 +1,78 @@
+import numpy as np
+
+def umeyama(src, dst, estimate_scale):
+    """Estimate N-D similarity transformation with or without scaling.
+    Parameters
+    ----------
+    src : (M, N) array
+        Source coordinates.
+    dst : (M, N) array
+        Destination coordinates.
+    estimate_scale : bool
+        Whether to estimate scaling factor.
+    Returns
+    -------
+    T : (N + 1, N + 1)
+        The homogeneous similarity transformation matrix. The matrix contains
+        NaN values only if the problem is not well-conditioned.
+    References
+    ----------
+    .. [1] "Least-squares estimation of transformation parameters between two
+            point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573
+    """
+
+    num = src.shape[0]
+    dim = src.shape[1]
+
+    # Compute mean of src and dst.
+    src_mean = src.mean(axis=0)
+    dst_mean = dst.mean(axis=0)
+
+    # Subtract mean from src and dst.
+    src_demean = src - src_mean
+    dst_demean = dst - dst_mean
+
+    # Eq. (38).
+    A = np.dot(dst_demean.T, src_demean) / num
+
+    # Eq. (39).
+    d = np.ones((dim,), dtype=np.double)
+    if np.linalg.det(A) < 0:
+        d[dim - 1] = -1
+
+    T = np.eye(dim + 1, dtype=np.double)
+    U, S, V = np.linalg.svd(A)
+    # print('A',A)
+    # print('V',V)
+    # print('U', U)
+    # print('S', S)
+    # print('T', T)
+
+    # Eq. (40) and (43).
+    rank = np.linalg.matrix_rank(A)
+    # print('rank', rank)
+    if rank == 0:
+        return np.nan * T
+    elif rank == dim - 1:
+        if np.linalg.det(U) * np.linalg.det(V) > 0:
+            T[:dim, :dim] = np.dot(U, V)
+        else:
+            s = d[dim - 1]
+            d[dim - 1] = -1
+            T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V))
+            d[dim - 1] = s
+    else:
+        T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T))
+
+    if estimate_scale:
+        # Eq. (41) and (42).
+        scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d)
+    else:
+        scale = 1.0
+    # print('scale', scale)
+    # print('dst_maen',dst_mean)
+    # print('src_mean', src_mean)
+    T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T)
+    T[:dim, :dim] *= scale
+
+    return T

project/hair_service_sd/core/utils/util.py

@@ -0,0 +1,287 @@
+import os
+import cv2
+import torch
+import logging
+import numpy as np
+# from utils.config import CONFIG
+# import torch.distributed as dist
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            os.makedirs(path)
+    else:
+        os.makedirs(paths)
+
+def make_dir(target_dir):
+    """
+    Create dir if not exists
+    """
+    if not os.path.exists(target_dir):
+        os.makedirs(target_dir)
+
+
+def print_network(model, name):
+    """
+    Print out the network information
+    """
+    logger = logging.getLogger("Logger")
+    num_params = 0
+    for p in model.parameters():
+        num_params += p.numel()
+
+    logger.info(model)
+    logger.info(name)
+    logger.info("Number of parameters: {}".format(num_params))
+
+
+def update_lr(lr, optimizer):
+    """
+    update learning rates
+    """
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def warmup_lr(init_lr, step, iter_num):
+    """
+    Warm up learning rate
+    """
+    return step/iter_num*init_lr
+
+
+def add_prefix_state_dict(state_dict, prefix="module"):
+    """
+    add prefix from the key of pretrained state dict for Data-Parallel
+    """
+    new_state_dict = {}
+    first_state_name = list(state_dict.keys())[0]
+    if not first_state_name.startswith(prefix):
+        for key, value in state_dict.items():
+            new_state_dict[prefix+"."+key] = state_dict[key].float()
+    else:
+        for key, value in state_dict.items():
+            new_state_dict[key] = state_dict[key].float()
+    return new_state_dict
+
+
+def remove_prefix_state_dict(state_dict, prefix="module"):
+    """
+    remove prefix from the key of pretrained state dict for Data-Parallel
+    """
+    new_state_dict = {}
+    first_state_name = list(state_dict.keys())[0]
+    if not first_state_name.startswith(prefix):
+        for key, value in state_dict.items():
+            new_state_dict[key] = state_dict[key].float()
+    else:
+        for key, value in state_dict.items():
+            new_state_dict[key[len(prefix)+1:]] = state_dict[key].float()
+    return new_state_dict
+
+#
+# def load_imagenet_pretrain(model, checkpoint_file):
+#     """
+#     Load imagenet pretrained resnet
+#     Add zeros channel to the first convolution layer
+#     Since we have the spectral normalization, we need to do a little more
+#     """
+#     checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu))
+#     state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
+#     for key, value in state_dict.items():
+#         state_dict[key] = state_dict[key].float()
+#
+#     logger = logging.getLogger("Logger")
+#     logger.debug("Imagenet pretrained keys:")
+#     logger.debug(state_dict.keys())
+#     logger.debug("Generator keys:")
+#     logger.debug(model.module.encoder.state_dict().keys())
+#     logger.debug("Intersection  keys:")
+#     logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys()))
+#
+#     weight_u = state_dict["conv1.module.weight_u"]
+#     weight_v = state_dict["conv1.module.weight_v"]
+#     weight_bar = state_dict["conv1.module.weight_bar"]
+#
+#     logger.debug("weight_v: {}".format(weight_v))
+#     logger.debug("weight_bar: {}".format(weight_bar.view(32, -1)))
+#     logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v))))
+#
+#     new_weight_v = torch.zeros(6, 3, 3).cuda()
+#     new_weight_bar = torch.zeros(32, 6, 3, 3).cuda()
+#
+#     new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3))
+#     new_weight_bar[:, :3, :, :].copy_(weight_bar)
+#
+#     logger.debug("new weight_v: {}".format(new_weight_v.view(-1)))
+#     logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1)))
+#     logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1)))))
+#
+#     state_dict["conv1.module.weight_v"] = new_weight_v.view(-1)
+#     state_dict["conv1.module.weight_bar"] = new_weight_bar
+#
+#     model.module.encoder.load_state_dict(state_dict, strict=False)
+
+
+def load_VGG_pretrain(model, checkpoint_file):
+    """
+    Load imagenet pretrained resnet
+    Add zeros channel to the first convolution layer
+    Since we have the spectral normalization, we need to do a little more
+    """
+    checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda())
+    backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
+
+    model.module.encoder.load_state_dict(backbone_state_dict, strict=False)
+
+
+def get_unknown_tensor(trimap):
+    """
+    get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor
+    """
+    # if CONFIG.model.trimap_channel == 3:
+    weight = trimap[:, 1:2, :, :].float()
+    # else:
+    #     weight = trimap.eq(1).float()
+    return weight
+
+
+def get_gaborfilter(angles):
+    """
+    generate gabor filter as the conv kernel
+    :param angles: number of different angles
+    """
+    gabor_filter = []
+    for angle in range(angles):
+        gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5))
+    gabor_filter = np.array(gabor_filter)
+    gabor_filter = np.expand_dims(gabor_filter, axis=1)
+    return gabor_filter.astype(np.float32)
+
+
+def get_gradfilter():
+    """
+    generate gradient filter as the conv kernel
+    """
+    grad_filter = []
+    grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
+    grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
+    grad_filter = np.array(grad_filter)
+    grad_filter = np.expand_dims(grad_filter, axis=1)
+    return grad_filter.astype(np.float32)
+
+
+# def reduce_tensor_dict(tensor_dict, mode='mean'):
+#     """
+#     average tensor dict over different GPUs
+#     """
+#     for key, tensor in tensor_dict.items():
+#         if tensor is not None:
+#             tensor_dict[key] = reduce_tensor(tensor, mode)
+#     return tensor_dict
+#
+#
+# def reduce_tensor(tensor, mode='mean'):
+#     """
+#     average tensor over different GPUs
+#     """
+#     rt = tensor.clone()
+#     dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+#     if mode == 'mean':
+#         rt /= CONFIG.world_size
+#     elif mode == 'sum':
+#         pass
+#     else:
+#         raise NotImplementedError("reduce mode can only be 'mean' or 'sum'")
+#     return rt
+
+def make_color_wheel():
+    # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros([ncols, 3])
+    col = 0
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
+    col += RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
+    colorwheel[col:col+YG, 1] = 255
+    col += YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
+    col += GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
+    colorwheel[col:col+CB, 2] = 255
+    col += CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
+    col += + BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+
+COLORWHEEL = make_color_wheel()
+
+def compute_color(u,v):
+    # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    h, w = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+    colorwheel = COLORWHEEL
+    # colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+    rad = np.sqrt(u**2+v**2)
+    a = np.arctan2(-v, -u) / np.pi
+    fk = (a+1) / 2 * (ncols - 1) + 1
+    k0 = np.floor(fk).astype(int)
+    k1 = k0 + 1
+    k1[k1 == ncols+1] = 1
+    f = fk - k0
+    for i in range(np.size(colorwheel,1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0-1] / 255
+        col1 = tmp[k1-1] / 255
+        col = (1-f) * col0 + f * col1
+        idx = rad <= 1
+        col[idx] = 1-rad[idx]*(1-col[idx])
+        notidx = np.logical_not(idx)
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
+    return img
+
+def flow_to_image(flow):
+    # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    maxrad = -1
+    u = flow[0, :, :]
+    v = flow[1, :, :]
+    rad = np.sqrt(u ** 2 + v ** 2)
+    maxrad = max(maxrad, np.max(rad))
+    u = u/(maxrad + np.finfo(float).eps)
+    v = v/(maxrad + np.finfo(float).eps)
+    img = compute_color(u, v)
+
+    return img
+
+
+if __name__ == "__main__":
+    import networks
+    logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
+                        datefmt='%m-%d %H:%M:%S')
+    G = networks.get_generator().cuda()
+    # load_imagenet_pretrain(G, CONFIG.model.imagenet_pretrain_path)
+    x = torch.randn(4,3,512,512).cuda()
+    y = torch.randn(4,3,512,512).cuda()
+    z = G(x, y)

project/hair_service_sd/core/utils/utils_3ddfa.py

@@ -0,0 +1,461 @@
+import numpy as np
+import math
+import scipy.io as scio
+import random
+from core import face3d
+from core.face3d import mesh
+from core.utils import params_3ddfa
+import cv2
+
+
+def RotationMatrix(angle_x, angle_y, angle_z):
+    phi = angle_x
+    gamma = angle_y
+    theta = angle_z
+    R_x = np.array([[1, 0, 0], [0, math.cos(phi), math.sin(phi)], [0, -math.sin(phi), math.cos(phi)]])
+    R_y = np.array([[math.cos(gamma), 0, -math.sin(gamma)], [0, 1, 0], [math.sin(gamma), 0, math.cos(gamma)]])
+    R_z = np.array([[math.cos(theta), math.sin(theta), 0], [-math.sin(theta), math.cos(theta), 0], [0, 0, 1]])
+    R = R_z @ R_x @ R_y
+    return R
+
+def process_uv(uv_coords, uv_h=256, uv_w=256):
+    uv_coords[:, 0] = uv_coords[:, 0] * (uv_w - 1)
+    uv_coords[:, 1] = uv_coords[:, 1] * (uv_h - 1)
+    uv_coords[:, 1] = uv_h - uv_coords[:, 1] - 1
+    uv_coords = np.hstack((uv_coords, np.zeros((uv_coords.shape[0], 1))))  # add z
+    return uv_coords
+
+def transform_params(params, tranform_mat, dst_img_height, src_img_height=256):
+    params_convert = np.zeros(76)
+    s_ori = np.sqrt(tranform_mat[:, 0:2].dot(np.transpose(tranform_mat[:, 0:2]))[0, 0])
+    s = s_ori * params[0]
+    roll = -np.arcsin(tranform_mat[0, 1] / s_ori)
+    txy_ori = np.ones([3, 1])
+    txy_ori[0] = params[4]
+    txy_ori[1] = src_img_height - 1 - params[5]
+    txy = tranform_mat.dot(txy_ori)
+    txy[1] = dst_img_height - 1 - txy[1]
+    params_convert[0] = s
+    params_convert[1:4] = params[1:4]
+    params_convert[3] += roll
+    params_convert[4] = txy[0]
+    params_convert[5] = txy[1]
+    params_convert[6:56] = params[6:56]
+    params_convert[56:] = params[56:]
+    return params_convert
+class FaceModel(object):
+
+    def __init__(self, img_size=384):
+        super(FaceModel, self).__init__()
+        self.init_status = False
+        from program_conf import ConfFactory
+        model_path = ConfFactory.getModelValue("model_dir")
+        # model_path = '/home/colomi/data/PycharmProjects/zao_service_test/models/model'
+        # model_path = '/home/colomi/Desktop/models'
+        face_model = scio.loadmat(model_path+'/face_3dfa/face_model.mat')
+        # face_model = scio.loadmat('data/face_model.mat')
+        self.mu_exp = face_model['mu_exp'].astype(np.float32)
+        self.mu_shape = face_model['mu_shape'].astype(np.float32)
+        self.w = face_model['w'].astype(np.float32)
+        self.w_exp = face_model['w_exp'].astype(np.float32)
+        self.sigma = face_model['sigma'].astype(np.float32).reshape((-1))
+        self.sigma_exp = face_model['sigma_exp'].astype(np.float32).reshape((-1))
+        self.tex = face_model['tex'].astype(np.float32).transpose((1, 0))
+        self.tri = face_model['tri'].transpose((1, 0)).astype(np.int32) - 1
+
+        keypoint834 = scio.loadmat(model_path+'/face_3dfa/keypt834.mat')
+        all_parrale = keypoint834['parrale_834']
+        all_iso_index = []
+        for line in all_parrale:
+            ddd = (line[0] - 1).reshape((-1))
+            all_iso_index.extend(((line[0] - 1).reshape((-1))).astype(np.int32).tolist())
+        tmp_mu = self.mu_exp + self.mu_shape
+        tmp_mu = tmp_mu.reshape((-1, 3))
+        # np.savetxt(r'H:\tmp\obj\tmp_mu.txt', tmp_mu, fmt='%f', delimiter=' ')
+        tmp_mu = tmp_mu[all_iso_index]
+        # np.savetxt(r'H:\tmp\obj\tmp_iso.txt', tmp_mu, fmt='%f', delimiter=' ')
+
+        # --load uv coords
+        uv_coords = face3d.load.load_uv_coords(model_path+'/face_3dfa/BFM_UV.mat')
+        uv_h = uv_w = img_size
+        self.image_h = self.image_w = img_size
+        self.uv_coords = process_uv(uv_coords, uv_h, uv_w)
+
+        trim_face = scio.loadmat(model_path+'/face_3dfa/trim_face.mat')
+        self.trim_idx = (trim_face['idx_face'].astype(np.int32) - 1).reshape((-1))
+        self.trim_tri = trim_face['tri_face'].transpose((1, 0)).astype(np.int32) - 1
+        self.trim_tri = self.trim_tri[:, ::-1]
+        keypoints834 = scio.loadmat(model_path+'/face_3dfa/keypt834.mat')
+        self.index_834 = (keypoints834['index_final'].astype(np.int32) - 1).reshape((-1))
+
+        self.mu_exp = self.mu_exp.reshape((-1, 3))[self.trim_idx, :].reshape((-1, 1))
+        self.mu_shape = self.mu_shape.reshape((-1, 3))[self.trim_idx, :].reshape((-1, 1))
+        self.w = self.w.reshape((-1, 3, 50))[self.trim_idx, :, :].reshape((-1, 50))
+        self.w_exp = self.w_exp.reshape((-1, 3, 20))[self.trim_idx, :, :].reshape((-1, 20))
+        self.tri = self.trim_tri
+        self.uv_coords = self.uv_coords[self.trim_idx]
+
+        self.mu = self.mu_exp + self.mu_shape
+
+        self.index30kTo137 = np.loadtxt(model_path+'/face_3dfa/index30kTo137.txt', dtype=np.int32)
+        self.init_status = True
+
+    def running(self):
+        return self.init_status
+
+    def draw_normal_depth_map_full_params(self, params, img_w, img_h, target_img=None, params_need_scale=False):
+        if params_need_scale:
+            f = params[0] / params_3ddfa.SCALE_F
+            phi = params[1] / params_3ddfa.SCALE_ROTATE
+            gamma = params[2] / params_3ddfa.SCALE_ROTATE
+            theta = params[3] / params_3ddfa.SCALE_ROTATE
+            t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET
+            alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis]
+            alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+        else:
+            f = params[0]
+            phi = params[1]
+            gamma = params[2]
+            theta = params[3]
+            t3d = np.array([params[4], params[5], 0])
+            alpha = params[6:56, np.newaxis]
+            alpha_exp = params[56:, np.newaxis]
+
+        express3d = self.w_exp @ alpha_exp
+        express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0))
+        shape3d = self.mu_shape + self.mu_exp + self.w @ alpha
+        shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0))
+        vertex3d = shape3d + express3d
+        R = RotationMatrix(phi, gamma, theta)
+        project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1))
+        project_vertex = project_vertex.transpose((1, 0))
+        project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1]
+        normal = face3d.mesh.light.get_normal(project_vertex, self.tri)
+        normal = (normal + 1) / 2
+
+        # np.savetxt('uv_coords.txt', self.uv_coords)
+        # np.savetxt('trim_tri.txt', self.trim_tri)
+        # np.savetxt('trim_tex.txt', trim_tex)
+        normal_map = mesh.render.render_colors(project_vertex, self.tri, normal, img_h, img_w, 3, BG=target_img)
+        # cv2.imwrite('uv_texture_map.png', uv_texture_map)
+
+        z = project_vertex[:, 2:]
+        z = z - np.min(z)
+        z = z / np.max(z)
+        depth_map = mesh.render.render_colors(project_vertex, self.tri, z, img_h, img_w, 1)
+
+        return normal_map, depth_map, project_vertex
+
+    def draw_normal_depth_map_full_params_with_bellus(self, params, bellus_shape, pts137_2d, img_w, img_h,
+                                                      target_img=None, params_need_scale=False):
+        if params_need_scale:
+            f = params[0] / params_3ddfa.SCALE_F
+            phi = params[1] / params_3ddfa.SCALE_ROTATE
+            gamma = params[2] / params_3ddfa.SCALE_ROTATE
+            theta = params[3] / params_3ddfa.SCALE_ROTATE
+            t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET
+            alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis]
+            alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+        else:
+            f = params[0]
+            phi = params[1]
+            gamma = params[2]
+            theta = params[3]
+            t3d = np.array([params[4], params[5], 0])
+            alpha = params[6:56, np.newaxis]
+            alpha_exp = params[56:, np.newaxis]
+
+        express3d = self.w_exp @ alpha_exp
+        express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0))
+
+        shape3d = bellus_shape
+
+        # shape3d = self.mu_shape + self.mu_exp + self.w @ alpha
+
+        shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0))
+        vertex3d = shape3d + express3d
+        R = RotationMatrix(phi, gamma, theta)
+        project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1))
+        project_vertex = project_vertex.transpose((1, 0))
+        project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1]
+
+        normal = face3d.mesh.light.get_normal(project_vertex, self.tri)
+        normal = (normal + 1) / 2
+
+        # np.savetxt('uv_coords.txt', self.uv_coords)
+        # np.savetxt('trim_tri.txt', self.trim_tri)
+        # np.savetxt('trim_tex.txt', trim_tex)
+
+        valid_2d_index = [list(range(0, 9)), list(range(14, 22)), [22, 42, 36, 29], list(range(64, 87)),
+                          list(range(88, 104)), list(range(105, 121))]
+        valid_2d_index = sum(valid_2d_index, [])
+
+        tmp2_ = target_img.copy()
+        tmp_ = target_img.copy()
+        for pt in pts137_2d[valid_2d_index].astype(np.int32):
+            cv2.circle(tmp_, (pt[0], pt[1]), 1, (0, 1, 0), 2)
+        # cv2.imshow('tmp_', tmp_)
+
+        index3d = self.index30kTo137[valid_2d_index]
+        # tmp_pts = project_vertex[index3d, :2]
+        # for pt in tmp_pts.astype(np.int32):
+        #     cv2.circle(tmp_, (pt[0], pt[1]), 1, (0, 0, 1), 2)
+        # cv2.imshow('tmp_', tmp_)
+
+        # Camera internals
+
+        focal_length = 384 * 1.8
+        center = (384 / 2, 384 / 2)
+        camera_matrix = np.array(
+            [[focal_length, 0, center[0]],
+             [0, focal_length, center[1]],
+             [0, 0, 1]], dtype="double"
+        )
+
+        express3d = self.w_exp @ alpha_exp
+        express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0))
+        shape3d = self.mu_shape + self.mu_exp + self.w @ alpha
+        shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0))
+        vertex3d = shape3d + express3d
+        R = RotationMatrix(phi, gamma, theta)
+        project_vertex_tmp = f * (R @ vertex3d) + t3d.reshape((3, 1))
+        project_vertex_tmp = project_vertex_tmp.transpose((1, 0))
+        project_vertex_tmp[:, 1] = img_h - 1 - project_vertex_tmp[:, 1]
+
+        dist_coeffs = np.zeros((4, 1))  # Assuming no lens distortion
+        # (success, rotation_vector, translation_vector) = cv2.solvePnP(project_vertex[index3d], pts137_2d[valid_2d_index], camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE)
+        (success, rotation_vector, translation_vector) = cv2.solvePnP(project_vertex[index3d],
+                                                                      project_vertex_tmp[index3d, :2], camera_matrix,
+                                                                      dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE)
+
+        (reproject, jacobian) = cv2.projectPoints(project_vertex[index3d], rotation_vector,
+                                                  translation_vector, camera_matrix, dist_coeffs)
+        reproject = reproject.squeeze()
+        for pt in reproject.astype(np.int32):
+            cv2.circle(tmp_, (pt[0], pt[1]), 1, (1, 0, 0), 2)
+        # cv2.imshow('tmp_', tmp_)
+
+        (reproject, jacobian) = cv2.projectPoints(project_vertex, rotation_vector,
+                                                  translation_vector, camera_matrix, dist_coeffs)
+        reproject = reproject.squeeze()
+        for pt in reproject.astype(np.int32):
+            cv2.circle(tmp2_, (pt[0], pt[1]), 1, (1, 0, 0), 1)
+        # cv2.imshow('tmp2_', tmp2_)
+
+        project_vertex[:, :2] = reproject
+
+        normal_map = mesh.render.render_colors(project_vertex, self.tri, normal, img_h, img_w, 3, BG=target_img)
+        # cv2.imwrite('uv_texture_map.png', uv_texture_map)
+
+        z = project_vertex[:, 2:]
+        z = z - np.min(z)
+        z = z / np.max(z)
+        depth_map = mesh.render.render_colors(project_vertex, self.tri, z, img_h, img_w, 1)
+
+        return normal_map, depth_map, project_vertex
+
+    def draw_normal_depth_map_137(self, user_params, movie_params, img_w, img_h, params_need_scale=False,
+                                  target_img=None, alpha_value=1):
+        if params_need_scale:
+            f_movie = movie_params[0] / params_3ddfa.SCALE_F
+            phi_movie = movie_params[1] / params_3ddfa.SCALE_ROTATE
+            gamma_movie = movie_params[2] / params_3ddfa.SCALE_ROTATE
+            theta = movie_params[3] / params_3ddfa.SCALE_ROTATE
+            t3d_movie = np.array([movie_params[4], movie_params[5], 0]) / params_3ddfa.SCALE_OFFSET
+            alpha_movie = (movie_params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis]
+            alpha_exp_movie = (movie_params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+            alpha_user = movie_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE + \
+                         (user_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE - movie_params[6:56,
+                                                                                     np.newaxis] / params_3ddfa.SCALE_SHAPE) * alpha_value
+            alpha_exp_user = (user_params[56:, np.newaxis] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+        else:
+            f_movie = movie_params[0]
+            phi_movie = movie_params[1]
+            gamma_movie = movie_params[2]
+            theta = movie_params[3]
+            t3d_movie = np.array([movie_params[4], movie_params[5], 0])
+            alpha_movie = movie_params[6:56, np.newaxis]
+            alpha_exp_movie = movie_params[56:, np.newaxis]
+            alpha_user = movie_params[6:56, np.newaxis] + (
+                        user_params[6:56, np.newaxis] - movie_params[6:56, np.newaxis]) * alpha_value
+            alpha_exp_user = user_params[56:, np.newaxis]
+
+        # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ')
+
+        express3d_movie = self.mu_exp + self.w_exp @ alpha_exp_movie
+        express3d_movie = np.reshape(express3d_movie, (-1, 3)).transpose((1, 0))
+        shape3d_user = self.mu_shape + self.w @ alpha_user
+        shape3d_user = np.reshape(shape3d_user, (-1, 3)).transpose((1, 0))
+        vertex3d_mix = shape3d_user + express3d_movie
+        R_movie = RotationMatrix(phi_movie, gamma_movie, theta)
+        project_vertex_mix = f_movie * (R_movie @ vertex3d_mix) + t3d_movie.reshape((3, 1))
+        project_vertex_mix = project_vertex_mix.transpose((1, 0))
+        project_vertex_mix[:, 1] = img_h - 1 - project_vertex_mix[:, 1]
+
+        project_vertex_normal = f_movie * (R_movie @ vertex3d_mix) + t3d_movie.reshape((3, 1))
+        project_vertex_normal = project_vertex_normal.transpose((1, 0))
+        project_vertex_normal[:, 1] = img_h - 1 - project_vertex_normal[:, 1]
+
+        normal = face3d.mesh.light.get_normal(project_vertex_normal, self.tri)
+        normal = (normal + 1) / 2
+        # normal_map = mesh.render.render_colors(project_vertex_normal, self.tri, normal, img_h, img_w, 3, BG=target_img)
+        # np.savetxt('uv_coords.txt', project_vertex_normal)
+        # np.savetxt('trim_tri.txt', self.tri)
+        # np.savetxt('trim_tex.txt', normal)
+        normal_map = mesh.render.render_colors(project_vertex_normal, self.tri, normal, img_h, img_w, 3, BG=target_img)
+        # cv2.imshow('normal_map', normal_map)
+        # cv2.waitKey()
+        # cv2.imwrite('normal_map.png', (normal_map * 255).astype(np.uint8))
+
+        z = project_vertex_normal[:, 2:]
+        z = z - np.min(z)
+        z = z / np.max(z)
+        gray_map = mesh.render.render_colors(project_vertex_normal, self.tri, z, img_h, img_w, 1)
+
+        # movie
+        shape3d_movie = self.mu_shape + self.w @ alpha_movie
+        shape3d_movie = np.reshape(shape3d_movie, (-1, 3)).transpose((1, 0))
+        vertex3d_movie = shape3d_movie + express3d_movie
+        project_vertex_movie = f_movie * (R_movie @ vertex3d_movie) + t3d_movie.reshape((3, 1))
+        project_vertex_movie = project_vertex_movie.transpose((1, 0))
+        project_vertex_movie[:, 1] = img_h - 1 - project_vertex_movie[:, 1]
+        point_usr_137_3DDFA = project_vertex_mix[self.index30kTo137, :2]
+        point_movie_137_3DDFA = project_vertex_movie[self.index30kTo137, :2]
+
+        return normal_map, gray_map, point_usr_137_3DDFA, point_movie_137_3DDFA
+
+    def get2Dpoint_137_orig(self, user_params, movie_params, img_w, img_h, params_need_scale=False, target_img=None,
+                            alpha_value=1):
+        if params_need_scale:
+            f_movie = movie_params[0] / params_3ddfa.SCALE_F
+            phi_movie = movie_params[1] / params_3ddfa.SCALE_ROTATE
+            gamma_movie = movie_params[2] / params_3ddfa.SCALE_ROTATE
+            theta = movie_params[3] / params_3ddfa.SCALE_ROTATE
+            t3d_movie = np.array([movie_params[4], movie_params[5], 0]) / params_3ddfa.SCALE_OFFSET
+            alpha_movie = (movie_params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis]
+            alpha_exp_movie = (movie_params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+            alpha_user = movie_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE + \
+                         (user_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE - movie_params[6:56,
+                                                                                     np.newaxis] / params_3ddfa.SCALE_SHAPE) * alpha_value
+        else:
+            f_movie = movie_params[0]
+            phi_movie = movie_params[1]
+            gamma_movie = movie_params[2]
+            theta = movie_params[3]
+            t3d_movie = np.array([movie_params[4], movie_params[5], 0])
+            alpha_movie = movie_params[6:56, np.newaxis]
+            alpha_exp_movie = movie_params[56:, np.newaxis]
+            alpha_user = movie_params[6:56, np.newaxis] + (
+                    user_params[6:56, np.newaxis] - movie_params[6:56, np.newaxis]) * alpha_value
+
+        # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ')
+
+        shape3d_user = self.mu_shape + self.w @ alpha_user
+        shape3d_user = np.reshape(shape3d_user, (-1, 3)).transpose((1, 0))
+        vertex3d_mix = shape3d_user
+
+        R_movie = RotationMatrix(phi_movie, gamma_movie, theta)
+        # project_vertex_mix = f_movie * (R_movie @  vertex3d_mix) + t3d_movie.reshape((3, 1))
+        project_vertex_mix = f_movie * (R_movie @ vertex3d_mix) + np.array([img_h / 2, img_w / 2, 0],
+                                                                           dtype=np.float32).reshape((3, 1))
+        project_vertex_mix = project_vertex_mix.transpose((1, 0))
+        project_vertex_mix[:, 1] = img_h - 1 - project_vertex_mix[:, 1]
+
+        # movie
+        shape3d_movie = self.mu_shape + self.w @ alpha_movie
+        shape3d_movie = np.reshape(shape3d_movie, (-1, 3)).transpose((1, 0))
+        vertex3d_movie = shape3d_movie
+        # project_vertex_movie = f_movie * (R_movie @  vertex3d_movie) + t3d_movie.reshape((3, 1))
+        project_vertex_movie = f_movie * (R_movie @ vertex3d_movie) + np.array([img_h / 2, img_w / 2, 0],
+                                                                               dtype=np.float32).reshape((3, 1))
+        project_vertex_movie = project_vertex_movie.transpose((1, 0))
+        project_vertex_movie[:, 1] = img_h - 1 - project_vertex_movie[:, 1]
+        point_usr_137_3DDFA_orig = project_vertex_mix[self.index30kTo137, :2]
+        point_movie_137_3DDFA_orig = project_vertex_movie[self.index30kTo137, :2]
+
+        return point_usr_137_3DDFA_orig, point_movie_137_3DDFA_orig
+
+    def draw_uv_map(self, img_target, params, params_need_scale=False, img_w=256, img_h=256, is_train=False):
+        if params_need_scale:
+            f = params[0] / params_3ddfa.SCALE_F
+            phi = params[1] / params_3ddfa.SCALE_ROTATE
+            gamma = params[2] / params_3ddfa.SCALE_ROTATE
+            theta = params[3] / params_3ddfa.SCALE_ROTATE
+            t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET
+            alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis]
+            alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis]
+        else:
+            f = params[0]
+            phi = params[1]
+            gamma = params[2]
+            theta = params[3]
+            t3d = np.array([params[4], params[5], 0])
+            alpha = params[6:56, np.newaxis]
+            alpha_exp = params[56:, np.newaxis]
+
+        if is_train:
+            scale_ratio = 0.7 + random.random() * 0.3
+            f = f * scale_ratio
+        else:
+            f = f * 0.85
+        # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ')
+
+        express3d = self.mu_exp + self.w_exp @ alpha_exp
+        express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0))
+        shape3d = self.mu_shape + self.w @ alpha
+        shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0))
+        vertex3d = shape3d + express3d
+        R = RotationMatrix(phi, gamma, theta)
+        project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1))
+        project_vertex = project_vertex.transpose((1, 0))
+        project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1]
+
+        trim_tex = np.zeros((project_vertex.shape[0], 3))
+        project_vertex_int = project_vertex.astype(np.int32)
+        project_vertex_int = np.clip(project_vertex_int, 0, img_w - 1)
+        for j in range(project_vertex_int.shape[0]):
+            tmp_ = img_target[project_vertex_int[j, 1], project_vertex_int[j, 0]]
+            trim_tex[j] = tmp_
+
+        # np.savetxt('uv_coords.txt', self.uv_coords)
+        # np.savetxt('trim_tri.txt', self.trim_tri)
+        # np.savetxt('trim_tex.txt', trim_tex)
+        uv_texture_map = mesh.render.render_colors(self.uv_coords, self.trim_tri, trim_tex, img_h, img_w, c=3).astype(
+            np.uint8)
+        # cv2.imwrite('uv_texture_map.png', uv_texture_map)
+        # uv_texture_map = cv2.cvtColor(uv_texture_map, cv2.COLOR_RGB2BGR)
+        return uv_texture_map
+
+    def compare_shape(self, pred_params, gt_params):
+        def parse_param_batch_noscale(param):
+            """Work for both numpy and tensor"""
+            N = param.shape[0]
+            f = param[:, 0]
+            R = np.zeros((N, 3, 3), dtype=np.float32)
+            for i in range(N):
+                R[i, :, :] = RotationMatrix(param[i, 1], param[i, 2], param[i, 3])
+            f = f.reshape((N, 1, 1))
+            p = f * R
+            offset = np.zeros((N, 3, 1), dtype=np.float32)
+            offset[:, :2, 0] = param[:, 4:6]
+            alpha_shp = param[:, 6:56].reshape((N, -1, 1))
+            alpha_exp = param[:, 56:].reshape((N, -1, 1))
+            return p, offset, alpha_shp, alpha_exp
+
+        gt_pred_params = gt_params.copy()
+        gt_pred_params[:, 6:56] = pred_params[:, 6:56]
+
+        pred_p, pred_offset, pred_alpha_shape, pred_alpha_exp = parse_param_batch_noscale(gt_pred_params)
+        gt_p, gt_offset, gt_alpha_shape, gt_alpha_exp = parse_param_batch_noscale(gt_params)
+
+        N = pred_params.shape[0]
+        gt_vertex = gt_p @ (self.mu + self.w @ gt_alpha_shape + self.w_exp @ gt_alpha_exp) \
+            .reshape((N, -1, 3)) \
+            .transpose((0, 2, 1)) + gt_offset
+        pred_vertex = pred_p @ (self.mu + self.w @ pred_alpha_shape + self.w_exp @ pred_alpha_exp) \
+            .reshape((N, -1, 3)) \
+            .transpose((0, 2, 1)) + pred_offset
+
+        diff = np.sqrt(np.sum((gt_vertex[:, :2, :] - pred_vertex[:, :2, :]) ** 2, axis=1))
+        loss = np.mean(diff)
+        return loss

project/hair_service_sd/core/utils/weight_init.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+import torch.nn as nn
+
+
+def c2_xavier_fill(module: nn.Module):
+    """
+    Initialize `module.weight` using the "XavierFill" implemented in Caffe2.
+    Also initializes `module.bias` to 0.
+
+    Args:
+        module (torch.nn.Module): module to initialize.
+    """
+    # Caffe2 implementation of XavierFill in fact
+    # corresponds to kaiming_uniform_ in PyTorch
+    nn.init.kaiming_uniform_(module.weight, a=1)
+    if module.bias is not None:
+        nn.init.constant_(module.bias, 0)
+
+
+def c2_msra_fill(module: nn.Module):
+    """
+    Initialize `module.weight` using the "MSRAFill" implemented in Caffe2.
+    Also initializes `module.bias` to 0.
+
+    Args:
+        module (torch.nn.Module): module to initialize.
+    """
+    nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+    if module.bias is not None:
+        nn.init.constant_(module.bias, 0)