add face_reconstruction model

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11527525

data/test/images/face_reconstruction.jpg

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

modelscope/metainfo.py

@@ -268,6 +268,7 @@ class Pipelines(object):
     image_object_detection_auto = 'yolox_image-object-detection-auto'
     hand_detection = 'yolox-pai_hand-detection'
     skin_retouching = 'unet-skin-retouching'
+    face_reconstruction = 'resnet50-face-reconstruction'
     tinynas_classification = 'tinynas-classification'
     easyrobust_classification = 'easyrobust-classification'
     tinynas_detection = 'tinynas-detection'

modelscope/models/cv/__init__.py

@@ -4,19 +4,19 @@
 from . import (action_recognition, animal_recognition, body_2d_keypoints,
                body_3d_keypoints, cartoon, cmdssl_video_embedding,
                crowd_counting, face_2d_keypoints, face_detection,
-               face_generation, human_wholebody_keypoint, image_classification,
-               image_color_enhance, image_colorization, image_defrcn_fewshot,
-               image_denoise, image_inpainting, image_instance_segmentation,
-               image_matching, image_mvs_depth_estimation,
-               image_panoptic_segmentation, image_portrait_enhancement,
-               image_probing_model, image_quality_assessment_mos,
-               image_reid_person, image_restoration,
-               image_semantic_segmentation, image_to_image_generation,
-               image_to_image_translation, language_guided_video_summarization,
-               movie_scene_segmentation, object_detection,
-               panorama_depth_estimation, pointcloud_sceneflow_estimation,
-               product_retrieval_embedding, realtime_object_detection,
-               referring_video_object_segmentation,
+               face_generation, face_reconstruction, human_wholebody_keypoint,
+               image_classification, image_color_enhance, image_colorization,
+               image_defrcn_fewshot, image_denoise, image_inpainting,
+               image_instance_segmentation, image_matching,
+               image_mvs_depth_estimation, image_panoptic_segmentation,
+               image_portrait_enhancement, image_probing_model,
+               image_quality_assessment_mos, image_reid_person,
+               image_restoration, image_semantic_segmentation,
+               image_to_image_generation, image_to_image_translation,
+               language_guided_video_summarization, movie_scene_segmentation,
+               object_detection, panorama_depth_estimation,
+               pointcloud_sceneflow_estimation, product_retrieval_embedding,
+               realtime_object_detection, referring_video_object_segmentation,
                robust_image_classification, salient_detection,
                shop_segmentation, super_resolution, video_frame_interpolation,
                video_object_segmentation, video_panoptic_segmentation,

modelscope/models/cv/face_reconstruction/models/bfm.py

@@ -0,0 +1,591 @@
+# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
+# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
+import os
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.io import loadmat
+
+from ..utils import read_obj, transferBFM09
+
+
+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 ParametricFaceModel:
+
+    def __init__(self,
+                 bfm_folder='./asset/BFM',
+                 recenter=True,
+                 camera_distance=10.,
+                 init_lit=np.array([0.8, 0, 0, 0, 0, 0, 0, 0, 0]),
+                 focal=1015.,
+                 center=112.,
+                 is_train=True,
+                 default_name='BFM_model_front.mat'):
+
+        if not os.path.isfile(os.path.join(bfm_folder, default_name)):
+            transferBFM09(bfm_folder)
+        model = loadmat(os.path.join(bfm_folder, default_name))
+        # mean face shape. [3*N,1]
+        self.mean_shape = model['meanshape'].astype(np.float32)
+
+        # identity basis. [3*N,80]
+        self.id_base = model['idBase'].astype(np.float32)
+
+        # expression basis. [3*N,64]
+        self.exp_base = model['exBase'].astype(np.float32)
+
+        # mean face texture. [3*N,1] (0-255)
+        self.mean_tex = model['meantex'].astype(np.float32)
+
+        # texture basis. [3*N,80]
+        self.tex_base = model['texBase'].astype(np.float32)
+
+        # face indices for each vertex that lies in. starts from 0. [N,8]
+        self.point_buf = model['point_buf'].astype(np.int64) - 1
+
+        # vertex indices for each face. starts from 0. [F,3]
+        self.face_buf = model['tri'].astype(np.int64) - 1
+
+        # vertex indices for 68 landmarks. starts from 0. [68,1]
+        self.keypoints = np.squeeze(model['keypoints']).astype(np.int64) - 1
+
+        self.mean_shape_ori = model['meanshape_ori'].astype(np.float32)
+        self.bfm_keep_inds = model['bfm_keep_inds'][0]
+        self.nose_reduced_part = model['nose_reduced_part'].reshape(
+            (1, -1)) - self.mean_shape
+        self.nonlinear_UVs = model['nonlinear_UVs']
+
+        if default_name == 'head_model_for_maas.mat':
+            self.ours_hair_area_inds = model['hair_area_inds'][0]
+
+            self.mean_tex = self.mean_tex.reshape(1, -1, 3)
+            mean_tex_keep = self.mean_tex[:, self.bfm_keep_inds]
+            self.mean_tex[:, :len(self.bfm_keep_inds)] = mean_tex_keep
+            self.mean_tex[:,
+                          len(self.bfm_keep_inds):] = np.array([200, 146,
+                                                                118])[None,
+                                                                      None]
+            self.mean_tex[:, self.ours_hair_area_inds] = 40.0
+            self.mean_tex = self.mean_tex.reshape(1, -1)
+            self.mean_tex = np.ascontiguousarray(self.mean_tex)
+
+            self.tex_base = self.tex_base.reshape(-1, 3, 80)
+            tex_base_keep = self.tex_base[self.bfm_keep_inds]
+            self.tex_base[:len(self.bfm_keep_inds)] = tex_base_keep
+            self.tex_base[len(self.bfm_keep_inds):] = 0.0
+            self.tex_base = self.tex_base.reshape(-1, 80)
+            self.tex_base = np.ascontiguousarray(self.tex_base)
+
+            self.point_buf = self.point_buf[:, :8] + 1
+
+            self.neck_adjust_part = model['neck_adjust_part'].reshape(
+                (1, -1)) - self.mean_shape
+            self.eyes_adjust_part = model['eyes_adjust_part'].reshape(
+                (1, -1)) - self.mean_shape
+
+            self.eye_corner_inds = model['eye_corner_inds'][0]
+            self.eye_corner_lines = model['eye_corner_lines']
+
+        if recenter:
+            mean_shape = self.mean_shape.reshape([-1, 3])
+            mean_shape_ori = self.mean_shape_ori.reshape([-1, 3])
+            mean_shape = mean_shape - np.mean(
+                mean_shape_ori[:35709, ...], axis=0, keepdims=True)
+            self.mean_shape = mean_shape.reshape([-1, 1])
+
+        self.center = center
+        self.persc_proj = perspective_projection(focal, self.center)
+        self.device = 'cpu'
+        self.camera_distance = camera_distance
+        self.SH = SH()
+        self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32)
+
+    def to(self, device):
+        self.device = device
+        for key, value in self.__dict__.items():
+            if type(value).__module__ == np.__name__:
+                setattr(self, key, torch.tensor(value).to(device))
+
+    def compute_shape(self,
+                      id_coeff,
+                      exp_coeff,
+                      nose_coeff=0.0,
+                      neck_coeff=0.0,
+                      eyes_coeff=0.0):
+        """
+        Return:
+            face_shape       -- torch.tensor, size (B, N, 3)
+
+        Parameters:
+            id_coeff         -- torch.tensor, size (B, 80), 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.id_base, id_coeff)
+        exp_part = torch.einsum('ij,aj->ai', self.exp_base, exp_coeff)
+        face_shape = id_part + exp_part + self.mean_shape.reshape([1, -1])
+
+        if nose_coeff != 0:
+            face_shape = face_shape + nose_coeff * self.nose_reduced_part
+        if neck_coeff != 0:
+            face_shape = face_shape + neck_coeff * self.neck_adjust_part
+        if eyes_coeff != 0 and self.eyes_adjust_part is not None:
+            face_shape = face_shape + eyes_coeff * self.eyes_adjust_part
+
+        return face_shape.reshape([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.tex_base,
+                                    tex_coeff) + self.mean_tex
+        if normalize:
+            face_texture = face_texture / 255.
+        return face_texture.reshape([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.device)],
+            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]
+        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)
+
+        y1 = a[0] * c[0] * torch.ones_like(face_norm[..., :1]).to(self.device)
+        y2 = -a[1] * c[1] * face_norm[..., 1:2]
+        y3 = a[1] * c[1] * face_norm[..., 2:]
+        y4 = -a[1] * c[1] * face_norm[..., :1]
+        y5 = a[2] * c[2] * face_norm[..., :1] * face_norm[..., 1:2]
+        y6 = -a[2] * c[2] * face_norm[..., 1:2] * face_norm[..., 2:]
+        y7 = 0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm[..., 2:]**2 - 1)
+        y8 = -a[2] * c[2] * face_norm[..., :1] * face_norm[..., 2:]
+        y9 = 0.5 * a[2] * c[2] * (
+            face_norm[..., :1]**2 - face_norm[..., 1:2]**2)
+        Y = torch.cat([y1, y2, y3, y4, y5, y6, y7, y8, y9], 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.device)
+        zeros = torch.zeros([batch_size, 1]).to(self.device)
+        x, y, z = angles[:, :1], angles[:, 1:2], angles[:, 2:],
+
+        value_list = [
+            ones, zeros, zeros, zeros,
+            torch.cos(x), -torch.sin(x), zeros,
+            torch.sin(x),
+            torch.cos(x)
+        ]
+        rot_x = torch.cat(value_list, dim=1).reshape([batch_size, 3, 3])
+
+        value_list = [
+            torch.cos(y), zeros,
+            torch.sin(y), zeros, ones, zeros, -torch.sin(y), zeros,
+            torch.cos(y)
+        ]
+        rot_y = torch.cat(value_list, dim=1).reshape([batch_size, 3, 3])
+
+        value_list = [
+            torch.cos(z), -torch.sin(z), zeros,
+            torch.sin(z),
+            torch.cos(z), zeros, zeros, zeros, ones
+        ]
+        rot_z = torch.cat(value_list, 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
+        face_proj = face_proj[..., :2] / face_proj[..., 2:]
+
+        return face_proj
+
+    def transform(self, face_shape, rot, trans):
+        """
+        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)
+            trans            -- torch.tensor, size (B, 3)
+        """
+        return face_shape @ rot + trans.unsqueeze(1)
+
+    def get_landmarks(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 split_coeff(self, coeffs):
+        """
+        Return:
+            coeffs_dict     -- a dict of torch.tensors
+
+        Parameters:
+            coeffs          -- torch.tensor, size (B, 256)
+        """
+        if type(coeffs) == dict and 'id' in coeffs:
+            return coeffs
+
+        id_coeffs = coeffs[:, :80]
+        exp_coeffs = coeffs[:, 80:144]
+        tex_coeffs = coeffs[:, 144:224]
+        angles = coeffs[:, 224:227]
+        gammas = coeffs[:, 227:254]
+        translations = coeffs[:, 254:]
+        return {
+            'id': id_coeffs,
+            'exp': exp_coeffs,
+            'tex': tex_coeffs,
+            'angle': angles,
+            'gamma': gammas,
+            'trans': translations
+        }
+
+    def merge_coeff(self, coeffs):
+        """
+        Return:
+            coeffs_dict     -- a dict of torch.tensors
+
+        Parameters:
+            coeffs          -- torch.tensor, size (B, 256)
+        """
+        names = ['id', 'exp', 'tex', 'angle', 'gamma', 'trans']
+        coeffs_merge = []
+        for name in names:
+            coeffs_merge.append(coeffs[name].detach())
+        coeffs_merge = torch.cat(coeffs_merge, dim=1)
+
+        return coeffs_merge
+
+    def compute_for_render(self, coeffs, coeffs_mvs=None):
+        """
+        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, 257)
+        """
+        if type(coeffs) == dict:
+            coef_dict = coeffs
+        elif type(coeffs) == torch.Tensor:
+            coef_dict = self.split_coeff(coeffs)
+
+        face_shape = self.compute_shape(
+            coef_dict['id'], coef_dict['exp'], nose_coeff=0.4, neck_coeff=0.6)
+
+        rotation = self.compute_rotation(coef_dict['angle'])
+
+        face_shape_transformed = self.transform(face_shape, rotation,
+                                                coef_dict['trans'])
+        face_vertex = self.to_camera(face_shape_transformed)
+        face_vertex_ori = self.to_camera(face_shape)
+
+        face_proj = self.to_image(face_vertex)
+        landmark = self.get_landmarks(face_proj)
+
+        face_texture = self.compute_texture(coef_dict['tex'])
+        face_norm = self.compute_norm(face_shape)
+        face_norm_roted = face_norm @ rotation
+        face_color = self.compute_color(face_texture, face_norm_roted,
+                                        coef_dict['gamma'])
+
+        if coeffs_mvs is not None:
+            mvs_face_shape = self.compute_shape(coeffs_mvs['id'],
+                                                coeffs_mvs['exp'])
+
+            mvs_face_shape_transformed = self.transform(
+                mvs_face_shape, rotation, coef_dict['trans'])
+            mvs_face_vertex = self.to_camera(mvs_face_shape_transformed)
+            return face_vertex, face_texture, face_color, landmark, mvs_face_vertex
+        else:
+            return face_vertex, face_texture, face_color, landmark, face_vertex_ori
+
+    def reverse_recenter(self, face_shape):
+        batch_size = face_shape.shape[0]
+        face_shape = face_shape.reshape([-1, 3])
+        mean_shape_ori = self.mean_shape_ori.reshape([-1, 3])
+        face_shape = face_shape + torch.mean(
+            mean_shape_ori[:35709, ...], dim=0, keepdim=True)
+        face_shape = face_shape.reshape([batch_size, -1, 3])
+        return face_shape
+
+    def add_nonlinear_offset_eyes(self, face_shape, shape_offset):
+        assert face_shape.shape[0] == 1 and shape_offset.shape[0] == 1
+        face_shape = face_shape[0]
+        shape_offset = shape_offset[0]
+
+        corner_inds = self.eye_corner_inds
+        lines = self.eye_corner_lines
+
+        corner_shape = face_shape[-625:, :]
+        corner_offset = shape_offset[corner_inds]
+        for i in range(len(lines)):
+            corner_shape[lines[i]] += corner_offset[i][None, ...]
+        face_shape[-625:, :] = corner_shape
+
+        l_eye_landmarks = [11540, 11541]
+        r_eye_landmarks = [4271, 4272]
+
+        l_eye_offset = torch.mean(
+            shape_offset[l_eye_landmarks], dim=0, keepdim=True)
+        face_shape[37082:37082 + 609] += l_eye_offset
+
+        r_eye_offset = torch.mean(
+            shape_offset[r_eye_landmarks], dim=0, keepdim=True)
+        face_shape[37082 + 609:37082 + 609 + 608] += r_eye_offset
+
+        face_shape = face_shape[None, ...]
+
+        return face_shape
+
+    def add_nonlinear_offset(self, face_shape, shape_offset_uv, UVs):
+        """
+
+        Args:
+            face_shape: torch.tensor, size (1, N, 3)
+            shape_offset_uv: torch.tensor, size (1, h, w, 3)
+            UVs: torch.tensor, size (N, 2)
+
+        Returns:
+
+        """
+        assert face_shape.shape[0] == 1 and shape_offset_uv.shape[0] == 1
+        face_shape = face_shape[0]
+        shape_offset_uv = shape_offset_uv[0]
+
+        h, w = shape_offset_uv.shape[:2]
+        UVs_coords = UVs.clone()
+        UVs_coords[:, 0] *= w
+        UVs_coords[:, 1] *= h
+        UVs_coords_int = torch.floor(UVs_coords)
+        UVs_coords_float = UVs_coords - UVs_coords_int
+        UVs_coords_int = UVs_coords_int.long()
+
+        shape_lt = shape_offset_uv[(h - 1
+                                    - UVs_coords_int[:, 1]).clamp(0, h - 1),
+                                   UVs_coords_int[:, 0].clamp(0, w - 1)]
+        shape_lb = shape_offset_uv[(h - UVs_coords_int[:, 1]).clamp(0, h - 1),
+                                   UVs_coords_int[:, 0].clamp(0, w - 1)]
+        shape_rt = shape_offset_uv[(h - 1
+                                    - UVs_coords_int[:, 1]).clamp(0, h - 1),
+                                   (UVs_coords_int[:, 0] + 1).clamp(0, w - 1)]
+        shape_rb = shape_offset_uv[(h - UVs_coords_int[:, 1]).clamp(0, h - 1),
+                                   (UVs_coords_int[:, 0] + 1).clamp(0, w - 1)]
+
+        value1 = shape_lt * (
+            1 - UVs_coords_float[:, :1]) * UVs_coords_float[:, 1:]
+        value2 = shape_lb * (1 - UVs_coords_float[:, :1]) * (
+            1 - UVs_coords_float[:, 1:])
+        value3 = shape_rt * UVs_coords_float[:, :1] * UVs_coords_float[:, 1:]
+        value4 = shape_rb * UVs_coords_float[:, :1] * (
+            1 - UVs_coords_float[:, 1:])
+        offset_shape = value1 + value2 + value3 + value4  # (N, 3)
+
+        face_shape = (face_shape + offset_shape)[None, ...]
+
+        return face_shape, offset_shape[None, ...]
+
+    def compute_for_render_train_nonlinear(self,
+                                           coeffs,
+                                           shape_offset_uv,
+                                           tex_offset_uv,
+                                           UVs,
+                                           reverse_recenter=True):
+        if type(coeffs) == dict:
+            coef_dict = coeffs
+        elif type(coeffs) == torch.Tensor:
+            coef_dict = self.split_coeff(coeffs)
+
+        face_shape = self.compute_shape(coef_dict['id'],
+                                        coef_dict['exp'])  # (1, n, 3)
+        if reverse_recenter:
+            face_shape_ori_noRecenter = self.reverse_recenter(
+                face_shape.clone())
+        else:
+            face_shape_ori_noRecenter = face_shape.clone()
+        face_vertex_ori = self.to_camera(face_shape_ori_noRecenter)
+
+        face_shape, shape_offset = self.add_nonlinear_offset(
+            face_shape, shape_offset_uv, UVs[:35709, :])  # (1, n, 3)
+        if reverse_recenter:
+            face_shape_offset_noRecenter = self.reverse_recenter(
+                face_shape.clone())
+        else:
+            face_shape_offset_noRecenter = face_shape.clone()
+        face_vertex_offset = self.to_camera(face_shape_offset_noRecenter)
+
+        rotation = self.compute_rotation(coef_dict['angle'])
+
+        face_shape_transformed = self.transform(face_shape, rotation,
+                                                coef_dict['trans'])
+        face_vertex = self.to_camera(face_shape_transformed)
+
+        face_proj = self.to_image(face_vertex)
+        landmark = self.get_landmarks(face_proj)
+
+        face_texture = self.compute_texture(coef_dict['tex'])  # (1, n, 3)
+        face_texture, texture_offset = self.add_nonlinear_offset(
+            face_texture, tex_offset_uv, UVs[:35709, :])
+        face_norm = self.compute_norm(face_shape)
+        face_norm_roted = face_norm @ rotation
+        face_color = self.compute_color(face_texture, face_norm_roted,
+                                        coef_dict['gamma'])
+
+        return face_vertex, face_texture, face_color, landmark, face_vertex_ori, face_vertex_offset, face_proj
+
+    def compute_for_render_nonlinear_full(self,
+                                          coeffs,
+                                          shape_offset_uv,
+                                          UVs,
+                                          nose_coeff=0.0,
+                                          eyes_coeff=0.0):
+        if type(coeffs) == dict:
+            coef_dict = coeffs
+        elif type(coeffs) == torch.Tensor:
+            coef_dict = self.split_coeff(coeffs)
+
+        face_shape = self.compute_shape(
+            coef_dict['id'],
+            coef_dict['exp'],
+            nose_coeff=nose_coeff,
+            neck_coeff=0.6,
+            eyes_coeff=eyes_coeff)  # (1, n, 3)
+        face_vertex_ori = self.to_camera(face_shape.clone())
+
+        face_shape[:, :35241, :], shape_offset = self.add_nonlinear_offset(
+            face_shape[:, :35241, :], shape_offset_uv,
+            UVs[:35709, :][self.bfm_keep_inds])
+        face_shape = self.add_nonlinear_offset_eyes(face_shape, shape_offset)
+        face_shape_noRecenter = self.reverse_recenter(face_shape.clone())
+        face_vertex_offset = self.to_camera(face_shape_noRecenter)
+
+        rotation = self.compute_rotation(coef_dict['angle'])
+
+        face_shape_transformed = self.transform(face_shape, rotation,
+                                                coef_dict['trans'])
+        face_vertex = self.to_camera(face_shape_transformed)
+
+        return face_vertex, face_vertex_ori, face_vertex_offset
+
+    def compute_for_render_train(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, 257)
+        """
+        if type(coeffs) == dict:
+            coef_dict = coeffs
+        elif type(coeffs) == torch.Tensor:
+            coef_dict = self.split_coeff(coeffs)
+
+        face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp'])
+        uv_geometry = self.render.world2uv(face_shape)
+
+        rotation = self.compute_rotation(coef_dict['angle'])
+
+        face_shape_transformed = self.transform(face_shape, rotation,
+                                                coef_dict['trans'])
+        face_vertex = self.to_camera(face_shape_transformed)
+
+        face_proj = self.to_image(face_vertex)
+        landmark = self.get_landmarks(face_proj)
+
+        face_texture = self.compute_texture(coef_dict['tex'])
+        face_norm = self.compute_norm(face_shape)
+        face_norm_roted = face_norm @ rotation
+        face_color = self.compute_color(face_texture, face_norm_roted,
+                                        coef_dict['gamma'])
+
+        return face_vertex, face_texture, face_color, landmark, uv_geometry

modelscope/models/cv/face_reconstruction/models/facelandmark/large_base_lmks_infer.py

@@ -0,0 +1,91 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import numpy as np
+import torch
+
+from .nets.large_base_lmks_net import LargeBaseLmksNet
+
+BASE_LANDMARK_NUM = 106
+INPUT_SIZE = 224
+ENLARGE_RATIO = 1.35
+
+
+class LargeBaseLmkInfer:
+
+    @staticmethod
+    def model_preload(model_path, use_gpu=True):
+        model = LargeBaseLmksNet(infer=False)
+        # using gpu
+        if use_gpu:
+            model = model.cuda()
+
+        checkpoint = []
+        if use_gpu:
+            checkpoint = torch.load(model_path, map_location='cuda')
+        else:
+            checkpoint = torch.load(model_path, map_location='cpu')
+
+        model.load_state_dict(
+            {
+                k.replace('module.', ''): v
+                for k, v in checkpoint['state_dict'].items()
+            },
+            strict=False)
+        model.eval()
+        return model
+
+    @staticmethod
+    def process_img(model, image, use_gpu=True):
+        img_resize = image
+
+        img_resize = (img_resize
+                      - [103.94, 116.78, 123.68]) / 255.0  # important
+        img_resize = img_resize.transpose([2, 0, 1])
+
+        if use_gpu:
+            img_resize = torch.from_numpy(img_resize).cuda()
+        else:
+            img_resize = torch.from_numpy(img_resize)
+
+        w_new = INPUT_SIZE
+        h_new = INPUT_SIZE
+        img_in = torch.zeros([1, 3, h_new, w_new], dtype=torch.float32)
+        if use_gpu:
+            img_in = img_in.cuda()
+
+        img_in[0, :] = img_resize
+
+        with torch.no_grad():
+            output = model(img_in)
+            output = output * INPUT_SIZE
+
+        if use_gpu:
+            output = output.cpu().numpy()
+        else:
+            output = output.numpy()
+
+        return output
+
+    @staticmethod
+    def smooth(cur_lmks, prev_lmks):
+        smooth_lmks = np.zeros((106, 2))
+
+        cur_rect_x1 = np.min(cur_lmks[:, 0])
+        cur_rect_x2 = np.max(cur_lmks[:, 0])
+
+        smooth_param = 60.0
+        factor = smooth_param / (cur_rect_x1 - cur_rect_x2)
+        for i in range(BASE_LANDMARK_NUM):
+            weightX = np.exp(factor * np.abs(cur_lmks[i][0] - prev_lmks[i][0]))
+            weightY = np.exp(factor * np.abs(cur_lmks[i][1] - prev_lmks[i][1]))
+
+            smooth_lmks[i][0] = (
+                1 - weightX) * cur_lmks[i][0] + weightX * prev_lmks[i][0]
+            smooth_lmks[i][1] = (
+                1 - weightY) * cur_lmks[i][1] + weightY * prev_lmks[i][1]
+
+        return smooth_lmks
+
+    @staticmethod
+    def infer_img(img, model, use_gpu=True):
+        lmks = LargeBaseLmkInfer.process_img(model, img, use_gpu)
+        return lmks

modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py

@@ -0,0 +1,430 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from modelscope.models.cv.skin_retouching.retinaface.predict_single import \
+    Model
+from ...utils import image_warp_grid1, spread_flow
+from .large_base_lmks_infer import LargeBaseLmkInfer
+
+INPUT_SIZE = 224
+ENLARGE_RATIO = 1.35
+
+
+def resize_on_long_side(img, long_side=800):
+    src_height = img.shape[0]
+    src_width = img.shape[1]
+
+    if src_height > src_width:
+        scale = long_side * 1.0 / src_height
+        _img = cv2.resize(
+            img, (int(src_width * scale), long_side),
+            interpolation=cv2.INTER_CUBIC)
+
+    else:
+        scale = long_side * 1.0 / src_width
+        _img = cv2.resize(
+            img, (long_side, int(src_height * scale)),
+            interpolation=cv2.INTER_CUBIC)
+
+    return _img, scale
+
+
+def draw_line(im, points, color, stroke_size=2, closed=False):
+    points = points.astype(np.int32)
+    for i in range(len(points) - 1):
+        cv2.line(im, tuple(points[i]), tuple(points[i + 1]), color,
+                 stroke_size)
+    if closed:
+        cv2.line(im, tuple(points[0]), tuple(points[-1]), color, stroke_size)
+
+
+def enlarged_bbox(bbox, img_width, img_height, enlarge_ratio=0.2):
+    '''
+    :param bbox: [xmin,ymin,xmax,ymax]
+    :return: bbox: [xmin,ymin,xmax,ymax]
+    '''
+
+    left = bbox[0]
+    top = bbox[1]
+
+    right = bbox[2]
+    bottom = bbox[3]
+
+    roi_width = right - left
+    roi_height = bottom - top
+
+    new_left = left - int(roi_width * enlarge_ratio)
+    new_left = 0 if new_left < 0 else new_left
+
+    new_top = top - int(roi_height * enlarge_ratio)
+    new_top = 0 if new_top < 0 else new_top
+
+    new_right = right + int(roi_width * enlarge_ratio)
+    new_right = img_width if new_right > img_width else new_right
+
+    new_bottom = bottom + int(roi_height * enlarge_ratio)
+    new_bottom = img_height if new_bottom > img_height else new_bottom
+
+    bbox = [new_left, new_top, new_right, new_bottom]
+
+    bbox = [int(x) for x in bbox]
+
+    return bbox
+
+
+class FaceInfo:
+
+    def __init__(self):
+        self.rect = np.asarray([0, 0, 0, 0])
+        self.points_array = np.zeros((106, 2))
+        self.eye_left = np.zeros((22, 2))
+        self.eye_right = np.zeros((22, 2))
+        self.eyebrow_left = np.zeros((13, 2))
+        self.eyebrow_right = np.zeros((13, 2))
+        self.lips = np.zeros((64, 2))
+
+
+class LargeModelInfer:
+
+    def __init__(self, ckpt, device='cuda'):
+        self.large_base_lmks_model = LargeBaseLmkInfer.model_preload(
+            ckpt,
+            device.lower() == 'cuda')
+        self.device = device.lower()
+        self.detector = Model(max_size=512, device=device)
+        detector_ckpt_name = 'retinaface_resnet50_2020-07-20_old_torch.pth'
+        state_dict = torch.load(
+            os.path.join(os.path.dirname(ckpt), detector_ckpt_name),
+            map_location='cpu')
+        self.detector.load_state_dict(state_dict)
+        self.detector.eval()
+
+    def infer(self, img_bgr):
+        landmarks = []
+
+        rgb_image = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        results = self.detector.predict_jsons(rgb_image)
+
+        boxes = []
+        for anno in results:
+            if anno['score'] == -1:
+                break
+            boxes.append({
+                'x1': anno['bbox'][0],
+                'y1': anno['bbox'][1],
+                'x2': anno['bbox'][2],
+                'y2': anno['bbox'][3]
+            })
+
+        for detect_result in boxes:
+            x1 = detect_result['x1']
+            y1 = detect_result['y1']
+            x2 = detect_result['x2']
+            y2 = detect_result['y2']
+
+            w = x2 - x1 + 1
+            h = y2 - y1 + 1
+
+            cx = (x2 + x1) / 2
+            cy = (y2 + y1) / 2
+
+            sz = max(h, w) * ENLARGE_RATIO
+
+            x1 = cx - sz / 2
+            y1 = cy - sz / 2
+            trans_x1 = x1
+            trans_y1 = y1
+            x2 = x1 + sz
+            y2 = y1 + sz
+
+            height, width, _ = rgb_image.shape
+            dx = max(0, -x1)
+            dy = max(0, -y1)
+            x1 = max(0, x1)
+            y1 = max(0, y1)
+
+            edx = max(0, x2 - width)
+            edy = max(0, y2 - height)
+            x2 = min(width, x2)
+            y2 = min(height, y2)
+
+            crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)]
+            if dx > 0 or dy > 0 or edx > 0 or edy > 0:
+                crop_img = cv2.copyMakeBorder(
+                    crop_img,
+                    int(dy),
+                    int(edy),
+                    int(dx),
+                    int(edx),
+                    cv2.BORDER_CONSTANT,
+                    value=(103.94, 116.78, 123.68))
+            crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE))
+
+            base_lmks = LargeBaseLmkInfer.infer_img(crop_img,
+                                                    self.large_base_lmks_model,
+                                                    self.device == 'cuda')
+
+            inv_scale = sz / INPUT_SIZE
+
+            affine_base_lmks = np.zeros((106, 2))
+            for idx in range(106):
+                affine_base_lmks[idx][
+                    0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1
+                affine_base_lmks[idx][
+                    1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1
+
+            x1 = np.min(affine_base_lmks[:, 0])
+            y1 = np.min(affine_base_lmks[:, 1])
+            x2 = np.max(affine_base_lmks[:, 0])
+            y2 = np.max(affine_base_lmks[:, 1])
+
+            w = x2 - x1 + 1
+            h = y2 - y1 + 1
+
+            cx = (x2 + x1) / 2
+            cy = (y2 + y1) / 2
+
+            sz = max(h, w) * ENLARGE_RATIO
+
+            x1 = cx - sz / 2
+            y1 = cy - sz / 2
+            trans_x1 = x1
+            trans_y1 = y1
+            x2 = x1 + sz
+            y2 = y1 + sz
+
+            height, width, _ = rgb_image.shape
+            dx = max(0, -x1)
+            dy = max(0, -y1)
+            x1 = max(0, x1)
+            y1 = max(0, y1)
+
+            edx = max(0, x2 - width)
+            edy = max(0, y2 - height)
+            x2 = min(width, x2)
+            y2 = min(height, y2)
+
+            crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)]
+            if dx > 0 or dy > 0 or edx > 0 or edy > 0:
+                crop_img = cv2.copyMakeBorder(
+                    crop_img,
+                    int(dy),
+                    int(edy),
+                    int(dx),
+                    int(edx),
+                    cv2.BORDER_CONSTANT,
+                    value=(103.94, 116.78, 123.68))
+            crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE))
+
+            base_lmks = LargeBaseLmkInfer.infer_img(
+                crop_img, self.large_base_lmks_model,
+                self.device.lower() == 'cuda')
+
+            inv_scale = sz / INPUT_SIZE
+
+            affine_base_lmks = np.zeros((106, 2))
+            for idx in range(106):
+                affine_base_lmks[idx][
+                    0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1
+                affine_base_lmks[idx][
+                    1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1
+
+            landmarks.append(affine_base_lmks)
+
+        return boxes, landmarks
+
+    def find_face_contour(self, image):
+
+        boxes, landmarks = self.infer(image)
+        landmarks = np.array(landmarks)
+
+        args = [[0, 33, False], [33, 38, False], [42, 47, False],
+                [51, 55, False], [57, 64, False], [66, 74, True],
+                [75, 83, True], [84, 96, True]]
+
+        roi_bboxs = []
+
+        for i in range(len(boxes)):
+            roi_bbox = enlarged_bbox([
+                boxes[i]['x1'], boxes[i]['y1'], boxes[i]['x2'], boxes[i]['y2']
+            ], image.shape[1], image.shape[0], 0.5)
+            roi_bbox = [int(x) for x in roi_bbox]
+            roi_bboxs.append(roi_bbox)
+
+        people_maps = []
+
+        for i in range(landmarks.shape[0]):
+            landmark = landmarks[i, :, :]
+            maps = []
+            whole_mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
+
+            roi_box = roi_bboxs[i]
+            roi_box_width = roi_box[2] - roi_box[0]
+            roi_box_height = roi_box[3] - roi_box[1]
+            short_side_length = roi_box_width if roi_box_width < roi_box_height else roi_box_height
+
+            line_width = short_side_length // 10
+
+            if line_width == 0:
+                line_width = 1
+
+            kernel_size = line_width * 2
+            gaussian_kernel = kernel_size if kernel_size % 2 == 1 else kernel_size + 1
+
+            for t, arg in enumerate(args):
+                mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
+                draw_line(mask, landmark[arg[0]:arg[1]], (255, 255, 255),
+                          line_width, arg[2])
+                mask = cv2.GaussianBlur(mask,
+                                        (gaussian_kernel, gaussian_kernel), 0)
+                if t >= 1:
+                    draw_line(whole_mask, landmark[arg[0]:arg[1]],
+                              (255, 255, 255), line_width * 2, arg[2])
+                maps.append(mask)
+            whole_mask = cv2.GaussianBlur(whole_mask,
+                                          (gaussian_kernel, gaussian_kernel),
+                                          0)
+            maps.append(whole_mask)
+            people_maps.append(maps)
+
+        return people_maps[0], boxes
+
+    def face2contour(self, image, stack_mode='column'):
+        '''
+
+        :param facer:
+        :param image:
+        :param stack_mode:
+        :return: final_maps: [map0, map1,....]
+                 roi_bboxs: [bbox0, bbox1, ...]
+        '''
+
+        boxes, landmarks = self.infer(image)
+        landmarks = np.array(landmarks)
+
+        args = [[0, 33, False], [33, 38, False], [42, 47, False],
+                [51, 55, False], [57, 64, False], [66, 74, True],
+                [75, 83, True], [84, 96, True]]
+
+        roi_bboxs = []
+
+        for i in range(len(boxes)):
+            roi_bbox = enlarged_bbox([
+                boxes[i]['x1'], boxes[i]['y1'], boxes[i]['x2'], boxes[i]['y2']
+            ], image.shape[1], image.shape[0], 0.5)
+            roi_bbox = [int(x) for x in roi_bbox]
+            roi_bboxs.append(roi_bbox)
+
+        people_maps = []
+
+        for i in range(landmarks.shape[0]):
+            landmark = landmarks[i, :, :]
+            maps = []
+            whole_mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
+
+            roi_box = roi_bboxs[i]
+            roi_box_width = roi_box[2] - roi_box[0]
+            roi_box_height = roi_box[3] - roi_box[1]
+            short_side_length = roi_box_width if roi_box_width < roi_box_height else roi_box_height
+
+            line_width = short_side_length // 50
+
+            if line_width == 0:
+                line_width = 1
+
+            kernel_size = line_width * 4
+            gaussian_kernel = kernel_size if kernel_size % 2 == 1 else kernel_size + 1
+
+            for arg in args:
+                mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
+                draw_line(mask, landmark[arg[0]:arg[1]], (255, 255, 255),
+                          line_width, arg[2])
+                mask = cv2.GaussianBlur(mask,
+                                        (gaussian_kernel, gaussian_kernel), 0)
+                draw_line(whole_mask, landmark[arg[0]:arg[1]], (255, 255, 255),
+                          line_width, arg[2])
+                maps.append(mask)
+            whole_mask = cv2.GaussianBlur(whole_mask,
+                                          (gaussian_kernel, gaussian_kernel),
+                                          0)
+            maps.append(whole_mask)
+            people_maps.append(maps)
+
+        if stack_mode == 'depth':
+            final_maps = []
+            for i, maps in enumerate(people_maps):
+                final_map = np.dstack(maps)
+                final_map = final_map[roi_bboxs[i][1]:roi_bboxs[i][3],
+                                      roi_bboxs[i][0]:roi_bboxs[i][2], :]
+                final_maps.append(final_map)
+            return final_maps, roi_bboxs
+
+        elif stack_mode == 'column':
+            final_maps = []
+            for i, maps in enumerate(people_maps):
+                joint_maps = [
+                    x[roi_bboxs[i][1]:roi_bboxs[i][3],
+                      roi_bboxs[i][0]:roi_bboxs[i][2]] for x in maps
+                ]
+                final_map = np.column_stack(joint_maps)
+                final_maps.append(final_map)
+            return final_maps, roi_bboxs
+
+    def fat_face(self, img, degree=0.1):
+
+        _img, scale = resize_on_long_side(img, 800)
+
+        contour_maps, boxes = self.find_face_contour(_img)
+
+        contour_map = contour_maps[0]
+
+        boxes = boxes[0]
+
+        Flow = np.zeros(
+            shape=(contour_map.shape[0], contour_map.shape[1], 2),
+            dtype=np.float32)
+
+        box_center = [(boxes['x1'] + boxes['x2']) / 2,
+                      (boxes['y1'] + boxes['y2']) / 2]
+
+        box_length = max(
+            abs(boxes['y1'] - boxes['y2']), abs(boxes['x1'] - boxes['x2']))
+
+        value_1 = 2 * (Flow.shape[0] - box_center[1] - 1)
+        value_2 = 2 * (Flow.shape[1] - box_center[0] - 1)
+        value_list = [
+            box_length * 2, 2 * (box_center[0] - 1), 2 * (box_center[1] - 1),
+            value_1, value_2
+        ]
+        flow_box_length = min(value_list)
+        flow_box_length = int(flow_box_length)
+
+        sf = spread_flow(100, flow_box_length * degree)
+        sf = cv2.resize(sf, (flow_box_length, flow_box_length))
+
+        Flow[int(box_center[1]
+                 - flow_box_length / 2):int(box_center[1]
+                                            + flow_box_length / 2),
+             int(box_center[0]
+                 - flow_box_length / 2):int(box_center[0]
+                                            + flow_box_length / 2)] = sf
+
+        Flow = Flow * np.dstack((contour_map, contour_map)) / 255.0
+
+        inter_face_maps = contour_maps[-1]
+
+        Flow = Flow * (1.0 - np.dstack(
+            (inter_face_maps, inter_face_maps)) / 255.0)
+
+        Flow = cv2.resize(Flow, (img.shape[1], img.shape[0]))
+
+        Flow = Flow / scale
+
+        pred, top_bound, bottom_bound, left_bound, right_bound = image_warp_grid1(
+            Flow[..., 0], Flow[..., 1], img, 1.0, [0, 0, 0, 0])
+
+        return pred

modelscope/models/cv/face_reconstruction/models/facelandmark/nets/large_base_lmks_net.py

@@ -0,0 +1,201 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+INPUT_SIZE = 224
+
+
+def constant_init(module, val, bias=0):
+    if hasattr(module, 'weight') and module.weight is not None:
+        nn.init.constant_(module.weight, val)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def kaiming_init(module,
+                 a=0,
+                 mode='fan_out',
+                 nonlinearity='relu',
+                 bias=0,
+                 distribution='normal'):
+    assert distribution in ['uniform', 'normal']
+    if distribution == 'uniform':
+        nn.init.kaiming_uniform_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    else:
+        nn.init.kaiming_normal_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def conv_bn(inp, oup, kernel, stride, padding=1):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel, stride, padding, bias=False),
+        nn.BatchNorm2d(oup), nn.PReLU(oup))
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup),
+        nn.ReLU(inplace=True))
+
+
+class InvertedResidual(nn.Module):
+
+    def __init__(self,
+                 inp,
+                 oup,
+                 stride,
+                 padding,
+                 use_res_connect,
+                 expand_ratio=6):
+
+        super(InvertedResidual, self).__init__()
+
+        self.stride = stride
+        assert stride in [1, 2]
+
+        self.use_res_connect = use_res_connect
+        hid_channels = inp * expand_ratio
+        self.conv = nn.Sequential(
+            nn.Conv2d(inp, hid_channels, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(hid_channels),
+            nn.PReLU(hid_channels),
+            nn.Conv2d(
+                hid_channels,
+                hid_channels,
+                3,
+                stride,
+                padding,
+                groups=hid_channels,
+                bias=False),
+            nn.BatchNorm2d(hid_channels),
+            nn.PReLU(hid_channels),
+            nn.Conv2d(hid_channels, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class SoftArgmax(nn.Module):
+
+    def __init__(self, beta: int = 1, infer=False):
+        if not 0.0 <= beta:
+            raise ValueError(f'Invalid beta: {beta}')
+        super().__init__()
+        self.beta = beta
+        self.infer = infer
+
+    def forward(self, heatmap: torch.Tensor) -> torch.Tensor:
+        heatmap = heatmap.mul(self.beta)
+        batch_size, num_channel, height, width = heatmap.size()
+        device: str = heatmap.device
+
+        if not self.infer:
+            softmax: torch.Tensor = F.softmax(
+                heatmap.view(batch_size, num_channel, height * width),
+                dim=2).view(batch_size, num_channel, height, width)
+
+            xx, yy = torch.meshgrid(list(map(torch.arange, [width, height])))
+
+            approx_x = (
+                softmax.mul(xx.float().to(device)).view(
+                    batch_size, num_channel,
+                    height * width).sum(2).unsqueeze(2))
+            approx_y = (
+                softmax.mul(yy.float().to(device)).view(
+                    batch_size, num_channel,
+                    height * width).sum(2).unsqueeze(2))
+
+            output = [approx_x / width, approx_y / height]
+            output = torch.cat(output, 2)
+            output = output.view(-1, output.size(1) * output.size(2))
+            return output
+        else:
+            softmax: torch.Tensor = F.softmax(
+                heatmap.view(batch_size, num_channel, height * width), dim=2)
+
+            return softmax
+
+
+class LargeBaseLmksNet(nn.Module):
+
+    def __init__(self, er=1.0, infer=False):
+
+        super(LargeBaseLmksNet, self).__init__()
+
+        self.infer = infer
+
+        self.block1 = conv_bn(3, int(64 * er), 3, 2, 1)
+        self.block2 = InvertedResidual(
+            int(64 * er), int(64 * er), 1, 1, False, 2)
+
+        self.block3 = InvertedResidual(
+            int(64 * er), int(64 * er), 2, 1, False, 2)
+        self.block4 = InvertedResidual(
+            int(64 * er), int(64 * er), 1, 1, True, 2)
+        self.block5 = InvertedResidual(
+            int(64 * er), int(64 * er), 1, 1, True, 2)
+        self.block6 = InvertedResidual(
+            int(64 * er), int(64 * er), 1, 1, True, 2)
+        self.block7 = InvertedResidual(
+            int(64 * er), int(64 * er), 1, 1, True, 2)
+
+        self.block8 = InvertedResidual(
+            int(64 * er), int(128 * er), 2, 1, False, 2)
+
+        self.block9 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, False, 4)
+        self.block10 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, True, 4)
+        self.block11 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, True, 4)
+        self.block12 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, True, 4)
+        self.block13 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, True, 4)
+        self.block14 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, True, 4)
+
+        self.block15 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, False, 2)  # [128, 14, 14]
+        self.block16 = InvertedResidual(
+            int(128 * er), int(128 * er), 2, 1, False, 2)
+        self.block17 = InvertedResidual(
+            int(128 * er), int(128 * er), 1, 1, False, 2)
+
+        self.block18 = conv_bn(int(128 * er), int(256 * er), 3, 1, 1)
+        self.block19 = nn.Conv2d(int(256 * er), 106, 3, 1, 1, bias=False)
+        self.softargmax = SoftArgmax(infer=infer)
+
+    def forward(self, x):  # x: 3, 224, 224
+
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        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)
+        x = self.softargmax(x)
+
+        return x

modelscope/models/cv/face_reconstruction/models/facelandmark/nets/large_eyeball_net.py

@@ -0,0 +1,160 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import torch.nn as nn
+
+FACE_PART_SIZE = 56
+
+
+class InvertedResidual(nn.Module):
+
+    def __init__(self,
+                 inp,
+                 oup,
+                 kernel_size,
+                 stride,
+                 padding,
+                 expand_ratio=2,
+                 use_connect=False,
+                 activation='relu'):
+        super(InvertedResidual, self).__init__()
+
+        hid_channels = int(inp * expand_ratio)
+        if activation == 'relu':
+            self.conv = nn.Sequential(
+                nn.Conv2d(inp, hid_channels, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hid_channels), nn.ReLU(inplace=True),
+                nn.Conv2d(
+                    hid_channels,
+                    hid_channels,
+                    kernel_size,
+                    stride,
+                    padding,
+                    groups=hid_channels,
+                    bias=False), nn.BatchNorm2d(hid_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(hid_channels, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup))
+        elif activation == 'prelu':
+            self.conv = nn.Sequential(
+                nn.Conv2d(inp, hid_channels, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hid_channels), nn.PReLU(hid_channels),
+                nn.Conv2d(
+                    hid_channels,
+                    hid_channels,
+                    kernel_size,
+                    stride,
+                    padding,
+                    groups=hid_channels,
+                    bias=False), nn.BatchNorm2d(hid_channels),
+                nn.PReLU(hid_channels),
+                nn.Conv2d(hid_channels, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup))
+        self.use_connect = use_connect
+
+    def forward(self, x):
+        if self.use_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class Residual(nn.Module):
+
+    def __init__(self,
+                 inp,
+                 oup,
+                 kernel_size,
+                 stride,
+                 padding,
+                 use_connect=False,
+                 activation='relu'):
+        super(Residual, self).__init__()
+
+        self.use_connect = use_connect
+
+        if activation == 'relu':
+            self.conv = nn.Sequential(
+                nn.Conv2d(
+                    inp,
+                    inp,
+                    kernel_size,
+                    stride,
+                    padding,
+                    groups=inp,
+                    bias=False), nn.BatchNorm2d(inp), nn.ReLU(inplace=True),
+                nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup),
+                nn.ReLU(inplace=True))
+        elif activation == 'prelu':
+            self.conv = nn.Sequential(
+                nn.Conv2d(
+                    inp,
+                    inp,
+                    kernel_size,
+                    stride,
+                    padding,
+                    groups=inp,
+                    bias=False), nn.BatchNorm2d(inp), nn.PReLU(inp),
+                nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup),
+                nn.PReLU(oup))
+
+    def forward(self, x):
+        if self.use_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+def conv_bn(inp, oup, kernel, stride, padding=1):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel, stride, padding, bias=False),
+        nn.BatchNorm2d(oup), nn.PReLU(oup))
+
+
+def conv_no_relu(inp, oup, kernel, stride, padding=1):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel, stride, padding, bias=False),
+        nn.BatchNorm2d(oup))
+
+
+class View(nn.Module):
+
+    def __init__(self, shape):
+        super(View, self).__init__()
+        self.shape = shape
+
+    def forward(self, x):
+        return x.view(*self.shape)
+
+
+class Softmax(nn.Module):
+
+    def __init__(self, dim):
+        super(Softmax, self).__init__()
+        self.softmax = nn.Softmax(dim)
+
+    def forward(self, x):
+        return self.softmax(x)
+
+
+class LargeEyeballNet(nn.Module):
+
+    def __init__(self):
+        super(LargeEyeballNet, self).__init__()
+
+        # v6/v7/v9
+        # iris : -1*2, 3, FACE_PART_SIZE, FACE_PART_SIZE
+        self.net = nn.Sequential(
+            conv_bn(3, 16, 3, 2, 0),
+            InvertedResidual(16, 16, 3, 1, 1, 2, True, activation='prelu'),
+            InvertedResidual(16, 32, 3, 2, 0, 2, False, activation='prelu'),
+            InvertedResidual(32, 32, 3, 1, 1, 2, True, activation='prelu'),
+            InvertedResidual(32, 64, 3, 2, 1, 2, False, activation='prelu'),
+            InvertedResidual(64, 64, 3, 1, 1, 2, True, activation='prelu'),
+            InvertedResidual(64, 64, 3, 2, 0, 2, False, activation='prelu'),
+            InvertedResidual(64, 64, 3, 1, 1, 2, True, activation='prelu'),
+            View((-1, 64 * 3 * 3, 1, 1)), conv_bn(64 * 3 * 3, 64, 1, 1, 0),
+            conv_no_relu(64, 40, 1, 1, 0), View((-1, 40)))
+
+    def forward(self, x):  # x: -1, 3, FACE_PART_SIZE, FACE_PART_SIZE
+        iris = self.net(x)
+
+        return iris

modelscope/models/cv/face_reconstruction/models/facerecon_model.py

@@ -0,0 +1,564 @@
+# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
+# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from modelscope.models import MODELS, TorchModel
+from modelscope.models.cv.face_reconstruction.models import opt
+from .. import utils
+from . import networks
+from .bfm import ParametricFaceModel
+from .losses import (CLIPLoss_relative, TVLoss, TVLoss_std, landmark_loss,
+                     perceptual_loss, photo_loss, points_loss_horizontal,
+                     reflectance_loss, reg_loss)
+from .nv_diffrast import MeshRenderer
+
+
+@MODELS.register_module('face-reconstruction', 'face_reconstruction')
+class FaceReconModel(TorchModel):
+
+    def __init__(self,
+                 model_dir,
+                 w_color=1.92,
+                 w_exp=0.8,
+                 w_gamma=10.0,
+                 w_id=1.0,
+                 w_lm=0.0016,
+                 w_reg=0.0003,
+                 w_tex=0.017,
+                 *args,
+                 **kwargs):
+        """The FaceReconModel is implemented based on Deep3DFaceRecon_pytorch, publicly available at
+        https://github.com/sicxu/Deep3DFaceRecon_pytorch
+
+        Args:
+            model_dir: the root directory of the model files
+            w_color: the weight of color loss
+            w_exp: the regularization weight of expression
+            w_gamma: the regularization weight of lighting
+            w_id: the regularization weight of identity
+            w_lm: the weight of landmark loss
+            w_reg: the weight of regularization loss
+            w_tex: the regularization weight of texture
+        """
+        super().__init__(model_dir, *args, **kwargs)
+
+        opt.bfm_folder = os.path.join(model_dir, 'assets')
+        self.opt = opt
+        self.w_color = w_color
+        self.w_exp = w_exp
+        self.w_gamma = w_gamma
+        self.w_id = w_id
+        self.w_lm = w_lm
+        self.w_reg = w_reg
+        self.w_tex = w_tex
+        self.device = torch.device('cpu')
+        self.isTrain = opt.isTrain
+        self.visual_names = ['output_vis']
+        self.model_names = ['net_recon']
+        self.parallel_names = self.model_names + ['renderer']
+
+        self.net_recon = networks.define_net_recon(
+            net_recon=opt.net_recon,
+            use_last_fc=opt.use_last_fc,
+            init_path=None)
+
+        self.facemodel = ParametricFaceModel(
+            bfm_folder=opt.bfm_folder,
+            camera_distance=opt.camera_d,
+            focal=opt.focal,
+            center=opt.center,
+            is_train=self.isTrain,
+            default_name=opt.bfm_model)
+
+        self.facemodel_front = ParametricFaceModel(
+            bfm_folder=opt.bfm_folder,
+            camera_distance=opt.camera_d,
+            focal=opt.focal,
+            center=opt.center,
+            is_train=self.isTrain,
+            default_name='face_model_for_maas.mat')
+
+        fov = 2 * np.arctan(opt.center / opt.focal) * 180 / np.pi
+        self.renderer = MeshRenderer(
+            rasterize_fov=fov,
+            znear=opt.z_near,
+            zfar=opt.z_far,
+            rasterize_size=int(2 * opt.center))
+
+        self.renderer_fitting = MeshRenderer(
+            rasterize_fov=fov,
+            znear=opt.z_near,
+            zfar=opt.z_far,
+            rasterize_size=int(2 * opt.center))
+
+        self.nonlinear_UVs = self.facemodel.nonlinear_UVs
+        self.nonlinear_UVs = torch.from_numpy(self.nonlinear_UVs).to(
+            torch.device('cuda'))
+
+        template_obj_path = os.path.join(opt.bfm_folder, 'template_mesh.obj')
+        self.template_mesh = utils.read_obj(template_obj_path)
+
+        self.input_imgs = []
+        self.input_img_hds = []
+        self.input_fat_img_hds = []
+        self.atten_masks = []
+        self.gt_lms = []
+        self.gt_lm_hds = []
+        self.trans_ms = []
+        self.img_names = []
+        self.face_masks = []
+        self.head_masks = []
+        self.input_imgs_coeff = []
+        self.gt_lms_coeff = []
+
+        self.loss_names = [
+            'all', 'feat', 'color', 'lm', 'reg', 'gamma', 'reflc'
+        ]
+
+        # loss func name: (compute_%s_loss) % loss_name
+        self.compute_feat_loss = perceptual_loss
+        self.comupte_color_loss = photo_loss
+        self.compute_lm_loss = landmark_loss
+        self.compute_reg_loss = reg_loss
+        self.compute_reflc_loss = reflectance_loss
+
+    def load_networks(self, load_path):
+        state_dict = torch.load(load_path, map_location=self.device)
+        print('loading the model from %s' % load_path)
+
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                net.load_state_dict(state_dict[name], strict=False)
+
+        if self.opt.phase != 'test':
+            if self.opt.continue_train:
+
+                try:
+                    for i, sched in enumerate(self.schedulers):
+                        sched.load_state_dict(state_dict['sched_%02d' % i])
+                except Exception as e:
+                    print(e)
+                    for i, sched in enumerate(self.schedulers):
+                        sched.last_epoch = self.opt.epoch_count - 1
+
+    def setup(self, checkpoint_path):
+        """Load and print networks; create schedulers
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        self.load_networks(checkpoint_path)
+
+    def parallelize(self, convert_sync_batchnorm=True):
+        if not self.opt.use_ddp:
+            for name in self.parallel_names:
+                if isinstance(name, str):
+                    module = getattr(self, name)
+                    setattr(self, name, module.to(self.device))
+        else:
+            for name in self.model_names:
+                if isinstance(name, str):
+                    module = getattr(self, name)
+                    if convert_sync_batchnorm:
+                        module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
+                            module)
+                    setattr(
+                        self, name,
+                        torch.nn.parallel.DistributedDataParallel(
+                            module.to(self.device),
+                            device_ids=[self.device.index],
+                            find_unused_parameters=True,
+                            broadcast_buffers=True))
+
+            # DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient.
+            for name in self.parallel_names:
+                if isinstance(name, str) and name not in self.model_names:
+                    module = getattr(self, name)
+                    setattr(self, name, module.to(self.device))
+
+        # put state_dict of optimizer to gpu device
+        if self.opt.phase != 'test':
+            if self.opt.continue_train:
+                for optim in self.optimizers:
+                    for state in optim.state.values():
+                        for k, v in state.items():
+                            if isinstance(v, torch.Tensor):
+                                state[k] = v.to(self.device)
+
+    def eval(self):
+        """Make models eval mode"""
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, name)
+                net.eval()
+
+    def set_render(self, image_res):
+        fov = 2 * np.arctan(self.opt.center / self.opt.focal) * 180 / np.pi
+        if image_res is None:
+            image_res = int(2 * self.opt.center)
+
+        self.renderer = MeshRenderer(
+            rasterize_fov=fov,
+            znear=self.opt.z_near,
+            zfar=self.opt.z_far,
+            rasterize_size=image_res)
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input: a dictionary that contains the data itself and its metadata information.
+        """
+        self.input_img = input['imgs'].to(self.device)
+        self.input_img_hd = input['imgs_hd'].to(
+            self.device) if 'imgs_hd' in input else None
+
+        if 'imgs_fat_hd' not in input or input['imgs_fat_hd'] is None:
+            self.input_fat_img_hd = self.input_img_hd
+        else:
+            self.input_fat_img_hd = input['imgs_fat_hd'].to(self.device)
+
+        self.atten_mask = input['msks'].to(
+            self.device) if 'msks' in input else None
+        self.gt_lm = input['lms'].to(self.device) if 'lms' in input else None
+        self.gt_lm_hd = input['lms_hd'].to(
+            self.device) if 'lms_hd' in input else None
+        self.trans_m = input['M'].to(self.device) if 'M' in input else None
+        self.image_paths = input['im_paths'] if 'im_paths' in input else None
+        self.img_name = input['img_name'] if 'img_name' in input else None
+        self.face_mask = input['face_mask'].to(
+            self.device) if 'face_mask' in input else None
+        self.head_mask = input['head_mask'].to(
+            self.device) if 'head_mask' in input else None
+        self.gt_normals = input['normals'].to(
+            self.device) if 'normals' in input else None
+        self.input_img_coeff = input['imgs_coeff'].to(
+            self.device) if 'imgs_coeff' in input else None
+        self.gt_lm_coeff = input['lms_coeff'].to(
+            self.device) if 'lms_coeff' in input else None
+
+    def get_edge_points_horizontal(self):
+        left_points = []
+        right_points = []
+        for i in range(self.face_mask.shape[2]):
+            inds = torch.where(self.face_mask[0, 0, i, :] > 0.5)  # 0.9
+            if len(inds[0]) > 0:  # i > 112 and len(inds[0]) > 0
+                left_points.append(int(inds[0][0]) + 1)
+                right_points.append(int(inds[0][-1]))
+            else:
+                left_points.append(0)
+                right_points.append(self.face_mask.shape[3] - 1)
+        self.left_points = torch.tensor(left_points).long().to(self.device)
+        self.right_points = torch.tensor(right_points).long().to(self.device)
+
+    def get_edge_points_vertical(self):
+        top_points = []
+        bottom_points = []
+        for i in range(self.face_mask.shape[3]):
+            inds = torch.where(self.face_mask[0, 0, :, i] > 0.9)
+            if len(inds[0]) > 0:
+                top_points.append(int(inds[0][0]))
+                bottom_points.append(int(inds[0][-1]))
+            else:
+                top_points.append(0)
+                bottom_points.append(self.face_mask.shape[2] - 1)
+        self.top_points = torch.tensor(top_points).long().to(self.device)
+        self.bottom_points = torch.tensor(bottom_points).long().to(self.device)
+
+    def blur_shape_offset_uv(self, global_blur=False, blur_size=3):
+        if self.edge_mask is not None:
+            shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu(
+            ).numpy()
+            shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur, (15, 15))
+            shape_offset_uv_blur = torch.from_numpy(
+                shape_offset_uv_blur).float().to(self.device)[None, ...]
+            value_1 = shape_offset_uv_blur * self.edge_mask[None, ..., None]
+            value_2 = self.shape_offset_uv * (
+                1 - self.edge_mask[None, ..., None])
+            self.shape_offset_uv = value_1 + value_2
+
+        self.shape_offset_uv = self.shape_offset_uv * self.fusion_mask[None,
+                                                                       ...,
+                                                                       None]
+
+        if global_blur and blur_size > 0:
+            shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu(
+            ).numpy()
+            shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur,
+                                            (blur_size, blur_size))
+            shape_offset_uv_blur = torch.from_numpy(
+                shape_offset_uv_blur).float().to(self.device)[None, ...]
+            self.shape_offset_uv = shape_offset_uv_blur
+
+    def get_fusion_mask(self):
+
+        h, w = self.shape_offset_uv.shape[1:3]
+        self.fusion_mask = torch.zeros((h, w)).to(self.device).float()
+        UVs_coords = self.nonlinear_UVs.clone()[:35709]
+        UVs_coords[:, 0] *= w
+        UVs_coords[:, 1] *= h
+        UVs_coords_int = torch.floor(UVs_coords)
+        UVs_coords_int = UVs_coords_int.long()
+
+        self.fusion_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:,
+                                                                      0]] = 1
+
+        # blur mask
+        self.fusion_mask = self.fusion_mask.cpu().numpy()
+        new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+        new_kernel2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8))
+        self.fusion_mask = cv2.dilate(self.fusion_mask, new_kernel1, 1)
+        self.fusion_mask = cv2.erode(self.fusion_mask, new_kernel2, 1)
+        self.fusion_mask = cv2.blur(self.fusion_mask, (17, 17))
+        self.fusion_mask = torch.from_numpy(self.fusion_mask).float().to(
+            self.device)
+
+    def get_edge_mask(self):
+
+        h, w = self.shape_offset_uv.shape[1:3]
+        self.edge_mask = torch.zeros((h, w)).to(self.device).float()
+        UVs_coords = self.nonlinear_UVs.clone()[self.edge_points_inds]
+        UVs_coords[:, 0] *= w
+        UVs_coords[:, 1] *= h
+        UVs_coords_int = torch.floor(UVs_coords)
+        UVs_coords_int = UVs_coords_int.long()
+
+        self.edge_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:, 0]] = 1
+
+        # blur mask
+        self.edge_mask = self.edge_mask.cpu().numpy()
+        new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8))
+        self.edge_mask = cv2.dilate(self.edge_mask, new_kernel1, 1)
+        self.edge_mask = cv2.blur(self.edge_mask, (5, 5))
+        self.edge_mask = torch.from_numpy(self.edge_mask).float().to(
+            self.device)
+
+    def fitting_nonlinear(self, coeff, debug=False, n_iters=100, out_dir=None):
+        output_coeff = coeff.detach().clone()
+
+        output_coeff = self.facemodel_front.split_coeff(output_coeff)
+        output_coeff['id'].requires_grad = True
+        output_coeff['exp'].requires_grad = True
+        output_coeff['tex'].requires_grad = True
+        output_coeff['angle'].requires_grad = True
+        output_coeff['gamma'].requires_grad = True
+        output_coeff['trans'].requires_grad = True
+
+        self.shape_offset_uv = torch.zeros(
+            (1, 300, 300, 3),
+            dtype=torch.float32).to(self.device)  # (1, 180, 256, 3)
+        self.shape_offset_uv.requires_grad = True
+
+        self.texture_offset_uv = torch.zeros(
+            (1, 300, 300, 3),
+            dtype=torch.float32).to(self.device)  # (1, 180, 256, 3)
+        self.texture_offset_uv.requires_grad = True
+
+        value_list = [
+            self.shape_offset_uv, self.texture_offset_uv, output_coeff['id'],
+            output_coeff['exp'], output_coeff['tex'], output_coeff['angle'],
+            output_coeff['gamma'], output_coeff['trans']
+        ]
+        optim = torch.optim.Adam(value_list, lr=1e-3)
+
+        self.get_edge_points_horizontal()
+        self.get_edge_points_vertical()
+
+        self.cur_iter = 0
+        for i in range(n_iters):  # 500
+            self.pred_vertex, _, self.pred_color, self.pred_lm, _, face_shape_offset, self.verts_proj = \
+                self.facemodel_front.compute_for_render_train_nonlinear(output_coeff, self.shape_offset_uv,
+                                                                        self.texture_offset_uv,
+                                                                        self.nonlinear_UVs[:35709, ...])
+            self.pred_mask, _, self.pred_face, self.occ = self.renderer_fitting(
+                self.pred_vertex,
+                self.facemodel_front.face_buf,
+                feat=self.pred_color)
+
+            self.pred_coeffs_dict = self.facemodel_front.split_coeff(
+                output_coeff)
+            self.compute_losses_fitting()
+            if debug and i % 10 == 0:
+                print('{}: total loss: {:.6f}'.format(i, self.loss_all.item()))
+
+            optim.zero_grad()
+            self.loss_all.backward()
+            optim.step()
+
+            self.cur_iter += 1
+
+        output_coeff = self.facemodel_front.merge_coeff(output_coeff)
+
+        self.get_edge_mask()
+        self.get_fusion_mask()
+        self.blur_shape_offset_uv()
+
+        self.pred_vertex, _, self.pred_color, self.pred_lm, _, face_shape_offset, self.verts_proj = \
+            self.facemodel_front.compute_for_render_train_nonlinear(output_coeff, self.shape_offset_uv,
+                                                                    self.texture_offset_uv,
+                                                                    self.nonlinear_UVs[:35709, ...])
+
+        if out_dir is not None:
+            input_img_numpy = 255. * (self.input_img).detach().cpu().permute(
+                0, 2, 3, 1).numpy()
+            input_img_numpy = np.squeeze(input_img_numpy)
+
+            output_vis = self.pred_face
+            output_vis_numpy_raw = 255. * output_vis.detach().cpu().permute(
+                0, 2, 3, 1).numpy()
+            output_vis_numpy_raw = np.squeeze(output_vis_numpy_raw)
+
+            output_vis_numpy = np.concatenate(
+                (input_img_numpy, output_vis_numpy_raw), axis=-2)
+
+            output_vis = np.squeeze(output_vis_numpy)
+            output_vis = output_vis[..., ::-1]  # rgb->bgr
+            output_face_mask = self.pred_mask.detach().cpu().permute(
+                0, 2, 3, 1).squeeze().numpy() * 255.0
+            output_vis = np.column_stack(
+                (output_vis, cv2.cvtColor(output_face_mask,
+                                          cv2.COLOR_GRAY2BGR)))
+            output_input_vis = output_vis[:, :224]
+            output_pred_vis = output_vis[:, 224:448]
+            output_mask_vis = output_vis[:, 448:]
+
+            face_mask_vis = 255. * self.face_mask.detach().cpu()[0, 0].numpy()
+
+            shape_offset_vis = self.shape_offset_uv.detach().cpu().numpy()[0]
+            shape_offset_vis = (shape_offset_vis - shape_offset_vis.min()) / (
+                shape_offset_vis.max() - shape_offset_vis.min()) * 255.0
+
+            cv2.imwrite(
+                os.path.join(out_dir, 'fitting_01_input.jpg'),
+                output_input_vis)
+            cv2.imwrite(
+                os.path.join(out_dir, 'fitting_02_pred.jpg'), output_pred_vis)
+            cv2.imwrite(
+                os.path.join(out_dir, 'fitting_03_mask.jpg'), output_mask_vis)
+            cv2.imwrite(
+                os.path.join(out_dir, 'fitting_04_facemask.jpg'),
+                face_mask_vis)
+            cv2.imwrite(
+                os.path.join(out_dir, 'fitting_05_shape_offset.jpg'),
+                shape_offset_vis)
+
+        recon_shape_offset = face_shape_offset
+        recon_shape_offset[..., -1] = 10 - recon_shape_offset[
+            ..., -1]  # from camera space to world space
+        recon_shape_offset = recon_shape_offset.detach().cpu().numpy()[0]
+
+        tri = self.facemodel_front.face_buf.cpu().numpy()
+        pred_color = self.pred_color.detach().cpu().numpy()[0].clip(0, 1)
+
+        output = {
+            'coeffs': output_coeff,
+            'face_vertices': recon_shape_offset,
+            'face_faces': tri + 1,
+            'face_colors': pred_color
+        }
+        return output
+
+    def forward(self, out_dir=None):
+        self.facemodel.to(self.device)
+        self.facemodel_front.to(self.device)
+        with torch.no_grad():
+
+            output_coeff = self.net_recon(self.input_img)
+
+        with torch.enable_grad():
+            output = self.fitting_nonlinear(
+                output_coeff, debug=True, out_dir=out_dir)
+
+        output_coeff = output['coeffs']
+        output_coeff = self.facemodel.split_coeff(output_coeff)
+        eye_coeffs = output_coeff['exp'][0, 16] + output_coeff['exp'][
+            0, 17] + output_coeff['exp'][0, 19]
+        if eye_coeffs > 1.0:
+            degree = 0.5
+        else:
+            degree = 1.0
+        output_coeff['exp'][0, 16] += 1 * degree
+        output_coeff['exp'][0, 17] += 1 * degree
+        output_coeff['exp'][0, 19] += 1.5 * degree
+        output_coeff = self.facemodel.merge_coeff(output_coeff)
+
+        self.pred_vertex, face_shape_ori, head_shape = \
+            self.facemodel.compute_for_render_nonlinear_full(output_coeff, self.shape_offset_uv.detach(),
+                                                             self.nonlinear_UVs, nose_coeff=0.1)
+
+        UVs_tensor = torch.tensor(
+            self.template_mesh['uvs'],
+            dtype=torch.float32)[None, ...].to(self.pred_vertex.device)
+        target_img = self.input_fat_img_hd.permute(0, 2, 3, 1)
+        with torch.enable_grad():
+            _, _, _, texture_map, _ = self.renderer.pred_shape_and_texture(
+                self.pred_vertex, self.facemodel.face_buf, UVs_tensor,
+                target_img)
+
+        recon_shape = head_shape
+        recon_shape[
+            ...,
+            -1] = 10 - recon_shape[..., -1]  # from camera space to world space
+        recon_shape = recon_shape.cpu().numpy()[0]
+        tri = self.facemodel.face_buf.cpu().numpy()
+        normals = utils.estimate_normals(recon_shape, tri)
+
+        output['head_vertices'] = recon_shape
+        output['head_faces'] = tri + 1
+        output['head_tex_map'] = texture_map
+        output['head_UVs'] = self.template_mesh['uvs']
+        output['head_faces_uv'] = self.template_mesh['faces_uv']
+        output['head_normals'] = normals
+
+        return output
+
+    def compute_losses_fitting(self):
+        face_mask = self.pred_mask
+
+        face_mask = face_mask.detach()
+        self.loss_color = self.w_color * self.comupte_color_loss(
+            self.pred_face, self.input_img, face_mask)  # 1.0
+
+        self.loss_color_nose = torch.tensor(0.0).float().to(self.device)
+
+        loss_reg, loss_gamma = self.compute_reg_loss(self.pred_coeffs_dict,
+                                                     self.w_id, self.w_exp,
+                                                     self.w_tex)
+        self.loss_reg = self.w_reg * loss_reg  # 1.0
+        self.loss_gamma = self.w_gamma * loss_gamma  # 1.0
+
+        self.loss_lm = self.w_lm * self.compute_lm_loss(
+            self.pred_lm, self.gt_lm) * 0.1  # 0.1
+
+        self.loss_smooth_offset = TVLoss()(self.shape_offset_uv.permute(
+            0, 3, 1, 2)) * 10000  # 10000
+
+        self.loss_reg_offset = torch.tensor(0.0).float().to(self.device)
+
+        self.loss_reg_textureOff = torch.mean(
+            torch.abs(self.texture_offset_uv)) * 10  # 10
+
+        self.loss_smooth_offset_std = TVLoss_std()(
+            self.shape_offset_uv.permute(0, 3, 1, 2)) * 50000  # 50000
+
+        self.loss_points_horizontal, self.edge_points_inds = points_loss_horizontal(
+            self.verts_proj, self.left_points, self.right_points)  # 20
+        self.loss_points_horizontal *= 20
+        self.loss_points_horizontal_jaw = torch.tensor(0.0).float().to(
+            self.device)
+        self.loss_points_vertical = torch.tensor(0.0).float().to(self.device)
+        self.loss_normals = torch.tensor(0.0).float().to(self.device)
+
+        self.loss_all = self.loss_color + self.loss_lm + self.loss_reg + self.loss_gamma + self.loss_smooth_offset
+        self.loss_all += self.loss_reg_offset + self.loss_smooth_offset_std + self.loss_points_horizontal
+        self.loss_all += self.loss_points_vertical + self.loss_reg_textureOff
+        self.loss_all += self.loss_color_nose + self.loss_normals + self.loss_points_horizontal_jaw
+
+        self.loss_mask = torch.tensor(0.0).float().to(self.device)

modelscope/models/cv/face_reconstruction/models/losses.py

@@ -0,0 +1,413 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import clip
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from kornia.geometry import warp_affine
+
+
+def resize_n_crop(image, M, dsize=112):
+    # image: (b, c, h, w)
+    # M   :  (b, 2, 3)
+    return warp_affine(image, M, dsize=(dsize, dsize))
+
+
+class CLIPLoss(torch.nn.Module):
+
+    def __init__(self):
+        super(CLIPLoss, self).__init__()
+        self.model, self.preprocess = clip.load('ViT-B/32', device='cuda')
+
+    def forward(self, image, text):
+        similarity = 1 - self.model(image, text)[0] / 100
+        return similarity
+
+
+class CLIPLoss_relative(torch.nn.Module):
+
+    def __init__(self):
+        super(CLIPLoss_relative, self).__init__()
+        self.model, self.preprocess = clip.load('ViT-B/32', device='cuda')
+
+    def forward(self, image, text, image_ori, text_ori):
+
+        image_features = self.model.encode_image(image)
+        text_features = self.model.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(
+            dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        image_features_ori = self.model.encode_image(image_ori)
+        text_features_ori = self.model.encode_text(text_ori)
+
+        # normalized features
+        image_features_ori = image_features_ori / image_features_ori.norm(
+            dim=1, keepdim=True)
+        text_features_ori = text_features_ori / text_features_ori.norm(
+            dim=1, keepdim=True)
+
+        delta_image = image_features - image_features_ori
+        delta_text = text_features - text_features_ori
+
+        loss = 1 - torch.sum(delta_image * delta_text) / (
+            torch.norm(delta_image) * torch.norm(delta_text))
+
+        return loss
+
+
+# perceptual level loss
+class PerceptualLoss(nn.Module):
+
+    def __init__(self, recog_net, input_size=112):
+        super(PerceptualLoss, self).__init__()
+        self.recog_net = recog_net
+        self.preprocess = lambda x: 2 * x - 1
+        self.input_size = input_size
+
+    def forward(self, imageA, imageB, M):
+        """
+        1 - cosine distance
+        Parameters:
+            imageA       --torch.tensor (B, 3, H, W), range (0, 1) , RGB order
+            imageB       --same as imageA
+        """
+
+        imageA = self.preprocess(resize_n_crop(imageA, M, self.input_size))
+        imageB = self.preprocess(resize_n_crop(imageB, M, self.input_size))
+
+        # freeze bn
+        self.recog_net.eval()
+
+        id_featureA = F.normalize(self.recog_net(imageA), dim=-1, p=2)
+        id_featureB = F.normalize(self.recog_net(imageB), dim=-1, p=2)
+        cosine_d = torch.sum(id_featureA * id_featureB, dim=-1)
+        return torch.sum(1 - cosine_d) / cosine_d.shape[0]
+
+
+def perceptual_loss(id_featureA, id_featureB):
+    cosine_d = torch.sum(id_featureA * id_featureB, dim=-1)
+    return torch.sum(1 - cosine_d) / cosine_d.shape[0]
+
+
+# image level loss
+def photo_loss(imageA, imageB, mask, eps=1e-6):
+    """
+    l2 norm (with sqrt, to ensure backward stabililty, use eps, otherwise Nan may occur)
+    Parameters:
+        imageA       --torch.tensor (B, 3, H, W), range (0, 1), RGB order
+        imageB       --same as imageA
+    """
+    loss = torch.sqrt(eps + torch.sum(
+        (imageA - imageB)**2, dim=1, keepdims=True)) * mask
+    loss = torch.sum(loss) / torch.max(
+        torch.sum(mask),
+        torch.tensor(1.0).to(mask.device))
+    return loss
+
+
+def landmark_loss(predict_lm, gt_lm, weight=None):
+    """
+    weighted mse loss
+    Parameters:
+        predict_lm    --torch.tensor (B, 68, 2)
+        gt_lm         --torch.tensor (B, 68, 2)
+        weight        --numpy.array (1, 68)
+    """
+    if not weight:
+        weight = np.ones([68])
+        weight[28:31] = 20
+        weight[-8:] = 20
+        weight = np.expand_dims(weight, 0)
+        weight = torch.tensor(weight).to(predict_lm.device)
+    loss = torch.sum((predict_lm - gt_lm)**2, dim=-1) * weight
+    loss = torch.sum(loss) / (predict_lm.shape[0] * predict_lm.shape[1])
+    return loss
+
+
+# regulization
+def reg_loss(coeffs_dict, w_id=1, w_exp=1, w_tex=1):
+    """
+    l2 norm without the sqrt, from yu's implementation (mse)
+    tf.nn.l2_loss https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss
+    Parameters:
+        coeffs_dict     -- a  dict of torch.tensors , keys: id, exp, tex, angle, gamma, trans
+
+    """
+    # coefficient regularization to ensure plausible 3d faces
+    value_1 = w_id * torch.sum(coeffs_dict['id']**2)
+    value_2 = w_exp * torch.sum(coeffs_dict['exp']**2)
+    value_3 = w_tex * torch.sum(coeffs_dict['tex']**2)
+    creg_loss = value_1 + value_2 + value_3
+    creg_loss = creg_loss / coeffs_dict['id'].shape[0]
+
+    # gamma regularization to ensure a nearly-monochromatic light
+    gamma = coeffs_dict['gamma'].reshape([-1, 3, 9])
+    gamma_mean = torch.mean(gamma, dim=1, keepdims=True)
+    gamma_loss = torch.mean((gamma - gamma_mean)**2)
+
+    return creg_loss, gamma_loss
+
+
+def reflectance_loss(texture, mask):
+    """
+    minimize texture variance (mse), albedo regularization to ensure an uniform skin albedo
+    Parameters:
+        texture       --torch.tensor, (B, N, 3)
+        mask          --torch.tensor, (N), 1 or 0
+
+    """
+    mask = mask.reshape([1, mask.shape[0], 1])
+    texture_mean = torch.sum(
+        mask * texture, dim=1, keepdims=True) / torch.sum(mask)
+    loss = torch.sum(((texture - texture_mean) * mask)**2) / (
+        texture.shape[0] * torch.sum(mask))
+    return loss
+
+
+def lm_3d_loss(pred_lm_3d, gt_lm_3d, mask):
+    loss = torch.abs(pred_lm_3d - gt_lm_3d)[mask, :]
+    loss = torch.mean(loss)
+    return loss
+
+
+class TVLoss(nn.Module):
+
+    def __init__(self, TVLoss_weight=1):
+        super(TVLoss, self).__init__()
+        self.TVLoss_weight = TVLoss_weight
+
+    def forward(self, x):
+        batch_size = x.size()[0]
+        h_x = x.size()[2]
+        w_x = x.size()[3]
+        count_h = self._tensor_size(x[:, :, 1:, :])
+        count_w = self._tensor_size(x[:, :, :, 1:])
+        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
+        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
+        return self.TVLoss_weight * 2 * (h_tv / count_h
+                                         + w_tv / count_w) / batch_size
+
+    def _tensor_size(self, t):
+        return t.size()[1] * t.size()[2] * t.size()[3]
+
+
+class TVLoss_std(nn.Module):
+
+    def __init__(self, TVLoss_weight=1):
+        super(TVLoss_std, self).__init__()
+        self.TVLoss_weight = TVLoss_weight
+
+    def forward(self, x):
+        batch_size = x.size()[0]
+        h_x = x.size()[2]
+        w_x = x.size()[3]
+        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2)
+        h_tv = ((h_tv - torch.mean(h_tv))**2).sum()
+        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2)
+        w_tv = ((w_tv - torch.mean(w_tv))**2).sum()
+        return self.TVLoss_weight * 2 * (h_tv + w_tv) / batch_size
+
+    def _tensor_size(self, t):
+        return t.size()[1] * t.size()[2] * t.size()[3]
+
+
+def photo_loss_sum(imageA, imageB, mask, eps=1e-6):
+    """
+    l2 norm (with sqrt, to ensure backward stabililty, use eps, otherwise Nan may occur)
+    Parameters:
+        imageA       --torch.tensor (B, 3, H, W), range (0, 1), RGB order
+        imageB       --same as imageA
+    """
+    loss = torch.sqrt(eps + torch.sum(
+        (imageA - imageB)**2, dim=1, keepdims=True)) * mask
+    loss = torch.sum(loss) / (
+        imageA.shape[0] * imageA.shape[2] * imageA.shape[3])
+    return loss
+
+
+def points_loss_horizontal(verts, left_points, right_points, width=224):
+    verts_int = torch.ceil(verts[0]).long().clamp(0, width - 1)  # (n, 2)
+    verts_left = left_points[width - 1 - verts_int[:, 1]].float()
+    verts_right = right_points[width - 1 - verts_int[:, 1]].float()
+    verts_x = verts[0, :, 0]
+    dist = (verts_left - verts_x) / width * (verts_right - verts_x) / width
+    dist /= torch.max(
+        torch.abs((verts_left - verts_x) / width),
+        torch.abs((verts_right - verts_x) / width))
+    edge_inds = torch.where(dist > 0)[0]
+    dist += 0.01
+    dist = torch.nn.functional.relu(dist).clone()
+    dist -= 0.01
+    dist = torch.abs(dist)
+    loss = torch.mean(dist)
+    return loss, edge_inds
+
+
+class LaplacianLoss(nn.Module):
+
+    def __init__(self):
+        super(LaplacianLoss, self).__init__()
+
+    def forward(self, x):
+        batch_size, slice_num = x.size()[:2]
+        z_x = x.size()[2]
+        h_x = x.size()[3]
+        w_x = x.size()[4]
+        count_z = self._tensor_size(x[:, :, 1:, :, :])
+        count_h = self._tensor_size(x[:, :, :, 1:, :])
+        count_w = self._tensor_size(x[:, :, :, :, 1:])
+        z_tv = torch.pow((x[:, :, 1:, :, :] - x[:, :, :z_x - 1, :, :]),
+                         2).sum()
+        h_tv = torch.pow((x[:, :, :, 1:, :] - x[:, :, :, :h_x - 1, :]),
+                         2).sum()
+        w_tv = torch.pow((x[:, :, :, :, 1:] - x[:, :, :, :, :w_x - 1]),
+                         2).sum()
+        return 2 * (z_tv / count_z + h_tv / count_h + w_tv / count_w) / (
+            batch_size * slice_num)
+
+    def _tensor_size(self, t):
+        return t.size()[2] * t.size()[3] * t.size()[4]
+
+
+class LaplacianLoss_L1(nn.Module):
+
+    def __init__(self):
+        super(LaplacianLoss_L1, self).__init__()
+
+    def forward(self, x):
+        batch_size, slice_num = x.size()[:2]
+        z_x = x.size()[2]
+        h_x = x.size()[3]
+        w_x = x.size()[4]
+        count_z = self._tensor_size(x[:, :, 1:, :, :])
+        count_h = self._tensor_size(x[:, :, :, 1:, :])
+        count_w = self._tensor_size(x[:, :, :, :, 1:])
+        z_tv = torch.abs((x[:, :, 1:, :, :] - x[:, :, :z_x - 1, :, :])).sum()
+        h_tv = torch.abs((x[:, :, :, 1:, :] - x[:, :, :, :h_x - 1, :])).sum()
+        w_tv = torch.abs((x[:, :, :, :, 1:] - x[:, :, :, :, :w_x - 1])).sum()
+        return 2 * (z_tv / count_z + h_tv / count_h + w_tv / count_w) / (
+            batch_size * slice_num)
+
+    def _tensor_size(self, t):
+        return t.size()[2] * t.size()[3] * t.size()[4]
+
+
+class GANLoss(nn.Module):
+
+    def __init__(self,
+                 gan_mode,
+                 target_real_label=1.0,
+                 target_fake_label=0.0,
+                 tensor=torch.FloatTensor):
+        super(GANLoss, self).__init__()
+        self.real_label = target_real_label
+        self.fake_label = target_fake_label
+        self.real_label_tensor = None
+        self.fake_label_tensor = None
+        self.zero_tensor = None
+        self.Tensor = tensor
+        self.gan_mode = gan_mode
+        if gan_mode == 'ls':
+            pass
+        elif gan_mode == 'original':
+            pass
+        elif gan_mode == 'w':
+            pass
+        elif gan_mode == 'hinge':
+            pass
+        else:
+            raise ValueError('Unexpected gan_mode {}'.format(gan_mode))
+
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            if self.real_label_tensor is None:
+                self.real_label_tensor = self.Tensor(1).fill_(self.real_label)
+                self.real_label_tensor.requires_grad_(False)
+            return self.real_label_tensor.expand_as(input)
+        else:
+            if self.fake_label_tensor is None:
+                self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)
+                self.fake_label_tensor.requires_grad_(False)
+            return self.fake_label_tensor.expand_as(input)
+
+    def get_zero_tensor(self, input):
+        if self.zero_tensor is None:
+            self.zero_tensor = self.Tensor(1).fill_(0)
+            self.zero_tensor.requires_grad_(False)
+        return self.zero_tensor.expand_as(input)
+
+    def loss(self, input, target_is_real, for_discriminator=True):
+        if self.gan_mode == 'original':  # cross entropy loss
+            target_tensor = self.get_target_tensor(input, target_is_real)
+            loss = F.binary_cross_entropy_with_logits(input, target_tensor)
+            return loss
+        elif self.gan_mode == 'ls':
+            target_tensor = self.get_target_tensor(input, target_is_real)
+            return F.mse_loss(input, target_tensor)
+        elif self.gan_mode == 'hinge':
+            if for_discriminator:
+                if target_is_real:
+                    minval = torch.min(input - 1, self.get_zero_tensor(input))
+                    loss = -torch.mean(minval)
+                else:
+                    minval = torch.min(-input - 1, self.get_zero_tensor(input))
+                    loss = -torch.mean(minval)
+            else:
+                assert target_is_real, "The generator's hinge loss must be aiming for real"
+                loss = -torch.mean(input)
+            return loss
+        else:
+            # wgan
+            if target_is_real:
+                return -input.mean()
+            else:
+                return input.mean()
+
+    def __call__(self, input, target_is_real, for_discriminator=True):
+        # computing loss is a bit complicated because |input| may not be
+        # a tensor, but list of tensors in case of multiscale discriminator
+        if isinstance(input, list):
+            loss = 0
+            for pred_i in input:
+                if isinstance(pred_i, list):
+                    pred_i = pred_i[-1]
+                loss_tensor = self.loss(pred_i, target_is_real,
+                                        for_discriminator)
+                bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)
+                new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)
+                loss += new_loss
+            return loss / len(input)
+        else:
+            return self.loss(input, target_is_real, for_discriminator)
+
+
+class BinaryDiceLoss(nn.Module):
+
+    def __init__(self, smooth=1, p=1, reduction='mean'):
+        super(BinaryDiceLoss, self).__init__()
+        self.smooth = smooth
+        self.p = p
+        self.reduction = reduction
+
+    def forward(self, predict, target):
+        assert predict.shape[0] == target.shape[
+            0], "predict & target batch size don't match"
+        predict = predict.contiguous().view(predict.shape[0], -1)
+        target = target.contiguous().view(target.shape[0], -1)
+
+        num = torch.sum(torch.mul(predict, target), dim=1)
+        den = torch.sum(predict + target, dim=1)
+
+        loss = 1 - (2 * num + self.smooth) / (den + self.smooth)
+
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        elif self.reduction == 'none':
+            return loss
+        else:
+            raise Exception('Unexpected reduction {}'.format(self.reduction))

modelscope/models/cv/face_reconstruction/models/networks.py

@@ -0,0 +1,577 @@
+# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
+# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
+import os
+from typing import Any, Callable, List, Optional, Type, Union
+
+import torch
+import torch.nn as nn
+from kornia.geometry import warp_affine
+from torch import Tensor
+from torch.optim import lr_scheduler
+
+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
+
+
+def resize_n_crop(image, M, dsize=112):
+    # image: (b, c, h, w)
+    # M   :  (b, 2, 3)
+    return warp_affine(image, M, dsize=(dsize, dsize))
+
+
+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 define_net_recon(net_recon, use_last_fc=False, init_path=None):
+    return ReconNetWrapper(
+        net_recon, use_last_fc=use_last_fc, init_path=init_path)
+
+
+def define_net_recon2(net_recon, use_last_fc=False, init_path=None):
+    return ReconNetWrapper2(
+        net_recon, use_last_fc=use_last_fc, init_path=init_path)
+
+
+class ReconNetWrapper(nn.Module):
+    fc_dim = 257
+
+    def __init__(self, net_recon, use_last_fc=False, init_path=None):
+        super(ReconNetWrapper, self).__init__()
+        self.use_last_fc = use_last_fc
+        if net_recon not in func_dict:
+            return NotImplementedError('network [%s] is not implemented',
+                                       net_recon)
+        func, last_dim = func_dict[net_recon]
+        backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_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 net_recon %s from %s' % (net_recon, init_path))
+        self.backbone = backbone
+        if not use_last_fc:
+            self.final_layers = nn.ModuleList([
+                conv1x1(last_dim, 80, bias=True),  # id layer
+                conv1x1(last_dim, 64, bias=True),  # exp layer
+                conv1x1(last_dim, 80, bias=True),  # tex layer
+                conv1x1(last_dim, 3, bias=True),  # angle layer
+                conv1x1(last_dim, 27, bias=True),  # gamma layer
+                conv1x1(last_dim, 2, bias=True),  # tx, ty
+                conv1x1(last_dim, 1, bias=True)  # tz
+            ])
+            for m in self.final_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.final_layers:
+                output.append(layer(x))
+            x = torch.flatten(torch.cat(output, dim=1), 1)
+        return x
+
+
+class ReconNetWrapper2(nn.Module):
+    fc_dim = 264
+
+    def __init__(self, net_recon, use_last_fc=False, init_path=None):
+        super(ReconNetWrapper2, self).__init__()
+        self.use_last_fc = use_last_fc
+        if net_recon not in func_dict:
+            return NotImplementedError('network [%s] is not implemented',
+                                       net_recon)
+        func, last_dim = func_dict[net_recon]
+        backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_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 net_recon %s from %s' % (net_recon, init_path))
+        self.backbone = backbone
+        if not use_last_fc:
+            self.final_layers2 = nn.ModuleList([
+                conv1x1(last_dim, 80, bias=True),  # id layer
+                conv1x1(last_dim, 51, bias=True),  # exp layer
+                conv1x1(last_dim, 100, bias=True),  # tex layer
+                conv1x1(last_dim, 3, bias=True),  # angle layer
+                conv1x1(last_dim, 27, bias=True),  # gamma layer
+                conv1x1(last_dim, 2, bias=True),  # tx, ty
+                conv1x1(last_dim, 1, bias=True)  # tz
+            ])
+            for m in self.final_layers2:
+                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.final_layers2:
+                output.append(layer(x))
+            x = torch.flatten(torch.cat(output, dim=1), 1)
+        return x
+
+
+# adapted from https://github.com/pytorch/vision/edit/master/torchvision/models/resnet.py
+__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-f37072fd.pth',
+    'resnet34':
+    'https://download.pytorch.org/models/resnet34-b627a593.pth',
+    'resnet50':
+    'https://download.pytorch.org/models/resnet50-0676ba61.pth',
+    'resnet101':
+    'https://download.pytorch.org/models/resnet101-63fe2227.pth',
+    'resnet152':
+    'https://download.pytorch.org/models/resnet152-394f9c45.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: 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 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 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 ResNet(nn.Module):
+
+    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(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 resnet18(pretrained: bool = False,
+             progress: bool = True,
+             **kwargs: Any) -> ResNet:
+    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,
+                   **kwargs)
+
+
+def resnet34(pretrained: bool = False,
+             progress: bool = True,
+             **kwargs: Any) -> ResNet:
+    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(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)
+
+
+def resnet101(pretrained: bool = False,
+              progress: bool = True,
+              **kwargs: Any) -> ResNet:
+    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, **kwargs)
+
+
+def resnet152(pretrained: bool = False,
+              progress: bool = True,
+              **kwargs: Any) -> ResNet:
+    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, **kwargs)
+
+
+def resnext50_32x4d(pretrained: bool = False,
+                    progress: bool = True,
+                    **kwargs: Any) -> ResNet:
+    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: bool = False,
+                     progress: bool = True,
+                     **kwargs: Any) -> ResNet:
+    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: bool = False,
+                    progress: bool = True,
+                    **kwargs: Any) -> ResNet:
+    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: bool = False,
+                     progress: bool = True,
+                     **kwargs: Any) -> ResNet:
+    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)
+
+
+func_dict = {'resnet18': (resnet18, 512), 'resnet50': (resnet50, 2048)}

modelscope/models/cv/face_reconstruction/models/nv_diffrast.py

@@ -0,0 +1,400 @@
+# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
+# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
+import warnings
+from typing import List
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .losses import TVLoss, TVLoss_std
+
+warnings.filterwarnings('ignore')
+
+
+def ndc_projection(x=0.1, n=1.0, f=50.0):
+    return np.array([[n / x, 0, 0, 0], [0, n / -x, 0, 0],
+                     [0, 0, -(f + n) / (f - n), -(2 * f * n) / (f - n)],
+                     [0, 0, -1, 0]]).astype(np.float32)
+
+
+def to_image(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)
+    """
+
+    focal = 1015.
+    center = 112.
+    persc_proj = np.array([focal, 0, center, 0, focal, center, 0, 0,
+                           1]).reshape([3, 3]).astype(np.float32).transpose()
+
+    persc_proj = torch.tensor(persc_proj).to(face_shape.device)
+
+    face_proj = face_shape @ persc_proj
+    face_proj = face_proj[..., :2] / face_proj[..., 2:]
+
+    return face_proj
+
+
+class MeshRenderer(nn.Module):
+
+    def __init__(self, rasterize_fov, znear=0.1, zfar=10, rasterize_size=224):
+        super(MeshRenderer, self).__init__()
+
+        x = np.tan(np.deg2rad(rasterize_fov * 0.5)) * znear
+        self.ndc_proj = torch.tensor(ndc_projection(
+            x=x, n=znear,
+            f=zfar)).matmul(torch.diag(torch.tensor([1., -1, -1, 1])))
+        self.rasterize_size = rasterize_size
+        self.glctx = None
+
+    def forward(self, vertex, tri, feat=None):
+        """
+        Return:
+            mask               -- torch.tensor, size (B, 1, H, W)
+            depth              -- torch.tensor, size (B, 1, H, W)
+            features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None
+
+        Parameters:
+            vertex          -- torch.tensor, size (B, N, 3)
+            tri             -- torch.tensor, size (B, M, 3) or (M, 3), triangles
+            feat(optional)  -- torch.tensor, size (B, C), features
+        """
+        device = vertex.device
+        rsize = int(self.rasterize_size)
+        ndc_proj = self.ndc_proj.to(device)
+        verts_proj = to_image(vertex)
+        # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
+        if vertex.shape[-1] == 3:
+            vertex = torch.cat(
+                [vertex, torch.ones([*vertex.shape[:2], 1]).to(device)],
+                dim=-1)
+            vertex[..., 1] = -vertex[..., 1]
+
+        vertex_ndc = vertex @ ndc_proj.t()
+        if self.glctx is None:
+            self.glctx = dr.RasterizeCudaContext(device=device)
+
+        ranges = None
+        if isinstance(tri, List) or len(tri.shape) == 3:
+            vum = vertex_ndc.shape[1]
+            fnum = torch.tensor([f.shape[0]
+                                 for f in tri]).unsqueeze(1).to(device)
+
+            print('fnum shape:{}'.format(fnum.shape))
+
+            fstartidx = torch.cumsum(fnum, dim=0) - fnum
+            ranges = torch.cat([fstartidx, fnum],
+                               axis=1).type(torch.int32).cpu()
+            for i in range(tri.shape[0]):
+                tri[i] = tri[i] + i * vum
+            vertex_ndc = torch.cat(vertex_ndc, dim=0)
+            tri = torch.cat(tri, dim=0)
+
+        # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
+        tri = tri.type(torch.int32).contiguous()
+        rast_out, _ = dr.rasterize(
+            self.glctx,
+            vertex_ndc.contiguous(),
+            tri,
+            resolution=[rsize, rsize],
+            ranges=ranges)
+
+        depth, _ = dr.interpolate(
+            vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(),
+            rast_out, tri)
+        depth = depth.permute(0, 3, 1, 2)
+        mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
+        depth = mask * depth
+
+        image = None
+
+        verts_x = verts_proj[0, :, 0]
+        verts_y = 224 - verts_proj[0, :, 1]
+        verts_int = torch.ceil(verts_proj[0]).long()  # (n, 2)
+        verts_xr_int = verts_int[:, 0].clamp(1, 224 - 1)
+        verts_yt_int = 224 - verts_int[:, 1].clamp(2, 224)
+        verts_right_float = verts_xr_int - verts_x
+        verts_left_float = 1 - verts_right_float
+        verts_top_float = verts_y - verts_yt_int
+        verts_bottom_float = 1 - verts_top_float
+
+        rast_lt = rast_out[0, verts_yt_int, verts_xr_int - 1, 3]
+        rast_lb = rast_out[0, verts_yt_int + 1, verts_xr_int - 1, 3]
+        rast_rt = rast_out[0, verts_yt_int, verts_xr_int, 3]
+        rast_rb = rast_out[0, verts_yt_int + 1, verts_xr_int, 3]
+
+        occ_feat = (rast_lt > 0) * 1.0 * (verts_left_float + verts_top_float) + \
+                   (rast_lb > 0) * 1.0 * (verts_left_float + verts_bottom_float) + \
+                   (rast_rt > 0) * 1.0 * (verts_right_float + verts_top_float) + \
+                   (rast_rb > 0) * 1.0 * (verts_right_float + verts_bottom_float)
+        occ_feat = occ_feat[None, :, None] / 4.0
+
+        occ, _ = dr.interpolate(occ_feat, rast_out, tri)
+        occ = occ.permute(0, 3, 1, 2)
+
+        if feat is not None:
+            image, _ = dr.interpolate(feat, rast_out, tri)
+            image = image.permute(0, 3, 1, 2)
+            image = mask * image
+
+        return mask, depth, image, occ
+
+    def render_uv_texture(self, vertex, tri, uv, uv_texture):
+        """
+        Return:
+            mask               -- torch.tensor, size (B, 1, H, W)
+            depth              -- torch.tensor, size (B, 1, H, W)
+            features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None
+
+        Parameters:
+            vertex          -- torch.tensor, size (B, N, 3)
+            tri             -- torch.tensor, size (M, 3), triangles
+            uv                -- torch.tensor, size (B,N, 2),  uv mapping
+            base_tex   -- torch.tensor, size (B,H,W,C)
+        """
+        device = vertex.device
+        rsize = int(self.rasterize_size)
+        ndc_proj = self.ndc_proj.to(device)
+        # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
+        if vertex.shape[-1] == 3:
+            vertex = torch.cat(
+                [vertex, torch.ones([*vertex.shape[:2], 1]).to(device)],
+                dim=-1)
+            vertex[..., 1] = -vertex[..., 1]
+
+        vertex_ndc = vertex @ ndc_proj.t()
+        if self.glctx is None:
+            self.glctx = dr.RasterizeCudaContext(device=device)
+
+        ranges = None
+        if isinstance(tri, List) or len(tri.shape) == 3:
+            vum = vertex_ndc.shape[1]
+            fnum = torch.tensor([f.shape[0]
+                                 for f in tri]).unsqueeze(1).to(device)
+
+            print('fnum shape:{}'.format(fnum.shape))
+
+            fstartidx = torch.cumsum(fnum, dim=0) - fnum
+            ranges = torch.cat([fstartidx, fnum],
+                               axis=1).type(torch.int32).cpu()
+            for i in range(tri.shape[0]):
+                tri[i] = tri[i] + i * vum
+            vertex_ndc = torch.cat(vertex_ndc, dim=0)
+            tri = torch.cat(tri, dim=0)
+
+        # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
+        tri = tri.type(torch.int32).contiguous()
+        rast_out, _ = dr.rasterize(
+            self.glctx,
+            vertex_ndc.contiguous(),
+            tri,
+            resolution=[rsize, rsize],
+            ranges=ranges)
+
+        depth, _ = dr.interpolate(
+            vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(),
+            rast_out, tri)
+        depth = depth.permute(0, 3, 1, 2)
+        mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
+        depth = mask * depth
+        uv[..., -1] = 1.0 - uv[..., -1]
+
+        rast_out, rast_db = dr.rasterize(
+            self.glctx,
+            vertex_ndc.contiguous(),
+            tri,
+            resolution=[rsize, rsize],
+            ranges=ranges)
+
+        interp_out, uv_da = dr.interpolate(
+            uv, rast_out, tri, rast_db, diff_attrs='all')
+
+        uv_texture = uv_texture.permute(0, 2, 3, 1).contiguous()
+        img = dr.texture(
+            uv_texture, interp_out, filter_mode='linear')  # , uv_da)
+        img = img * torch.clamp(rast_out[..., -1:], 0,
+                                1)  # Mask out background.
+
+        tex_map = uv_texture[0].detach().cpu().numpy()[..., ::-1] * 255.0
+
+        image = img.permute(0, 3, 1, 2)
+
+        return mask, depth, image, tex_map
+
+    def pred_shape_and_texture(self,
+                               vertex,
+                               tri,
+                               uv,
+                               target_img,
+                               base_tex=None):
+        """
+        Return:
+            mask               -- torch.tensor, size (B, 1, H, W)
+            depth              -- torch.tensor, size (B, 1, H, W)
+            features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None
+
+        Parameters:
+            vertex          -- torch.tensor, size (B, N, 3)
+            tri             -- torch.tensor, size (B, M, 3) or (M, 3), triangles
+            uv                -- torch.tensor, size (B,N, 2),  uv mapping
+            base_tex   -- torch.tensor, size (B,H,W,C)
+        """
+        vertex = torch.cat([vertex[:, :35241, :], vertex[:, 37082:, :]],
+                           dim=1)  # BFM front
+        tri = torch.cat([tri[:69732, :], tri[73936:, ]], dim=0)
+        uv = torch.cat([uv[:, :35241, :], uv[:, 37082:, :]], dim=1)
+        tri[69732:, :] = tri[69732:, :] - (37082 - 35241)
+
+        device = vertex.device
+        rsize = int(self.rasterize_size)
+        ndc_proj = self.ndc_proj.to(device)
+        # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
+        if vertex.shape[-1] == 3:
+            vertex = torch.cat(
+                [vertex, torch.ones([*vertex.shape[:2], 1]).to(device)],
+                dim=-1)
+            vertex[..., 1] = -vertex[..., 1]
+
+        vertex_ndc = vertex @ ndc_proj.t()
+        if self.glctx is None:
+            self.glctx = dr.RasterizeCudaContext(device=device)
+
+        ranges = None
+        if isinstance(tri, List) or len(tri.shape) == 3:
+            vum = vertex_ndc.shape[1]
+            fnum = torch.tensor([f.shape[0]
+                                 for f in tri]).unsqueeze(1).to(device)
+
+            fstartidx = torch.cumsum(fnum, dim=0) - fnum
+            ranges = torch.cat([fstartidx, fnum],
+                               axis=1).type(torch.int32).cpu()
+            for i in range(tri.shape[0]):
+                tri[i] = tri[i] + i * vum
+            vertex_ndc = torch.cat(vertex_ndc, dim=0)
+            tri = torch.cat(tri, dim=0)
+
+        # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
+        tri = tri.type(torch.int32).contiguous()
+        rast_out, _ = dr.rasterize(
+            self.glctx,
+            vertex_ndc.contiguous(),
+            tri,
+            resolution=[rsize, rsize],
+            ranges=ranges)
+
+        depth, _ = dr.interpolate(
+            vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(),
+            rast_out, tri)
+        depth = depth.permute(0, 3, 1, 2)
+        mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
+        depth = mask * depth
+        uv[..., -1] = 1.0 - uv[..., -1]
+
+        rast_out, rast_db = dr.rasterize(
+            self.glctx,
+            vertex_ndc.contiguous(),
+            tri,
+            resolution=[rsize, rsize],
+            ranges=ranges)
+
+        interp_out, uv_da = dr.interpolate(
+            uv, rast_out, tri, rast_db, diff_attrs='all')
+
+        mask_3c = mask.permute(0, 2, 3, 1)
+        mask_3c = torch.cat((mask_3c, mask_3c, mask_3c), dim=-1)
+        maskout_img = mask_3c * target_img
+        mean_color = torch.sum(maskout_img, dim=(1, 2))
+        valid_pixel_count = torch.sum(mask)
+
+        mean_color = mean_color / valid_pixel_count
+
+        tex = torch.zeros((1, 128 * 5 // 4, 128, 3), dtype=torch.float32)
+        tex[:, :, :, 0] = mean_color[0, 0]
+        tex[:, :, :, 1] = mean_color[0, 1]
+        tex[:, :, :, 2] = mean_color[0, 2]
+
+        tex = tex.cuda()
+
+        tex_mask = torch.zeros((1, 2048 * 5 // 4, 2048, 3),
+                               dtype=torch.float32)
+        tex_mask[:, :, :, 1] = 1.0
+        tex_mask = tex_mask.cuda()
+        tex_mask.requires_grad = True
+        tex_mask = tex_mask.contiguous()
+
+        criterionTV = TVLoss()
+
+        if base_tex is not None:
+            base_tex = base_tex.cuda()
+
+        for tex_resolution in [64, 128, 256, 512, 1024, 2048]:
+            tex = tex.detach()
+            tex = tex.permute(0, 3, 1, 2)
+            tex = F.interpolate(tex, (tex_resolution * 5 // 4, tex_resolution))
+            tex = tex.permute(0, 2, 3, 1).contiguous()
+
+            if base_tex is not None:
+                _base_tex = base_tex.permute(0, 3, 1, 2)
+                _base_tex = F.interpolate(
+                    _base_tex, (tex_resolution * 5 // 4, tex_resolution))
+                _base_tex = _base_tex.permute(0, 2, 3, 1).contiguous()
+                tex += _base_tex
+
+            tex.requires_grad = True
+            optim = torch.optim.Adam([tex], lr=1e-2)
+
+            texture_opt_iters = 100
+
+            if tex_resolution == 2048:
+                optim_mask = torch.optim.Adam([tex_mask], lr=1e-2)
+
+            for i in range(int(texture_opt_iters)):
+
+                if tex_resolution == 2048:
+                    optim_mask.zero_grad()
+                    rendered = dr.texture(
+                        tex_mask, interp_out, filter_mode='linear')  # , uv_da)
+                    rendered = rendered * torch.clamp(
+                        rast_out[..., -1:], 0, 1)  # Mask out background.
+                    tex_loss = torch.mean((target_img - rendered)**2)
+
+                    tex_loss.backward()
+                    optim_mask.step()
+
+                optim.zero_grad()
+
+                img = dr.texture(
+                    tex, interp_out, filter_mode='linear')  # , uv_da)
+                img = img * torch.clamp(rast_out[..., -1:], 0,
+                                        1)  # Mask out background.
+                recon_loss = torch.mean((target_img - img)**2)
+
+                if tex_resolution < 2048:
+                    tv_loss = criterionTV(tex.permute(0, 3, 1, 2))
+
+                    total_loss = recon_loss + tv_loss * 0.01
+                else:
+
+                    total_loss = recon_loss
+
+                total_loss.backward()
+                optim.step()
+
+        tex_map = tex[0].detach().cpu().numpy()[..., ::-1] * 255.0
+
+        image = img.permute(0, 3, 1, 2)
+
+        tex_mask = tex_mask[0].detach().cpu().numpy() * 255.0
+        tex_mask = np.where(tex_mask[..., 1] > 250, 1.0, 0.0) * np.where(
+            tex_mask[..., 0] < 10, 1.0, 0) * np.where(tex_mask[..., 2] < 10,
+                                                      1.0, 0)
+        tex_mask = 1.0 - tex_mask
+
+        return mask, depth, image, tex_map, tex_mask

modelscope/models/cv/face_reconstruction/models/opt.py

@@ -0,0 +1,13 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+bfm_folder = ''
+bfm_model = 'head_model_for_maas.mat'
+camera_d = 10.0
+center = 112.0
+focal = 1015.0
+isTrain = False
+net_recon = 'resnet50'
+phase = 'test'
+use_ddp = False
+use_last_fc = False
+z_far = 15.0
+z_near = 5.0

modelscope/models/cv/face_reconstruction/utils.py

@@ -0,0 +1,752 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import argparse
+import math
+import os
+import os.path as osp
+from array import array
+
+import cv2
+import numba
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from scipy.io import loadmat, savemat
+
+
+def img_value_rescale(img, old_range: list, new_range: list):
+    assert len(old_range) == 2
+    assert len(new_range) == 2
+    img = (img - old_range[0]) / (old_range[1] - old_range[0]) * (
+        new_range[1] - new_range[0]) + new_range[0]
+    return img
+
+
+def resize_on_long_side(img, long_side=800):
+    src_height = img.shape[0]
+    src_width = img.shape[1]
+
+    if src_height > src_width:
+        scale = long_side * 1.0 / src_height
+        _img = cv2.resize(
+            img, (int(src_width * scale), long_side),
+            interpolation=cv2.INTER_CUBIC)
+
+    else:
+        scale = long_side * 1.0 / src_width
+        _img = cv2.resize(
+            img, (long_side, int(src_height * scale)),
+            interpolation=cv2.INTER_CUBIC)
+
+    return _img, scale
+
+
+def get_mg_layer(src, gt, skin_mask=None):
+    """
+    src, gt shape: [h, w, 3] value: [0, 1]
+    return: mg, shape: [h, w, 1] value: [0, 1]
+    """
+    mg = (src * src - gt + 1e-10) / (2 * src * src - 2 * src + 2e-10)
+    mg[mg < 0] = 0.5
+    mg[mg > 1] = 0.5
+
+    diff_abs = np.abs(gt - src)
+    mg[diff_abs < (1 / 255.0)] = 0.5
+
+    if skin_mask is not None:
+        mg[skin_mask == 0] = 0.5
+
+    return mg
+
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def spread_flow(length, spread_ratio=2):
+    Flow = np.zeros(shape=(length, length, 2), dtype=np.float32)
+    mag = np.zeros(shape=(length, length), dtype=np.float32)
+
+    radius = length * 0.5
+    for h in range(Flow.shape[0]):
+        for w in range(Flow.shape[1]):
+
+            if (h - length // 2)**2 + (w - length // 2)**2 <= radius**2:
+                Flow[h, w, 0] = -(w - length // 2)
+                Flow[h, w, 1] = -(h - length // 2)
+
+                distance = np.sqrt((w - length // 2)**2 + (h - length // 2)**2)
+
+                if distance <= radius / 2.0:
+                    mag[h, w] = 2.0 / radius * distance
+                else:
+                    mag[h, w] = -2.0 / radius * distance + 2.0
+
+    _, ang = cv2.cartToPolar(Flow[..., 0] + 1e-8, Flow[..., 1] + 1e-8)
+
+    mag *= spread_ratio
+
+    x, y = cv2.polarToCart(mag, ang, angleInDegrees=False)
+    Flow = np.dstack((x, y))
+
+    return Flow
+
+
+@numba.jit(nopython=True, parallel=True)
+def bilinear_interp(x, y, v11, v12, v21, v22):
+    t = 0.2
+
+    if x < t and y < t:
+        return v11
+    elif x < t and y > 1 - t:
+        return v12
+    elif x > 1 - t and y < t:
+        return v21
+    elif x > 1 - t and y > 1 - t:
+        return v22
+    else:
+        result = (v11 * (1 - y) + v12 * y) * (1 - x) + \
+                 (v21 * (1 - y) + v22 * y) * x
+        if result < 0:
+            result = 0
+
+        if result > 255:
+            result = 255
+        return result
+
+
+@numba.jit(nopython=True, parallel=True)
+def image_warp_grid1(rDx, rDy, oriImg, transRatio, pads):
+    # assert oriImg.dtype == np.uint8
+    srcW = oriImg.shape[1]
+    srcH = oriImg.shape[0]
+
+    padTop, padBottom, padLeft, padRight = pads
+
+    left_bound = padLeft + 1
+    right_bound = srcW - padRight
+    bottom_bound = srcH - padBottom
+    top_bound = padTop + 1
+
+    newImg = oriImg.copy()
+
+    for i in range(srcH):
+        for j in range(srcW):
+            _i = i
+            _j = j
+
+            deltaX = rDx[_i, _j]
+            deltaY = rDy[_i, _j]
+
+            if abs(deltaX) < 0.2 and abs(deltaY) < 0.2:
+                continue
+
+            nx = _j + deltaX * transRatio
+            ny = _i + deltaY * transRatio
+
+            if nx >= srcW - padRight:
+                if nx > srcW - 1:
+                    nx = srcW - 1
+
+                if _j < right_bound:
+                    right_bound = _j
+
+            if ny >= srcH - padBottom:
+                if ny > srcH - 1:
+                    ny = srcH - 1
+
+                if _i < bottom_bound:
+                    bottom_bound = _i
+
+            if nx < padLeft:
+                if nx < 0:
+                    nx = 0
+
+                if _j + 1 > left_bound:
+                    left_bound = _j + 1
+
+            if ny < padTop:
+                if ny < 0:
+                    ny = 0
+
+                if _i + 1 > top_bound:
+                    top_bound = _i + 1
+
+            nxi = int(math.floor(nx))
+            nyi = int(math.floor(ny))
+            nxi1 = int(math.ceil(nx))
+            nyi1 = int(math.ceil(ny))
+
+            if nxi < 0:
+                nxi = 0
+            if nxi > oriImg.shape[1] - 1:
+                nxi = oriImg.shape[1] - 1
+
+            if nxi1 < 0:
+                nxi1 = 0
+            if nxi1 > oriImg.shape[1] - 1:
+                nxi1 = oriImg.shape[1] - 1
+
+            if nyi < 0:
+                nyi = 0
+            if nyi > oriImg.shape[0] - 1:
+                nyi = oriImg.shape[0] - 1
+
+            if nyi1 < 0:
+                nyi1 = 0
+            if nyi1 > oriImg.shape[0] - 1:
+                nyi1 = oriImg.shape[0] - 1
+
+            for ll in range(3):
+                newImg[_i, _j,
+                       ll] = bilinear_interp(ny - nyi, nx - nxi,
+                                             oriImg[nyi, nxi,
+                                                    ll], oriImg[nyi, nxi1, ll],
+                                             oriImg[nyi1, nxi,
+                                                    ll], oriImg[nyi1, nxi1,
+                                                                ll])
+
+    return newImg, top_bound, bottom_bound, left_bound, right_bound
+
+
+def warp(x, flow, mode='bilinear', padding_mode='zeros', coff=0.1):
+    """
+
+    Args:
+        x: [n, c, h, w]
+        flow: [n, h, w, 2]
+        mode:
+        padding_mode:
+        coff:
+
+    Returns:
+
+    """
+    n, c, h, w = x.size()
+    yv, xv = torch.meshgrid([torch.arange(h), torch.arange(w)])
+    xv = xv.float() / (w - 1) * 2.0 - 1
+    yv = yv.float() / (h - 1) * 2.0 - 1
+    '''
+    grid[0,:,:,0] =
+    -1, .....1
+    -1, .....1
+    -1, .....1
+
+    grid[0,:,:,1] =
+    -1,  -1, -1
+     ;        ;
+     1,   1,  1
+
+    '''
+
+    if torch.cuda.is_available():
+        grid = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)),
+                         -1).unsqueeze(0).cuda()
+    else:
+        grid = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)), -1).unsqueeze(0)
+    grid_x = grid + 2 * flow * coff
+    warp_x = F.grid_sample(x, grid_x, mode=mode, padding_mode=padding_mode)
+    return warp_x
+
+
+# load expression basis
+def LoadExpBasis(bfm_folder='asset/BFM'):
+    n_vertex = 53215
+    Expbin = open(osp.join(bfm_folder, 'Exp_Pca.bin'), 'rb')
+    exp_dim = array('i')
+    exp_dim.fromfile(Expbin, 1)
+    expMU = array('f')
+    expPC = array('f')
+    expMU.fromfile(Expbin, 3 * n_vertex)
+    expPC.fromfile(Expbin, 3 * exp_dim[0] * n_vertex)
+    Expbin.close()
+
+    expPC = np.array(expPC)
+    expPC = np.reshape(expPC, [exp_dim[0], -1])
+    expPC = np.transpose(expPC)
+
+    expEV = np.loadtxt(osp.join(bfm_folder, 'std_exp.txt'))
+
+    return expPC, expEV
+
+
+# transfer original BFM09 to our face model
+def transferBFM09(bfm_folder='BFM'):
+    print('Transfer BFM09 to BFM_model_front......')
+    original_BFM = loadmat(osp.join(bfm_folder, '01_MorphableModel.mat'))
+    shapePC = original_BFM['shapePC']  # shape basis, 160470*199
+    shapeEV = original_BFM['shapeEV']  # corresponding eigen value, 199*1
+    shapeMU = original_BFM['shapeMU']  # mean face, 160470*1
+    texPC = original_BFM['texPC']  # texture basis, 160470*199
+    texEV = original_BFM['texEV']  # eigen value, 199*1
+    texMU = original_BFM['texMU']  # mean texture, 160470*1
+
+    expPC, expEV = LoadExpBasis()
+
+    # transfer BFM09 to our face model
+
+    idBase = shapePC * np.reshape(shapeEV, [-1, 199])
+    idBase = idBase / 1e5  # unify the scale to decimeter
+    idBase = idBase[:, :80]  # use only first 80 basis
+
+    exBase = expPC * np.reshape(expEV, [-1, 79])
+    exBase = exBase / 1e5  # unify the scale to decimeter
+    exBase = exBase[:, :64]  # use only first 64 basis
+
+    texBase = texPC * np.reshape(texEV, [-1, 199])
+    texBase = texBase[:, :80]  # use only first 80 basis
+
+    # our face model is cropped along face landmarks and contains only 35709 vertex.
+    # original BFM09 contains 53490 vertex, and expression basis provided by Guo et al. contains 53215 vertex.
+    # thus we select corresponding vertex to get our face model.
+
+    index_exp = loadmat(osp.join(bfm_folder, 'BFM_front_idx.mat'))
+    index_exp = index_exp['idx'].astype(
+        np.int32) - 1  # starts from 0 (to 53215)
+
+    index_shape = loadmat(osp.join(bfm_folder, 'BFM_exp_idx.mat'))
+    index_shape = index_shape['trimIndex'].astype(
+        np.int32) - 1  # starts from 0 (to 53490)
+    index_shape = index_shape[index_exp]
+
+    idBase = np.reshape(idBase, [-1, 3, 80])
+    idBase = idBase[index_shape, :, :]
+    idBase = np.reshape(idBase, [-1, 80])
+
+    texBase = np.reshape(texBase, [-1, 3, 80])
+    texBase = texBase[index_shape, :, :]
+    texBase = np.reshape(texBase, [-1, 80])
+
+    exBase = np.reshape(exBase, [-1, 3, 64])
+    exBase = exBase[index_exp, :, :]
+    exBase = np.reshape(exBase, [-1, 64])
+
+    meanshape = np.reshape(shapeMU, [-1, 3]) / 1e5
+    meanshape = meanshape[index_shape, :]
+    meanshape = np.reshape(meanshape, [1, -1])
+
+    meantex = np.reshape(texMU, [-1, 3])
+    meantex = meantex[index_shape, :]
+    meantex = np.reshape(meantex, [1, -1])
+
+    # other info contains triangles, region used for computing photometric loss,
+    # region used for skin texture regularization, and 68 landmarks index etc.
+    other_info = loadmat(osp.join(bfm_folder, 'facemodel_info.mat'))
+    frontmask2_idx = other_info['frontmask2_idx']
+    skinmask = other_info['skinmask']
+    keypoints = other_info['keypoints']
+    point_buf = other_info['point_buf']
+    tri = other_info['tri']
+    tri_mask2 = other_info['tri_mask2']
+
+    # save our face model
+    savemat(
+        osp.join(bfm_folder, 'BFM_model_front.mat'), {
+            'meanshape': meanshape,
+            'meantex': meantex,
+            'idBase': idBase,
+            'exBase': exBase,
+            'texBase': texBase,
+            'tri': tri,
+            'point_buf': point_buf,
+            'tri_mask2': tri_mask2,
+            'keypoints': keypoints,
+            'frontmask2_idx': frontmask2_idx,
+            'skinmask': skinmask
+        })
+
+
+# load landmarks for standard face, which is used for image preprocessing
+def load_lm3d(bfm_folder):
+
+    Lm3D = loadmat(osp.join(bfm_folder, 'similarity_Lm3D_all.mat'))
+    Lm3D = Lm3D['lm']
+
+    # calculate 5 facial landmarks using 68 landmarks
+    lm_idx = np.array([31, 37, 40, 43, 46, 49, 55]) - 1
+    value_list = [
+        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], :]
+    ]
+    Lm3D = np.stack(value_list, axis=0)
+    Lm3D = Lm3D[[1, 2, 0, 3, 4], :]
+
+    return Lm3D
+
+
+def write_obj(save_path, mesh):
+    save_dir = os.path.dirname(save_path)
+    save_name = os.path.splitext(os.path.basename(save_path))[0]
+
+    if 'texture_map' in mesh:
+        cv2.imwrite(
+            os.path.join(save_dir, save_name + '.jpg'), mesh['texture_map'])
+
+        with open(os.path.join(save_dir, save_name + '.mtl'), 'w') as wf:
+            wf.write('# Created by ModelScope\n')
+            wf.write('newmtl material_0\n')
+            wf.write('Ka 1.000000 0.000000 0.000000\n')
+            wf.write('Kd 1.000000 1.000000 1.000000\n')
+            wf.write('Ks 0.000000 0.000000 0.000000\n')
+            wf.write('Tr 0.000000\n')
+            wf.write('illum 0\n')
+            wf.write('Ns 0.000000\n')
+            wf.write('map_Kd {}\n'.format(save_name + '.jpg'))
+
+    with open(save_path, 'w') as wf:
+        if 'texture_map' in mesh:
+            wf.write('# Create by ModelScope\n')
+            wf.write('mtllib ./{}.mtl\n'.format(save_name))
+
+        if 'colors' in mesh:
+            for i, v in enumerate(mesh['vertices']):
+                wf.write('v {} {} {} {} {} {}\n'.format(
+                    v[0], v[1], v[2], mesh['colors'][i][0],
+                    mesh['colors'][i][1], mesh['colors'][i][2]))
+        else:
+            for v in mesh['vertices']:
+                wf.write('v {} {} {}\n'.format(v[0], v[1], v[2]))
+
+        if 'UVs' in mesh:
+            for uv in mesh['UVs']:
+                wf.write('vt {} {}\n'.format(uv[0], uv[1]))
+
+        if 'normals' in mesh:
+            for vn in mesh['normals']:
+                wf.write('vn {} {} {}\n'.format(vn[0], vn[1], vn[2]))
+
+        if 'faces' in mesh:
+            for ind, face in enumerate(mesh['faces']):
+                if 'faces_uv' in mesh or 'faces_normal' in mesh:
+                    if 'faces_uv' in mesh:
+                        face_uv = mesh['faces_uv'][ind]
+                    else:
+                        face_uv = face
+                    if 'faces_normal' in mesh:
+                        face_normal = mesh['faces_normal'][ind]
+                    else:
+                        face_normal = face
+                    row = 'f ' + ' '.join([
+                        '{}/{}/{}'.format(face[i], face_uv[i], face_normal[i])
+                        for i in range(len(face))
+                    ]) + '\n'
+                else:
+                    row = 'f ' + ' '.join(
+                        ['{}'.format(face[i])
+                         for i in range(len(face))]) + '\n'
+                wf.write(row)
+
+
+def read_obj(obj_path, print_shape=True):
+    with open(obj_path, 'r') as f:
+        bfm_lines = f.readlines()
+
+    vertices = []
+    faces = []
+    uvs = []
+    vns = []
+    faces_uv = []
+    faces_normal = []
+    max_face_length = 0
+    for line in bfm_lines:
+        if line[:2] == 'v ':
+            vertex = [
+                float(a) for a in line.strip().split(' ')[1:] if len(a) > 0
+            ]
+            vertices.append(vertex)
+
+        if line[:2] == 'f ':
+            items = line.strip().split(' ')[1:]
+            face = [int(a.split('/')[0]) for a in items if len(a) > 0]
+            max_face_length = max(max_face_length, len(face))
+            if len(faces) > 0 and len(face) != len(faces[0]):
+                continue
+            faces.append(face)
+
+            if '/' in items[0] and len(items[0].split('/')[1]) > 0:
+                face_uv = [int(a.split('/')[1]) for a in items if len(a) > 0]
+                faces_uv.append(face_uv)
+
+            if '/' in items[0] and len(items[0].split('/')) >= 3 and len(
+                    items[0].split('/')[2]) > 0:
+                face_normal = [
+                    int(a.split('/')[2]) for a in items if len(a) > 0
+                ]
+                faces_normal.append(face_normal)
+
+        if line[:3] == 'vt ':
+            items = line.strip().split(' ')[1:]
+            uv = [float(a) for a in items if len(a) > 0]
+            uvs.append(uv)
+
+        if line[:3] == 'vn ':
+            items = line.strip().split(' ')[1:]
+            vn = [float(a) for a in items if len(a) > 0]
+            vns.append(vn)
+
+    vertices = np.array(vertices).astype(np.float32)
+    if max_face_length <= 3:
+        faces = np.array(faces).astype(np.int32)
+
+    if vertices.shape[1] == 3:
+        mesh = {
+            'vertices': vertices,
+            'faces': faces,
+        }
+    else:
+        mesh = {
+            'vertices': vertices[:, :3],
+            'colors': vertices[:, 3:],
+            'faces': faces,
+        }
+
+    if len(uvs) > 0:
+        uvs = np.array(uvs).astype(np.float32)
+        mesh['uvs'] = uvs
+
+    if len(vns) > 0:
+        vns = np.array(vns).astype(np.float32)
+        mesh['vns'] = vns
+
+    if len(faces_uv) > 0:
+        if max_face_length <= 3:
+            faces_uv = np.array(faces_uv).astype(np.int32)
+        mesh['faces_uv'] = faces_uv
+
+    if len(faces_normal) > 0:
+        if max_face_length <= 3:
+            faces_normal = np.array(faces_normal).astype(np.int32)
+        mesh['faces_normal'] = faces_normal
+
+    return mesh
+
+
+# calculating least square problem for image alignment
+def POS(xp, x):
+    npts = xp.shape[1]
+
+    A = np.zeros([2 * npts, 8])
+
+    A[0:2 * npts - 1:2, 0:3] = x.transpose()
+    A[0:2 * npts - 1:2, 3] = 1
+
+    A[1:2 * npts:2, 4:7] = x.transpose()
+    A[1:2 * npts:2, 7] = 1
+
+    b = np.reshape(xp.transpose(), [2 * npts, 1])
+
+    k, _, _, _ = np.linalg.lstsq(A, b)
+
+    R1 = k[0:3]
+    R2 = k[4:7]
+    sTx = k[3]
+    sTy = k[7]
+    s = (np.linalg.norm(R1) + np.linalg.norm(R2)) / 2
+    t = np.stack([sTx, sTy], axis=0)
+
+    return t, s
+
+
+# bounding box for 68 landmark detection
+def BBRegression(points, params):
+    w1 = params['W1']
+    b1 = params['B1']
+    w2 = params['W2']
+    b2 = params['B2']
+    data = points.copy()
+    data = data.reshape([5, 2])
+    data_mean = np.mean(data, axis=0)
+    x_mean = data_mean[0]
+    y_mean = data_mean[1]
+    data[:, 0] = data[:, 0] - x_mean
+    data[:, 1] = data[:, 1] - y_mean
+
+    rms = np.sqrt(np.sum(data**2) / 5)
+    data = data / rms
+    data = data.reshape([1, 10])
+    data = np.transpose(data)
+    inputs = np.matmul(w1, data) + b1
+    inputs = 2 / (1 + np.exp(-2 * inputs)) - 1
+    inputs = np.matmul(w2, inputs) + b2
+    inputs = np.transpose(inputs)
+    x = inputs[:, 0] * rms + x_mean
+    y = inputs[:, 1] * rms + y_mean
+    w = 224 / inputs[:, 2] * rms
+    rects = [x, y, w, w]
+    return np.array(rects).reshape([4])
+
+
+# utils for landmark detection
+def img_padding(img, box):
+    success = True
+    bbox = box.copy()
+    res = np.zeros([2 * img.shape[0], 2 * img.shape[1], 3])
+    res[img.shape[0] // 2:img.shape[0] + img.shape[0] // 2,
+        img.shape[1] // 2:img.shape[1] + img.shape[1] // 2] = img
+
+    bbox[0] = bbox[0] + img.shape[1] // 2
+    bbox[1] = bbox[1] + img.shape[0] // 2
+    if bbox[0] < 0 or bbox[1] < 0:
+        success = False
+    return res, bbox, success
+
+
+# utils for landmark detection
+def crop(img, bbox):
+    padded_img, padded_bbox, flag = img_padding(img, bbox)
+    if flag:
+        crop_img = padded_img[padded_bbox[1]:padded_bbox[1] + padded_bbox[3],
+                              padded_bbox[0]:padded_bbox[0] + padded_bbox[2]]
+        crop_img = cv2.resize(
+            crop_img.astype(np.uint8), (224, 224),
+            interpolation=cv2.INTER_CUBIC)
+        scale = 224 / padded_bbox[3]
+        return crop_img, scale
+    else:
+        return padded_img, 0
+
+
+# utils for landmark detection
+def scale_trans(img, lm, t, s):
+    imgw = img.shape[1]
+    imgh = img.shape[0]
+    M_s = np.array(
+        [[1, 0, -t[0] + imgw // 2 + 0.5], [0, 1, -imgh // 2 + t[1]]],
+        dtype=np.float32)
+    img = cv2.warpAffine(img, M_s, (imgw, imgh))
+    w = int(imgw / s * 100)
+    h = int(imgh / s * 100)
+    img = cv2.resize(img, (w, h))
+    lm = np.stack([lm[:, 0] - t[0] + imgw // 2, lm[:, 1] - t[1] + imgh // 2],
+                  axis=1) / s * 100
+
+    left = w // 2 - 112
+    up = h // 2 - 112
+    bbox = [left, up, 224, 224]
+    cropped_img, scale2 = crop(img, bbox)
+    assert (scale2 != 0)
+    t1 = np.array([bbox[0], bbox[1]])
+
+    # back to raw img s * crop + s * t1 + t2
+    t1 = np.array([w // 2 - 112, h // 2 - 112])
+    scale = s / 100
+    t2 = np.array([t[0] - imgw / 2, t[1] - imgh / 2])
+    inv = (scale / scale2, scale * t1 + t2.reshape([2]))
+    return cropped_img, inv
+
+
+# utils for landmark detection
+def align_for_lm(img, five_points, params):
+    five_points = np.array(five_points).reshape([1, 10])
+    bbox = BBRegression(five_points, params)
+    assert (bbox[2] != 0)
+    bbox = np.round(bbox).astype(np.int32)
+    crop_img, scale = crop(img, bbox)
+    return crop_img, scale, bbox
+
+
+# resize and crop images for face reconstruction
+def resize_n_crop_img(img, lm, t, s, target_size=224., mask=None):
+    w0, h0 = img.size
+    w = (w0 * s).astype(np.int32)
+    h = (h0 * s).astype(np.int32)
+    left = (w / 2 - target_size / 2 + float(
+        (t[0] - w0 / 2) * s)).astype(np.int32)
+    right = left + target_size
+    up = (h / 2 - target_size / 2 + float(
+        (h0 / 2 - t[1]) * s)).astype(np.int32)
+    below = up + target_size
+
+    new_img = img.resize((w, h), resample=Image.BICUBIC)
+    new_img = new_img.crop((left, up, right, below))
+
+    if mask is not None:
+        mask = mask.resize((w, h), resample=Image.BICUBIC)
+        mask = mask.crop((left, up, right, below))
+
+    new_lm = np.stack([lm[:, 0] - t[0] + w0 / 2, lm[:, 1] - t[1] + h0 / 2],
+                      axis=1) * s
+    new_lm = new_lm - np.reshape(
+        np.array([(w / 2 - target_size / 2),
+                  (h / 2 - target_size / 2)]), [1, 2])
+
+    return new_img, new_lm, mask
+
+
+# utils for face reconstruction
+def extract_5p(lm):
+    lm_idx = np.array([31, 37, 40, 43, 46, 49, 55]) - 1
+    value_list = [
+        lm[lm_idx[0], :],
+        np.mean(lm[lm_idx[[1, 2]], :], 0),
+        np.mean(lm[lm_idx[[3, 4]], :], 0), lm[lm_idx[5], :], lm[lm_idx[6], :]
+    ]
+    lm5p = np.stack(value_list, axis=0)
+    lm5p = lm5p[[1, 2, 0, 3, 4], :]
+    return lm5p
+
+
+# utils for face reconstruction
+def align_img(img, lm, lm3D, mask=None, target_size=224., rescale_factor=102.):
+    """
+    Return:
+        transparams        --numpy.array  (raw_W, raw_H, scale, tx, ty)
+        img_new            --PIL.Image  (target_size, target_size, 3)
+        lm_new             --numpy.array  (68, 2), y direction is opposite to v direction
+        mask_new           --PIL.Image  (target_size, target_size)
+
+    Parameters:
+        img                --PIL.Image  (raw_H, raw_W, 3)
+        lm                 --numpy.array  (68, 2), y direction is opposite to v direction
+        lm3D               --numpy.array  (5, 3)
+        mask               --PIL.Image  (raw_H, raw_W, 3)
+    """
+
+    w0, h0 = img.size
+    if lm.shape[0] != 5:
+        lm5p = extract_5p(lm)
+    else:
+        lm5p = lm
+
+    # calculate translation and scale factors using 5 facial landmarks and standard landmarks of a 3D face
+    t, s = POS(lm5p.transpose(), lm3D.transpose())
+    s = rescale_factor / s
+
+    # processing the image
+    img_new, lm_new, mask_new = resize_n_crop_img(
+        img, lm, t, s, target_size=target_size, mask=mask)
+    trans_params = np.array([w0, h0, s, t[0], t[1]])
+
+    return trans_params, img_new, lm_new, mask_new
+
+
+def normalize_v3(arr):
+    ''' Normalize a numpy array of 3 component vectors shape=(n,3) '''
+    lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2)[:, None]
+    arr /= lens
+    return arr
+
+
+def estimate_normals(vertices, faces):
+    norm = np.zeros(vertices.shape, dtype=vertices.dtype)
+    tris = vertices[faces]
+    n = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0])
+    n[(n[:, 0] == 0) * (n[:, 1] == 0) * (n[:, 2] == 0)] = [0, 0, 1.0]
+    n = normalize_v3(n)
+    for i in range(3):
+        for j in range(faces.shape[0]):
+            norm[faces[j, i]] += n[j]
+
+    inds = (norm[:, 0] == 0) * (norm[:, 1] == 0) * (norm[:, 2] == 0)
+    norm[inds] = [0, 0, 1.0]
+    result = normalize_v3(norm)
+    return result

modelscope/outputs/outputs.py

@@ -390,6 +390,21 @@ TASK_OUTPUTS = {
     Tasks.body_3d_keypoints:
     [OutputKeys.KEYPOINTS, OutputKeys.TIMESTAMPS, OutputKeys.OUTPUT_VIDEO],
 
+    # 3D face reconstruction result for single sample
+    # {
+    #     "output": {
+    #         "vertices": np.array with shape(n, 3),
+    #         "faces": np.array with shape(n, 3),
+    #         "faces_uv": np.array with shape(n, 3),
+    #         "faces_normal": np.array with shape(n, 3),
+    #         "colors": np.array with shape(n, 3),
+    #         "UVs": np.array with shape(n, 2),
+    #         "normals": np.array with shape(n, 3),
+    #         "texture_map": np.array with shape(h, w, 3),
+    #     }
+    # }
+    Tasks.face_reconstruction: [OutputKeys.OUTPUT],
+
     # 2D hand keypoints result for single sample
     # {
     #     "keypoints": [

modelscope/pipeline_inputs.py

@@ -68,6 +68,8 @@ TASK_INPUTS = {
     InputType.IMAGE,
     Tasks.face_recognition:
     InputType.IMAGE,
+    Tasks.face_reconstruction:
+    InputType.IMAGE,
     Tasks.human_detection:
     InputType.IMAGE,
     Tasks.face_image_generation:

modelscope/pipelines/cv/__init__.py

@@ -51,6 +51,7 @@ if TYPE_CHECKING:
     from .license_plate_detection_pipeline import LicensePlateDetectionPipeline
     from .table_recognition_pipeline import TableRecognitionPipeline
     from .skin_retouching_pipeline import SkinRetouchingPipeline
+    from .face_reconstruction_pipeline import FaceReconstructionPipeline
     from .tinynas_classification_pipeline import TinynasClassificationPipeline
     from .video_category_pipeline import VideoCategoryPipeline
     from .virtual_try_on_pipeline import VirtualTryonPipeline
@@ -150,6 +151,7 @@ else:
         'license_plate_detection_pipeline': ['LicensePlateDetectionPipeline'],
         'table_recognition_pipeline': ['TableRecognitionPipeline'],
         'skin_retouching_pipeline': ['SkinRetouchingPipeline'],
+        'face_reconstruction_pipeline': ['FaceReconstructionPipeline'],
         'tinynas_classification_pipeline': ['TinynasClassificationPipeline'],
         'video_category_pipeline': ['VideoCategoryPipeline'],
         'virtual_try_on_pipeline': ['VirtualTryonPipeline'],

modelscope/pipelines/cv/face_reconstruction_pipeline.py

@@ -0,0 +1,370 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+from typing import Any, Dict
+
+import cv2
+import face_alignment
+import numpy as np
+import PIL.Image
+import tensorflow as tf
+import torch
+from scipy.io import loadmat, savemat
+
+from modelscope.metainfo import Pipelines
+from modelscope.models import Model
+from modelscope.models.cv.face_reconstruction.models.facelandmark.large_model_infer import \
+    LargeModelInfer
+from modelscope.models.cv.face_reconstruction.utils import (align_for_lm,
+                                                            align_img,
+                                                            load_lm3d,
+                                                            read_obj,
+                                                            write_obj)
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.base import Input, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors import LoadImage
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.device import create_device, device_placement
+from modelscope.utils.logger import get_logger
+
+if tf.__version__ >= '2.0':
+    tf = tf.compat.v1
+    tf.disable_eager_execution()
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(
+    Tasks.face_reconstruction, module_name=Pipelines.face_reconstruction)
+class FaceReconstructionPipeline(Pipeline):
+
+    def __init__(self, model: str, device: str):
+        """The inference pipeline for face reconstruction task.
+
+        Args:
+            model (`str` or `Model` or module instance): A model instance or a model local dir
+                or a model id in the model hub.
+            device ('str'): device str, should be either cpu, cuda, gpu, gpu:X or cuda:X.
+
+        Example:
+            >>> from modelscope.pipelines import pipeline
+            >>> test_image = 'data/test/images/face_reconstruction.jpg'
+            >>> pipeline_faceRecon = pipeline('face-reconstruction',
+                model='damo/cv_resnet50_face-reconstruction')
+            >>> result = pipeline_faceRecon(test_image)
+            >>> write_obj('result_face_reconstruction.obj', result[OutputKeys.OUTPUT])
+        """
+        super().__init__(model=model, device=device)
+
+        model_root = model
+        bfm_folder = os.path.join(model_root, 'assets')
+        checkpoint_path = os.path.join(model_root, ModelFile.TORCH_MODEL_FILE)
+
+        self.face_mark_model = LargeModelInfer(
+            os.path.join(model_root, 'large_base_net.pth'), device='cuda')
+
+        device = torch.device(0)
+        torch.cuda.set_device(device)
+        self.model.setup(checkpoint_path)
+        self.model.device = device
+        self.model.parallelize()
+        self.model.eval()
+        self.model.set_render(image_res=1024)
+
+        save_ckpt_dir = os.path.join(
+            os.path.expanduser('~'), '.cache/torch/hub/checkpoints')
+        if not os.path.exists(save_ckpt_dir):
+            os.makedirs(save_ckpt_dir)
+        shutil.copy(
+            os.path.join(model_root, 'face_alignment', 's3fd-619a316812.pth'),
+            save_ckpt_dir)
+        shutil.copy(
+            os.path.join(model_root, 'face_alignment',
+                         '3DFAN4-4a694010b9.zip'), save_ckpt_dir)
+        shutil.copy(
+            os.path.join(model_root, 'face_alignment', 'depth-6c4283c0e0.zip'),
+            save_ckpt_dir)
+        self.lm_sess = face_alignment.FaceAlignment(
+            face_alignment.LandmarksType._3D, flip_input=False)
+
+        config = tf.ConfigProto(allow_soft_placement=True)
+        config.gpu_options.per_process_gpu_memory_fraction = 0.2
+        config.gpu_options.allow_growth = True
+        g1 = tf.Graph()
+        self.face_sess = tf.Session(graph=g1, config=config)
+        with self.face_sess.as_default():
+            with g1.as_default():
+                with tf.gfile.FastGFile(
+                        os.path.join(model_root, 'segment_face.pb'),
+                        'rb') as f:
+                    graph_def = tf.GraphDef()
+                    graph_def.ParseFromString(f.read())
+                    self.face_sess.graph.as_default()
+                    tf.import_graph_def(graph_def, name='')
+                    self.face_sess.run(tf.global_variables_initializer())
+
+        self.tex_size = 4096
+
+        self.bald_tex_bg = cv2.imread(
+            '{}/assets/template_texture.jpg'.format(model_root)).astype(
+                np.float32)
+
+        front_mask = cv2.imread(
+            '{}/assets/face_mask.jpg'.format(model_root)).astype(
+                np.float32) / 255
+        front_mask = cv2.resize(front_mask, (1024, 1024))
+        front_mask = cv2.resize(front_mask, (0, 0), fx=0.1, fy=0.1)
+        front_mask = cv2.erode(front_mask,
+                               np.ones(shape=(7, 7), dtype=np.float32))
+        front_mask = cv2.GaussianBlur(front_mask, (13, 13), 0)
+        self.front_mask = cv2.resize(front_mask,
+                                     (self.tex_size, self.tex_size))
+        self.binary_front_mask = self.front_mask.copy()
+        self.binary_front_mask[(self.front_mask < 0.3)
+                               + (self.front_mask > 0.7)] = 0
+        self.binary_front_mask[self.binary_front_mask != 0] = 1.0
+        self.binary_front_mask_ = self.binary_front_mask.copy()
+        self.binary_front_mask = np.zeros((4096 + 1024, 4096, 3),
+                                          dtype=np.float32)
+        self.binary_front_mask[:4096, :] = self.binary_front_mask_
+        self.front_mask_ = self.front_mask.copy()
+        self.front_mask = np.zeros((4096 + 1024, 4096, 3), dtype=np.float32)
+        self.front_mask[:4096, :] = self.front_mask_
+
+        l_eye_mask = cv2.imread(
+            '{}/assets/l_eye_mask.png'.format(model_root))[:, :, :1] / 255.0
+        l_eye_mask = cv2.erode(l_eye_mask,
+                               np.ones(shape=(5, 5), dtype=np.float32))
+        self.l_eye_mask = cv2.GaussianBlur(l_eye_mask, (7, 7), 0)[..., None]
+        self.l_eye_binary_mask = self.l_eye_mask.copy()
+        self.l_eye_binary_mask[(self.l_eye_mask < 0.3)
+                               + (self.l_eye_mask > 0.7)] = 0
+        self.l_eye_binary_mask[self.l_eye_binary_mask != 0] = 1.0
+
+        r_eye_mask = cv2.imread(
+            '{}/assets/r_eye_mask.png'.format(model_root))[:, :, :1] / 255.0
+        r_eye_mask = cv2.dilate(r_eye_mask,
+                                np.ones(shape=(7, 7), dtype=np.float32))
+        self.r_eye_mask = cv2.GaussianBlur(r_eye_mask, (7, 7), 0)[..., None]
+        self.r_eye_binary_mask = self.r_eye_mask.copy()
+        self.r_eye_binary_mask[(self.r_eye_mask < 0.3)
+                               + (self.r_eye_mask > 0.7)] = 0
+        self.r_eye_binary_mask[self.r_eye_binary_mask != 0] = 1.0
+
+        self.lm3d_std = load_lm3d(bfm_folder)
+        self.align_params = loadmat(
+            '{}/assets/BBRegressorParam_r.mat'.format(model_root))
+
+        device = create_device(self.device_name)
+        self.device = device
+
+    def preprocess(self, input: Input) -> Dict[str, Any]:
+        img = LoadImage.convert_to_ndarray(input)
+        if len(img.shape) == 2:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        img = img.astype(np.float)
+        result = {'img': img}
+        return result
+
+    def read_data(self,
+                  img,
+                  lm,
+                  lm3d_std,
+                  to_tensor=True,
+                  image_res=1024,
+                  img_fat=None):
+        # to RGB
+        im = PIL.Image.fromarray(img[..., ::-1])
+        W, H = im.size
+        lm[:, -1] = H - 1 - lm[:, -1]
+
+        im_lr_coeff, lm_lr_coeff = None, None
+        head_mask = None
+
+        _, im_lr, lm_lr, mask_lr_head = align_img(
+            im, lm, lm3d_std, mask=head_mask)
+        _, im_hd, lm_hd, _ = align_img(
+            im,
+            lm,
+            lm3d_std,
+            target_size=image_res,
+            rescale_factor=102.0 * image_res / 224)
+
+        mask_lr = self.face_sess.run(
+            self.face_sess.graph.get_tensor_by_name('output_alpha:0'),
+            feed_dict={'input_image:0': np.array(im_lr)})
+
+        if img_fat is not None:
+            assert img_fat.shape == img.shape
+            im_fat = PIL.Image.fromarray(img_fat[..., ::-1])
+
+            _, im_hd, _, _ = align_img(
+                im_fat,
+                lm,
+                lm3d_std,
+                target_size=image_res,
+                rescale_factor=102.0 * image_res / 224)
+
+        im_hd = np.array(im_hd).astype(np.float32)
+
+        if to_tensor:
+            im_lr = torch.tensor(
+                np.array(im_lr) / 255.,
+                dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
+            im_hd = torch.tensor(
+                np.array(im_hd) / 255.,
+                dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
+            mask_lr = torch.tensor(
+                np.array(mask_lr) / 255., dtype=torch.float32)[None,
+                                                               None, :, :]
+            mask_lr_head = torch.tensor(
+                np.array(mask_lr_head) / 255., dtype=torch.float32)[
+                    None, None, :, :] if mask_lr_head is not None else None
+            lm_lr = torch.tensor(lm_lr).unsqueeze(0)
+            lm_hd = torch.tensor(lm_hd).unsqueeze(0)
+        return im_lr, lm_lr, im_hd, lm_hd, mask_lr, mask_lr_head, im_lr_coeff, lm_lr_coeff
+
+    def prepare_data(self, img, lm_sess, five_points=None):
+        input_img, scale, bbox = align_for_lm(
+            img, five_points,
+            self.align_params)  # align for 68 landmark detection
+
+        if scale == 0:
+            return None
+
+        # detect landmarks
+        input_img = np.reshape(input_img, [1, 224, 224, 3]).astype(np.float32)
+
+        input_img = input_img[0, :, :, ::-1]
+        landmark = lm_sess.get_landmarks_from_image(input_img)[0]
+
+        landmark = landmark[:, :2] / scale
+        landmark[:, 0] = landmark[:, 0] + bbox[0]
+        landmark[:, 1] = landmark[:, 1] + bbox[1]
+
+        return landmark
+
+    def blend_eye_corner(self, tex_map, template_tex):
+        tex_map = tex_map.astype(np.float32)
+
+        x1 = int(288 * 4096 / 758)
+        y1 = int(235 * 4096 / 758)
+        w = int(90 * 4096 / 758)
+        h = int(50 * 4096 / 758)
+        template_tex_l = template_tex[y1:y1 + h, x1:x1 + w]
+        pred_tex_l = tex_map[y1:y1 + h, x1:x1 + w]
+        pred_tex_l_mean_rgb = np.sum(
+            pred_tex_l * self.l_eye_binary_mask, axis=(0, 1))
+        template_tex_l_mean_rgb = np.sum(
+            template_tex_l * self.l_eye_binary_mask, axis=(0, 1))
+        for ch in range(3):
+            template_tex_l[:, :, ch] *= pred_tex_l_mean_rgb[
+                ch] / template_tex_l_mean_rgb[ch]
+        pred_tex_l = pred_tex_l * (
+            1 - self.l_eye_mask) + template_tex_l * self.l_eye_mask
+
+        x2 = 4096 - x1 - w
+        y2 = y1
+        template_tex_r = template_tex[y2:y2 + h, x2:x2 + w]
+        pred_tex_r = tex_map[y2:y2 + h, x2:x2 + w]
+        pred_tex_r_mean_rgb = np.sum(
+            pred_tex_r * self.r_eye_binary_mask, axis=(0, 1))
+        template_tex_r_mean_rgb = np.sum(
+            template_tex_r * self.r_eye_binary_mask, axis=(0, 1))
+        for ch in range(3):
+            template_tex_r[:, :, ch] *= pred_tex_r_mean_rgb[
+                ch] / template_tex_r_mean_rgb[ch]
+        pred_tex_r = pred_tex_r * (
+            1 - self.r_eye_mask) + template_tex_r * self.r_eye_mask
+
+        tex_map[y1:y1 + h, x1:x1 + w] = pred_tex_l
+        tex_map[y2:y2 + h, x2:x2 + w] = pred_tex_r
+
+        return tex_map
+
+    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        rgb_image = input['img'].cpu().numpy().astype(np.uint8)
+
+        bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
+
+        img = bgr_image
+        # preprocess
+        flag = 0
+        box, results = self.face_mark_model.infer(img)
+        if results is None or np.array(results).shape[0] == 0:
+            flag = 1  # no face
+            return flag, {}
+
+        fatbgr = self.face_mark_model.fat_face(img, degree=0.02)
+
+        landmarks = []
+        results = results[0]
+        for idx in [74, 83, 54, 84, 90]:
+            landmarks.append([results[idx][0], results[idx][1]])
+        landmarks = np.array(landmarks)
+
+        landmarks = self.prepare_data(img, self.lm_sess, five_points=landmarks)
+
+        im_tensor, lm_tensor, im_hd_tensor, lm_hd_tensor, mask, _, _, _ = self.read_data(
+            img, landmarks, self.lm3d_std, image_res=1024, img_fat=fatbgr)
+        data = {
+            'imgs': im_tensor,
+            'imgs_hd': im_hd_tensor,
+            'lms': lm_tensor,
+            'lms_hd': lm_hd_tensor,
+            'face_mask': mask,
+            'img_name': 'temp',
+        }
+        self.model.set_input(data)  # unpack data from data loader
+
+        # reconstruct
+        out_dir = None
+        output = self.model(out_dir=out_dir)  # run inference
+
+        # process texture map
+        tex_map = output['head_tex_map'].astype(np.float32)
+        tex_map = cv2.resize(tex_map, (self.tex_size, self.tex_size + 1024))
+        bg_mean_rgb = np.sum(
+            self.bald_tex_bg * self.binary_front_mask, axis=(0, 1))
+        pred_tex_mean_rgb = np.sum(
+            tex_map * self.binary_front_mask, axis=(0, 1)) * 1.05
+        mid_mean_rgb = bg_mean_rgb * 0.8 + pred_tex_mean_rgb * 0.2
+        tex_map += (
+            (mid_mean_rgb - pred_tex_mean_rgb)
+            / np.sum(self.binary_front_mask, axis=(0, 1)))[None, None] * 0.5
+        pred_tex_mean_rgb = np.sum(
+            tex_map * self.binary_front_mask, axis=(0, 1)) * 1.05
+        _bald_tex_bg = self.bald_tex_bg.copy()
+        for ch in range(3):
+            _bald_tex_bg[:, :, ch] *= pred_tex_mean_rgb[ch] / bg_mean_rgb[ch]
+        tex_map = _bald_tex_bg * (
+            1. - self.front_mask) + tex_map * self.front_mask
+        tex_map = tex_map * 1.05
+        tex_map = self.blend_eye_corner(tex_map, self.bald_tex_bg)
+
+        # export mesh
+        results = {
+            'vertices': output['head_vertices'],
+            'faces': output['head_faces'],
+            'UVs': output['head_UVs'],
+            'faces_uv': output['head_faces_uv'],
+            'normals': output['head_normals'],
+            'texture_map': tex_map,
+        }
+
+        if out_dir is not None:
+            face_mesh = {
+                'vertices': output['face_vertices'],
+                'faces': output['face_faces'],
+                'colors': output['face_colors'],
+            }
+
+            write_obj(os.path.join(out_dir, 'face.obj'), face_mesh)
+            write_obj(os.path.join(out_dir, 'head.obj'), results)
+
+        return {OutputKeys.OUTPUT: results}
+
+    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        return inputs

modelscope/utils/constant.py

@@ -124,6 +124,9 @@ class CVTasks(object):
     # domain specific object detection
     domain_specific_object_detection = 'domain-specific-object-detection'
 
+    # 3d reconstruction
+    face_reconstruction = 'face-reconstruction'
+
     # image quality assessment mos
     image_quality_assessment_mos = 'image-quality-assessment-mos'
     # motion generation

requirements/cv.txt

@@ -6,6 +6,7 @@ clip>=1.0
 ddpm_guided_diffusion
 easydict
 easyrobust
+face_alignment>=1.3.5
 fairscale>=0.4.1
 fastai>=1.0.51
 ffmpeg>=1.4
@@ -41,6 +42,7 @@ tensorflow-estimator>=1.15.1
 tf_slim
 timm>=0.4.9
 torchmetrics>=0.6.2
+torchsummary>=1.5.1
 torchvision
 ujson
 utils

tests/pipelines/test_face_reconstruction.py

@@ -0,0 +1,52 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path as osp
+import sys
+import unittest
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.models.cv.face_reconstruction.utils import write_obj
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.base import Pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+sys.path.append('.')
+
+
+class FaceReconstructionTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.face_reconstruction
+        self.model_id = 'damo/cv_resnet50_face-reconstruction'
+        self.test_image = 'data/test/images/face_reconstruction.jpg'
+
+    def pipeline_inference(self, pipeline: Pipeline, input_location: str):
+        result = pipeline(input_location)
+        mesh = result[OutputKeys.OUTPUT]
+        write_obj('result_face_reconstruction.obj', mesh)
+        print(
+            f'Output written to {osp.abspath("result_face_reconstruction.obj")}'
+        )
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_run_by_direct_model_download(self):
+        model_dir = snapshot_download(self.model_id)
+        face_reconstruction = pipeline(
+            Tasks.face_reconstruction, model=model_dir)
+        self.pipeline_inference(face_reconstruction, self.test_image)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_modelhub(self):
+        face_reconstruction = pipeline(
+            Tasks.face_reconstruction, model=self.model_id)
+        self.pipeline_inference(face_reconstruction, self.test_image)
+
+    @unittest.skip('demo compatibility test is only enabled on a needed-basis')
+    def test_demo_compatibility(self):
+        self.compatibility_check()
+
+
+if __name__ == '__main__':
+    unittest.main()