diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 750b7aa0..207a4003 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -366,6 +366,8 @@ class Pipelines(object): hand_detection = 'yolox-pai_hand-detection' skin_retouching = 'unet-skin-retouching' face_reconstruction = 'resnet50-face-reconstruction' + head_reconstruction = 'HRN-head-reconstruction' + text_to_head = 'HRN-text-to-head' tinynas_classification = 'tinynas-classification' easyrobust_classification = 'easyrobust-classification' tinynas_detection = 'tinynas-detection' diff --git a/modelscope/models/cv/head_reconstruction/__init__.py b/modelscope/models/cv/head_reconstruction/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/head_reconstruction/models/__init__.py b/modelscope/models/cv/head_reconstruction/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/head_reconstruction/models/bfm.py b/modelscope/models/cv/head_reconstruction/models/bfm.py new file mode 100644 index 00000000..d0aebd8e --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/bfm.py @@ -0,0 +1,673 @@ +# 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 modelscope.models.cv.face_reconstruction.utils import read_obj + + +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, + assets_root='assets', + 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'): + + model = loadmat(os.path.join(assets_root, '3dmm/BFM', default_name)) + model_bfm_front = loadmat( + os.path.join(assets_root, '3dmm/BFM/BFM_model_front.mat')) + self.mean_shape_ori = model_bfm_front['meanshape'].astype(np.float32) + # mean face shape. [3*N,1] + self.mean_shape = model['meanshape'].astype(np.float32) # (1, 107127) + + # identity basis. [3*N,80] + self.id_base = model['idBase'].astype(np.float32) # (107127, 80) + + # expression basis. [3*N,64] + self.exp_base = model['exBase'].astype(np.float32) # (107127, 64) + + # mean face texture. [3*N,1] (0-255) + self.mean_tex = model['meantex'].astype(np.float32) # (1, 107127) + + # texture basis. [3*N,80] + self.tex_base = model['texBase'].astype(np.float32) # (107127, 80) + + self.bfm_keep_inds = np.load( + os.path.join(assets_root, '3dmm/inds/bfm_keep_inds.npy')) + + self.ours_hair_area_inds = np.load( + os.path.join(assets_root, '3dmm/inds/ours_hair_area_inds.npy')) + + if default_name == 'ourRefineFull_model.mat': + 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) + + # face indices for each vertex that lies in. starts from 0. [N,8] + self.point_buf = model['point_buf'].astype(np.int64) - 1 # (35709, 8) + + # vertex indices for each face. starts from 0. [F,3] + self.face_buf = model['tri'].astype(np.int64) - 1 # (70789, 3) + + # vertex indices for 68 landmarks. starts from 0. [68,1] + self.keypoints = np.squeeze(model['keypoints']).astype(np.int64) - 1 + + if default_name == 'ourRefineFull_model.mat': + self.keypoints = np.load( + os.path.join( + assets_root, + '3dmm/inds/our_refine0223_basis_withoutEyes_withUV_keypoints_inds.npy' + )).astype(np.int64) + self.point_buf = self.point_buf[:, :8] + 1 + + if is_train: + # vertex indices for small face region to compute photometric error. starts from 0. + self.front_mask = np.squeeze(model['frontmask2_idx']).astype( + np.int64) - 1 + # vertex indices for each face from small face region. starts from 0. [f,3] + self.front_face_buf = model['tri_mask2'].astype(np.int64) - 1 + # vertex indices for pre-defined skin region to compute reflectance loss + self.skin_mask = np.squeeze(model['skinmask']) + + if default_name == 'ourRefineFull_model.mat': + nose_reduced_mesh = read_obj( + os.path.join(assets_root, + '3dmm/adjust_part/our_full/145_nose.obj')) + self.nose_reduced_part = nose_reduced_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + neck_mesh = read_obj( + os.path.join(assets_root, + '3dmm/adjust_part/our_full/154_neck.obj')) + self.neck_adjust_part = neck_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + eyes_mesh = read_obj( + os.path.join( + assets_root, + '3dmm/adjust_part/our_full/our_mean_adjust_eyes.obj')) + self.eyes_adjust_part = eyes_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + self.neck_slim_part = None + self.neck_stretch_part = None + elif default_name == 'ourRefineBFMEye0504_model.mat': + nose_reduced_mesh = read_obj( + os.path.join(assets_root, + '3dmm/adjust_part/our_full_bfmEyes/145_nose.obj')) + self.nose_reduced_part = nose_reduced_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + neck_mesh = read_obj( + os.path.join(assets_root, + '3dmm/adjust_part/our_full_bfmEyes/146_neck.obj')) + self.neck_adjust_part = neck_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + self.eyes_adjust_part = None + + neck_slim_mesh = read_obj( + os.path.join( + assets_root, + '3dmm/adjust_part/our_full_bfmEyes/147_neckSlim2.obj')) + self.neck_slim_part = neck_slim_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + + neck_stretch_mesh = read_obj( + os.path.join( + assets_root, + '3dmm/adjust_part/our_full_bfmEyes/148_neckLength.obj')) + self.neck_stretch_part = neck_stretch_mesh['vertices'].reshape( + (1, -1)) - self.mean_shape + else: + self.nose_reduced_part = None + + self.neck_adjust_part = None + self.eyes_adjust_part = None + self.neck_slim_part = None + self.neck_stretch_part = None + + 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]) + + eye_corner_inds = np.load( + os.path.join(assets_root, '3dmm/inds/eye_corner_inds.npy')) + self.eye_corner_inds = torch.from_numpy(eye_corner_inds).long() + eye_lines = np.load( + os.path.join(assets_root, '3dmm/inds/eye_corner_lines.npy')) + self.eye_lines = torch.from_numpy(eye_lines).long() + + self.center = center + self.persc_proj = perspective_projection(focal, self.center) + 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, + neckSlim_coeff=0.0, + neckStretch_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 + if neckSlim_coeff != 0 and self.neck_slim_part is not None: + face_shape = face_shape + neckSlim_coeff * self.neck_slim_part + if neckStretch_coeff != 0 and self.neck_stretch_part is not None: + + neck_stretch_part = self.neck_stretch_part.reshape(1, -1, 3) + neck_stretch_part_top = neck_stretch_part[0, 37476, 1] + neck_stretch_part_bottom = neck_stretch_part[0, 37357, 1] + neck_stretch_height = neck_stretch_part_top - neck_stretch_part_bottom + + face_shape_ = face_shape.reshape(1, -1, 3) + face_shape_top = face_shape_[0, 37476, 1] + face_shape_bottom = face_shape_[0, 37357, 1] + face_shape_height = face_shape_top - face_shape_bottom + + target_neck_height = 0.72 # top ind 37476, bottom ind 37357 + neckStretch_coeff = (target_neck_height + - face_shape_height) / neck_stretch_height + + face_shape = face_shape + neckStretch_coeff * self.neck_stretch_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 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_shape = face_shape[-625:, :] + corner_offset = shape_offset[self.eye_corner_inds] + for i in range(len(self.eye_lines)): + corner_shape[self.eye_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)] # (N, 3) + 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)] + + value_1 = shape_lt * ( + 1 - UVs_coords_float[:, :1]) * UVs_coords_float[:, 1:] + value_2 = shape_lb * (1 - UVs_coords_float[:, :1]) * ( + 1 - UVs_coords_float[:, 1:]) + value_3 = shape_rt * UVs_coords_float[:, :1] * UVs_coords_float[:, 1:] + value_4 = shape_rb * UVs_coords_float[:, :1] * ( + 1 - UVs_coords_float[:, 1:]) + + offset_shape = value_1 + value_2 + value_3 + value_4 # (B, N, 3) + + face_shape = (face_shape + offset_shape)[None, ...] + + return face_shape, offset_shape[None, ...] + + def compute_for_render_head_fitting(self, + coeffs, + shape_offset_uv, + texture_offset_uv, + shape_offset_uv_head, + texture_offset_uv_head, + UVs, + reverse_recenter=True, + get_eyes=False, + get_neck=False, + nose_coeff=0.0, + neck_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=neck_coeff, + eyes_coeff=eyes_coeff) # (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[:, :35241, :], shape_offset = self.add_nonlinear_offset( + face_shape[:, :35241, :], shape_offset_uv, + UVs[:35709, ...][self.bfm_keep_inds]) # (1, n, 3) + if get_eyes: + face_shape = self.add_nonlinear_offset_eyes( + face_shape, shape_offset) + if get_neck: + face_shape[:, 35241:37082, ...], _ = self.add_nonlinear_offset( + face_shape[:, 35241:37082, ...], shape_offset_uv_head, + UVs[35709:, ...]) # (1, n, 3) + else: + face_shape[:, self.ours_hair_area_inds, + ...], _ = self.add_nonlinear_offset( + face_shape[:, self.ours_hair_area_inds, + ...], shape_offset_uv_head, + UVs[self.ours_hair_area_inds + (35709 - 35241), + ...]) # (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[:, :35241, :], texture_offset = self.add_nonlinear_offset( + face_texture[:, :35241, :], texture_offset_uv, + UVs[:35709, ...][self.bfm_keep_inds]) + face_texture[:, 35241:37082, :], _ = self.add_nonlinear_offset( + face_texture[:, 35241:37082, :], texture_offset_uv_head, + 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_head(self, + coeffs, + shape_offset_uv, + texture_offset_uv, + shape_offset_uv_head, + texture_offset_uv_head, + UVs, + reverse_recenter=True, + nose_coeff=0.0, + neck_coeff=0.0, + eyes_coeff=0.0, + neckSlim_coeff=0.0, + neckStretch_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=neck_coeff, + eyes_coeff=eyes_coeff, + neckSlim_coeff=neckSlim_coeff, + neckStretch_coeff=neckStretch_coeff) # (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[:, :35709, :], shape_offset = self.add_nonlinear_offset( + face_shape[:, :35709, :], shape_offset_uv, UVs[:35709, + ...]) # (1, n, 3) + face_shape[:, 35709:, + ...], _ = self.add_nonlinear_offset(face_shape[:, 35709:, + ...], + shape_offset_uv_head, + 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[:, :35709, :], texture_offset = self.add_nonlinear_offset( + face_texture[:, :35709, :], texture_offset_uv, UVs[:35709, ...]) + face_texture[:, 35709:, :], _ = self.add_nonlinear_offset( + face_texture[:, 35709:, :], texture_offset_uv_head, 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 diff --git a/modelscope/models/cv/head_reconstruction/models/head_segmentation.py b/modelscope/models/cv/head_reconstruction/models/head_segmentation.py new file mode 100644 index 00000000..ef784a7d --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/head_segmentation.py @@ -0,0 +1,196 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os + +import cv2 +import json +import numpy as np +import tensorflow as tf + +if tf.__version__ >= '2.0': + tf = tf.compat.v1 + tf.disable_eager_execution() + + +class HeadSegmentor(): + + def __init__(self, model_root): + """The HeadSegmentor is implemented based on https://arxiv.org/abs/2004.04955 + Args: + model_root: the root directory of the model files + """ + self.sess = self.load_sess( + os.path.join(model_root, 'head_segmentation', + 'Matting_headparser_6_18.pb')) + self.sess_detect = self.load_sess( + os.path.join(model_root, 'head_segmentation', 'face_detect.pb')) + self.sess_face = self.load_sess( + os.path.join(model_root, 'head_segmentation', 'segment_face.pb')) + + def load_sess(self, model_path): + config = tf.ConfigProto(allow_soft_placement=True) + config.gpu_options.allow_growth = True + sess = tf.Session(config=config) + with tf.gfile.FastGFile(model_path, 'rb') as f: + graph_def = tf.GraphDef() + graph_def.ParseFromString(f.read()) + sess.graph.as_default() + tf.import_graph_def(graph_def, name='') + sess.run(tf.global_variables_initializer()) + return sess + + def process(self, image): + """ image: bgr + """ + + h, w, c = image.shape + faceRects = self.detect_face(image) + face_num = len(faceRects) + all_head_alpha = [] + all_face_mask = [] + for i in range(face_num): + y1 = faceRects[i][0] + y2 = faceRects[i][1] + x1 = faceRects[i][2] + x2 = faceRects[i][3] + pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box( + y1, y2, x1, x2, 0.15, 0.15, 0.15, 0.15, h, w) + temp_img = image.copy() + roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2] + output_alpha = self.sess_face.run( + self.sess_face.graph.get_tensor_by_name('output_alpha_face:0'), + feed_dict={'input_image_face:0': roi_img[:, :, ::-1]}) + face_mask = np.zeros((h, w, 3)) + face_mask[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha + all_face_mask.append(face_mask) + cv2.imwrite(str(i) + 'face.jpg', face_mask) + cv2.imwrite(str(i) + 'face_roi.jpg', roi_img) + + for i in range(face_num): + y1 = faceRects[i][0] + y2 = faceRects[i][1] + x1 = faceRects[i][2] + x2 = faceRects[i][3] + pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box( + y1, y2, x1, x2, 1.47, 1.47, 1.3, 2.0, h, w) + temp_img = image.copy() + for j in range(face_num): + y1 = faceRects[j][0] + y2 = faceRects[j][1] + x1 = faceRects[j][2] + x2 = faceRects[j][3] + small_y1, small_y2, small_x1, small_x2 = self.pad_box( + y1, y2, x1, x2, -0.1, -0.1, -0.1, -0.1, h, w) + small_width = small_x2 - small_x1 + small_height = small_y2 - small_y1 + if (small_x1 < 0 or small_y1 < 0 or small_width < 3 + or small_height < 3 or small_x2 > w or small_y2 > h): + continue + # if(i!=j): + # temp_img[small_y1:small_y2,small_x1:small_x2]=0 + if (i != j): + temp_img = temp_img * (1.0 - all_face_mask[j] / 255.0) + + roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2] + output_alpha = self.sess.run( + self.sess.graph.get_tensor_by_name('output_alpha:0'), + feed_dict={'input_image:0': roi_img[:, :, ::-1]}) + head_alpha = np.zeros((h, w)) + head_alpha[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha[:, :, 0] + if np.sum(head_alpha) > 255 * w * h * 0.01 * 0.01: + all_head_alpha.append(head_alpha) + + head_num = len(all_head_alpha) + head_elements = [] + if head_num == 0: + return head_elements + + for i in range(head_num): + head_alpha = all_head_alpha[i] + head_elements.append(head_alpha) + + return head_elements + + def pad_box(self, y1, y2, x1, x2, left_ratio, right_ratio, top_ratio, + bottom_ratio, h, w): + box_w = x2 - x1 + box_h = y2 - y1 + pad_y1 = np.maximum(np.int32(y1 - top_ratio * box_h), 0) + pad_y2 = np.minimum(np.int32(y2 + bottom_ratio * box_h), h - 1) + pad_x1 = np.maximum(np.int32(x1 - left_ratio * box_w), 0) + pad_x2 = np.minimum(np.int32(x2 + right_ratio * box_w), w - 1) + return pad_y1, pad_y2, pad_x1, pad_x2 + + def detect_face(self, img): + h, w, c = img.shape + input_img = cv2.resize(img[:, :, ::-1], (512, 512)) + boxes, scores, num_detections = self.sess_detect.run( + [ + self.sess_detect.graph.get_tensor_by_name('tower_0/boxes:0'), + self.sess_detect.graph.get_tensor_by_name('tower_0/scores:0'), + self.sess_detect.graph.get_tensor_by_name( + 'tower_0/num_detections:0') + ], + feed_dict={ + 'tower_0/images:0': input_img[np.newaxis], + 'training_flag:0': False + }) + faceRects = [] + for i in range(num_detections[0]): + if scores[0, i] < 0.5: + continue + y1 = np.int32(boxes[0, i, 0] * h) + x1 = np.int32(boxes[0, i, 1] * w) + y2 = np.int32(boxes[0, i, 2] * h) + x2 = np.int32(boxes[0, i, 3] * w) + if x2 <= x1 + 3 or y2 <= y1 + 3: + continue + faceRects.append((y1, y2, x1, x2, y2 - y1, x2 - x1)) + sorted(faceRects, key=lambda x: x[4] * x[5], reverse=True) + return faceRects + + def generate_json(self, status_code, status_msg, ori_url, result_element, + track_id): + data = {} + data['originUri'] = ori_url + data['elements'] = result_element + data['statusCode'] = status_code + data['statusMessage'] = status_msg + data['requestId'] = track_id + return json.dumps(data) + + def get_box(self, alpha): + h, w = alpha.shape + start_h = 0 + end_h = 0 + start_w = 0 + end_w = 0 + for i in range(0, h, 3): + line = alpha[i, :] + if np.max(line) >= 1: + start_h = i + break + + for i in range(0, w, 3): + line = alpha[:, i] + if np.max(line) >= 1: + start_w = i + break + + for i in range(0, h, 3): + i = h - 1 - i + line = alpha[i, :] + if np.max(line) >= 1: + end_h = i + if end_h < h - 1: + end_h = end_h + 1 + break + for i in range(0, w, 3): + i = w - 1 - i + line = alpha[:, i] + if np.max(line) >= 1: + end_w = i + if end_w < w - 1: + end_w = end_w + 1 + break + + return start_h, start_w, end_h, end_w diff --git a/modelscope/models/cv/head_reconstruction/models/headrecon_model.py b/modelscope/models/cv/head_reconstruction/models/headrecon_model.py new file mode 100644 index 00000000..e515421c --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/headrecon_model.py @@ -0,0 +1,564 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +from collections import OrderedDict + +import cv2 +import numpy as np +import torch + +from modelscope.models import MODELS, TorchModel +from modelscope.models.cv.face_reconstruction.utils import (estimate_normals, + read_obj) +from . import networks, opt +from .bfm import ParametricFaceModel +from .losses import (BinaryDiceLoss, TVLoss, TVLoss_std, landmark_loss, + perceptual_loss, photo_loss, points_loss_horizontal, + reflectance_loss, reg_loss) +from .nv_diffrast import MeshRenderer + + +@MODELS.register_module('head-reconstruction', 'head_reconstruction') +class HeadReconModel(TorchModel): + + def __init__(self, model_dir, *args, **kwargs): + """The HeadReconModel is implemented based on HRN, publicly available at + https://github.com/youngLBW/HRN + + Args: + model_dir: the root directory of the model files + """ + super().__init__(model_dir, *args, **kwargs) + + self.model_dir = model_dir + opt.bfm_folder = os.path.join(model_dir, 'assets') + self.opt = opt + self.isTrain = opt.isTrain + self.visual_names = ['output_vis'] + self.model_names = ['net_recon'] + self.parallel_names = self.model_names + [ + 'renderer', 'renderer_fitting' + ] + + # networks + self.net_recon = networks.define_net_recon( + net_recon=opt.net_recon, + use_last_fc=opt.use_last_fc, + init_path=None) + + # assets + self.headmodel = ParametricFaceModel( + assets_root=opt.bfm_folder, + camera_distance=opt.camera_d, + focal=opt.focal, + center=opt.center, + is_train=self.isTrain, + default_name='ourRefineBFMEye0504_model.mat') + + self.headmodel_for_fitting = ParametricFaceModel( + assets_root=opt.bfm_folder, + camera_distance=opt.camera_d, + focal=opt.focal, + center=opt.center, + is_train=self.isTrain, + default_name='ourRefineFull_model.mat') + + # renderer + 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)) + + template_obj_path = os.path.join( + model_dir, + 'assets/3dmm/template_mesh/template_ourFull_bfmEyes.obj') + self.template_output_mesh = read_obj(template_obj_path) + + self.nonlinear_UVs = self.template_output_mesh['uvs'] + self.nonlinear_UVs = torch.from_numpy(self.nonlinear_UVs) + + self.jaw_edge_mask = cv2.imread( + os.path.join(model_dir, + 'assets/texture/jaw_edge_mask2.png'))[..., 0].astype( + np.float32) / 255.0 + self.jaw_edge_mask = cv2.resize(self.jaw_edge_mask, (300, 300))[..., + None] + + 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' + ] + + 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 + + if opt.isTrain: + self.optimizer = torch.optim.Adam( + self.net_recon.parameters(), lr=opt.lr) + self.optimizers = [self.optimizer] + self.parallel_names += ['net_recog'] + + def set_device(self, device): + self.device = device + self.net_recon = self.net_recon.to(self.device) + self.headmodel.to(self.device) + self.headmodel_for_fitting.to(self.device) + self.nonlinear_UVs = self.nonlinear_UVs.to(self.device) + + 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) + + 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=1024): + 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 check_head_pose(self, coeffs): + pi = 3.14 + if coeffs[0, 225] > pi / 6 or coeffs[0, 225] < -pi / 6: + return False + elif coeffs[0, 224] > pi / 6 or coeffs[0, 224] < -pi / 6: + return False + elif coeffs[0, 226] > pi / 6 or coeffs[0, 226] < -pi / 6: + return False + else: + return True + + def get_fusion_mask(self, keep_forehead=True): + self.without_forehead_inds = torch.from_numpy( + np.load( + os.path.join(self.model_dir, + 'assets/3dmm/inds/bfm_withou_forehead_inds.npy')) + ).long().to(self.device) + + h, w = self.shape_offset_uv.shape[1:3] + self.fusion_mask = torch.zeros((h, w)).to(self.device).float() + if keep_forehead: + UVs_coords = self.nonlinear_UVs.clone()[:35709][ + self.without_forehead_inds] + else: + 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 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, ...] + self.shape_offset_uv = shape_offset_uv_blur * self.edge_mask[ + None, ..., None] + self.shape_offset_uv * ( + 1 - self.edge_mask[None, ..., None]) + + 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 blur_offset_edge(self): + shape_offset_uv = self.shape_offset_uv[0].detach().cpu().numpy() + shape_offset_uv_head = self.shape_offset_uv_head[0].detach().cpu( + ).numpy() + shape_offset_uv_head = cv2.resize(shape_offset_uv_head, (300, 300)) + shape_offset_uv_head = shape_offset_uv_head * ( + 1 - self.jaw_edge_mask) + shape_offset_uv * self.jaw_edge_mask + shape_offset_uv_head = cv2.resize(shape_offset_uv_head, (100, 100)) + + self.shape_offset_uv_head = torch.from_numpy( + shape_offset_uv_head).float().to(self.device)[None, ...] + + def fitting_nonlinear(self, coeff, n_iters=250): + output_coeff = coeff.detach().clone() + + output_coeff = self.headmodel_for_fitting.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) + self.shape_offset_uv.requires_grad = True + + self.texture_offset_uv = torch.zeros( + (1, 300, 300, 3), dtype=torch.float32).to(self.device) + self.texture_offset_uv.requires_grad = True + + self.shape_offset_uv_head = torch.zeros( + (1, 100, 100, 3), dtype=torch.float32).to(self.device) + self.shape_offset_uv_head.requires_grad = True + + self.texture_offset_uv_head = torch.zeros( + (1, 100, 100, 3), dtype=torch.float32).to(self.device) + self.texture_offset_uv_head.requires_grad = True + + head_face_inds = np.load( + os.path.join(self.model_dir, + 'assets/3dmm/inds/ours_head_face_inds.npy')) + head_face_inds = torch.from_numpy(head_face_inds).to(self.device) + head_faces = self.headmodel_for_fitting.face_buf[head_face_inds] + + # print('before fitting', output_coeff) + + opt_parameters = [ + self.shape_offset_uv, self.texture_offset_uv, + self.shape_offset_uv_head, self.texture_offset_uv_head, + output_coeff['id'], output_coeff['exp'], output_coeff['tex'], + output_coeff['gamma'] + ] + optim = torch.optim.Adam(opt_parameters, lr=1e-3) + + optim_pose = torch.optim.Adam([output_coeff['trans']], lr=1e-1) + + self.get_edge_points_horizontal() + + for i in range(n_iters): # 500 + self.pred_vertex_head, self.pred_tex, self.pred_color_head, self.pred_lm, face_shape, \ + face_shape_offset, self.verts_proj_head = \ + self.headmodel_for_fitting.compute_for_render_head_fitting(output_coeff, self.shape_offset_uv, + self.texture_offset_uv, + self.shape_offset_uv_head, + self.texture_offset_uv_head, + self.nonlinear_UVs) + self.pred_vertex = self.pred_vertex_head[:, :35241] + self.pred_color = self.pred_color_head[:, :35241] + self.verts_proj = self.verts_proj_head[:, :35241] + self.pred_mask_head, _, self.pred_head, self.occ_head = self.renderer_fitting( + self.pred_vertex_head, head_faces, feat=self.pred_color_head) + self.pred_mask, _, self.pred_face, self.occ_face = self.renderer_fitting( + self.pred_vertex, + self.headmodel_for_fitting.face_buf[:69732], + feat=self.pred_color) + + self.pred_coeffs_dict = self.headmodel_for_fitting.split_coeff( + output_coeff) + self.compute_losses_fitting() + + if i < 150: + optim_pose.zero_grad() + (self.loss_lm + self.loss_color * 0.1).backward() + optim_pose.step() + else: + optim.zero_grad() + self.loss_all.backward() + optim.step() + + output_coeff = self.headmodel_for_fitting.merge_coeff(output_coeff) + + self.get_edge_mask() + self.get_fusion_mask(keep_forehead=False) + self.blur_shape_offset_uv(global_blur=True) + self.blur_offset_edge() + return output_coeff + + def forward(self): + with torch.no_grad(): + output_coeff = self.net_recon(self.input_img_coeff) + + if not self.check_head_pose(output_coeff): + return None + + with torch.enable_grad(): + output_coeff = self.fitting_nonlinear(output_coeff) + + output_coeff = self.headmodel.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 + # degree = 0.5 + 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.headmodel.merge_coeff(output_coeff) + + self.pred_vertex, _, _, _, face_shape_ori, face_shape, _ = \ + self.headmodel.compute_for_render_head(output_coeff, + self.shape_offset_uv.detach(), + self.texture_offset_uv.detach(), + self.shape_offset_uv_head.detach() * 0, + self.texture_offset_uv_head.detach(), + self.nonlinear_UVs, + nose_coeff=0.1, + neck_coeff=0.3, + neckSlim_coeff=0.5, + neckStretch_coeff=0.5) + + UVs = np.array(self.template_output_mesh['uvs']) + UVs_tensor = torch.tensor(UVs, dtype=torch.float32) + UVs_tensor = torch.unsqueeze(UVs_tensor, 0).to(self.pred_vertex.device) + + target_img = self.input_fat_img_hd + target_img = target_img.permute(0, 2, 3, 1) + face_buf = self.headmodel.face_buf + # get texture map + with torch.enable_grad(): + pred_mask, _, pred_face, texture_map, texture_mask = self.renderer.pred_shape_and_texture( + self.pred_vertex, face_buf, UVs_tensor, target_img, None) + self.pred_coeffs_dict = self.headmodel.split_coeff(output_coeff) + + recon_shape = face_shape # get reconstructed shape, [1, 35709, 3] + recon_shape[ + ..., + -1] = 10 - recon_shape[..., -1] # from camera space to world space + recon_shape = recon_shape.cpu().numpy()[0] + tri = self.headmodel.face_buf.cpu().numpy() + + output = {} + output['flag'] = 0 + + output['tex_map'] = texture_map + output['tex_mask'] = texture_mask * 255.0 + ''' + coeffs + { + 'id': id_coeffs, + 'exp': exp_coeffs, + 'tex': tex_coeffs, + 'angle': angles, + 'gamma': gammas, + 'trans': translations + } + ''' + output['coeffs'] = self.pred_coeffs_dict + + normals = estimate_normals(recon_shape, tri) + + output['vertices'] = recon_shape + output['triangles'] = tri + output['uvs'] = UVs + output['faces_uv'] = self.template_output_mesh['faces_uv'] + output['normals'] = normals + + return output + + 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 compute_losses_fitting(self): + face_mask = self.pred_mask + face_mask = face_mask.detach() + self.loss_color = self.opt.w_color * self.comupte_color_loss( + self.pred_face, self.input_img, face_mask) # 1.0 + + loss_reg, loss_gamma = self.compute_reg_loss( + self.pred_coeffs_dict, + w_id=self.opt.w_id, + w_exp=self.opt.w_exp, + w_tex=self.opt.w_tex) + self.loss_reg = self.opt.w_reg * loss_reg # 1.0 + self.loss_gamma = self.opt.w_gamma * loss_gamma # 1.0 + + self.loss_lm = self.opt.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_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_all = self.loss_color + self.loss_lm + self.loss_reg + self.loss_gamma + self.loss_all += self.loss_smooth_offset + self.loss_smooth_offset_std + self.loss_reg_textureOff + self.loss_all += self.loss_points_horizontal + + head_mask = self.pred_mask_head + head_mask = head_mask.detach() + self.loss_color_head = self.opt.w_color * self.comupte_color_loss( + self.pred_head, self.input_img, head_mask) # 1.0 + self.loss_smooth_offset_head = TVLoss()( + self.shape_offset_uv_head.permute(0, 3, 1, 2)) * 100 # 10000 + self.loss_smooth_offset_std_head = TVLoss_std()( + self.shape_offset_uv_head.permute(0, 3, 1, 2)) * 500 # 50000 + self.loss_mask = BinaryDiceLoss()(self.occ_head, self.head_mask) * 20 + + self.loss_all += self.loss_mask + self.loss_color_head + self.loss_all += self.loss_smooth_offset_head + self.loss_smooth_offset_std_head diff --git a/modelscope/models/cv/head_reconstruction/models/losses.py b/modelscope/models/cv/head_reconstruction/models/losses.py new file mode 100644 index 00000000..6d4af4e8 --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/losses.py @@ -0,0 +1,367 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +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)) + + +# 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)) diff --git a/modelscope/models/cv/head_reconstruction/models/networks.py b/modelscope/models/cv/head_reconstruction/models/networks.py new file mode 100644 index 00000000..1eb5770b --- /dev/null +++ b/modelscope/models/cv/head_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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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" `_. + + 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)} diff --git a/modelscope/models/cv/head_reconstruction/models/nv_diffrast.py b/modelscope/models/cv/head_reconstruction/models/nv_diffrast.py new file mode 100644 index 00000000..f765c6c7 --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/nv_diffrast.py @@ -0,0 +1,414 @@ +# 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 +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: + if nvdiffrast.__version__ == '0.2.7': + self.glctx = dr.RasterizeGLContext(device=device) + else: + 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: + if nvdiffrast.__version__ == '0.2.7': + self.glctx = dr.RasterizeGLContext(device=device) + else: + 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. + + image = img.permute(0, 3, 1, 2) + + return mask, depth, image + + 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) + """ + uv = uv.clone() + + 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: + if nvdiffrast.__version__ == '0.2.7': + self.glctx = dr.RasterizeGLContext(device=device) + else: + self.glctx = dr.RasterizeCudaContext(device=device) + # print("create glctx on device cuda:%d" % device.index) + + # print('vertex_ndc shape:{}'.format(vertex_ndc.shape)) # Size([1, 35709, 4]) + # print('tri shape:{}'.format(tri.shape)) # Size([70789, 3]) + + 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') + + 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, int(128), 128, 3), dtype=torch.float32) + # tex = torch.zeros((1, 128, 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, int(2048), 2048, 3), dtype=torch.float32) + # tex_mask = torch.zeros((1, 2048, 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, (int(tex_resolution), tex_resolution)) + # tex = F.interpolate(tex, (tex_resolution, 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, (int(tex_resolution), tex_resolution)) + # _base_tex = F.interpolate(_base_tex, (tex_resolution, 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 = 200 + + 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 diff --git a/modelscope/models/cv/head_reconstruction/models/opt.py b/modelscope/models/cv/head_reconstruction/models/opt.py new file mode 100644 index 00000000..d6be64b3 --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/opt.py @@ -0,0 +1,21 @@ +# 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 +lr = 0.001 +w_color = 1.92 +w_reg = 3.0e-4 +w_gamma = 10.0 +w_lm = 1.6e-3 +w_id = 1.0 +w_exp = 0.8 +w_tex = 1.7e-2 diff --git a/modelscope/models/cv/head_reconstruction/models/tex_processor.py b/modelscope/models/cv/head_reconstruction/models/tex_processor.py new file mode 100644 index 00000000..e202ab7d --- /dev/null +++ b/modelscope/models/cv/head_reconstruction/models/tex_processor.py @@ -0,0 +1,145 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os + +import cv2 +import numpy as np + + +def get_fade_out_mask(length, start_value, end_value, fade_start_ratio, + fade_end_ratio): + fade_start_ind = int(length * fade_start_ratio) + fade_end_ind = int(length * fade_end_ratio) + + left_part = np.array([start_value] * fade_start_ind) + fade_part = np.linspace(start_value, end_value, + fade_end_ind - fade_start_ind) + len_right = length - len(left_part) - len(fade_part) + right_part = np.array([end_value] * len_right) + + fade_out_mask = np.concatenate([left_part, fade_part, right_part], axis=0) + return fade_out_mask + + +class TexProcesser(): + + def __init__(self, model_root): + + self.tex_size = 4096 + + self.bald_tex_bg = cv2.imread( + os.path.join(model_root, + 'assets/texture/template_bald_tex_2.jpg')).astype( + np.float32) + self.hair_tex_bg = cv2.imread( + os.path.join(model_root, + 'assets/texture/template_withHair_tex.jpg')).astype( + np.float32) + + self.hair_mask = cv2.imread( + os.path.join(model_root, + 'assets/texture/hair_mask_male.png'))[..., 0].astype( + np.float32) / 255.0 + self.hair_mask = cv2.resize(self.hair_mask, (4096, 4096 + 1024)) + + front_mask = cv2.imread( + os.path.join(model_root, + 'assets/texture/face_mask_singleview.jpg')).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_[:int(4096 * 375 / 950)] = 0 + self.binary_front_mask_[int(4096 * 600 / 950):] = 0 + 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_ + + self.fg_mask = cv2.imread( + os.path.join(model_root, + 'assets/texture/fg_mask.png'))[..., 0].astype( + np.float32) / 255.0 + self.fg_mask = cv2.resize(self.fg_mask, (256, 256)) + self.fg_mask = cv2.dilate(self.fg_mask, + np.ones(shape=(13, 13), dtype=np.float32)) + self.fg_mask = cv2.blur(self.fg_mask, (27, 27), 0) + self.fg_mask = cv2.resize(self.fg_mask, (4096, 4096 + 1024)) + self.fg_mask = self.fg_mask[..., None] + + self.cheek_mask = cv2.imread( + os.path.join(model_root, + 'assets/texture/cheek_area_mask.png'))[..., 0].astype( + np.float32) / 255.0 + self.cheek_mask = cv2.resize(self.cheek_mask, (4096, 4096 + 1024)) + self.cheek_mask = self.cheek_mask[..., None] + + self.bald_tex_bg = self.bald_tex_bg[:4096] + self.hair_tex_bg = self.hair_tex_bg[:4096] + self.fg_mask = self.fg_mask[:4096] + self.hair_mask = self.hair_mask[:4096] + self.front_mask = self.front_mask[:4096] + self.binary_front_mask = self.binary_front_mask[:4096] + self.front_mask_ = self.front_mask_[:4096] + + self.cheek_mask_left = self.cheek_mask[:4096] + self.cheek_mask_right = self.cheek_mask[:4096].copy()[:, ::-1] + + def post_process_texture(self, tex_map, hair_tex=True): + tex_map = cv2.resize(tex_map, (self.tex_size, self.tex_size)) + + # if hair_tex is true and there is a dark side, use the mirror texture + if hair_tex: + left_cheek_light_mean = np.mean( + tex_map[self.cheek_mask_left[..., 0] == 1.0]) + right_cheek_light_mean = np.mean( + tex_map[self.cheek_mask_right[..., 0] == 1.0]) + + tex_map_flip = tex_map[:, ::-1, :] + w = tex_map.shape[1] + half_w = w // 2 + if left_cheek_light_mean > right_cheek_light_mean * 1.5: + tex_map[:, half_w:, :] = tex_map_flip[:, half_w:, :] + elif right_cheek_light_mean > left_cheek_light_mean * 2: + tex_map[:, :half_w, :] = tex_map_flip[:, :half_w, :] + + # change the color of template texture + bg_mean_rgb = np.mean( + self.bald_tex_bg[self.binary_front_mask[..., 0] == 1.0], + axis=0)[None, None] + pred_tex_mean_rgb = np.mean( + tex_map[self.binary_front_mask[..., 0] == 1.0], axis=0)[None, + None] * 1.1 + _bald_tex_bg = self.bald_tex_bg.copy() + _bald_tex_bg = self.bald_tex_bg + (pred_tex_mean_rgb - bg_mean_rgb) + + if hair_tex: + # inpaint hair + tex_gray = cv2.cvtColor( + tex_map.astype(np.uint8), + cv2.COLOR_BGR2GRAY).astype(np.float32) + hair_mask = (self.hair_mask == 1.0) * (tex_gray < 120) + hair_bgr = np.mean(tex_map[hair_mask, :], axis=0) * 0.5 + if hair_bgr is None: + hair_bgr = 20.0 + _bald_tex_bg[self.hair_mask == 1.0] = hair_bgr + + # fuse + tex_map = _bald_tex_bg * (1. + - self.fg_mask) + tex_map * self.fg_mask + else: + # fuse + tex_map = _bald_tex_bg * ( + 1. - self.front_mask) + tex_map * self.front_mask + + return tex_map diff --git a/modelscope/models/cv/text_to_head/__init__.py b/modelscope/models/cv/text_to_head/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/text_to_head/text_to_head_model.py b/modelscope/models/cv/text_to_head/text_to_head_model.py new file mode 100644 index 00000000..ed09b8fa --- /dev/null +++ b/modelscope/models/cv/text_to_head/text_to_head_model.py @@ -0,0 +1,55 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +from collections import OrderedDict + +import cv2 +import numpy as np +import torch +from diffusers import (ControlNetModel, DDIMScheduler, + StableDiffusionControlNetPipeline) +from diffusers.utils import load_image + +from modelscope.models import MODELS, TorchModel + + +@MODELS.register_module('text-to-head', 'text_to_head') +class TextToHeadModel(TorchModel): + + def __init__(self, model_dir, *args, **kwargs): + """The HeadReconModel is implemented based on HRN, publicly available at + https://github.com/youngLBW/HRN + + Args: + model_dir: the root directory of the model files + """ + super().__init__(model_dir, *args, **kwargs) + + self.model_dir = model_dir + + base_model_path = os.path.join(model_dir, 'base_model') + controlnet_path = os.path.join(model_dir, 'control_net') + + controlnet = ControlNetModel.from_pretrained( + controlnet_path, torch_dtype=torch.float16) + self.face_gen_pipeline = StableDiffusionControlNetPipeline.from_pretrained( + base_model_path, controlnet=controlnet, torch_dtype=torch.float16) + self.face_gen_pipeline.scheduler = DDIMScheduler.from_config( + self.face_gen_pipeline.scheduler.config) + self.face_gen_pipeline.enable_model_cpu_offload() + + self.add_prompt = ', 4K, good looking face, epic realistic, Sony a7, sharp, ' \ + 'skin detail pores, soft light, uniform illumination' + self.neg_prompt = 'ugly, cross eye, bangs, teeth, glasses, hat, dark, shadow' + + control_pose_path = os.path.join(self.model_dir, 'control_pose.jpg') + self.control_pose = load_image(control_pose_path) + + def forward(self, input): + prompt = input['text'] + self.add_prompt + image = self.face_gen_pipeline( + prompt, + negative_prompt=self.neg_prompt, + image=self.control_pose, + num_inference_steps=20).images[0] # PIL.Image + + return image diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index f1ae964c..10fc6e7f 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -861,6 +861,41 @@ TASK_OUTPUTS = { # } Tasks.face_reconstruction: [OutputKeys.OUTPUT], + # 3D head reconstruction result for single sample + # { + # "output_obj": io.BytesIO, + # "output_img": np.array with shape(h, w, 3), + # "output": { + # "mesh": { + # "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), + # "UVs": np.array with shape(n, 2), + # "normals": np.array with shape(n, 3), + # }, + # } + # } + Tasks.head_reconstruction: [OutputKeys.OUTPUT], + + # text to head result for text input + # { + # "output_obj": io.BytesIO, + # "output_img": np.array with shape(h, w, 3), + # "output": { + # "mesh": { + # "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), + # "UVs": np.array with shape(n, 2), + # "normals": np.array with shape(n, 3), + # }, + # }, + # "image": np.array with shape(h, w, 3), + # } + Tasks.text_to_head: [OutputKeys.OUTPUT], + # 3D human reconstruction result for single sample # { # "output": { diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index 3be03682..6997504b 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -126,6 +126,10 @@ TASK_INPUTS = { InputType.IMAGE, Tasks.face_reconstruction: InputType.IMAGE, + Tasks.head_reconstruction: + InputType.IMAGE, + Tasks.text_to_head: + InputType.TEXT, Tasks.human_detection: InputType.IMAGE, Tasks.face_image_generation: diff --git a/modelscope/pipelines/cv/head_reconstruction_pipeline.py b/modelscope/pipelines/cv/head_reconstruction_pipeline.py new file mode 100644 index 00000000..03dcc5b1 --- /dev/null +++ b/modelscope/pipelines/cv/head_reconstruction_pipeline.py @@ -0,0 +1,607 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import io +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.cv.face_reconstruction.models.facelandmark.large_base_lmks_infer import \ + LargeBaseLmkInfer +from modelscope.models.cv.face_reconstruction.utils import ( + POS, align_for_lm, draw_line, enlarged_bbox, extract_5p, image_warp_grid1, + load_lm3d, mesh_to_string, read_obj, resize_n_crop_img, + resize_on_long_side, spread_flow, write_obj) +from modelscope.models.cv.head_reconstruction.models.head_segmentation import \ + HeadSegmentor +from modelscope.models.cv.head_reconstruction.models.tex_processor import \ + TexProcesser +from modelscope.models.cv.skin_retouching.retinaface.predict_single import \ + Model +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.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.device import create_device, device_placement +from modelscope.utils.logger import get_logger + +try: + from torch.hub import get_dir +except BaseException: + from torch.hub import _get_torch_home as get_dir + +if tf.__version__ >= '2.0': + tf = tf.compat.v1 + tf.disable_eager_execution() + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.head_reconstruction, module_name=Pipelines.head_reconstruction) +class HeadReconstructionPipeline(Pipeline): + + def __init__(self, model: str, device: str, hair_tex=False): + """The inference pipeline for head 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_headRecon = pipeline('head-reconstruction', + model='damo/cv_HRN_head-reconstruction') + >>> result = pipeline_headRecon(test_image) + >>> mesh = result[OutputKeys.OUTPUT]['mesh'] + >>> texture_map = result[OutputKeys.OUTPUT_IMG] + >>> mesh['texture_map'] = texture_map + >>> write_obj('head_reconstruction.obj', mesh) + """ + 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) + + config_path = os.path.join(model_root, ModelFile.CONFIGURATION) + logger.info(f'loading config from {config_path}') + self.cfg = Config.from_file(config_path) + + self.hair_tex = hair_tex + + if 'gpu' in device: + self.device_name_ = 'cuda' + else: + self.device_name_ = device + self.device_name_ = self.device_name_.lower() + lmks_cpkt_path = os.path.join(model_root, 'large_base_net.pth') + self.large_base_lmks_model = LargeBaseLmkInfer.model_preload( + lmks_cpkt_path, self.device_name_ == 'cuda') + self.detector = Model(max_size=512, device=self.device_name_) + detector_ckpt_name = 'retinaface_resnet50_2020-07-20_old_torch.pth' + state_dict = torch.load( + os.path.join(os.path.dirname(lmks_cpkt_path), detector_ckpt_name), + map_location='cpu') + self.detector.load_state_dict(state_dict) + self.detector.eval() + + device = torch.device(self.device_name_) + self.model.set_device(device) + self.model.setup(checkpoint_path) + self.model.parallelize() + self.model.eval() + self.model.set_render() + + hub_dir = get_dir() + save_ckpt_dir = os.path.join(hub_dir, '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.THREE_D, + flip_input=False) # face_alignment.LandmarksType._3D + + 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.head_segmentor = HeadSegmentor(model_root=model_root) + + self.tex_processor = TexProcesser(model_root=model_root) + + 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]: + if isinstance(input, str): + img = LoadImage.convert_to_ndarray(input) + if len(img.shape) == 2: + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + img = img.astype(float) + else: + img = input.astype(float) + result = {'img': img} + return result + + def align_img(self, + 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][0], t[1][0]]) + + return trans_params, img_new, lm_new, mask_new + + def read_data(self, + img, + lm, + lm3d_std, + to_tensor=True, + image_res=1024, + img_fat=None, + head_mask=None, + rescale_factor=75.0): + # to RGB + im = PIL.Image.fromarray(img[..., ::-1]) + W, H = im.size + lm[:, -1] = H - 1 - lm[:, -1] + + head_mask = PIL.Image.fromarray(head_mask) + im_fat = PIL.Image.fromarray(img_fat[..., ::-1]) + + _, im_lr_coeff, lm_lr_coeff, _ = self.align_img(im, lm, lm3d_std) + _, im_lr, lm_lr, mask_lr_head = self.align_img( + im, lm, lm3d_std, mask=head_mask, rescale_factor=rescale_factor) + _, im_hd, lm_hd, _ = self.align_img( + im_fat, + lm, + lm3d_std, + target_size=image_res, + rescale_factor=rescale_factor * 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 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) + im_lr_coeff = torch.tensor( + np.array(im_lr_coeff) / 255., + dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) + lm_lr_coeff = torch.tensor(lm_lr_coeff).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 infer_lmks(self, img_bgr): + INPUT_SIZE = 224 + ENLARGE_RATIO = 1.35 + + 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_name_ == '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_name_.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_lmks(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 fat_face(self, img, degree=0.04): + + _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 + + 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 + + if img.shape[0] > 2000 or img.shape[1] > 2000: + img, _ = resize_on_long_side(img, 1500) + + box, results = self.infer_lmks(img) + + if results is None or np.array(results).shape[0] == 0: + return {} + + fatbgr = self.fat_face(img) + + 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) + + head_mask = self.head_segmentor.process(img)[0] + + im_tensor, lm_tensor, im_hd_tensor, lm_hd_tensor, mask, head_mask, im_co, lm_co = self.read_data( + img, landmarks, self.lm3d_std, img_fat=fatbgr, head_mask=head_mask) + + data = { + 'imgs': im_tensor, + 'imgs_hd': im_hd_tensor, + 'lms': lm_tensor, + 'lms_hd': lm_hd_tensor, + 'face_mask': mask, + 'head_mask': head_mask, + 'imgs_coeff': im_co, + 'lms_coeff': lm_co, + } + self.model.set_input(data) # unpack data from data loader + + output = self.model() # run inference + + assert output is not None + + tex_map = output['tex_map'].astype(np.float32) + + # post-process texture map + tex_map = self.tex_processor.post_process_texture( + tex_map, hair_tex=self.hair_tex) + + head_mesh = { + 'vertices': output['vertices'], + 'faces': output['triangles'] + 1, + 'UVs': output['uvs'], + 'faces_uv': output['faces_uv'], + 'normals': output['normals'], + 'texture_map': tex_map + } + + results = { + 'mesh': head_mesh, + } + + return { + OutputKeys.OUTPUT_OBJ: None, + OutputKeys.OUTPUT_IMG: tex_map, + OutputKeys.OUTPUT: results + } + + def postprocess(self, inputs, **kwargs) -> Dict[str, Any]: + render = kwargs.get('render', False) + output_obj = inputs[OutputKeys.OUTPUT_OBJ] + texture_map = inputs[OutputKeys.OUTPUT_IMG] + results = inputs[OutputKeys.OUTPUT] + + if render: + output_obj = io.BytesIO() + mesh_str = mesh_to_string(results['mesh']) + mesh_bytes = mesh_str.encode(encoding='utf-8') + output_obj.write(mesh_bytes) + + result = { + OutputKeys.OUTPUT_OBJ: output_obj, + OutputKeys.OUTPUT_IMG: texture_map, + OutputKeys.OUTPUT: None if render else results, + } + return result diff --git a/modelscope/pipelines/cv/text_to_head_pipeline.py b/modelscope/pipelines/cv/text_to_head_pipeline.py new file mode 100644 index 00000000..0558e4fc --- /dev/null +++ b/modelscope/pipelines/cv/text_to_head_pipeline.py @@ -0,0 +1,91 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import io +import os +import shutil +from typing import Any, Dict + +import numpy as np + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.face_reconstruction.utils import ( + align_for_lm, align_img, draw_line, enlarged_bbox, image_warp_grid1, + load_lm3d, mesh_to_string, read_obj, resize_on_long_side, spread_flow, + 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 + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.text_to_head, module_name=Pipelines.text_to_head) +class TextToHeadPipeline(Pipeline): + + def __init__(self, model: str, device: str, hair_tex=True): + """The inference pipeline for text-to-head 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 + >>> from modelscope.models.cv.face_reconstruction.utils import write_obj + >>> test_prompt = "a clown with red nose" + >>> pipeline_textToHead = pipeline('text-to-head', + model='damo/cv_HRN_text-to-head') + >>> result = pipeline_textToHead(test_prompt) + >>> mesh = result[OutputKeys.OUTPUT]['mesh'] + >>> texture_map = result[OutputKeys.OUTPUT_IMG] + >>> mesh['texture_map'] = texture_map + >>> write_obj('text_to_head.obj', mesh) + """ + super().__init__(model=model, device=device) + + self.hair_tex = hair_tex + + head_recon_model_id = 'damo/cv_HRN_head-reconstruction' + self.head_reconstructor = pipeline( + Tasks.head_reconstruction, + model=head_recon_model_id, + model_revision='v0.1', + hair_tex=hair_tex) + + def preprocess(self, input: Input) -> Dict[str, Any]: + result = {'text': input} + return result + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + image = self.model(input) + image = np.array(image) + + results = self.head_reconstructor(image) + results['image'] = image + return results + + def postprocess(self, inputs, **kwargs) -> Dict[str, Any]: + render = kwargs.get('render', False) + output_obj = inputs[OutputKeys.OUTPUT_OBJ] + texture_map = inputs[OutputKeys.OUTPUT_IMG] + results = inputs[OutputKeys.OUTPUT] + + if render: + output_obj = io.BytesIO() + mesh_str = mesh_to_string(results['mesh']) + mesh_bytes = mesh_str.encode(encoding='utf-8') + output_obj.write(mesh_bytes) + + result = { + OutputKeys.OUTPUT_OBJ: output_obj, + OutputKeys.OUTPUT_IMG: texture_map, + OutputKeys.OUTPUT: None if render else results, + 'image': inputs['image'] + } + return result diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index fb315f4b..52b854b1 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -148,6 +148,8 @@ class CVTasks(object): # 3d face reconstruction face_reconstruction = 'face-reconstruction' + head_reconstruction = 'head-reconstruction' + text_to_head = 'text-to-head' # 3d human reconstruction human_reconstruction = 'human-reconstruction' diff --git a/tests/pipelines/test_head_reconstruction.py b/tests/pipelines/test_head_reconstruction.py new file mode 100644 index 00000000..d94dd922 --- /dev/null +++ b/tests/pipelines/test_head_reconstruction.py @@ -0,0 +1,60 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import io +import os +import os.path as osp +import sys +import unittest + +import cv2 + +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.test_utils import test_level + +sys.path.append('.') + + +class HeadReconstructionTest(unittest.TestCase): + + def setUp(self) -> None: + self.task = Tasks.head_reconstruction + self.model_id = 'damo/cv_HRN_head-reconstruction' + self.test_image = 'data/test/images/face_reconstruction.jpg' + + def save_results(self, result, save_root): + os.makedirs(save_root, exist_ok=True) + + # export obj and texture + mesh = result[OutputKeys.OUTPUT]['mesh'] + texture_map = result[OutputKeys.OUTPUT_IMG] + mesh['texture_map'] = texture_map + write_obj(os.path.join(save_root, 'head_recon_result.obj'), mesh) + + print(f'Output written to {osp.abspath(save_root)}') + + def pipeline_inference(self, pipeline: Pipeline, input_location: str): + result = pipeline(input_location) + self.save_results(result, './head_reconstruction_results') + + @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, revision='v0.2') + head_reconstruction = pipeline( + Tasks.head_reconstruction, model=model_dir) + self.pipeline_inference(head_reconstruction, self.test_image) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_modelhub(self): + head_reconstruction = pipeline( + Tasks.head_reconstruction, + model=self.model_id, + model_revision='v0.2') + self.pipeline_inference(head_reconstruction, self.test_image) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_text_to_head.py b/tests/pipelines/test_text_to_head.py new file mode 100644 index 00000000..4f081f70 --- /dev/null +++ b/tests/pipelines/test_text_to_head.py @@ -0,0 +1,62 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import io +import os +import os.path as osp +import sys +import unittest + +import cv2 + +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.test_utils import test_level + +sys.path.append('.') + + +class TextToHeadTest(unittest.TestCase): + + def setUp(self) -> None: + self.task = Tasks.text_to_head + self.model_id = 'damo/cv_HRN_text-to-head' + self.test_prompt = 'a clown with red nose' + + def save_results(self, result, save_root): + os.makedirs(save_root, exist_ok=True) + + # export obj and texture + mesh = result[OutputKeys.OUTPUT]['mesh'] + texture_map = result[OutputKeys.OUTPUT_IMG] + mesh['texture_map'] = texture_map + write_obj(os.path.join(save_root, 'text_to_head_result.obj'), mesh) + + image = result['image'] + image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + cv2.imwrite( + os.path.join(save_root, 'text_to_head_image.jpg'), image_bgr) + + print(f'Output written to {osp.abspath(save_root)}') + + def pipeline_inference(self, pipeline: Pipeline, prompt: str): + result = pipeline(prompt) + self.save_results(result, './text_to_head_results') + + @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, revision='v0.1') + text_to_head = pipeline(Tasks.text_to_head, model=model_dir) + self.pipeline_inference(text_to_head, self.test_prompt) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_modelhub(self): + face_reconstruction = pipeline( + Tasks.text_to_head, model=self.model_id, model_revision='v0.1') + self.pipeline_inference(face_reconstruction, self.test_prompt) + + +if __name__ == '__main__': + unittest.main()