initial

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mat filter=lfs diff=lfs merge=lfs -text

pre_trained/unet_denoise/checkpoint-70000/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNetDenoise2DConditionModel",
+  "_diffusers_version": "0.25.1",
+  "_name_or_path": "/scratch/project_462000772/wmengting/diffusion_models/stable-diffusion-v1-5",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 12,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

pre_trained/unet_denoise/checkpoint-70000/diffusion_pytorch_model.safetensors

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

pre_trained/unet_id/checkpoint-70000/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNetID2DConditionModel",
+  "_diffusers_version": "0.25.1",
+  "_name_or_path": "/scratch/project_462000772/wmengting/diffusion_models/stable-diffusion-v1-5",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

pre_trained/unet_id/checkpoint-70000/diffusion_pytorch_model.safetensors

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

utils/checkpoints/net_seg_res18.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:041ab78a4f8f756cd7e93df0d2840d03162e46b9c463c144f7fdf3ee3e6c4233
+size 57429148

utils/checkpoints/third_party/BFM_model_front.mat

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

utils/checkpoints/third_party/d3dfr_res50_nofc.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52c54b90304a06c16b6813910c26faff1a907d4f8bd69a71ad4ecff43b41a090
+size 96449126

utils/third_party/d3dfr/bfm.py

@@ -0,0 +1,473 @@
+import os
+import torch
+# import torch.nn as nn
+from scipy.io import loadmat
+import numpy as np
+import torch.nn.functional as F
+
+# CURRENT_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def perspective_projection(focal, center):
+    # return p.T (N, 3) @ (3, 3)
+    return np.array([
+        focal, 0, center,
+        0, focal, center,
+        0, 0, 1
+    ]).reshape([3, 3]).astype(np.float32).transpose()
+
+
+class SH:
+    def __init__(self):
+        self.a = [np.pi, 2 * np.pi / np.sqrt(3.), 2 * np.pi / np.sqrt(8.)]
+        self.c = [1/np.sqrt(4 * np.pi), np.sqrt(3.) / np.sqrt(4 * np.pi), 3 * np.sqrt(5.) / np.sqrt(12 * np.pi)]
+
+
+class BFM(torch.nn.Module):
+    # BFM 3D face model
+    def __init__(self,
+                 recenter=True,
+                 camera_distance=10.,
+                 init_lit=np.array([0.8, 0, 0, 0, 0, 0, 0, 0, 0]),
+                 focal=1015.,
+                 image_size=224,
+                 bfm_model_path='pretrained/BFM_model_front.mat'
+                 ):
+        super().__init__()
+        model = loadmat(bfm_model_path)
+        # self.bfm_uv = loadmat(os.path.join(CURRENT_PATH, 'BFM/BFM_UV.mat'))
+        # print(model.keys())
+        # mean face shape. [3*N,1]
+        # self.meanshape = torch.from_numpy(model['meanshape'])
+        self.register_buffer('meanshape', torch.from_numpy(model['meanshape']).float())
+
+        if recenter:
+            meanshape = self.meanshape.view(-1, 3)
+            meanshape = meanshape - torch.mean(meanshape, dim=0, keepdim=True)
+            self.meanshape = meanshape.view(-1, 1)
+
+        # identity basis. [3*N,80]
+        # self.idBase = torch.from_numpy(model['idBase'])
+        self.register_buffer('idBase', torch.from_numpy(model['idBase']).float())
+        # self.idBase = nn.Parameter(torch.from_numpy(model['idBase']).float())
+        # self.exBase = torch.from_numpy(model['exBase'].astype(
+        #     np.float32))    # expression basis. [3*N,64]
+        self.register_buffer('exBase', torch.from_numpy(model['exBase']).float())
+        # self.exBase = nn.Parameter(torch.from_numpy(model['exBase']).float())
+        # mean face texture. [3*N,1] (0-255)
+        # self.meantex = torch.from_numpy(model['meantex'])
+        self.register_buffer('meantex', torch.from_numpy(model['meantex']).float())
+        # texture basis. [3*N,80]
+        # self.texBase = torch.from_numpy(model['texBase'])
+        self.register_buffer('texBase', torch.from_numpy(model['texBase']).float())
+        # self.texBase = nn.Parameter(torch.from_numpy(model['texBase']).float())
+
+        # triangle indices for each vertex that lies in. starts from 0. [N,8]
+        self.register_buffer('point_buf', torch.from_numpy(model['point_buf']).long()-1)
+        # self.point_buf = model['point_buf'].astype(np.int32)
+        # vertex indices in each triangle. starts from 0. [F,3]
+        self.register_buffer('face_buf', torch.from_numpy(model['tri']).long()-1)
+        # self.tri = model['tri'].astype(np.int32)
+        # vertex indices of 68 facial landmarks. starts from 0. [68]
+        self.register_buffer('keypoints', torch.from_numpy(model['keypoints']).long().view(68)-1)
+        # self.keypoints = model['keypoints'].astype(np.int32)[0]
+        # print(self.keypoints)
+        # print('keypoints', self.keypoints)
+
+        # vertex indices for small face region to compute photometric error. starts from 0.
+        # self.front_mask = np.squeeze(model['frontmask2_idx']).astype(np.int64) - 1
+        self.register_buffer('front_mask', torch.from_numpy(np.squeeze(model['frontmask2_idx'])).long()-1)
+        # vertex indices for each face from small face region. starts from 0. [f,3]
+        # self.front_face_buf = model['tri_mask2'].astype(np.int64) - 1
+        self.register_buffer('front_face_buf', torch.from_numpy(np.squeeze(model['tri_mask2'])).long() - 1)
+        # vertex indices for pre-defined skin region to compute reflectance loss
+        # self.skin_mask = np.squeeze(model['skinmask'])
+        self.register_buffer('skin_mask', torch.from_numpy(np.squeeze(model['skinmask'])))
+
+
+        # keypoints_222 = []
+        # with open(os.path.join(CURRENT_PATH, 'BFM/D3DFR_222.txt'), 'r') as f:
+        #     for line in f.readlines():
+        #         idx = int(line.strip())
+        #         keypoints_222.append(max(idx, 0))
+        # self.register_buffer('keypoints_222', torch.from_numpy(np.array(keypoints_222)).long())
+
+        # (1) right eye outer corner, (2) right eye inner corner, (3) left eye inner corner, (4) left eye outer corner,
+        # (5) nose bottom, (6) right mouth corner, (7) left mouth corner
+        self.register_buffer('keypoints_7', self.keypoints[[36, 39, 42, 45, 33, 48, 54]])
+
+        # self.persc_proj = torch.from_numpy(perspective_projection(focal, center)).float()
+        self.register_buffer('persc_proj', torch.from_numpy(perspective_projection(focal, image_size/2)))
+        self.camera_distance = camera_distance
+        self.image_size = image_size
+        self.SH = SH()
+        # self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32)
+        self.register_buffer('init_lit', torch.from_numpy(init_lit.reshape([1, 1, -1]).astype(np.float32)))
+
+        # (1) right eye outer corner, (2) right eye inner corner, (3) left eye inner corner, (4) left eye outer corner,
+        # (5) nose bottom, (6) right mouth corner, (7) left mouth corner
+        # print(self.keypoints[[36, 39, 42, 45, 33, 48, 54]])
+
+        # Lm3D = loadmat(os.path.join(CURRENT_PATH, 'BFM/similarity_Lm3D_all.mat'))
+        # Lm3D = Lm3D['lm']
+        # # print(Lm3D)
+        #
+        # # calculate 5 facial landmarks using 68 landmarks
+        # lm_idx = np.array([31, 37, 40, 43, 46, 49, 55]) - 1
+        # Lm3D = np.stack([Lm3D[lm_idx[0], :], np.mean(Lm3D[lm_idx[[1, 2]], :], 0), np.mean(
+        #     Lm3D[lm_idx[[3, 4]], :], 0), Lm3D[lm_idx[5], :], Lm3D[lm_idx[6], :]], axis=0)
+        # Lm3D = Lm3D[[1, 2, 0, 3, 4], :]
+        # self.Lm3D = Lm3D
+        # print(Lm3D.shape)
+
+    def split_coeff(self, coeff):
+        # input: coeff with shape [1,258]
+        id_coeff = coeff[:, 0:80]        # identity(shape) coeff of dim 80
+        ex_coeff = coeff[:, 80:144]      # expression coeff of dim 64
+        tex_coeff = coeff[:, 144:224]    # texture(albedo) coeff of dim 80
+        gamma = coeff[:, 227:254]        # lighting coeff for 3 channel SH function of dim 27
+        angles = coeff[:, 224:227]       # ruler angles(x,y,z) for rotation of dim 3
+        translation = coeff[:, 254:257]  # translation coeff of dim 3
+
+        return id_coeff, ex_coeff, tex_coeff, gamma, angles, translation
+
+    def split_coeff_orderly(self, coeff):
+        # input: coeff with shape [1,258]
+        id_coeff = coeff[:, 0:80]        # identity(shape) coeff of dim 80
+        ex_coeff = coeff[:, 80:144]      # expression coeff of dim 64
+        tex_coeff = coeff[:, 144:224]    # texture(albedo) coeff of dim 80
+        angles = coeff[:, 224:227]  # ruler angles(x,y,z) for rotation of dim 3
+        gamma = coeff[:, 227:254]        # lighting coeff for 3 channel SH function of dim 27
+        translation = coeff[:, 254:257]  # translation coeff of dim 3
+
+        return id_coeff, ex_coeff, tex_coeff, angles, gamma, translation
+
+    def compute_exp_deform(self, exp_coeff):
+        exp_part = torch.einsum('ij,aj->ai', self.exBase, exp_coeff)
+        return exp_part
+
+    def compute_id_deform(self, id_coeff):
+        id_part = torch.einsum('ij,aj->ai', self.idBase, id_coeff)
+        return id_part
+
+    def compute_shape_from_coeff(self, coeff):
+        id_coeff = coeff[:, 0:80]
+        ex_coeff = coeff[:, 80:144]
+        batch_size = coeff.shape[0]
+        id_part = torch.einsum('ij,aj->ai', self.idBase, id_coeff)  #B, n
+        exp_part = torch.einsum('ij,aj->ai', self.exBase, ex_coeff) #B, n
+        face_shape = id_part + exp_part + self.meanshape.view(1, -1)
+        return face_shape.view(batch_size, -1, 3)
+
+    def compute_shape(self, id_coeff, exp_coeff):
+        """
+        Return:
+            face_shape       -- torch.tensor, size (B, N, 3)
+        Parameters:
+            id_coeff         -- torch.tensor, size (B, 80), identity coeffs
+            id_relative_scale  -- torch.tensor, size (B, 1), identity coeffs
+            exp_coeff        -- torch.tensor, size (B, 64), expression coeffs
+        """
+        batch_size = id_coeff.shape[0]
+        id_part = torch.einsum('ij,aj->ai', self.idBase, id_coeff)  #B, n
+        exp_part = torch.einsum('ij,aj->ai', self.exBase, exp_coeff) #B, n
+        face_shape = id_part + exp_part + self.meanshape.view(1, -1)
+        return face_shape.view(batch_size, -1, 3)
+
+    def compute_texture(self, tex_coeff, normalize=True):
+        """
+        Return:
+            face_texture     -- torch.tensor, size (B, N, 3), in RGB order, range (0, 1.)
+        Parameters:
+            tex_coeff        -- torch.tensor, size (B, 80)
+        """
+        batch_size = tex_coeff.shape[0]
+        face_texture = torch.einsum('ij,aj->ai', self.texBase, tex_coeff) + self.meantex
+        if normalize:
+            face_texture = face_texture / 255.
+        return face_texture.view(batch_size, -1, 3)
+
+    def compute_norm(self, face_shape):
+        """
+        Return:
+            vertex_norm      -- torch.tensor, size (B, N, 3)
+        Parameters:
+            face_shape       -- torch.tensor, size (B, N, 3)
+        """
+
+        v1 = face_shape[:, self.face_buf[:, 0]]
+        v2 = face_shape[:, self.face_buf[:, 1]]
+        v3 = face_shape[:, self.face_buf[:, 2]]
+        e1 = v1 - v2
+        e2 = v2 - v3
+        face_norm = torch.cross(e1, e2, dim=-1)
+        face_norm = F.normalize(face_norm, dim=-1, p=2)
+        face_norm = torch.cat([face_norm, torch.zeros(face_norm.shape[0], 1, 3).to(self.meanshape)], dim=1)
+
+        vertex_norm = torch.sum(face_norm[:, self.point_buf], dim=2)
+        vertex_norm = F.normalize(vertex_norm, dim=-1, p=2)
+        return vertex_norm
+
+    def compute_color(self, face_texture, face_norm, gamma):
+        """
+        Return:
+            face_color       -- torch.tensor, size (B, N, 3), range (0, 1.)
+        Parameters:
+            face_texture     -- torch.tensor, size (B, N, 3), from texture model, range (0, 1.)
+            face_norm        -- torch.tensor, size (B, N, 3), rotated face normal
+            gamma            -- torch.tensor, size (B, 27), SH coeffs
+        """
+        batch_size = gamma.shape[0]
+        v_num = face_texture.shape[1]
+        a, c = self.SH.a, self.SH.c
+        gamma = gamma.reshape([batch_size, 3, 9])
+        gamma = gamma + self.init_lit
+        gamma = gamma.permute(0, 2, 1)
+        Y = torch.cat([
+            a[0] * c[0] * torch.ones_like(face_norm[..., :1]).to(self.meanshape),
+            -a[1] * c[1] * face_norm[..., 1:2],
+            a[1] * c[1] * face_norm[..., 2:],
+            -a[1] * c[1] * face_norm[..., :1],
+            a[2] * c[2] * face_norm[..., :1] * face_norm[..., 1:2],
+            -a[2] * c[2] * face_norm[..., 1:2] * face_norm[..., 2:],
+            0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm[..., 2:] ** 2 - 1),
+            -a[2] * c[2] * face_norm[..., :1] * face_norm[..., 2:],
+            0.5 * a[2] * c[2] * (face_norm[..., :1] ** 2 - face_norm[..., 1:2] ** 2)
+        ], dim=-1)
+        r = Y @ gamma[..., :1]
+        g = Y @ gamma[..., 1:2]
+        b = Y @ gamma[..., 2:]
+        face_color = torch.cat([r, g, b], dim=-1) * face_texture
+        return face_color
+
+    def compute_rotation(self, angles):
+        """
+        Return:
+            rot              -- torch.tensor, size (B, 3, 3) pts @ trans_mat
+        Parameters:
+            angles           -- torch.tensor, size (B, 3), radian
+        """
+
+        batch_size = angles.shape[0]
+        ones = torch.ones([batch_size, 1]).to(self.meanshape)
+        zeros = torch.zeros([batch_size, 1]).to(self.meanshape)
+        x, y, z = angles[:, :1], angles[:, 1:2], angles[:, 2:],
+
+        rot_x = torch.cat([
+            ones, zeros, zeros,
+            zeros, torch.cos(x), -torch.sin(x),
+            zeros, torch.sin(x), torch.cos(x)
+        ], dim=1).reshape([batch_size, 3, 3])
+
+        rot_y = torch.cat([
+            torch.cos(y), zeros, torch.sin(y),
+            zeros, ones, zeros,
+            -torch.sin(y), zeros, torch.cos(y)
+        ], dim=1).reshape([batch_size, 3, 3])
+
+        rot_z = torch.cat([
+            torch.cos(z), -torch.sin(z), zeros,
+            torch.sin(z), torch.cos(z), zeros,
+            zeros, zeros, ones
+        ], dim=1).reshape([batch_size, 3, 3])
+
+        rot = rot_z @ rot_y @ rot_x
+        return rot.permute(0, 2, 1)
+
+    def to_camera(self, face_shape):
+        face_shape[..., -1] = self.camera_distance - face_shape[..., -1]
+        return face_shape
+
+    def to_image(self, face_shape):
+        """
+        Return:
+            face_proj        -- torch.tensor, size (B, N, 2), y direction is opposite to v direction
+        Parameters:
+            face_shape       -- torch.tensor, size (B, N, 3)
+        """
+        # to image_plane
+        face_proj = face_shape @ self.persc_proj
+
+        # print(face_proj.shape)
+        face_proj = face_proj[..., :2] / face_proj[..., 2:]
+
+        return face_proj
+
+    def rotate(self, face_shape, rot):
+        """
+        Return:
+            face_shape       -- torch.tensor, size (B, N, 3) pts @ rot + trans
+        Parameters:
+            face_shape       -- torch.tensor, size (B, N, 3)
+            rot              -- torch.tensor, size (B, 3, 3)
+
+        """
+        return face_shape @ rot
+
+    def get_landmarks7(self, face_proj):
+        """
+        Return:
+            face_lms         -- torch.tensor, size (B, 68, 2)
+        Parameters:
+            face_proj       -- torch.tensor, size (B, N, 2)
+        """
+        return face_proj[:, self.keypoints_7, :]
+
+    def get_landmarks68(self, face_proj):
+        """
+        Return:
+            face_lms         -- torch.tensor, size (B, 68, 2)
+        Parameters:
+            face_proj       -- torch.tensor, size (B, N, 2)
+        """
+        return face_proj[:, self.keypoints, :]
+
+    def get_landmarks222(self, face_proj):
+        """
+        Return:
+            face_lms         -- torch.tensor, size (B, 68, 2)
+        Parameters:
+            face_proj       -- torch.tensor, size (B, N, 2)
+        """
+        return face_proj[:, self.keypoints_222, :]
+
+    def compute_for_render(self, coeffs):
+        """
+        Return:
+            face_vertex     -- torch.tensor, size (B, N, 3), in camera coordinate
+            face_color      -- torch.tensor, size (B, N, 3), in RGB order
+            landmark        -- torch.tensor, size (B, 68, 2), y direction is opposite to v direction
+        Parameters:
+            coeffs          -- torch.tensor, size (B, 258)
+        """
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeffs)
+        # id_relative_scale = id_relative_scale.clamp(0.9,1.1)
+        face_shape = self.compute_shape(id_coeff, ex_coeff)
+        # face_shape_noexp = self.compute_shape(id_coeff, torch.zeros_like(ex_coeff))
+        # print(face_shape.size())
+        rotation = self.compute_rotation(angles)
+        # print('rotation')
+
+        face_shape_rotated = self.rotate(face_shape, rotation)
+        face_shape_transformed = face_shape_rotated + translation.unsqueeze(1)
+        face_vertex = self.to_camera(face_shape_transformed)
+        face_proj = self.to_image(face_vertex)
+
+        # face_shape_transformed_noexp = self.transform(face_shape_noexp, rotation, translation, scale_xyz)
+        # face_vertex_noexp = self.to_camera(face_shape_transformed_noexp)
+
+        landmark68 = self.get_landmarks68(face_proj)
+        # landmark_face = face_proj[:,self.front_mask[::32], :]
+        landmark68[:, :, 1] = self.image_size - 1 - landmark68[:, :, 1]
+
+        face_texture = self.compute_texture(tex_coeff)
+        face_norm_roted = self.compute_norm(face_shape_rotated)
+        # face_norm_roted = face_norm @ rotation
+        face_color = self.compute_color(face_texture, face_norm_roted, gamma)
+
+        # face_norm_noexp = self.compute_norm(face_shape_noexp)
+        # face_norm_noexp_roted = face_norm_noexp @ rotation
+        # face_color_noexp = self.compute_color(face_texture, face_norm_noexp_roted, gamma)
+
+        return face_shape, face_vertex, face_color, face_texture, landmark68
+    
+    def get_lm68(self, coeffs):
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeffs)
+        ex_coeff = torch.zeros_like(ex_coeff)
+        # id_relative_scale = id_relative_scale.clamp(0.9,1.1)
+        face_shape = self.compute_shape(id_coeff, ex_coeff)
+        # face_shape_noexp = self.compute_shape(id_coeff, torch.zeros_like(ex_coeff))
+        # print(face_shape.size())
+        rotation = self.compute_rotation(angles)
+        # print('rotation')
+
+        face_shape_rotated = self.rotate(face_shape, rotation)
+        face_shape_transformed = face_shape_rotated + translation.unsqueeze(1)
+        face_vertex = self.to_camera(face_shape_transformed)
+        face_proj = self.to_image(face_vertex)
+
+        landmark68 = self.get_landmarks68(face_proj)
+        # landmark_face = face_proj[:,self.front_mask[::32], :]
+        landmark68[:, :, 1] = self.image_size - 1 - landmark68[:, :, 1]
+        return landmark68, ex_coeff
+
+    def get_coeffs(self, coeffs):
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeffs)
+        return id_coeff, ex_coeff, tex_coeff, gamma, angles, translation
+
+    def get_vertex(self, coeffs):
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeffs)
+        # id_relative_scale = id_relative_scale.clamp(0.9,1.1)
+        face_shape = self.compute_shape(id_coeff, ex_coeff)
+        # face_shape_noexp = self.compute_shape(id_coeff, torch.zeros_like(ex_coeff))
+        # print(face_shape.size())
+        rotation = self.compute_rotation(angles)
+        # print('rotation')
+
+        face_shape_rotated = self.rotate(face_shape, rotation)
+        face_shape_transformed = face_shape_rotated + translation.unsqueeze(1)
+        face_vertex = self.to_camera(face_shape_transformed)
+        face_proj = self.to_image(face_vertex)
+
+        return face_proj
+        
+
+    def forward(self, coeffs):
+        face_shape, face_vertex, face_color, face_texture, landmark68 = self.compute_for_render(coeffs)
+        return face_shape, face_vertex, face_color, face_texture, landmark68
+
+
+    def save_obj(self, coeff, obj_name):
+        # The image size is 224 * 224
+        # face reconstruction with coeff and BFM model
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeff)
+
+        # compute face shape
+        face_shape = self.compute_shape(id_coeff, ex_coeff).cpu().detach().numpy()[0]
+        face_tri = self.face_buf.cpu().numpy()
+
+        with open(obj_name, 'w') as fobj:
+            for i in range(face_shape.shape[0]):
+                fobj.write(
+                    'v ' + str(face_shape[i][0]) + ' ' + str(face_shape[i][1]) + ' ' + str(face_shape[i][2]) + '\n')
+
+            # start from 1
+            for i in range(face_tri.shape[0]):
+                fobj.write('f ' + str(face_tri[i][0] + 1) + ' ' + str(face_tri[i][1] + 1) + ' ' + str(
+                    face_tri[i][2] + 1) + '\n')
+
+        # lm7 = face_shape[[2215,  5828, 10455, 14066,  8204,  5522, 10795], :]
+        # with open(obj_name[:-3]+'txt', 'w') as f:
+        #     for point in lm7:
+        #         f.write('{} {} {}\n'.format(point[0], point[1], point[2]))
+
+    def save_neutral_obj(self, coeff, obj_name):
+        # The image size is 224 * 224
+        # face reconstruction with coeff and BFM model
+        id_coeff, ex_coeff, tex_coeff, gamma, angles, translation = self.split_coeff(coeff)
+
+        # compute face shape
+        face_shape = self.compute_shape(id_coeff, ex_coeff*0).cpu().numpy()[0]
+        face_tri = self.face_buf.cpu().numpy()
+
+        with open(obj_name, 'w') as fobj:
+            for i in range(face_shape.shape[0]):
+                fobj.write(
+                    'v ' + str(face_shape[i][0]) + ' ' + str(face_shape[i][1]) + ' ' + str(face_shape[i][2]) + '\n')
+
+            # start from 1
+            for i in range(face_tri.shape[0]):
+                fobj.write('f ' + str(face_tri[i][0] + 1) + ' ' + str(face_tri[i][1] + 1) + ' ' + str(
+                    face_tri[i][2] + 1) + '\n')
+
+        # lm7 = face_shape[[2215,  5828, 10455, 14066,  8204,  5522, 10795], :]
+        # with open(obj_name[:-3]+'txt', 'w') as f:
+        #     for point in lm7:
+        #         f.write('{} {} {}\n'.format(point[0], point[1], point[2]))
+
+    # def clip(self, g_ratio=0.1, t_ratio=0.1):
+    #     self.idBase.data = torch.minimum(torch.maximum(self.idBase_org * (1 - g_ratio), self.idBase.data), self.idBase_org * (1 + g_ratio))
+    #     self.exBase.data = self.exBase_org #torch.minimum(torch.maximum(self.exBase_org * (1 - 0.001), self.exBase.data), self.exBase_org * (1 + 0.001))
+    #     self.texBase.data = torch.minimum(torch.maximum(self.texBase_org * (1 - t_ratio), self.texBase.data), self.texBase_org * (1 + t_ratio))
+
+

utils/third_party/d3dfr_res50_nofc.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62e7e6a6bc4e16fb567643182ccabf55f8222746269cce3392109a6c592babc1
+size 288887131

utils/third_party/insightface_backbone_conv.py

@@ -0,0 +1,237 @@
+import os
+import torch
+from torch import nn
+
+__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200', 'getarcface']
+
+
+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 IBasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(IBasicBlock, self).__init__()
+        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")
+        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.prelu = nn.PReLU(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.prelu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        return out
+
+
+class IResNet(nn.Module):
+    fc_scale = 7 * 7
+    def __init__(self,
+                 block, layers, dropout=0, num_features=512, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
+        super(IResNet, self).__init__()
+        self.fp16 = fp16
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            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=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+        self.prelu = nn.PReLU(self.inplanes)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        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.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
+        self.dropout = nn.Dropout(p=dropout, inplace=True)
+        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
+        self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+        nn.init.constant_(self.features.weight, 1.0)
+        self.features.weight.requires_grad = False
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, 0, 0.1)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, IBasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        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),
+                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+            )
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups,
+                  self.base_width, previous_dilation))
+        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))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x, return_id512=False):
+        
+        bz = x.shape[0]
+        # with torch.cuda.amp.autocast(self.fp16):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.prelu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        if not return_id512:
+            return x.view(bz,512,-1).permute(0,2,1).contiguous()
+        else:
+            x = self.bn2(x)
+            x = torch.flatten(x, 1)
+                # x = self.dropout(x)
+            # x = self.fc(x.float() if self.fp16 else x)
+            x = self.fc(x)
+            x = self.features(x)
+            return x
+        
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = IResNet(block, layers, **kwargs)
+    if pretrained:
+        raise ValueError()
+    return model
+
+
+def iresnet18(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet34(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet50(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet100(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet200(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
+                    progress, **kwargs)
+
+
+def getarcface(pretrained=None):
+    model = iresnet100().eval()
+    for param in model.parameters():
+        param.requires_grad=False
+
+    if pretrained is not None and os.path.exists(pretrained):
+        info = model.load_state_dict(torch.load(pretrained))
+        print(info)
+    return model
+
+
+if __name__=='__main__':
+    ckpt = 'pretrained/insightface_glint360k.pth'
+    arcface = iresnet100().eval()
+    info = arcface.load_state_dict(torch.load(ckpt))
+    print(info)
+    
+    id = arcface(torch.randn(1,3,128,128))
+    print(id.shape)
+    
+    # import cv2
+    # import numpy as np
+    # im1_crop256 = cv2.imread('happy.jpg')
+    # im2_crop256 = cv2.imread('angry.jpg')
+    
+    # im1_crop112 = cv2.resize(im1_crop256, (128,128))[0:112,8:120,:]
+    # im2_crop112 = cv2.resize(im2_crop256, (128,128))[0:112,8:120,:]
+    
+    # cv2.imwrite('1_112.jpg', im1_crop112)
+    # cv2.imwrite('2_112.jpg', im2_crop112)
+    
+    # # [-1,1] rgb
+    # im1_crop112_tensor = torch.from_numpy(im1_crop112[:,:,[2,1,0]].transpose(2, 0, 1).astype(np.float32)).unsqueeze(0)/127.5-1
+    # im2_crop112_tensor = torch.from_numpy(im2_crop112[:,:,[2,1,0]].transpose(2, 0, 1).astype(np.float32)).unsqueeze(0)/127.5-1
+    
+    # im1_id = arcface(im1_crop112_tensor)
+    # im2_id = arcface(im2_crop112_tensor)
+    
+    # loss_cos = torch.mean(1-torch.cosine_similarity(im1_id, im2_id, dim=1))
+    
+    # print(loss_cos)

utils/third_party/model_parsing.py

@@ -0,0 +1,381 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torch.utils.model_zoo as modelzoo
+
+resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_chan, out_chan, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_chan, out_chan, stride)
+        self.bn1 = nn.BatchNorm2d(out_chan)
+        self.conv2 = conv3x3(out_chan, out_chan)
+        self.bn2 = nn.BatchNorm2d(out_chan)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        if in_chan != out_chan or stride != 1:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_chan, out_chan,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_chan),
+                )
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = F.relu(self.bn1(residual))
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+
+        shortcut = x
+        if self.downsample is not None:
+            shortcut = self.downsample(x)
+
+        out = shortcut + residual
+        out = self.relu(out)
+        return out
+
+
+def create_layer_basic(in_chan, out_chan, bnum, stride=1):
+    layers = [BasicBlock(in_chan, out_chan, stride=stride)]
+    for i in range(bnum-1):
+        layers.append(BasicBlock(out_chan, out_chan, stride=1))
+    return nn.Sequential(*layers)
+
+
+class Resnet18(nn.Module):
+    def __init__(self):
+        super(Resnet18, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
+        self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
+        self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
+        self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(self.bn1(x))
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        feat8 = self.layer2(x) # 1/8
+        feat16 = self.layer3(feat8) # 1/16
+        feat32 = self.layer4(feat16) # 1/32
+        return feat8, feat16, feat32
+
+    def init_weight(self):
+        if os.path.isfile('./checkpoints/third_party/resnet18-5c106cde.pth'):
+            state_dict = torch.load('./checkpoints/third_party/resnet18-5c106cde.pth')
+        else:
+            state_dict = modelzoo.load_url(resnet18_url)
+        self_state_dict = self.state_dict()
+        for k, v in state_dict.items():
+            if 'fc' in k: continue
+            self_state_dict.update({k: v})
+        self.load_state_dict(self_state_dict)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module,  nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+
+class ConvBNReLU(nn.Module):
+    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_chan,
+                out_chan,
+                kernel_size = ks,
+                stride = stride,
+                padding = padding,
+                bias = False)
+        self.bn = nn.BatchNorm2d(out_chan)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = F.relu(self.bn(x))
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+class BiSeNetOutput(nn.Module):
+    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
+        super(BiSeNetOutput, self).__init__()
+        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.conv_out(x)
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class AttentionRefinementModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(AttentionRefinementModule, self).__init__()
+        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
+        self.bn_atten = nn.BatchNorm2d(out_chan)
+        self.sigmoid_atten = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv(x)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv_atten(atten)
+        atten = self.bn_atten(atten)
+        atten = self.sigmoid_atten(atten)
+        out = torch.mul(feat, atten)
+        return out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class ContextPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(ContextPath, self).__init__()
+        self.resnet = Resnet18()
+        self.arm16 = AttentionRefinementModule(256, 128)
+        self.arm32 = AttentionRefinementModule(512, 128)
+        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
+
+        self.init_weight()
+
+    def forward(self, x):
+        H0, W0 = x.size()[2:]
+        feat8, feat16, feat32 = self.resnet(x)
+        H8, W8 = feat8.size()[2:]
+        H16, W16 = feat16.size()[2:]
+        H32, W32 = feat32.size()[2:]
+
+        avg = F.avg_pool2d(feat32, feat32.size()[2:])
+        avg = self.conv_avg(avg)
+        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
+
+        feat32_arm = self.arm32(feat32)
+        feat32_sum = feat32_arm + avg_up
+        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
+        feat32_up = self.conv_head32(feat32_up)
+
+        feat16_arm = self.arm16(feat16)
+        feat16_sum = feat16_arm + feat32_up
+        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
+        feat16_up = self.conv_head16(feat16_up)
+
+        return feat8, feat16_up, feat32_up  # x8, x8, x16
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+### This is not used, since I replace this with the resnet feature with the same size
+class SpatialPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SpatialPath, self).__init__()
+        self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
+        self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv1(x)
+        feat = self.conv2(feat)
+        feat = self.conv3(feat)
+        feat = self.conv_out(feat)
+        return feat
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class FeatureFusionModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(FeatureFusionModule, self).__init__()
+        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+        self.conv1 = nn.Conv2d(out_chan,
+                out_chan//4,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.conv2 = nn.Conv2d(out_chan//4,
+                out_chan,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, fsp, fcp):
+        fcat = torch.cat([fsp, fcp], dim=1)
+        feat = self.convblk(fcat)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv1(atten)
+        atten = self.relu(atten)
+        atten = self.conv2(atten)
+        atten = self.sigmoid(atten)
+        feat_atten = torch.mul(feat, atten)
+        feat_out = feat_atten + feat
+        return feat_out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class BiSeNet(nn.Module):
+    def __init__(self, n_classes, *args, **kwargs):
+        super(BiSeNet, self).__init__()
+        self.cp = ContextPath()
+        ## here self.sp is deleted
+        self.ffm = FeatureFusionModule(256, 256)
+        self.conv_out = BiSeNetOutput(256, 256, n_classes)
+        self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
+        self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
+        self.init_weight()
+
+    def forward(self, x):
+        H, W = x.size()[2:]
+        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
+        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
+        feat_fuse = self.ffm(feat_sp, feat_cp8)
+
+        feat_out = self.conv_out(feat_fuse)
+        feat_out16 = self.conv_out16(feat_cp8)
+        feat_out32 = self.conv_out32(feat_cp16)
+
+        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
+        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
+        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
+        return feat_out, feat_out16, feat_out32
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
+        for name, child in self.named_children():
+            child_wd_params, child_nowd_params = child.get_params()
+            if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
+                lr_mul_wd_params += child_wd_params
+                lr_mul_nowd_params += child_nowd_params
+            else:
+                wd_params += child_wd_params
+                nowd_params += child_nowd_params
+        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
+
+
+def get_face_parsing(save_pth = 'third_party/pretrained/79999_iter.pth'):
+    net = BiSeNet(n_classes=19)
+    net.load_state_dict(torch.load(save_pth))
+    return net
+
+
+if __name__ == "__main__":
+    net = BiSeNet(19)
+    net.cuda()
+    net.eval()
+    in_ten = torch.randn(16, 3, 640, 480).cuda()
+    out, out16, out32 = net(in_ten)
+    print(out.shape)
+
+    net.get_params()

utils/third_party/model_resnet_d3dfr.py

@@ -0,0 +1,554 @@
+import os
+import torch.nn as nn
+from torch.nn import Linear
+from torch.nn import Conv2d
+from torch.nn import BatchNorm1d
+from torch.nn import BatchNorm2d
+from torch.nn import ReLU
+from torch.nn import Dropout
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+from torch.nn import MaxPool2d
+from torch.nn import Sequential
+from torch.nn import Module
+import torch
+from torch import Tensor
+from typing import Type, Any, Callable, Union, List, Optional
+
+
+model_urls = {
+    'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',
+}
+
+def filter_state_dict(state_dict, remove_name='fc'):
+    new_state_dict = {}
+    for key in state_dict:
+        if remove_name in key:
+            continue
+        new_state_dict[key] = state_dict[key]
+    return new_state_dict
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """ 3x3 convolution with padding
+    """
+    return Conv2d(in_planes,
+                  out_planes,
+                  kernel_size=3,
+                  stride=stride,
+                  padding=1,
+                  bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1, bias=False):
+    """ 1x1 convolution
+    """
+    return Conv2d(in_planes,
+                  out_planes,
+                  kernel_size=1,
+                  stride=stride,
+                  bias=bias)
+
+def conv3x3_(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+    """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: int, out_planes: int, stride: int = 1, bias: bool = False) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=bias)
+
+
+class Bottleneck(Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = BatchNorm2d(planes * self.expansion)
+        self.relu = 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 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: int = 4
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            groups: int = 1,
+            base_width: int = 64,
+            dilation: int = 1,
+            norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> 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: Tensor) -> Tensor:
+        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 BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            groups: int = 1,
+            base_width: int = 64,
+            dilation: int = 1,
+            norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> 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: Tensor) -> Tensor:
+        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 ResNet(Module):
+    """ ResNet backbone
+    """
+    def __init__(self, input_size, block, layers, zero_init_residual=True):
+        """ Args:
+            input_size: input_size of backbone
+            block: block function
+            layers: layers in each block
+        """
+        super(ResNet, self).__init__()
+        assert input_size[0] in [112, 224], \
+            "input_size should be [112, 112] or [224, 224]"
+        self.inplanes = 64
+        self.conv1 = Conv2d(3, 64,
+                            kernel_size=7,
+                            stride=2,
+                            padding=3,
+                            bias=False)
+        self.bn1 = BatchNorm2d(64)
+        self.relu = ReLU(inplace=True)
+        self.maxpool = 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.bn_o1 = BatchNorm2d(2048)
+        self.dropout = Dropout()
+        if input_size[0] == 112:
+            self.fc = Linear(2048 * 4 * 4, 512)
+        else:
+            self.fc = Linear(2048 * 7 * 7, 512)
+        self.bn_o2 = BatchNorm1d(512)
+
+        # initialize_weights(self.modules)
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                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 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.bn_o1(x)
+        x = self.dropout(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        x = self.bn_o2(x)
+
+        return x
+
+
+class resNet(nn.Module): # ori resnet
+
+    def __init__(
+            self,
+            block_: Type[Union[BasicBlock, Bottleneck_]],
+            layers: List[int],
+            num_classes: int = 1000,
+            zero_init_residual: bool = False,
+            use_last_fc: bool = False,
+            groups: int = 1,
+            width_per_group: int = 64,
+            replace_stride_with_dilation: Optional[List[bool]] = None,
+            norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> 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.use_last_fc = use_last_fc
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        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))
+
+        if self.use_last_fc:
+            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)  # type: ignore[arg-type]
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
+
+    def _make_layer(self, block_: Type[Union[BasicBlock, Bottleneck_]], planes: int, blocks: int,
+                    stride: int = 1, dilate: bool = False) -> nn.Sequential:
+        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: Tensor) -> Tensor:
+        # 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)
+        if self.use_last_fc:
+            x = torch.flatten(x, 1)
+            x = self.fc(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+def ResNet_50(input_size, **kwargs):
+    """ Constructs a ResNet-50 model.
+    """
+    model = ResNet(input_size, Bottleneck, [3, 4, 6, 3], **kwargs)
+
+    return model
+
+
+class ResNet50_nofc(Module):
+    """ ResNet backbone
+    """
+    def __init__(self, input_size, output_dim, use_last_fc=False, init_path=None):
+        """ Args:
+            input_size: input_size of backbone
+            block: block function
+            layers: layers in each block
+        """
+        super(ResNet50_nofc, self).__init__()
+        assert input_size[0] in [112, 224, 256], \
+            "input_size should be [112, 112] or [224, 224]"
+        func, last_dim = func_dict['resnet50']
+        self.use_last_fc=use_last_fc
+        backbone = func(use_last_fc=use_last_fc, num_classes=output_dim)
+        if init_path and os.path.isfile(init_path):
+            state_dict = filter_state_dict(torch.load(init_path, map_location='cpu'))
+            backbone.load_state_dict(state_dict)
+            print("Loading init recon %s from %s"%('resnet50', init_path))
+        self.backbone = backbone
+        if not use_last_fc:
+            self.fianl_layers = nn.ModuleList([
+                conv1x1(last_dim, 80, bias=True), # id
+                conv1x1(last_dim, 64, bias=True), # exp
+                conv1x1(last_dim, 80, bias=True), # tex
+                conv1x1(last_dim, 3, bias=True), # angle
+                conv1x1(last_dim, 27, bias=True), # gamma
+                conv1x1(last_dim, 2, bias=True), # tx, ty
+                conv1x1(last_dim, 1, bias=True), # tz
+                conv1x1(last_dim, 4, bias=True) # pupil
+            ])
+            for m in self.fianl_layers:
+                nn.init.constant_(m.weight, 0.)
+                nn.init.constant_(m.bias, 0.)
+
+
+    def forward(self, x):
+        x = self.backbone(x)
+        if not self.use_last_fc:
+            output = []
+            for layer in self.fianl_layers:
+                output.append(layer(x))
+            x = torch.flatten(torch.cat(output, dim=1), 1)
+        return x
+
+
+def _resnet(
+        arch: str,
+        block: Type[Union[BasicBlock, Bottleneck_]],
+        layers: List[int],
+        pretrained: bool,
+        progress: bool,
+        **kwargs: Any
+) -> ResNet:
+    model = resNet(block, layers, **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch],
+                                              progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> resNet:
+    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,
+                   **kwargs)
+
+
+func_dict = {
+    'resnet50': (resnet50, 2048),
+}
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+    
+
+def fuse(conv, bn):
+    w = conv.weight
+    mean = bn.running_mean
+    var_sqrt = torch.sqrt(bn.running_var + bn.eps)
+
+    beta = bn.weight
+    gamma = bn.bias
+
+    if conv.bias is not None:
+        b = conv.bias
+    else:
+        b = mean.new_zeros(mean.shape)
+
+    w = w * (beta / var_sqrt).reshape([conv.out_channels, 1, 1, 1])
+    b = (b - mean) / var_sqrt * beta + gamma
+
+    fused_conv = nn.Conv2d(
+        conv.in_channels,
+        conv.out_channels,
+        conv.kernel_size,
+        conv.stride,
+        conv.padding,
+        conv.dilation,
+        conv.groups,
+        bias=True,
+        padding_mode=conv.padding_mode
+    )
+    fused_conv.weight = nn.Parameter(w)
+    fused_conv.bias = nn.Parameter(b)
+    return fused_conv
+
+
+def fuse_module(m):
+    children = list(m.named_children())
+    conv = None
+    conv_name = None
+    for name, child in children:
+        if isinstance(child, nn.BatchNorm2d) and conv:
+            bc = fuse(conv, child)
+            m._modules[conv_name] = bc
+            m._modules[name] = Identity()
+            conv = None
+        elif isinstance(child, nn.Conv2d):
+            conv = child
+            conv_name = name
+        else:
+            fuse_module(child)
+
+
+def getd3dfr_res50(pretrained="./d3dfr_res50_nofc.pth"):
+    model = ResNet50_nofc([256, 256], 257+4, use_last_fc=False)
+    for param in model.parameters():
+        param.requires_grad=False
+    if pretrained is not None and os.path.exists(pretrained):
+        checkpoint_no_module = {}
+        checkpoint = torch.load(pretrained, map_location=lambda storage, loc: storage)
+        for k, v in checkpoint.items():
+            if k.startswith('module'):
+                k = k[7:]
+            checkpoint_no_module[k] = v
+        info = model.load_state_dict(checkpoint_no_module, strict=False)
+
+        print(pretrained, info)
+    model = model.eval()
+    fuse_module(model)
+    return model
+if __name__ == '__main__':
+    model = getd3dfr_res50()

utils/third_party_files/79999_iter.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:468e13ca13a9b43cc0881a9f99083a430e9c0a38abd935431d1c28ee94b26567
+size 53289463

utils/third_party_files/BFM_model_front.mat

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

utils/third_party_files/d3dfr_res50_nofc.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52c54b90304a06c16b6813910c26faff1a907d4f8bd69a71ad4ecff43b41a090
+size 96449126

utils/third_party_files/insightface_glint360k.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f631718e783448b41631e15073bdc622eaeef56509bbad4e5085f23bd32db83
+size 261223796

utils/third_party_files/models/antelopev2/1k3d68.onnx

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

utils/third_party_files/models/antelopev2/2d106det.onnx

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

utils/third_party_files/models/antelopev2/antelopev2.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7353a5fdca5a90e11d2792e0236032b2fe42adc1ea23eaef5cf8c8b57e7e9393
+size 360743453

utils/third_party_files/models/antelopev2/genderage.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fde69b1c810857b88c64a335084f1c3fe8f01246c9a191b48c7bb756d6652fb
+size 1322532

utils/third_party_files/models/antelopev2/glintr100.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ab1d6435d639628a6f3e5008dd4f929edf4c4124b1a7169e1048f9fef534cdf
+size 260665334

utils/third_party_files/models/antelopev2/scrfd_10g_bnkps.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5838f7fe053675b1c7a08b633df49e7af5495cee0493c7dcf6697200b85b5b91
+size 16923827

utils/third_party_files/resnet18-5c106cde.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c106cde386e87d4033832f2996f5493238eda96ccf559d1d62760c4de0613f8
+size 46827520