From 9ec8f14f6f98211b159e9627506feb5d9fb91677 Mon Sep 17 00:00:00 2001 From: "biwen.lbw" Date: Thu, 6 Apr 2023 14:03:27 +0800 Subject: [PATCH] update face reconstruction to HRN(CVPR2023) Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12161605 * update to HRN * update test image * update results format * delete useless code & change header * add comment * reorganize code --- .../cv/face_reconstruction/models/bfm.py | 663 ++++++++----- .../models/de_retouching_module.py | 47 + .../models/facelandmark/large_model_infer.py | 430 --------- .../models/facerecon_model.py | 894 ++++++++++-------- .../cv/face_reconstruction/models/losses.py | 46 - .../face_reconstruction/models/nv_diffrast.py | 186 +--- .../cv/face_reconstruction/models/opt.py | 1 + .../models/pix2pix/__init__.py | 0 .../models/pix2pix/networks.py | 848 +++++++++++++++++ .../models/pix2pix/pix2pix_model.py | 125 +++ .../models/pix2pix/pix2pix_options.py | 26 + .../cv/face_reconstruction/models/renderer.py | 325 +++++++ .../cv/face_reconstruction/models/unet.py | 173 ++++ .../models/cv/face_reconstruction/utils.py | 266 +++++- modelscope/outputs/outputs.py | 22 +- .../cv/face_reconstruction_pipeline.py | 541 +++++++---- tests/pipelines/test_face_reconstruction.py | 39 +- 17 files changed, 3135 insertions(+), 1497 deletions(-) create mode 100644 modelscope/models/cv/face_reconstruction/models/de_retouching_module.py delete mode 100644 modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py create mode 100644 modelscope/models/cv/face_reconstruction/models/pix2pix/__init__.py create mode 100644 modelscope/models/cv/face_reconstruction/models/pix2pix/networks.py create mode 100644 modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_model.py create mode 100644 modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_options.py create mode 100755 modelscope/models/cv/face_reconstruction/models/renderer.py create mode 100644 modelscope/models/cv/face_reconstruction/models/unet.py diff --git a/modelscope/models/cv/face_reconstruction/models/bfm.py b/modelscope/models/cv/face_reconstruction/models/bfm.py index be4455bf..bd2b7c25 100644 --- a/modelscope/models/cv/face_reconstruction/models/bfm.py +++ b/modelscope/models/cv/face_reconstruction/models/bfm.py @@ -1,5 +1,6 @@ # 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 @@ -7,7 +8,9 @@ import torch import torch.nn.functional as F from scipy.io import loadmat +from .. import utils from ..utils import read_obj, transferBFM09 +from .renderer import SRenderY, set_rasterizer def perspective_projection(focal, center): @@ -30,7 +33,7 @@ class SH: class ParametricFaceModel: def __init__(self, - bfm_folder='./asset/BFM', + assets_folder='assets', recenter=True, camera_distance=10., init_lit=np.array([0.8, 0, 0, 0, 0, 0, 0, 0, 0]), @@ -39,80 +42,53 @@ class ParametricFaceModel: is_train=True, default_name='BFM_model_front.mat'): - if not os.path.isfile(os.path.join(bfm_folder, default_name)): - transferBFM09(bfm_folder) - model = loadmat(os.path.join(bfm_folder, default_name)) + if not os.path.isfile(os.path.join(assets_folder, default_name)): + transferBFM09(assets_folder) + model = loadmat(os.path.join(assets_folder, default_name)) # mean face shape. [3*N,1] - self.mean_shape = model['meanshape'].astype(np.float32) + self.mean_shape = model['meanshape'].astype(np.float32) # (1, 107127) # identity basis. [3*N,80] - self.id_base = model['idBase'].astype(np.float32) + self.id_base = model['idBase'].astype(np.float32) # (107127, 80) # expression basis. [3*N,64] - self.exp_base = model['exBase'].astype(np.float32) + 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) + self.mean_tex = model['meantex'].astype(np.float32) # (1, 107127) # texture basis. [3*N,80] - self.tex_base = model['texBase'].astype(np.float32) + self.tex_base = model['texBase'].astype(np.float32) # (107127, 80) + + mean_tex_uv_path = os.path.join(assets_folder, 'bfm_tex_mean2.npy') + tex_base_uv_path = os.path.join(assets_folder, 'bfm_texmap_base2.npy') + self.mean_tex_uv = np.load(mean_tex_uv_path) + self.mean_tex_uv = self.mean_tex_uv.reshape((1, -1)) + self.tex_base_uv = np.load(tex_base_uv_path) + self.tex_base_uv = self.tex_base_uv.reshape((-1, 80)) + set_rasterizer() + template_obj_path = os.path.join(assets_folder, 'template_bfm.obj') + uvcoords_path = os.path.join(assets_folder, 'bfm_uvs2.npy') + self.render = SRenderY( + 224, template_obj_path, uvcoords_path, uv_size=256) # face indices for each vertex that lies in. starts from 0. [N,8] - self.point_buf = model['point_buf'].astype(np.int64) - 1 + 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 + 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 - self.mean_shape_ori = model['meanshape_ori'].astype(np.float32) - self.bfm_keep_inds = model['bfm_keep_inds'][0] - self.nose_reduced_part = model['nose_reduced_part'].reshape( - (1, -1)) - self.mean_shape - self.nonlinear_UVs = model['nonlinear_UVs'] - - if default_name == 'head_model_for_maas.mat': - self.ours_hair_area_inds = model['hair_area_inds'][0] - - self.mean_tex = self.mean_tex.reshape(1, -1, 3) - mean_tex_keep = self.mean_tex[:, self.bfm_keep_inds] - self.mean_tex[:, :len(self.bfm_keep_inds)] = mean_tex_keep - self.mean_tex[:, - len(self.bfm_keep_inds):] = np.array([200, 146, - 118])[None, - None] - self.mean_tex[:, self.ours_hair_area_inds] = 40.0 - self.mean_tex = self.mean_tex.reshape(1, -1) - self.mean_tex = np.ascontiguousarray(self.mean_tex) - - self.tex_base = self.tex_base.reshape(-1, 3, 80) - tex_base_keep = self.tex_base[self.bfm_keep_inds] - self.tex_base[:len(self.bfm_keep_inds)] = tex_base_keep - self.tex_base[len(self.bfm_keep_inds):] = 0.0 - self.tex_base = self.tex_base.reshape(-1, 80) - self.tex_base = np.ascontiguousarray(self.tex_base) - - self.point_buf = self.point_buf[:, :8] + 1 - - self.neck_adjust_part = model['neck_adjust_part'].reshape( - (1, -1)) - self.mean_shape - self.eyes_adjust_part = model['eyes_adjust_part'].reshape( - (1, -1)) - self.mean_shape - - self.eye_corner_inds = model['eye_corner_inds'][0] - self.eye_corner_lines = model['eye_corner_lines'] - if recenter: mean_shape = self.mean_shape.reshape([-1, 3]) - mean_shape_ori = self.mean_shape_ori.reshape([-1, 3]) mean_shape = mean_shape - np.mean( - mean_shape_ori[:35709, ...], axis=0, keepdims=True) + mean_shape[:35709, ...], axis=0, keepdims=True) self.mean_shape = mean_shape.reshape([-1, 1]) self.center = center self.persc_proj = perspective_projection(focal, self.center) - self.device = 'cpu' self.camera_distance = camera_distance self.SH = SH() self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32) @@ -122,13 +98,9 @@ class ParametricFaceModel: for key, value in self.__dict__.items(): if type(value).__module__ == np.__name__: setattr(self, key, torch.tensor(value).to(device)) + self.render = self.render.to(device) - def compute_shape(self, - id_coeff, - exp_coeff, - nose_coeff=0.0, - neck_coeff=0.0, - eyes_coeff=0.0): + def compute_shape(self, id_coeff, exp_coeff): """ Return: face_shape -- torch.tensor, size (B, N, 3) @@ -142,16 +114,9 @@ class ParametricFaceModel: exp_part = torch.einsum('ij,aj->ai', self.exp_base, exp_coeff) face_shape = id_part + exp_part + self.mean_shape.reshape([1, -1]) - if nose_coeff != 0: - face_shape = face_shape + nose_coeff * self.nose_reduced_part - if neck_coeff != 0: - face_shape = face_shape + neck_coeff * self.neck_adjust_part - if eyes_coeff != 0 and self.eyes_adjust_part is not None: - face_shape = face_shape + eyes_coeff * self.eyes_adjust_part - return face_shape.reshape([batch_size, -1, 3]) - def compute_texture(self, tex_coeff, normalize=True): + def compute_albedo(self, tex_coeff, normalize=True): """ Return: face_texture -- torch.tensor, size (B, N, 3), in RGB order, range (0, 1.) @@ -166,6 +131,21 @@ class ParametricFaceModel: face_texture = face_texture / 255. return face_texture.reshape([batch_size, -1, 3]) + def compute_albedo_map(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_uv, + tex_coeff) + self.mean_tex_uv + if normalize: + face_texture = face_texture / 255. + return face_texture.reshape([batch_size, 512, 512, 3]) + def compute_norm(self, face_shape): """ Return: @@ -224,6 +204,101 @@ class ParametricFaceModel: face_color = torch.cat([r, g, b], dim=-1) * face_texture return face_color + def compute_color_map(self, face_texture_uv, face_norm_uv, gamma): + """ + Return: + face_color -- torch.tensor, size (B, N, 3), range (0, 1.) + + Parameters: + face_texture_uv -- torch.tensor, (B, 3, 256, 256) + face_norm_uv -- torch.tensor, (B, 3, 256, 256) + gamma -- torch.tensor, size (B, 27), SH coeffs + """ + face_texture_uv = face_texture_uv.permute( + 0, 2, 3, 1).contiguous() # (B, 256, 256, 3) + face_norm_uv = face_norm_uv.permute(0, 2, 3, + 1).contiguous() # (B, 256, 256, 3) + size = face_texture_uv.shape[1] + + 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) # (B, 9, 3) + y1 = a[0] * c[0] * torch.ones_like(face_norm_uv[..., :1]).to( + self.device) + y2 = -a[1] * c[1] * face_norm_uv[..., 1:2] + y3 = a[1] * c[1] * face_norm_uv[..., 2:] + y4 = -a[1] * c[1] * face_norm_uv[..., :1] + y5 = a[2] * c[2] * face_norm_uv[..., :1] * face_norm_uv[..., 1:2] + y6 = -a[2] * c[2] * face_norm_uv[..., 1:2] * face_norm_uv[..., 2:] + y7 = 0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm_uv[..., 2:]**2 + - 1) + y8 = -a[2] * c[2] * face_norm_uv[..., :1] * face_norm_uv[..., 2:] + y9 = 0.5 * a[2] * c[2] * ( + face_norm_uv[..., :1]**2 - face_norm_uv[..., 1:2]**2) + Y = torch.cat([y1, y2, y3, y4, y5, y6, y7, y8, y9], + dim=-1) # (B, 256, 256, 9) + Y = Y.reshape(batch_size, -1, 9) + r = Y @ gamma[..., :1] + g = Y @ gamma[..., 1:2] + b = Y @ gamma[..., 2:] + face_color = torch.cat([r, g, b], dim=-1).reshape( + batch_size, size, size, 3) * face_texture_uv # (B, 256, 256, 3) + face_color = face_color.permute(0, 3, 1, + 2).contiguous() # (B, 3, 256, 256) + return face_color + + def displacement2normal(self, uv_z, coarse_verts, coarse_normals): + ''' Convert displacement map into detail normal map + ''' + batch_size = uv_z.shape[0] + uv_coarse_vertices = self.render.world2uv(coarse_verts) + uv_coarse_normals = self.render.world2uv(coarse_normals) + + uv_detail_vertices = uv_coarse_vertices + dense_vertices = uv_detail_vertices.permute(0, 2, 3, 1).reshape( + [batch_size, -1, 3]) + uv_detail_normals = utils.vertex_normals( + dense_vertices, self.render.dense_faces.expand(batch_size, -1, -1)) + uv_detail_normals = uv_detail_normals.reshape([ + batch_size, uv_coarse_vertices.shape[2], + uv_coarse_vertices.shape[3], 3 + ]).permute(0, 3, 1, 2) + uv_detail_normals[:, :2, ...] = -uv_detail_normals[:, :2, ...] + offset = uv_coarse_normals - uv_detail_normals + + uv_detail_vertices = uv_coarse_vertices + uv_z * uv_coarse_normals + dense_vertices = uv_detail_vertices.permute(0, 2, 3, 1).reshape( + [batch_size, -1, 3]) + uv_detail_normals = utils.vertex_normals( + dense_vertices, self.render.dense_faces.expand(batch_size, -1, -1)) + uv_detail_normals = uv_detail_normals.reshape([ + batch_size, uv_coarse_vertices.shape[2], + uv_coarse_vertices.shape[3], 3 + ]).permute(0, 3, 1, 2) + uv_detail_normals[:, :2, ...] = -uv_detail_normals[:, :2, ...] + uv_detail_normals = uv_detail_normals + offset + return uv_detail_normals + + def compute_color_with_displacement(self, face_texture_uv, verts, normals, + displacement_uv, 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 + """ + uv_detail_normals = self.displacement2normal( + displacement_uv, verts, normals + ) # verts: (B, n, 3), ops['normals']: (B, n, 3), uv_detail_normals: (B, 3, 256, 256) + uv_texture = self.compute_color_map(face_texture_uv, uv_detail_normals, + gamma) + return uv_texture + def compute_rotation(self, angles): """ Return: @@ -309,7 +384,7 @@ class ParametricFaceModel: coeffs_dict -- a dict of torch.tensors Parameters: - coeffs -- torch.tensor, size (B, 256) + coeffs -- torch.tensor, size (B, 257) """ if type(coeffs) == dict and 'id' in coeffs: return coeffs @@ -345,49 +420,272 @@ class ParametricFaceModel: return coeffs_merge - def compute_for_render(self, coeffs, coeffs_mvs=None): - """ - Return: - face_vertex -- torch.tensor, size (B, N, 3), in camera coordinate - face_color -- torch.tensor, size (B, N, 3), in RGB order - landmark -- torch.tensor, size (B, 68, 2), y direction is opposite to v direction - Parameters: - coeffs -- torch.tensor, size (B, 257) - """ + def get_texture_map(self, face_vertex, input_img): + batch_size = input_img.shape[0] + h, w = input_img.shape[2:] + face_vertex_uv = self.render.world2uv(face_vertex) # (B, 3, 256, 256) + face_vertex_uv = face_vertex_uv.reshape(batch_size, 3, -1).permute( + 0, 2, 1) # (B, N, 3), N: 256*256 + face_vertex_uv_proj = self.to_image( + face_vertex_uv) # (B, N, 2) , project to image (size=224) + face_vertex_uv_proj[..., 0] *= w / 224 + face_vertex_uv_proj[..., 1] *= h / 224 + face_vertex_uv_proj[torch.isnan(face_vertex_uv_proj)] = 0 + + face_vertex_uv_proj[..., -1] = h - 1 - face_vertex_uv_proj[..., -1] + + input_img = input_img.permute(0, 2, 3, 1) # (B, h, w, 3) + + face_vertex_uv_proj_int = torch.floor(face_vertex_uv_proj) + face_vertex_uv_proj_float = face_vertex_uv_proj - face_vertex_uv_proj_int # (B, N, 2) + face_vertex_uv_proj_float = face_vertex_uv_proj_float.reshape( + -1, 2) # (B * N, 2) + face_vertex_uv_proj_int = face_vertex_uv_proj_int.long() # (B, N, 2) + + batch_indices = torch.arange(0, batch_size)[:, None, None].repeat( + 1, face_vertex_uv_proj_int.shape[1], + 1).long().to(face_vertex_uv_proj_int.device) # (B, N, 1) + indices = torch.cat([face_vertex_uv_proj_int, batch_indices], dim=2) + indices = indices.reshape(-1, 3) # (B * N, 3) + + face_vertex_uv_proj_lt = input_img[indices[:, 2], indices[:, 1].clamp( + 0, h - 1), indices[:, 0].clamp(0, w - 1)] # (B * N, 3) + face_vertex_uv_proj_lb = input_img[indices[:, 2], + (indices[:, 1] + 1).clamp(0, h - 1), + indices[:, 0].clamp(0, w - 1)] + face_vertex_uv_proj_rt = input_img[indices[:, 2], + indices[:, 1].clamp(0, h - 1), + (indices[:, 0] + 1).clamp(0, w - 1)] + face_vertex_uv_proj_rb = input_img[indices[:, 2], + (indices[:, 1] + 1).clamp(0, h - 1), + (indices[:, 0] + 1).clamp(0, w - 1)] + + value_1 = face_vertex_uv_proj_lt * ( + 1 - face_vertex_uv_proj_float[:, :1] + ) * face_vertex_uv_proj_float[:, 1:] + value_2 = face_vertex_uv_proj_lb * ( + 1 - face_vertex_uv_proj_float[:, :1]) * ( + 1 - face_vertex_uv_proj_float[:, 1:]) + value_3 = face_vertex_uv_proj_rt * face_vertex_uv_proj_float[:, : + 1] * face_vertex_uv_proj_float[:, + 1:] + value_4 = face_vertex_uv_proj_rb * face_vertex_uv_proj_float[:, :1] * ( + 1 - face_vertex_uv_proj_float[:, 1:]) + + texture_map = value_1 + value_2 + value_3 + value_4 # (B * N, 3) + + texture_map = texture_map.reshape(batch_size, self.render.uv_size, + self.render.uv_size, + -1) # (B, 256, 256, 3) + + return texture_map + + def compute_for_render(self, coeffs): if type(coeffs) == dict: coef_dict = coeffs elif type(coeffs) == torch.Tensor: coef_dict = self.split_coeff(coeffs) - face_shape = self.compute_shape( - coef_dict['id'], coef_dict['exp'], nose_coeff=0.4, neck_coeff=0.6) + face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp']) rotation = self.compute_rotation(coef_dict['angle']) face_shape_transformed = self.transform(face_shape, rotation, coef_dict['trans']) - face_vertex = self.to_camera(face_shape_transformed) - face_vertex_ori = self.to_camera(face_shape) + face_vertex = self.to_camera(face_shape_transformed.clone()) + face_vertex_noTrans = self.to_camera(face_shape.clone()) face_proj = self.to_image(face_vertex) landmark = self.get_landmarks(face_proj) - face_texture = self.compute_texture(coef_dict['tex']) face_norm = self.compute_norm(face_shape) face_norm_roted = face_norm @ rotation - face_color = self.compute_color(face_texture, face_norm_roted, - coef_dict['gamma']) - if coeffs_mvs is not None: - mvs_face_shape = self.compute_shape(coeffs_mvs['id'], - coeffs_mvs['exp']) + face_albedo_map = self.compute_albedo_map( + coef_dict['tex']) # (1, 512, 512, 3) + face_albedo_map = face_albedo_map.permute(0, 3, 1, 2) + face_albedo_map = torch.nn.functional.interpolate( + face_albedo_map, [self.render.uv_size, self.render.uv_size], + mode='bilinear') + face_norm_roted_uv = self.render.world2uv(face_norm_roted) + face_color_map = self.compute_color_map(face_albedo_map, + face_norm_roted_uv, + coef_dict['gamma']) - mvs_face_shape_transformed = self.transform( - mvs_face_shape, rotation, coef_dict['trans']) - mvs_face_vertex = self.to_camera(mvs_face_shape_transformed) - return face_vertex, face_texture, face_color, landmark, mvs_face_vertex - else: - return face_vertex, face_texture, face_color, landmark, face_vertex_ori + position_map = self.render.world2uv(face_shape) + + return face_vertex, face_albedo_map, face_color_map, landmark, face_vertex_noTrans, position_map + + def recolor_texture(self, shaded_texture): + rgb_mean = torch.mean(shaded_texture, dim=(2, 3), keepdim=True) + target_mean = torch.ones_like(rgb_mean) * 0.6 + shaded_texture = shaded_texture * target_mean / rgb_mean + return shaded_texture + + def compute_for_render_hierarchical_mid(self, + coeffs, + deformation_map, + UVs, + visualize=False, + de_retouched_albedo_map=None): + 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']) # (B, n, 3) + face_shape_base = face_shape.clone() + + face_shape, shape_offset = self.add_nonlinear_offset( + face_shape, deformation_map, UVs) # (B, n, 3) + + 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.clone()) + + face_proj = self.to_image(face_vertex) + landmark = self.get_landmarks(face_proj) + + face_norm = self.compute_norm(face_shape) + face_norm_roted = face_norm @ rotation + + face_albedo_map = self.compute_albedo_map( + coef_dict['tex']) # (B, 512, 512, 3) + face_albedo_map = face_albedo_map.permute(0, 3, 1, 2) + face_albedo_map = torch.nn.functional.interpolate( + face_albedo_map, [self.render.uv_size, self.render.uv_size], + mode='bilinear') + face_norm_roted_uv = self.render.world2uv(face_norm_roted) + + albedo_for_render = face_albedo_map if de_retouched_albedo_map is None else de_retouched_albedo_map + face_color_map = self.compute_color_map(albedo_for_render, + face_norm_roted_uv, + coef_dict['gamma']) + + extra_results = None + if visualize: + extra_results = {} + extra_results['tex_mid_color'] = face_color_map + + face_shape_transformed_base = self.transform( + face_shape_base, rotation, coef_dict['trans']) + face_vertex_base = self.to_camera( + face_shape_transformed_base.clone()) + extra_results['pred_vertex_base'] = face_vertex_base + face_norm_base = self.compute_norm(face_shape_base) + face_norm_roted_base = face_norm_base @ rotation + + batch_size = albedo_for_render.shape[0] + size = albedo_for_render.shape[2] + gray_tex = torch.ones((batch_size, 3, size, size), + dtype=torch.float32).to(self.device) * 0.8 + zero_displacement = torch.zeros( + (batch_size, 1, size, size), + dtype=torch.float32).to(self.device) + + tex_mid_gray = self.compute_color_with_displacement( + gray_tex.detach(), face_shape_transformed, face_norm_roted, + zero_displacement, coef_dict['gamma']) + tex_mid_gray = self.recolor_texture(tex_mid_gray) + extra_results['tex_mid_gray'] = tex_mid_gray + + tex_base_color = self.compute_color_with_displacement( + albedo_for_render, face_shape_transformed_base, + face_norm_roted_base, zero_displacement, coef_dict['gamma']) + extra_results['tex_base_color'] = tex_base_color + + tex_base_gray = self.compute_color_with_displacement( + gray_tex.detach(), face_shape_transformed_base, + face_norm_roted_base, zero_displacement, coef_dict['gamma']) + tex_base_gray = self.recolor_texture(tex_base_gray) + extra_results['tex_base_gray'] = tex_base_gray + + # to export rotate video + init_angle = torch.zeros_like(coef_dict['angle']).to( + coef_dict['angle'].device) + + pi = 3.14 + n_frame = 30 + y_angles = torch.linspace(-pi / 6, pi / 6, steps=n_frame).float() + + extra_results['face_shape_transformed_list'] = [] + extra_results['face_norm_roted_list'] = [] + extra_results['face_vertex_list'] = [] + for y_angle in y_angles: + cur_angle = init_angle.clone() + cur_angle[:, 1] = y_angle + cur_angle[:, 0] = pi / 36 + cur_rotation = self.compute_rotation(cur_angle) + cur_face_shape_transformed = self.transform( + face_shape, cur_rotation, coef_dict['trans'] * 0) + cur_face_norm_roted = face_norm @ cur_rotation + cur_face_vertex = self.to_camera( + cur_face_shape_transformed.clone()) + + extra_results['face_shape_transformed_list'].append( + cur_face_shape_transformed) + extra_results['face_norm_roted_list'].append( + cur_face_norm_roted) + extra_results['face_vertex_list'].append(cur_face_vertex) + + return (face_vertex, face_color_map, landmark, face_proj, + face_albedo_map, face_shape_transformed, face_norm_roted, + extra_results) + + def compute_for_render_hierarchical_high(self, + coeffs, + displacement_uv, + face_albedo_map, + face_shape_transformed, + face_norm_roted, + extra_results=None): + if type(coeffs) == dict: + coef_dict = coeffs + elif type(coeffs) == torch.Tensor: + coef_dict = self.split_coeff(coeffs) + + face_color_map = self.compute_color_with_displacement( + face_albedo_map, face_shape_transformed, face_norm_roted, + displacement_uv, coef_dict['gamma']) + + if extra_results is not None: + extra_results['tex_high_color'] = face_color_map + + batch_size = face_albedo_map.shape[0] + size = face_albedo_map.shape[2] + gray_tex = torch.ones((batch_size, 3, size, size), + dtype=torch.float32).to(self.device) * 0.8 + + tex_high_gray = self.compute_color_with_displacement( + gray_tex.detach(), face_shape_transformed, face_norm_roted, + displacement_uv, coef_dict['gamma']) + tex_high_gray = self.recolor_texture(tex_high_gray) + extra_results['tex_high_gray'] = tex_high_gray + + if 'face_shape_transformed_list' in extra_results: + extra_results['tex_high_gray_list'] = [] + extra_results['tex_high_color_list'] = [] + for i in range( + len(extra_results['face_shape_transformed_list'])): + tex_high_gray_i = self.compute_color_with_displacement( + gray_tex.detach(), + extra_results['face_shape_transformed_list'][i], + extra_results['face_norm_roted_list'][i], + displacement_uv, coef_dict['gamma']) + extra_results['tex_high_gray_list'].append(tex_high_gray_i) + + tex_high_color_i = self.compute_color_with_displacement( + face_albedo_map.detach(), + extra_results['face_shape_transformed_list'][i], + extra_results['face_norm_roted_list'][i], + displacement_uv, coef_dict['gamma']) + extra_results['tex_high_color_list'].append( + tex_high_color_i) + + return face_color_map, extra_results def reverse_recenter(self, face_shape): batch_size = face_shape.shape[0] @@ -398,51 +696,18 @@ class ParametricFaceModel: face_shape = face_shape.reshape([batch_size, -1, 3]) return face_shape - def add_nonlinear_offset_eyes(self, face_shape, shape_offset): - assert face_shape.shape[0] == 1 and shape_offset.shape[0] == 1 - face_shape = face_shape[0] - shape_offset = shape_offset[0] - - corner_inds = self.eye_corner_inds - lines = self.eye_corner_lines - - corner_shape = face_shape[-625:, :] - corner_offset = shape_offset[corner_inds] - for i in range(len(lines)): - corner_shape[lines[i]] += corner_offset[i][None, ...] - face_shape[-625:, :] = corner_shape - - l_eye_landmarks = [11540, 11541] - r_eye_landmarks = [4271, 4272] - - l_eye_offset = torch.mean( - shape_offset[l_eye_landmarks], dim=0, keepdim=True) - face_shape[37082:37082 + 609] += l_eye_offset - - r_eye_offset = torch.mean( - shape_offset[r_eye_landmarks], dim=0, keepdim=True) - face_shape[37082 + 609:37082 + 609 + 608] += r_eye_offset - - face_shape = face_shape[None, ...] - - return face_shape - def add_nonlinear_offset(self, face_shape, shape_offset_uv, UVs): """ Args: - face_shape: torch.tensor, size (1, N, 3) - shape_offset_uv: torch.tensor, size (1, h, w, 3) + face_shape: torch.tensor, size (B, N, 3) + shape_offset_uv: torch.tensor, size (B, 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] + h, w = shape_offset_uv.shape[1:3] UVs_coords = UVs.clone() UVs_coords[:, 0] *= w UVs_coords[:, 1] *= h @@ -450,142 +715,28 @@ class ParametricFaceModel: 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), + shape_lt = shape_offset_uv[:, (h - 1 - UVs_coords_int[:, 1]).clamp( + 0, h - 1), UVs_coords_int[:, 0].clamp(0, w - 1)] # (B, 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_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), + 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), + shape_rb = shape_offset_uv[:, + (h - UVs_coords_int[:, 1]).clamp(0, h - 1), (UVs_coords_int[:, 0] + 1).clamp(0, w - 1)] - value1 = shape_lt * ( + value_1 = shape_lt * ( 1 - UVs_coords_float[:, :1]) * UVs_coords_float[:, 1:] - value2 = shape_lb * (1 - UVs_coords_float[:, :1]) * ( + value_2 = shape_lb * (1 - UVs_coords_float[:, :1]) * ( 1 - UVs_coords_float[:, 1:]) - value3 = shape_rt * UVs_coords_float[:, :1] * UVs_coords_float[:, 1:] - value4 = shape_rb * UVs_coords_float[:, :1] * ( + 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 = value1 + value2 + value3 + value4 # (N, 3) - face_shape = (face_shape + offset_shape)[None, ...] + offset_shape = value_1 + value_2 + value_3 + value_4 # (B, N, 3) - return face_shape, offset_shape[None, ...] + face_shape = face_shape + offset_shape - def compute_for_render_train_nonlinear(self, - coeffs, - shape_offset_uv, - tex_offset_uv, - UVs, - reverse_recenter=True): - if type(coeffs) == dict: - coef_dict = coeffs - elif type(coeffs) == torch.Tensor: - coef_dict = self.split_coeff(coeffs) - - face_shape = self.compute_shape(coef_dict['id'], - coef_dict['exp']) # (1, n, 3) - if reverse_recenter: - face_shape_ori_noRecenter = self.reverse_recenter( - face_shape.clone()) - else: - face_shape_ori_noRecenter = face_shape.clone() - face_vertex_ori = self.to_camera(face_shape_ori_noRecenter) - - face_shape, shape_offset = self.add_nonlinear_offset( - face_shape, shape_offset_uv, UVs[:35709, :]) # (1, n, 3) - if reverse_recenter: - face_shape_offset_noRecenter = self.reverse_recenter( - face_shape.clone()) - else: - face_shape_offset_noRecenter = face_shape.clone() - face_vertex_offset = self.to_camera(face_shape_offset_noRecenter) - - rotation = self.compute_rotation(coef_dict['angle']) - - face_shape_transformed = self.transform(face_shape, rotation, - coef_dict['trans']) - face_vertex = self.to_camera(face_shape_transformed) - - face_proj = self.to_image(face_vertex) - landmark = self.get_landmarks(face_proj) - - face_texture = self.compute_texture(coef_dict['tex']) # (1, n, 3) - face_texture, texture_offset = self.add_nonlinear_offset( - face_texture, tex_offset_uv, UVs[:35709, :]) - face_norm = self.compute_norm(face_shape) - face_norm_roted = face_norm @ rotation - face_color = self.compute_color(face_texture, face_norm_roted, - coef_dict['gamma']) - - return face_vertex, face_texture, face_color, landmark, face_vertex_ori, face_vertex_offset, face_proj - - def compute_for_render_nonlinear_full(self, - coeffs, - shape_offset_uv, - UVs, - nose_coeff=0.0, - eyes_coeff=0.0): - if type(coeffs) == dict: - coef_dict = coeffs - elif type(coeffs) == torch.Tensor: - coef_dict = self.split_coeff(coeffs) - - face_shape = self.compute_shape( - coef_dict['id'], - coef_dict['exp'], - nose_coeff=nose_coeff, - neck_coeff=0.6, - eyes_coeff=eyes_coeff) # (1, n, 3) - face_vertex_ori = self.to_camera(face_shape.clone()) - - face_shape[:, :35241, :], shape_offset = self.add_nonlinear_offset( - face_shape[:, :35241, :], shape_offset_uv, - UVs[:35709, :][self.bfm_keep_inds]) - face_shape = self.add_nonlinear_offset_eyes(face_shape, shape_offset) - face_shape_noRecenter = self.reverse_recenter(face_shape.clone()) - face_vertex_offset = self.to_camera(face_shape_noRecenter) - - rotation = self.compute_rotation(coef_dict['angle']) - - face_shape_transformed = self.transform(face_shape, rotation, - coef_dict['trans']) - face_vertex = self.to_camera(face_shape_transformed) - - return face_vertex, face_vertex_ori, face_vertex_offset - - def compute_for_render_train(self, coeffs): - """ - Return: - face_vertex -- torch.tensor, size (B, N, 3), in camera coordinate - face_color -- torch.tensor, size (B, N, 3), in RGB order - landmark -- torch.tensor, size (B, 68, 2), y direction is opposite to v direction - Parameters: - coeffs -- torch.tensor, size (B, 257) - """ - if type(coeffs) == dict: - coef_dict = coeffs - elif type(coeffs) == torch.Tensor: - coef_dict = self.split_coeff(coeffs) - - face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp']) - uv_geometry = self.render.world2uv(face_shape) - - rotation = self.compute_rotation(coef_dict['angle']) - - face_shape_transformed = self.transform(face_shape, rotation, - coef_dict['trans']) - face_vertex = self.to_camera(face_shape_transformed) - - face_proj = self.to_image(face_vertex) - landmark = self.get_landmarks(face_proj) - - face_texture = self.compute_texture(coef_dict['tex']) - face_norm = self.compute_norm(face_shape) - face_norm_roted = face_norm @ rotation - face_color = self.compute_color(face_texture, face_norm_roted, - coef_dict['gamma']) - - return face_vertex, face_texture, face_color, landmark, uv_geometry + return face_shape, offset_shape diff --git a/modelscope/models/cv/face_reconstruction/models/de_retouching_module.py b/modelscope/models/cv/face_reconstruction/models/de_retouching_module.py new file mode 100644 index 00000000..92711f6e --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/de_retouching_module.py @@ -0,0 +1,47 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch + +from .unet import UNet + + +class DeRetouchingModule(): + + def __init__(self, model_path): + self.retouching_network = UNet(3, 3).to('cuda') + self.retouching_network.load_state_dict( + torch.load(model_path, map_location='cpu')['generator']) + self.retouching_network.eval() + + def run(self, face_albedo_map, texture_map): + """ + + :param face_albedo_map: tensor, (1, 3, 256, 256), 0~1, rgb + :param texture_map: tensor, (1, 3, 256, 256), -1~1, rgb + :return: + """ + h, w = texture_map.shape[2:] + retouch_input = torch.nn.functional.interpolate( + texture_map, (512, 512), mode='bilinear') + + # predict blend layer + blend_layer = self.retouching_network(retouch_input) # value: 0~1 + blend_layer = torch.nn.functional.interpolate( + blend_layer, (h, w), mode='bilinear') + + # retouch texture map + tex = (texture_map + 1.0) / 2 + retouched_tex = (1 - 2 * blend_layer + ) * tex * tex + 2 * blend_layer * tex # value: 0~1 + + # de-retouch albedo map + A_0 = face_albedo_map + T_0 = retouched_tex + T_ = tex + + phi = 1e6 * T_0 * T_0 * T_0 + 1e-6 + B = (T_ + 1 / phi) / (T_0 + 1 / phi) + A_ = A_0 * B + + de_retouched_albedo = A_ + + return de_retouched_albedo diff --git a/modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py b/modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py deleted file mode 100644 index c61a0d13..00000000 --- a/modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py +++ /dev/null @@ -1,430 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import os - -import cv2 -import numpy as np -import torch - -from modelscope.models.cv.skin_retouching.retinaface.predict_single import \ - Model -from ...utils import image_warp_grid1, spread_flow -from .large_base_lmks_infer import LargeBaseLmkInfer - -INPUT_SIZE = 224 -ENLARGE_RATIO = 1.35 - - -def resize_on_long_side(img, long_side=800): - src_height = img.shape[0] - src_width = img.shape[1] - - if src_height > src_width: - scale = long_side * 1.0 / src_height - _img = cv2.resize( - img, (int(src_width * scale), long_side), - interpolation=cv2.INTER_CUBIC) - - else: - scale = long_side * 1.0 / src_width - _img = cv2.resize( - img, (long_side, int(src_height * scale)), - interpolation=cv2.INTER_CUBIC) - - return _img, scale - - -def draw_line(im, points, color, stroke_size=2, closed=False): - points = points.astype(np.int32) - for i in range(len(points) - 1): - cv2.line(im, tuple(points[i]), tuple(points[i + 1]), color, - stroke_size) - if closed: - cv2.line(im, tuple(points[0]), tuple(points[-1]), color, stroke_size) - - -def enlarged_bbox(bbox, img_width, img_height, enlarge_ratio=0.2): - ''' - :param bbox: [xmin,ymin,xmax,ymax] - :return: bbox: [xmin,ymin,xmax,ymax] - ''' - - left = bbox[0] - top = bbox[1] - - right = bbox[2] - bottom = bbox[3] - - roi_width = right - left - roi_height = bottom - top - - new_left = left - int(roi_width * enlarge_ratio) - new_left = 0 if new_left < 0 else new_left - - new_top = top - int(roi_height * enlarge_ratio) - new_top = 0 if new_top < 0 else new_top - - new_right = right + int(roi_width * enlarge_ratio) - new_right = img_width if new_right > img_width else new_right - - new_bottom = bottom + int(roi_height * enlarge_ratio) - new_bottom = img_height if new_bottom > img_height else new_bottom - - bbox = [new_left, new_top, new_right, new_bottom] - - bbox = [int(x) for x in bbox] - - return bbox - - -class FaceInfo: - - def __init__(self): - self.rect = np.asarray([0, 0, 0, 0]) - self.points_array = np.zeros((106, 2)) - self.eye_left = np.zeros((22, 2)) - self.eye_right = np.zeros((22, 2)) - self.eyebrow_left = np.zeros((13, 2)) - self.eyebrow_right = np.zeros((13, 2)) - self.lips = np.zeros((64, 2)) - - -class LargeModelInfer: - - def __init__(self, ckpt, device='cuda'): - self.large_base_lmks_model = LargeBaseLmkInfer.model_preload( - ckpt, - device.lower() == 'cuda') - self.device = device.lower() - self.detector = Model(max_size=512, device=device) - detector_ckpt_name = 'retinaface_resnet50_2020-07-20_old_torch.pth' - state_dict = torch.load( - os.path.join(os.path.dirname(ckpt), detector_ckpt_name), - map_location='cpu') - self.detector.load_state_dict(state_dict) - self.detector.eval() - - def infer(self, img_bgr): - landmarks = [] - - rgb_image = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) - results = self.detector.predict_jsons(rgb_image) - - boxes = [] - for anno in results: - if anno['score'] == -1: - break - boxes.append({ - 'x1': anno['bbox'][0], - 'y1': anno['bbox'][1], - 'x2': anno['bbox'][2], - 'y2': anno['bbox'][3] - }) - - for detect_result in boxes: - x1 = detect_result['x1'] - y1 = detect_result['y1'] - x2 = detect_result['x2'] - y2 = detect_result['y2'] - - w = x2 - x1 + 1 - h = y2 - y1 + 1 - - cx = (x2 + x1) / 2 - cy = (y2 + y1) / 2 - - sz = max(h, w) * ENLARGE_RATIO - - x1 = cx - sz / 2 - y1 = cy - sz / 2 - trans_x1 = x1 - trans_y1 = y1 - x2 = x1 + sz - y2 = y1 + sz - - height, width, _ = rgb_image.shape - dx = max(0, -x1) - dy = max(0, -y1) - x1 = max(0, x1) - y1 = max(0, y1) - - edx = max(0, x2 - width) - edy = max(0, y2 - height) - x2 = min(width, x2) - y2 = min(height, y2) - - crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)] - if dx > 0 or dy > 0 or edx > 0 or edy > 0: - crop_img = cv2.copyMakeBorder( - crop_img, - int(dy), - int(edy), - int(dx), - int(edx), - cv2.BORDER_CONSTANT, - value=(103.94, 116.78, 123.68)) - crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE)) - - base_lmks = LargeBaseLmkInfer.infer_img(crop_img, - self.large_base_lmks_model, - self.device == 'cuda') - - inv_scale = sz / INPUT_SIZE - - affine_base_lmks = np.zeros((106, 2)) - for idx in range(106): - affine_base_lmks[idx][ - 0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1 - affine_base_lmks[idx][ - 1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1 - - x1 = np.min(affine_base_lmks[:, 0]) - y1 = np.min(affine_base_lmks[:, 1]) - x2 = np.max(affine_base_lmks[:, 0]) - y2 = np.max(affine_base_lmks[:, 1]) - - w = x2 - x1 + 1 - h = y2 - y1 + 1 - - cx = (x2 + x1) / 2 - cy = (y2 + y1) / 2 - - sz = max(h, w) * ENLARGE_RATIO - - x1 = cx - sz / 2 - y1 = cy - sz / 2 - trans_x1 = x1 - trans_y1 = y1 - x2 = x1 + sz - y2 = y1 + sz - - height, width, _ = rgb_image.shape - dx = max(0, -x1) - dy = max(0, -y1) - x1 = max(0, x1) - y1 = max(0, y1) - - edx = max(0, x2 - width) - edy = max(0, y2 - height) - x2 = min(width, x2) - y2 = min(height, y2) - - crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)] - if dx > 0 or dy > 0 or edx > 0 or edy > 0: - crop_img = cv2.copyMakeBorder( - crop_img, - int(dy), - int(edy), - int(dx), - int(edx), - cv2.BORDER_CONSTANT, - value=(103.94, 116.78, 123.68)) - crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE)) - - base_lmks = LargeBaseLmkInfer.infer_img( - crop_img, self.large_base_lmks_model, - self.device.lower() == 'cuda') - - inv_scale = sz / INPUT_SIZE - - affine_base_lmks = np.zeros((106, 2)) - for idx in range(106): - affine_base_lmks[idx][ - 0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1 - affine_base_lmks[idx][ - 1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1 - - landmarks.append(affine_base_lmks) - - return boxes, landmarks - - def find_face_contour(self, image): - - boxes, landmarks = self.infer(image) - landmarks = np.array(landmarks) - - args = [[0, 33, False], [33, 38, False], [42, 47, False], - [51, 55, False], [57, 64, False], [66, 74, True], - [75, 83, True], [84, 96, True]] - - roi_bboxs = [] - - for i in range(len(boxes)): - roi_bbox = enlarged_bbox([ - boxes[i]['x1'], boxes[i]['y1'], boxes[i]['x2'], boxes[i]['y2'] - ], image.shape[1], image.shape[0], 0.5) - roi_bbox = [int(x) for x in roi_bbox] - roi_bboxs.append(roi_bbox) - - people_maps = [] - - for i in range(landmarks.shape[0]): - landmark = landmarks[i, :, :] - maps = [] - whole_mask = np.zeros((image.shape[0], image.shape[1]), np.uint8) - - roi_box = roi_bboxs[i] - roi_box_width = roi_box[2] - roi_box[0] - roi_box_height = roi_box[3] - roi_box[1] - short_side_length = roi_box_width if roi_box_width < roi_box_height else roi_box_height - - line_width = short_side_length // 10 - - if line_width == 0: - line_width = 1 - - kernel_size = line_width * 2 - gaussian_kernel = kernel_size if kernel_size % 2 == 1 else kernel_size + 1 - - for t, arg in enumerate(args): - mask = np.zeros((image.shape[0], image.shape[1]), np.uint8) - draw_line(mask, landmark[arg[0]:arg[1]], (255, 255, 255), - line_width, arg[2]) - mask = cv2.GaussianBlur(mask, - (gaussian_kernel, gaussian_kernel), 0) - if t >= 1: - draw_line(whole_mask, landmark[arg[0]:arg[1]], - (255, 255, 255), line_width * 2, arg[2]) - maps.append(mask) - whole_mask = cv2.GaussianBlur(whole_mask, - (gaussian_kernel, gaussian_kernel), - 0) - maps.append(whole_mask) - people_maps.append(maps) - - return people_maps[0], boxes - - def face2contour(self, image, stack_mode='column'): - ''' - - :param facer: - :param image: - :param stack_mode: - :return: final_maps: [map0, map1,....] - roi_bboxs: [bbox0, bbox1, ...] - ''' - - boxes, landmarks = self.infer(image) - landmarks = np.array(landmarks) - - args = [[0, 33, False], [33, 38, False], [42, 47, False], - [51, 55, False], [57, 64, False], [66, 74, True], - [75, 83, True], [84, 96, True]] - - roi_bboxs = [] - - for i in range(len(boxes)): - roi_bbox = enlarged_bbox([ - boxes[i]['x1'], boxes[i]['y1'], boxes[i]['x2'], boxes[i]['y2'] - ], image.shape[1], image.shape[0], 0.5) - roi_bbox = [int(x) for x in roi_bbox] - roi_bboxs.append(roi_bbox) - - people_maps = [] - - for i in range(landmarks.shape[0]): - landmark = landmarks[i, :, :] - maps = [] - whole_mask = np.zeros((image.shape[0], image.shape[1]), np.uint8) - - roi_box = roi_bboxs[i] - roi_box_width = roi_box[2] - roi_box[0] - roi_box_height = roi_box[3] - roi_box[1] - short_side_length = roi_box_width if roi_box_width < roi_box_height else roi_box_height - - line_width = short_side_length // 50 - - if line_width == 0: - line_width = 1 - - kernel_size = line_width * 4 - gaussian_kernel = kernel_size if kernel_size % 2 == 1 else kernel_size + 1 - - for arg in args: - mask = np.zeros((image.shape[0], image.shape[1]), np.uint8) - draw_line(mask, landmark[arg[0]:arg[1]], (255, 255, 255), - line_width, arg[2]) - mask = cv2.GaussianBlur(mask, - (gaussian_kernel, gaussian_kernel), 0) - draw_line(whole_mask, landmark[arg[0]:arg[1]], (255, 255, 255), - line_width, arg[2]) - maps.append(mask) - whole_mask = cv2.GaussianBlur(whole_mask, - (gaussian_kernel, gaussian_kernel), - 0) - maps.append(whole_mask) - people_maps.append(maps) - - if stack_mode == 'depth': - final_maps = [] - for i, maps in enumerate(people_maps): - final_map = np.dstack(maps) - final_map = final_map[roi_bboxs[i][1]:roi_bboxs[i][3], - roi_bboxs[i][0]:roi_bboxs[i][2], :] - final_maps.append(final_map) - return final_maps, roi_bboxs - - elif stack_mode == 'column': - final_maps = [] - for i, maps in enumerate(people_maps): - joint_maps = [ - x[roi_bboxs[i][1]:roi_bboxs[i][3], - roi_bboxs[i][0]:roi_bboxs[i][2]] for x in maps - ] - final_map = np.column_stack(joint_maps) - final_maps.append(final_map) - return final_maps, roi_bboxs - - def fat_face(self, img, degree=0.1): - - _img, scale = resize_on_long_side(img, 800) - - contour_maps, boxes = self.find_face_contour(_img) - - contour_map = contour_maps[0] - - boxes = boxes[0] - - Flow = np.zeros( - shape=(contour_map.shape[0], contour_map.shape[1], 2), - dtype=np.float32) - - box_center = [(boxes['x1'] + boxes['x2']) / 2, - (boxes['y1'] + boxes['y2']) / 2] - - box_length = max( - abs(boxes['y1'] - boxes['y2']), abs(boxes['x1'] - boxes['x2'])) - - value_1 = 2 * (Flow.shape[0] - box_center[1] - 1) - value_2 = 2 * (Flow.shape[1] - box_center[0] - 1) - value_list = [ - box_length * 2, 2 * (box_center[0] - 1), 2 * (box_center[1] - 1), - value_1, value_2 - ] - flow_box_length = min(value_list) - flow_box_length = int(flow_box_length) - - sf = spread_flow(100, flow_box_length * degree) - sf = cv2.resize(sf, (flow_box_length, flow_box_length)) - - Flow[int(box_center[1] - - flow_box_length / 2):int(box_center[1] - + flow_box_length / 2), - int(box_center[0] - - flow_box_length / 2):int(box_center[0] - + flow_box_length / 2)] = sf - - Flow = Flow * np.dstack((contour_map, contour_map)) / 255.0 - - inter_face_maps = contour_maps[-1] - - Flow = Flow * (1.0 - np.dstack( - (inter_face_maps, inter_face_maps)) / 255.0) - - Flow = cv2.resize(Flow, (img.shape[1], img.shape[0])) - - Flow = Flow / scale - - pred, top_bound, bottom_bound, left_bound, right_bound = image_warp_grid1( - Flow[..., 0], Flow[..., 1], img, 1.0, [0, 0, 0, 0]) - - return pred diff --git a/modelscope/models/cv/face_reconstruction/models/facerecon_model.py b/modelscope/models/cv/face_reconstruction/models/facerecon_model.py index c94cf1f9..d753b416 100644 --- a/modelscope/models/cv/face_reconstruction/models/facerecon_model.py +++ b/modelscope/models/cv/face_reconstruction/models/facerecon_model.py @@ -1,6 +1,6 @@ -# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch, -# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch +# Copyright (c) Alibaba, Inc. and its affiliates. import os +from collections import OrderedDict import cv2 import numpy as np @@ -11,10 +11,11 @@ from modelscope.models.cv.face_reconstruction.models import opt from .. import utils from . import networks from .bfm import ParametricFaceModel -from .losses import (CLIPLoss_relative, TVLoss, TVLoss_std, landmark_loss, - perceptual_loss, photo_loss, points_loss_horizontal, - reflectance_loss, reg_loss) +from .de_retouching_module import DeRetouchingModule +from .losses import TVLoss, photo_loss from .nv_diffrast import MeshRenderer +from .pix2pix.pix2pix_model import Pix2PixModel +from .pix2pix.pix2pix_options import Pix2PixOptions @MODELS.register_module('face-reconstruction', 'face_reconstruction') @@ -23,12 +24,8 @@ class FaceReconModel(TorchModel): def __init__(self, model_dir, w_color=1.92, - w_exp=0.8, - w_gamma=10.0, - w_id=1.0, - w_lm=0.0016, - w_reg=0.0003, - w_tex=0.017, + tex_iters=100, + w_tex_smooth=10, *args, **kwargs): """The FaceReconModel is implemented based on Deep3DFaceRecon_pytorch, publicly available at @@ -37,94 +34,86 @@ class FaceReconModel(TorchModel): Args: model_dir: the root directory of the model files w_color: the weight of color loss - w_exp: the regularization weight of expression - w_gamma: the regularization weight of lighting - w_id: the regularization weight of identity - w_lm: the weight of landmark loss - w_reg: the weight of regularization loss - w_tex: the regularization weight of texture """ super().__init__(model_dir, *args, **kwargs) opt.bfm_folder = os.path.join(model_dir, 'assets') self.opt = opt self.w_color = w_color - self.w_exp = w_exp - self.w_gamma = w_gamma - self.w_id = w_id - self.w_lm = w_lm - self.w_reg = w_reg - self.w_tex = w_tex - self.device = torch.device('cpu') + self.w_tex_smooth = w_tex_smooth + self.tex_iters = tex_iters self.isTrain = opt.isTrain self.visual_names = ['output_vis'] - self.model_names = ['net_recon'] - self.parallel_names = self.model_names + ['renderer'] + self.model_names = ['net_recon', 'mid_net', 'high_net'] + self.parallel_names = self.model_names + [ + 'renderer', 'renderer_high_res' + ] + # networks self.net_recon = networks.define_net_recon( net_recon=opt.net_recon, use_last_fc=opt.use_last_fc, init_path=None) - self.facemodel = ParametricFaceModel( - bfm_folder=opt.bfm_folder, - camera_distance=opt.camera_d, - focal=opt.focal, - center=opt.center, - is_train=self.isTrain, - default_name=opt.bfm_model) + de_retouching_model_path = os.path.join(model_dir, + 'de_retouching_model.pth') + self.de_retouching_module = DeRetouchingModule( + de_retouching_model_path) + self.mid_opt = Pix2PixOptions() + self.mid_opt.input_nc = 6 + self.mid_opt.output_nc = 3 + self.mid_opt.name = 'mid_net' + self.mid_net = Pix2PixModel(self.mid_opt).netG + + self.high_opt = Pix2PixOptions() + self.high_opt.input_nc = 9 + self.high_opt.output_nc = 1 + self.high_opt.name = 'high_net' + self.high_net = Pix2PixModel(self.high_opt).netG + + self.alpha = (torch.ones(1, dtype=torch.float32) * 0.01) + self.alpha.requires_grad = True + self.beta = (torch.ones(1, dtype=torch.float32) * 0.01) + self.beta.requires_grad = True + + # assets self.facemodel_front = ParametricFaceModel( - bfm_folder=opt.bfm_folder, + assets_folder=opt.bfm_folder, camera_distance=opt.camera_d, focal=opt.focal, center=opt.center, is_train=self.isTrain, - default_name='face_model_for_maas.mat') + default_name='BFM_model_front.mat') + bfm_uvs_path = os.path.join(opt.bfm_folder, 'bfm_uvs2.npy') + self.bfm_UVs = np.load(bfm_uvs_path) + self.bfm_UVs = torch.from_numpy(self.bfm_UVs).float() + + # renderer fov = 2 * np.arctan(opt.center / opt.focal) * 180 / np.pi + self.renderer_high_res = MeshRenderer( + rasterize_fov=fov, + znear=opt.z_near, + zfar=opt.z_far, + rasterize_size=int(2 * opt.center)) + self.renderer = MeshRenderer( rasterize_fov=fov, znear=opt.z_near, zfar=opt.z_far, rasterize_size=int(2 * opt.center)) - self.renderer_fitting = MeshRenderer( - rasterize_fov=fov, - znear=opt.z_near, - zfar=opt.z_far, - rasterize_size=int(2 * opt.center)) - - self.nonlinear_UVs = self.facemodel.nonlinear_UVs - self.nonlinear_UVs = torch.from_numpy(self.nonlinear_UVs).to( - torch.device('cuda')) - - template_obj_path = os.path.join(opt.bfm_folder, 'template_mesh.obj') - self.template_mesh = utils.read_obj(template_obj_path) - - self.input_imgs = [] - self.input_img_hds = [] - self.input_fat_img_hds = [] - self.atten_masks = [] - self.gt_lms = [] - self.gt_lm_hds = [] - self.trans_ms = [] - self.img_names = [] - self.face_masks = [] - self.head_masks = [] - self.input_imgs_coeff = [] - self.gt_lms_coeff = [] - - self.loss_names = [ - 'all', 'feat', 'color', 'lm', 'reg', 'gamma', 'reflc' - ] - - # loss func name: (compute_%s_loss) % loss_name - self.compute_feat_loss = perceptual_loss self.comupte_color_loss = photo_loss - self.compute_lm_loss = landmark_loss - self.compute_reg_loss = reg_loss - self.compute_reflc_loss = reflectance_loss + + def set_device(self, device): + self.device = device + self.mid_net = self.mid_net.to(self.device) + self.high_net = self.high_net.to(self.device) + self.alpha = self.alpha.to(self.device) + self.beta = self.beta.to(self.device) + self.bfm_UVs = self.bfm_UVs.to(self.device) + self.facemodel_front.to(self.device) def load_networks(self, load_path): state_dict = torch.load(load_path, map_location=self.device) @@ -137,6 +126,9 @@ class FaceReconModel(TorchModel): net = net.module net.load_state_dict(state_dict[name], strict=False) + self.alpha = state_dict['alpha'] + self.beta = state_dict['beta'] + if self.opt.phase != 'test': if self.opt.continue_train: @@ -199,18 +191,21 @@ class FaceReconModel(TorchModel): net = getattr(self, name) net.eval() + self.alpha.requires_grad = False + self.beta.requires_grad = False + def set_render(self, image_res): fov = 2 * np.arctan(self.opt.center / self.opt.focal) * 180 / np.pi if image_res is None: image_res = int(2 * self.opt.center) - self.renderer = MeshRenderer( + self.renderer_high_res = MeshRenderer( rasterize_fov=fov, znear=self.opt.z_near, zfar=self.opt.z_far, rasterize_size=image_res) - def set_input(self, input): + def set_input_base(self, input): """Unpack input data from the dataloader and perform necessary pre-processing steps. Parameters: @@ -219,346 +214,477 @@ class FaceReconModel(TorchModel): self.input_img = input['imgs'].to(self.device) self.input_img_hd = input['imgs_hd'].to( self.device) if 'imgs_hd' in input else None - - if 'imgs_fat_hd' not in input or input['imgs_fat_hd'] is None: - self.input_fat_img_hd = self.input_img_hd - else: - self.input_fat_img_hd = input['imgs_fat_hd'].to(self.device) - - self.atten_mask = input['msks'].to( - self.device) if 'msks' in input else None self.gt_lm = input['lms'].to(self.device) if 'lms' in input else None self.gt_lm_hd = input['lms_hd'].to( self.device) if 'lms_hd' in input else None - self.trans_m = input['M'].to(self.device) if 'M' in input else None - self.image_paths = input['im_paths'] if 'im_paths' in input else None - self.img_name = input['img_name'] if 'img_name' in input else None self.face_mask = input['face_mask'].to( self.device) if 'face_mask' in input else None - self.head_mask = input['head_mask'].to( - self.device) if 'head_mask' in input else None - self.gt_normals = input['normals'].to( - self.device) if 'normals' in input else None - self.input_img_coeff = input['imgs_coeff'].to( - self.device) if 'imgs_coeff' in input else None - self.gt_lm_coeff = input['lms_coeff'].to( - self.device) if 'lms_coeff' in input else None - def get_edge_points_horizontal(self): - left_points = [] - right_points = [] - for i in range(self.face_mask.shape[2]): - inds = torch.where(self.face_mask[0, 0, i, :] > 0.5) # 0.9 - if len(inds[0]) > 0: # i > 112 and len(inds[0]) > 0 - left_points.append(int(inds[0][0]) + 1) - right_points.append(int(inds[0][-1])) - else: - left_points.append(0) - right_points.append(self.face_mask.shape[3] - 1) - self.left_points = torch.tensor(left_points).long().to(self.device) - self.right_points = torch.tensor(right_points).long().to(self.device) + def set_input_hrn(self, input): + """Unpack input data from the dataloader and perform necessary pre-processing steps. - def get_edge_points_vertical(self): - top_points = [] - bottom_points = [] - for i in range(self.face_mask.shape[3]): - inds = torch.where(self.face_mask[0, 0, :, i] > 0.9) - if len(inds[0]) > 0: - top_points.append(int(inds[0][0])) - bottom_points.append(int(inds[0][-1])) - else: - top_points.append(0) - bottom_points.append(self.face_mask.shape[2] - 1) - self.top_points = torch.tensor(top_points).long().to(self.device) - self.bottom_points = torch.tensor(bottom_points).long().to(self.device) - - def blur_shape_offset_uv(self, global_blur=False, blur_size=3): - if self.edge_mask is not None: - shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu( - ).numpy() - shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur, (15, 15)) - shape_offset_uv_blur = torch.from_numpy( - shape_offset_uv_blur).float().to(self.device)[None, ...] - value_1 = shape_offset_uv_blur * self.edge_mask[None, ..., None] - value_2 = self.shape_offset_uv * ( - 1 - self.edge_mask[None, ..., None]) - self.shape_offset_uv = value_1 + value_2 - - self.shape_offset_uv = self.shape_offset_uv * self.fusion_mask[None, - ..., - None] - - if global_blur and blur_size > 0: - shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu( - ).numpy() - shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur, - (blur_size, blur_size)) - shape_offset_uv_blur = torch.from_numpy( - shape_offset_uv_blur).float().to(self.device)[None, ...] - self.shape_offset_uv = shape_offset_uv_blur - - def get_fusion_mask(self): - - h, w = self.shape_offset_uv.shape[1:3] - self.fusion_mask = torch.zeros((h, w)).to(self.device).float() - UVs_coords = self.nonlinear_UVs.clone()[:35709] - UVs_coords[:, 0] *= w - UVs_coords[:, 1] *= h - UVs_coords_int = torch.floor(UVs_coords) - UVs_coords_int = UVs_coords_int.long() - - self.fusion_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:, - 0]] = 1 - - # blur mask - self.fusion_mask = self.fusion_mask.cpu().numpy() - new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) - new_kernel2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8)) - self.fusion_mask = cv2.dilate(self.fusion_mask, new_kernel1, 1) - self.fusion_mask = cv2.erode(self.fusion_mask, new_kernel2, 1) - self.fusion_mask = cv2.blur(self.fusion_mask, (17, 17)) - self.fusion_mask = torch.from_numpy(self.fusion_mask).float().to( + Parameters: + input: a dictionary that contains the data itself and its metadata information. + """ + self.input_img = input['input_img'].to(self.device) + self.input_img_for_tex = input['input_img_for_tex'].to(self.device) + self.input_img_hd = input['input_img_hd'].to(self.device) + self.face_mask = input['face_mask'].to(self.device) + self.gt_lm = input['gt_lm'].to(self.device) + self.coeffs = input['coeffs'].to(self.device) + self.position_map = input['position_map'].to(self.device) + self.texture_map = input['texture_map'].to(self.device) + self.tex_valid_mask = input['tex_valid_mask'].to(self.device) + self.de_retouched_albedo_map = input['de_retouched_albedo_map'].to( self.device) - def get_edge_mask(self): - - h, w = self.shape_offset_uv.shape[1:3] - self.edge_mask = torch.zeros((h, w)).to(self.device).float() - UVs_coords = self.nonlinear_UVs.clone()[self.edge_points_inds] - UVs_coords[:, 0] *= w - UVs_coords[:, 1] *= h - UVs_coords_int = torch.floor(UVs_coords) - UVs_coords_int = UVs_coords_int.long() - - self.edge_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:, 0]] = 1 - - # blur mask - self.edge_mask = self.edge_mask.cpu().numpy() - new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8)) - self.edge_mask = cv2.dilate(self.edge_mask, new_kernel1, 1) - self.edge_mask = cv2.blur(self.edge_mask, (5, 5)) - self.edge_mask = torch.from_numpy(self.edge_mask).float().to( - self.device) - - def fitting_nonlinear(self, coeff, debug=False, n_iters=100, out_dir=None): - output_coeff = coeff.detach().clone() - - output_coeff = self.facemodel_front.split_coeff(output_coeff) - output_coeff['id'].requires_grad = True - output_coeff['exp'].requires_grad = True - output_coeff['tex'].requires_grad = True - output_coeff['angle'].requires_grad = True - output_coeff['gamma'].requires_grad = True - output_coeff['trans'].requires_grad = True - - self.shape_offset_uv = torch.zeros( - (1, 300, 300, 3), - dtype=torch.float32).to(self.device) # (1, 180, 256, 3) - self.shape_offset_uv.requires_grad = True - - self.texture_offset_uv = torch.zeros( - (1, 300, 300, 3), - dtype=torch.float32).to(self.device) # (1, 180, 256, 3) - self.texture_offset_uv.requires_grad = True - - value_list = [ - self.shape_offset_uv, self.texture_offset_uv, output_coeff['id'], - output_coeff['exp'], output_coeff['tex'], output_coeff['angle'], - output_coeff['gamma'], output_coeff['trans'] - ] - optim = torch.optim.Adam(value_list, lr=1e-3) - - self.get_edge_points_horizontal() - self.get_edge_points_vertical() - - self.cur_iter = 0 - for i in range(n_iters): # 500 - self.pred_vertex, _, self.pred_color, self.pred_lm, _, face_shape_offset, self.verts_proj = \ - self.facemodel_front.compute_for_render_train_nonlinear(output_coeff, self.shape_offset_uv, - self.texture_offset_uv, - self.nonlinear_UVs[:35709, ...]) - self.pred_mask, _, self.pred_face, self.occ = self.renderer_fitting( - self.pred_vertex, - self.facemodel_front.face_buf, - feat=self.pred_color) - - self.pred_coeffs_dict = self.facemodel_front.split_coeff( - output_coeff) - self.compute_losses_fitting() - if debug and i % 10 == 0: - print('{}: total loss: {:.6f}'.format(i, self.loss_all.item())) - - optim.zero_grad() - self.loss_all.backward() - optim.step() - - self.cur_iter += 1 - - output_coeff = self.facemodel_front.merge_coeff(output_coeff) - - self.get_edge_mask() - self.get_fusion_mask() - self.blur_shape_offset_uv() - - self.pred_vertex, _, self.pred_color, self.pred_lm, _, face_shape_offset, self.verts_proj = \ - self.facemodel_front.compute_for_render_train_nonlinear(output_coeff, self.shape_offset_uv, - self.texture_offset_uv, - self.nonlinear_UVs[:35709, ...]) - - if out_dir is not None: - input_img_numpy = 255. * (self.input_img).detach().cpu().permute( - 0, 2, 3, 1).numpy() - input_img_numpy = np.squeeze(input_img_numpy) - - output_vis = self.pred_face - output_vis_numpy_raw = 255. * output_vis.detach().cpu().permute( - 0, 2, 3, 1).numpy() - output_vis_numpy_raw = np.squeeze(output_vis_numpy_raw) - - output_vis_numpy = np.concatenate( - (input_img_numpy, output_vis_numpy_raw), axis=-2) - - output_vis = np.squeeze(output_vis_numpy) - output_vis = output_vis[..., ::-1] # rgb->bgr - output_face_mask = self.pred_mask.detach().cpu().permute( - 0, 2, 3, 1).squeeze().numpy() * 255.0 - output_vis = np.column_stack( - (output_vis, cv2.cvtColor(output_face_mask, - cv2.COLOR_GRAY2BGR))) - output_input_vis = output_vis[:, :224] - output_pred_vis = output_vis[:, 224:448] - output_mask_vis = output_vis[:, 448:] - - face_mask_vis = 255. * self.face_mask.detach().cpu()[0, 0].numpy() - - shape_offset_vis = self.shape_offset_uv.detach().cpu().numpy()[0] - shape_offset_vis = (shape_offset_vis - shape_offset_vis.min()) / ( - shape_offset_vis.max() - shape_offset_vis.min()) * 255.0 - - cv2.imwrite( - os.path.join(out_dir, 'fitting_01_input.jpg'), - output_input_vis) - cv2.imwrite( - os.path.join(out_dir, 'fitting_02_pred.jpg'), output_pred_vis) - cv2.imwrite( - os.path.join(out_dir, 'fitting_03_mask.jpg'), output_mask_vis) - cv2.imwrite( - os.path.join(out_dir, 'fitting_04_facemask.jpg'), - face_mask_vis) - cv2.imwrite( - os.path.join(out_dir, 'fitting_05_shape_offset.jpg'), - shape_offset_vis) - - recon_shape_offset = face_shape_offset - recon_shape_offset[..., -1] = 10 - recon_shape_offset[ - ..., -1] # from camera space to world space - recon_shape_offset = recon_shape_offset.detach().cpu().numpy()[0] - - tri = self.facemodel_front.face_buf.cpu().numpy() - pred_color = self.pred_color.detach().cpu().numpy()[0].clip(0, 1) - - output = { - 'coeffs': output_coeff, - 'face_vertices': recon_shape_offset, - 'face_faces': tri + 1, - 'face_colors': pred_color - } - return output - - def forward(self, out_dir=None): - self.facemodel.to(self.device) - self.facemodel_front.to(self.device) + def predict_results_base(self): + # predict low-frequency coefficients with torch.no_grad(): - output_coeff = self.net_recon(self.input_img) - with torch.enable_grad(): - output = self.fitting_nonlinear( - output_coeff, debug=True, out_dir=out_dir) + # 3DMM + face_vertex, face_albedo_map, face_color_map, landmark, face_vertex_noTrans, position_map = \ + self.facemodel_front.compute_for_render(output_coeff) - output_coeff = output['coeffs'] - output_coeff = self.facemodel.split_coeff(output_coeff) - eye_coeffs = output_coeff['exp'][0, 16] + output_coeff['exp'][ - 0, 17] + output_coeff['exp'][0, 19] - if eye_coeffs > 1.0: - degree = 0.5 - else: - degree = 1.0 - output_coeff['exp'][0, 16] += 1 * degree - output_coeff['exp'][0, 17] += 1 * degree - output_coeff['exp'][0, 19] += 1.5 * degree - output_coeff = self.facemodel.merge_coeff(output_coeff) + # get texture map + texture_map = self.facemodel_front.get_texture_map( + face_vertex, self.input_img_hd) - self.pred_vertex, face_shape_ori, head_shape = \ - self.facemodel.compute_for_render_nonlinear_full(output_coeff, self.shape_offset_uv.detach(), - self.nonlinear_UVs, nose_coeff=0.1) + # de-retouch + texture_map_input_high = texture_map.permute( + 0, 3, 1, 2).detach() # (1, 3, 256, 256) + texture_map_input_high = (texture_map_input_high - 0.5) * 2 + de_retouched_face_albedo_map = self.de_retouching_module.run( + face_albedo_map, texture_map_input_high) - UVs_tensor = torch.tensor( - self.template_mesh['uvs'], - dtype=torch.float32)[None, ...].to(self.pred_vertex.device) - target_img = self.input_fat_img_hd.permute(0, 2, 3, 1) - with torch.enable_grad(): - _, _, _, texture_map, _ = self.renderer.pred_shape_and_texture( - self.pred_vertex, self.facemodel.face_buf, UVs_tensor, - target_img) + # get valid texture mask to deal with occlusion + valid_mask = self.facemodel_front.get_texture_map( + face_vertex, self.face_mask) # (256, 256, 1) + valid_mask = valid_mask.permute(0, 3, 1, + 2).detach() # (1, 1, 256, 256) - recon_shape = head_shape - recon_shape[ - ..., - -1] = 10 - recon_shape[..., -1] # from camera space to world space - recon_shape = recon_shape.cpu().numpy()[0] - tri = self.facemodel.face_buf.cpu().numpy() - normals = utils.estimate_normals(recon_shape, tri) + # render + pred_mask, _, pred_face = self.renderer.render_uv_texture( + face_vertex, self.facemodel_front.face_buf, self.bfm_UVs.clone(), + face_color_map) - output['head_vertices'] = recon_shape - output['head_faces'] = tri + 1 - output['head_tex_map'] = texture_map - output['head_UVs'] = self.template_mesh['uvs'] - output['head_faces_uv'] = self.template_mesh['faces_uv'] - output['head_normals'] = normals + input_img_numpy = 255. * (self.input_img).detach().cpu().permute( + 0, 2, 3, 1).numpy() + input_img_numpy = np.squeeze(input_img_numpy) + output_vis = pred_face * pred_mask + (1 - pred_mask) * self.input_img + output_vis_numpy_raw = 255. * output_vis.detach().cpu().permute( + 0, 2, 3, 1).numpy() + output_vis_numpy_raw = np.squeeze(output_vis_numpy_raw) + output_vis_numpy = np.concatenate( + (input_img_numpy, output_vis_numpy_raw), axis=-2) + output_vis = np.squeeze(output_vis_numpy) + output_vis = output_vis[..., ::-1] # rgb->bgr + output_face_mask = pred_mask.detach().cpu().permute( + 0, 2, 3, 1).squeeze().numpy() * 255.0 + output_vis = np.column_stack( + (output_vis, cv2.cvtColor(output_face_mask, cv2.COLOR_GRAY2BGR))) + output_input_vis = output_vis[:, :224] + output_pred_vis = output_vis[:, 224:448] + + input_img_hd = 255. * (self.input_img_hd).detach().cpu().permute( + 0, 2, 3, 1).numpy()[..., ::-1] + input_img_hd = np.squeeze(input_img_hd) + + # from camera space to world space + recon_vertices = face_vertex # [1, 35709, 3] + recon_vertices[..., -1] = 10 - recon_vertices[..., -1] + + recon_shape = face_vertex_noTrans # [1, 35709, 3] + recon_shape[..., -1] = 10 - recon_shape[..., -1] + + tri = self.facemodel_front.face_buf + + # output + output = {} + output['coeffs'] = output_coeff.detach() # [B, 257] + output['vertices'] = recon_vertices.detach() # [B, 35709, 3] + output['vertices_noTrans'] = recon_shape.detach() # [B, 35709, 3] + output['triangles'] = tri.detach() # [n_faces, 3], start from 0 + output['UVs'] = self.bfm_UVs.detach() # [35709, 2] + output['texture_map'] = texture_map.detach() # [B, h, w, 3], RGB + output['albedo_map'] = face_albedo_map.detach() # [B, 3, h, w] + output['color_map'] = face_color_map.detach() # [B, 3, h, w] + output['position_map'] = position_map.detach() # [B, 3, h, w] + + output['input_face'] = output_input_vis + output['pred_face'] = output_pred_vis + output['input_face_hd'] = input_img_hd + + output['input_img'] = self.input_img + output['input_img_hd'] = self.input_img_hd + output['gt_lm'] = self.gt_lm + output['face_mask'] = self.face_mask + output['tex_valid_mask'] = valid_mask + + output[ + 'de_retouched_albedo_map'] = de_retouched_face_albedo_map.detach() return output - def compute_losses_fitting(self): - face_mask = self.pred_mask + def forward(self, visualize=False): + self.bfm_UVs = self.bfm_UVs - face_mask = face_mask.detach() - self.loss_color = self.w_color * self.comupte_color_loss( - self.pred_face, self.input_img, face_mask) # 1.0 + # get valid mask to deal with occlusion + tex_valid_mask = self.tex_valid_mask # (B, 1, 256, 256) + if visualize: + tex_valid_mask = self.smooth_valid_mask(tex_valid_mask) + tex_valid_mask_mid = torch.nn.functional.interpolate( + tex_valid_mask, (64, 64), mode='bilinear') - self.loss_color_nose = torch.tensor(0.0).float().to(self.device) + # mid frequency + texture_map_input = self.texture_map.permute(0, 3, 1, + 2).to(self.device) + texture_map_input = (texture_map_input - 0.5) * 2 + texture_map_input_mid = torch.nn.functional.interpolate( + texture_map_input, (64, 64), mode='bilinear') + position_map_input_mid = torch.nn.functional.interpolate( + self.position_map, (64, 64), mode='bilinear') + input_mid = torch.cat([position_map_input_mid, texture_map_input_mid], + dim=1) + self.deformation_map = self.mid_net( + input_mid) * 0.1 * self.alpha # ori * 0.1 * self.alpha + self.deformation_map = self.deformation_map * tex_valid_mask_mid + self.deformation_map = self.deformation_map.permute(0, 2, 3, 1) - loss_reg, loss_gamma = self.compute_reg_loss(self.pred_coeffs_dict, - self.w_id, self.w_exp, - self.w_tex) - self.loss_reg = self.w_reg * loss_reg # 1.0 - self.loss_gamma = self.w_gamma * loss_gamma # 1.0 + # render + self.pred_vertex, self.pred_color, self.pred_lm, self.verts_proj, self.face_albedo_map, \ + face_shape_transformed, face_norm_roted, self.extra_results = \ + self.facemodel_front.compute_for_render_hierarchical_mid(self.coeffs, self.deformation_map, self.bfm_UVs, + visualize=visualize, + de_retouched_albedo_map= + self.de_retouched_albedo_map) + self.pred_mask, _, self.pred_face_mid = self.renderer.render_uv_texture( + self.pred_vertex, self.facemodel_front.face_buf, + self.bfm_UVs.clone(), self.pred_color) + self.deformation_map = self.deformation_map.permute(0, 3, 1, 2) - self.loss_lm = self.w_lm * self.compute_lm_loss( - self.pred_lm, self.gt_lm) * 0.1 # 0.1 + # get re-aligned texture + texture_map_input_high = self.facemodel_front.get_texture_map( + self.pred_vertex, self.input_img_hd) # (1, 256, 256, 3) + texture_map_input_high = texture_map_input_high.permute( + 0, 3, 1, 2).detach() # (1, 3, 256, 256) + texture_map_input_high = (texture_map_input_high - 0.5) * 2 - self.loss_smooth_offset = TVLoss()(self.shape_offset_uv.permute( - 0, 3, 1, 2)) * 10000 # 10000 + # high frequency + position_map_input_high = torch.nn.functional.interpolate( + self.position_map, (256, 256), mode='bilinear') + deformation_map_input_high = torch.nn.functional.interpolate( + self.deformation_map, (256, 256), mode='bilinear') + value = [ + position_map_input_high, texture_map_input_high, + deformation_map_input_high + ] + input_high = torch.cat(value, dim=1) + self.displacement_map = self.high_net( + input_high) * 0.1 * self.beta # ori * 0.1 * self.alpha + self.displacement_map = self.displacement_map * tex_valid_mask - self.loss_reg_offset = torch.tensor(0.0).float().to(self.device) + # render + self.pred_color_high, self.extra_results = self.facemodel_front.compute_for_render_hierarchical_high( + self.coeffs, + self.displacement_map, + self.de_retouched_albedo_map, + face_shape_transformed, + face_norm_roted, + extra_results=self.extra_results) + _, _, self.pred_face_high = self.renderer.render_uv_texture( + self.pred_vertex, self.facemodel_front.face_buf, + self.bfm_UVs.clone(), self.pred_color_high) - self.loss_reg_textureOff = torch.mean( - torch.abs(self.texture_offset_uv)) * 10 # 10 + self.pred_coeffs_dict = self.facemodel_front.split_coeff(self.coeffs) - self.loss_smooth_offset_std = TVLoss_std()( - self.shape_offset_uv.permute(0, 3, 1, 2)) * 50000 # 50000 + if visualize: + # high + self.extra_results['pred_mask_high'] = self.pred_mask + self.extra_results['pred_face_high_color'] = self.pred_face_high + _, _, self.extra_results[ + 'pred_face_high_gray'] = self.renderer.render_uv_texture( + self.pred_vertex, self.facemodel_front.face_buf, + self.bfm_UVs.clone(), self.extra_results['tex_high_gray']) - self.loss_points_horizontal, self.edge_points_inds = points_loss_horizontal( - self.verts_proj, self.left_points, self.right_points) # 20 - self.loss_points_horizontal *= 20 - self.loss_points_horizontal_jaw = torch.tensor(0.0).float().to( - self.device) - self.loss_points_vertical = torch.tensor(0.0).float().to(self.device) - self.loss_normals = torch.tensor(0.0).float().to(self.device) + # mid + self.extra_results['pred_mask_mid'] = self.pred_mask + self.extra_results['pred_face_mid_color'] = self.pred_face_mid + _, _, self.extra_results[ + 'pred_face_mid_gray'] = self.renderer.render_uv_texture( + self.pred_vertex, self.facemodel_front.face_buf, + self.bfm_UVs.clone(), self.extra_results['tex_mid_gray']) - self.loss_all = self.loss_color + self.loss_lm + self.loss_reg + self.loss_gamma + self.loss_smooth_offset - self.loss_all += self.loss_reg_offset + self.loss_smooth_offset_std + self.loss_points_horizontal - self.loss_all += self.loss_points_vertical + self.loss_reg_textureOff - self.loss_all += self.loss_color_nose + self.loss_normals + self.loss_points_horizontal_jaw + # base + self.extra_results['pred_mask_base'], _, self.extra_results[ + 'pred_face_base_color'] = self.renderer.render_uv_texture( + self.extra_results['pred_vertex_base'], + self.facemodel_front.face_buf, self.bfm_UVs.clone(), + self.extra_results['tex_base_color']) + _, _, self.extra_results[ + 'pred_face_base_gray'] = self.renderer.render_uv_texture( + self.extra_results['pred_vertex_base'], + self.facemodel_front.face_buf, self.bfm_UVs.clone(), + self.extra_results['tex_base_gray']) - self.loss_mask = torch.tensor(0.0).float().to(self.device) + # fit texture + with torch.enable_grad(): + texture_offset = torch.zeros( + (1, 3, 256, 256), dtype=torch.float32).to(self.device) + texture_offset.requires_grad = True + + optim = torch.optim.Adam([texture_offset], lr=1e-2) + + for i in range(self.tex_iters): + pred_color_high = self.pred_color_high.detach( + ) + texture_offset + _, _, pred_face_high = self.renderer.render_uv_texture( + self.pred_vertex.detach(), + self.facemodel_front.face_buf, self.bfm_UVs.clone(), + pred_color_high) + + loss_color_high = self.w_color * self.comupte_color_loss( + pred_face_high, self.input_img_for_tex, + self.pred_mask.detach()) + loss_smooth = TVLoss()(texture_offset) * self.w_tex_smooth + loss_all = loss_color_high + loss_smooth + optim.zero_grad() + loss_all.backward() + optim.step() + + self.pred_color_high = (self.pred_color_high + + texture_offset).detach() + + # for video + self.extra_results['pred_face_high_gray_list'] = [] + self.extra_results['pred_face_high_color_list'] = [] + for i in range(len(self.extra_results['tex_high_gray_list'])): + _, _, pred_face_high_gray_i = self.renderer_high_res.render_uv_texture( + self.extra_results['face_vertex_list'][i], + self.facemodel_front.face_buf, self.bfm_UVs.clone(), + self.extra_results['tex_high_gray_list'][i]) + self.extra_results['pred_face_high_gray_list'].append( + pred_face_high_gray_i) + + _, _, pred_face_high_color_i = self.renderer_high_res.render_uv_texture( + self.extra_results['face_vertex_list'][i], + self.facemodel_front.face_buf, self.bfm_UVs.clone(), + self.pred_color_high) + self.extra_results['pred_face_high_color_list'].append( + pred_face_high_color_i) + + def get_edge_points_horizontal(self): + left_points_list = [] + right_points_list = [] + for k in range(self.face_mask.shape[0]): + left_points = [] + right_points = [] + for i in range(self.face_mask.shape[2]): + inds = torch.where(self.face_mask[k, 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) + left_points_list.append( + torch.tensor(left_points).long().to(self.device)) + right_points_list.append( + torch.tensor(right_points).long().to(self.device)) + self.left_points = torch.stack( + left_points_list, dim=0).long().to(self.device) + self.right_points = torch.stack( + right_points_list, dim=0).long().to(self.device) + + def smooth_valid_mask(self, tex_valid_mask): + """ + + :param tex_valid_mask: torch.tensor, (B, 1, 256, 256), value: 0~1 + :return: + """ + batch_size = tex_valid_mask.shape[0] + tex_valid_mask_ = tex_valid_mask.detach().cpu().numpy() + result_list = [] + for i in range(batch_size): + mask = tex_valid_mask_[i, 0] + mask = cv2.erode(mask, np.ones(shape=(3, 3), dtype=np.float32)) + mask = cv2.blur(mask, (11, 11), 0) + result_list.append( + torch.from_numpy(mask)[None].float().to(tex_valid_mask.device)) + smoothed_mask = torch.stack(result_list, dim=0) + return smoothed_mask + + def compute_visuals_hrn(self): + with torch.no_grad(): + + input_img_vis = 255. * self.input_img.detach().cpu().permute( + 0, 2, 3, 1).numpy() + output_vis_mid = self.pred_face_mid * self.pred_mask + ( + 1 - self.pred_mask) * self.input_img + output_vis_mid = 255. * output_vis_mid.detach().cpu().permute( + 0, 2, 3, 1).numpy() + output_vis_high = self.pred_face_high * self.pred_mask + ( + 1 - self.pred_mask) * self.input_img + output_vis_high = 255. * output_vis_high.detach().cpu().permute( + 0, 2, 3, 1).numpy() + + deformation_map_vis = torch.nn.functional.interpolate( + self.deformation_map, + input_img_vis.shape[1:3], + mode='bilinear').permute(0, 2, 3, 1) + deformation_map_vis = ( + deformation_map_vis - deformation_map_vis.min()) / ( + deformation_map_vis.max() - deformation_map_vis.min()) + deformation_map_vis = 255. * deformation_map_vis.detach().cpu( + ).numpy() + + displacement_map_vis = torch.nn.functional.interpolate( + self.displacement_map, + input_img_vis.shape[1:3], + mode='bilinear').permute(0, 2, 3, 1) + displacement_vis = ( + displacement_map_vis - displacement_map_vis.min()) / ( + displacement_map_vis.max() - displacement_map_vis.min()) + displacement_vis = 255. * displacement_vis.detach().cpu().numpy() + displacement_vis = np.concatenate( + [displacement_vis, displacement_vis, displacement_vis], + axis=-1) + + face_albedo_map_vis = torch.nn.functional.interpolate( + self.face_albedo_map, + input_img_vis.shape[1:3], + mode='bilinear').permute(0, 2, 3, 1) + face_albedo_map_vis = 255. * face_albedo_map_vis.detach().cpu( + ).numpy() + + de_retouched_face_albedo_map_vis = torch.nn.functional.interpolate( + self.de_retouched_albedo_map, + input_img_vis.shape[1:3], + mode='bilinear').permute(0, 2, 3, 1) + de_retouched_face_albedo_map_vis = 255. * de_retouched_face_albedo_map_vis.detach( + ).cpu().numpy() + + if self.extra_results is not None: + pred_mask_base = self.extra_results['pred_mask_base'] + output_vis_base = self.extra_results[ + 'pred_face_base_color'] * pred_mask_base + ( + 1 - pred_mask_base) * self.input_img + output_vis_base = 255. * output_vis_base.detach().cpu( + ).permute(0, 2, 3, 1).numpy() + + output_vis_base_gray = self.extra_results[ + 'pred_face_base_gray'] * pred_mask_base + ( + 1 - pred_mask_base) * self.input_img + output_vis_base_gray = 255. * output_vis_base_gray.detach( + ).cpu().permute(0, 2, 3, 1).numpy() + output_vis_mid_gray = self.extra_results[ + 'pred_face_mid_gray'] * self.pred_mask + ( + 1 - self.pred_mask) * self.input_img + output_vis_mid_gray = 255. * output_vis_mid_gray.detach().cpu( + ).permute(0, 2, 3, 1).numpy() + output_vis_high_gray = self.extra_results[ + 'pred_face_high_gray'] * self.pred_mask + ( + 1 - self.pred_mask) * self.input_img + output_vis_high_gray = 255. * output_vis_high_gray.detach( + ).cpu().permute(0, 2, 3, 1).numpy() + + output_vis_numpy = np.concatenate( + (input_img_vis, output_vis_high, output_vis_base_gray, + output_vis_mid_gray, output_vis_high_gray, + deformation_map_vis, displacement_vis, + face_albedo_map_vis, de_retouched_face_albedo_map_vis), + axis=-2) + + elif self.gt_lm is not None: + gt_lm_numpy = self.gt_lm.cpu().numpy() + pred_lm_numpy = self.pred_lm.detach().cpu().numpy() + output_vis_high_lm = utils.draw_landmarks( + output_vis_high, gt_lm_numpy, 'b') + output_vis_high_lm = utils.draw_landmarks( + output_vis_high_lm, pred_lm_numpy, 'r') + + output_vis_numpy = np.concatenate( + (input_img_vis, output_vis_mid, output_vis_high, + output_vis_high_lm, deformation_map_vis, + displacement_vis), + axis=-2) + else: + output_vis_numpy = np.concatenate( + (input_img_vis, output_vis_mid, output_vis_high, + deformation_map_vis, displacement_vis), + axis=-2) + + self.output_vis = torch.tensor( + output_vis_numpy / 255., + dtype=torch.float32).permute(0, 3, 1, 2).to(self.device) + + def get_current_visuals(self): + """Return visualization images. train.py will display these images with visdom, and save the images to a HTML""" + visual_ret = OrderedDict() + for name in self.visual_names: + if isinstance(name, str): + visual_ret[name] = getattr(self, name)[:, :3, ...] + return visual_ret + + def save_results_hrn(self): + self.compute_visuals_hrn() + results = self.get_current_visuals() + + hrn_output_vis_batch = (255.0 * results['output_vis']).permute( + 0, 2, 3, 1).detach().cpu().numpy()[..., ::-1] + + vertices_batch = self.pred_vertex.detach( + ) # get reconstructed shape, [1, 35709, 3] + vertices_batch[..., -1] = 10 - vertices_batch[ + ..., -1] # from camera space to world space + vertices_batch = vertices_batch.cpu().numpy() + + texture_map_batch = (255.0 * self.pred_color_high).permute( + 0, 2, 3, 1).detach().cpu().numpy()[..., ::-1] + + output = {} + + output['vis_image'] = hrn_output_vis_batch[0] + + texture_map = texture_map_batch[0] + vertices = vertices_batch[0] + normals = utils.estimate_normals( + vertices, + self.facemodel_front.face_buf.cpu().numpy()) + face_mesh = { + 'vertices': vertices, + 'faces': self.facemodel_front.face_buf.cpu().numpy() + 1, + 'UVs': self.bfm_UVs.detach().cpu().numpy(), + 'faces_uv': self.facemodel_front.face_buf.cpu().numpy() + 1, + 'normals': normals, + 'faces_normal': self.facemodel_front.face_buf.cpu().numpy() + 1, + } + output['face_mesh'] = face_mesh + output['texture_map'] = texture_map + + frame_list = [] + if 'pred_face_high_gray_list' in self.extra_results: + input_img_vis = 255. * self.input_img_hd.detach().cpu().permute( + 0, 2, 3, 1).numpy()[0] + for j in range( + len(self.extra_results['pred_face_high_gray_list'])): + pred_face_high_gray_j = self.extra_results[ + 'pred_face_high_gray_list'][j][0, ...] + pred_face_high_gray_j = 255. * pred_face_high_gray_j.detach( + ).cpu().permute(1, 2, 0).numpy() + pred_face_high_color_j = self.extra_results[ + 'pred_face_high_color_list'][j][0, ...] + pred_face_high_color_j = 255. * pred_face_high_color_j.detach( + ).cpu().permute(1, 2, 0).numpy() + + value = [ + input_img_vis, pred_face_high_gray_j, + pred_face_high_color_j + ] + vis_j = np.concatenate(value, axis=1) + + frame_list.append(vis_j.clip(0, 255).astype(np.uint8)) + output['frame_list'] = frame_list + + return output diff --git a/modelscope/models/cv/face_reconstruction/models/losses.py b/modelscope/models/cv/face_reconstruction/models/losses.py index 7a73d61e..6d4af4e8 100644 --- a/modelscope/models/cv/face_reconstruction/models/losses.py +++ b/modelscope/models/cv/face_reconstruction/models/losses.py @@ -1,5 +1,4 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -import clip import numpy as np import torch import torch.nn as nn @@ -13,51 +12,6 @@ def resize_n_crop(image, M, dsize=112): return warp_affine(image, M, dsize=(dsize, dsize)) -class CLIPLoss(torch.nn.Module): - - def __init__(self): - super(CLIPLoss, self).__init__() - self.model, self.preprocess = clip.load('ViT-B/32', device='cuda') - - def forward(self, image, text): - similarity = 1 - self.model(image, text)[0] / 100 - return similarity - - -class CLIPLoss_relative(torch.nn.Module): - - def __init__(self): - super(CLIPLoss_relative, self).__init__() - self.model, self.preprocess = clip.load('ViT-B/32', device='cuda') - - def forward(self, image, text, image_ori, text_ori): - - image_features = self.model.encode_image(image) - text_features = self.model.encode_text(text) - - # normalized features - image_features = image_features / image_features.norm( - dim=1, keepdim=True) - text_features = text_features / text_features.norm(dim=1, keepdim=True) - - image_features_ori = self.model.encode_image(image_ori) - text_features_ori = self.model.encode_text(text_ori) - - # normalized features - image_features_ori = image_features_ori / image_features_ori.norm( - dim=1, keepdim=True) - text_features_ori = text_features_ori / text_features_ori.norm( - dim=1, keepdim=True) - - delta_image = image_features - image_features_ori - delta_text = text_features - text_features_ori - - loss = 1 - torch.sum(delta_image * delta_text) / ( - torch.norm(delta_image) * torch.norm(delta_text)) - - return loss - - # perceptual level loss class PerceptualLoss(nn.Module): diff --git a/modelscope/models/cv/face_reconstruction/models/nv_diffrast.py b/modelscope/models/cv/face_reconstruction/models/nv_diffrast.py index f17246e5..348c8112 100644 --- a/modelscope/models/cv/face_reconstruction/models/nv_diffrast.py +++ b/modelscope/models/cv/face_reconstruction/models/nv_diffrast.py @@ -4,6 +4,7 @@ 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 @@ -79,7 +80,10 @@ class MeshRenderer(nn.Module): vertex_ndc = vertex @ ndc_proj.t() if self.glctx is None: - self.glctx = dr.RasterizeCudaContext(device=device) + 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: @@ -171,7 +175,10 @@ class MeshRenderer(nn.Module): vertex_ndc = vertex @ ndc_proj.t() if self.glctx is None: - self.glctx = dr.RasterizeCudaContext(device=device) + 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: @@ -222,179 +229,6 @@ class MeshRenderer(nn.Module): img = img * torch.clamp(rast_out[..., -1:], 0, 1) # Mask out background. - tex_map = uv_texture[0].detach().cpu().numpy()[..., ::-1] * 255.0 - image = img.permute(0, 3, 1, 2) - return mask, depth, image, tex_map - - def pred_shape_and_texture(self, - vertex, - tri, - uv, - target_img, - base_tex=None): - """ - Return: - mask -- torch.tensor, size (B, 1, H, W) - depth -- torch.tensor, size (B, 1, H, W) - features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None - - Parameters: - vertex -- torch.tensor, size (B, N, 3) - tri -- torch.tensor, size (B, M, 3) or (M, 3), triangles - uv -- torch.tensor, size (B,N, 2), uv mapping - base_tex -- torch.tensor, size (B,H,W,C) - """ - vertex = torch.cat([vertex[:, :35241, :], vertex[:, 37082:, :]], - dim=1) # BFM front - tri = torch.cat([tri[:69732, :], tri[73936:, ]], dim=0) - uv = torch.cat([uv[:, :35241, :], uv[:, 37082:, :]], dim=1) - tri[69732:, :] = tri[69732:, :] - (37082 - 35241) - - device = vertex.device - rsize = int(self.rasterize_size) - ndc_proj = self.ndc_proj.to(device) - # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v - if vertex.shape[-1] == 3: - vertex = torch.cat( - [vertex, torch.ones([*vertex.shape[:2], 1]).to(device)], - dim=-1) - vertex[..., 1] = -vertex[..., 1] - - vertex_ndc = vertex @ ndc_proj.t() - if self.glctx is None: - self.glctx = dr.RasterizeCudaContext(device=device) - - ranges = None - if isinstance(tri, List) or len(tri.shape) == 3: - vum = vertex_ndc.shape[1] - fnum = torch.tensor([f.shape[0] - for f in tri]).unsqueeze(1).to(device) - - fstartidx = torch.cumsum(fnum, dim=0) - fnum - ranges = torch.cat([fstartidx, fnum], - axis=1).type(torch.int32).cpu() - for i in range(tri.shape[0]): - tri[i] = tri[i] + i * vum - vertex_ndc = torch.cat(vertex_ndc, dim=0) - tri = torch.cat(tri, dim=0) - - # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3] - tri = tri.type(torch.int32).contiguous() - rast_out, _ = dr.rasterize( - self.glctx, - vertex_ndc.contiguous(), - tri, - resolution=[rsize, rsize], - ranges=ranges) - - depth, _ = dr.interpolate( - vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(), - rast_out, tri) - depth = depth.permute(0, 3, 1, 2) - mask = (rast_out[..., 3] > 0).float().unsqueeze(1) - depth = mask * depth - uv[..., -1] = 1.0 - uv[..., -1] - - rast_out, rast_db = dr.rasterize( - self.glctx, - vertex_ndc.contiguous(), - tri, - resolution=[rsize, rsize], - ranges=ranges) - - interp_out, uv_da = dr.interpolate( - uv, rast_out, tri, rast_db, diff_attrs='all') - - mask_3c = mask.permute(0, 2, 3, 1) - mask_3c = torch.cat((mask_3c, mask_3c, mask_3c), dim=-1) - maskout_img = mask_3c * target_img - mean_color = torch.sum(maskout_img, dim=(1, 2)) - valid_pixel_count = torch.sum(mask) - - mean_color = mean_color / valid_pixel_count - - tex = torch.zeros((1, 128 * 5 // 4, 128, 3), dtype=torch.float32) - tex[:, :, :, 0] = mean_color[0, 0] - tex[:, :, :, 1] = mean_color[0, 1] - tex[:, :, :, 2] = mean_color[0, 2] - - tex = tex.cuda() - - tex_mask = torch.zeros((1, 2048 * 5 // 4, 2048, 3), - dtype=torch.float32) - tex_mask[:, :, :, 1] = 1.0 - tex_mask = tex_mask.cuda() - tex_mask.requires_grad = True - tex_mask = tex_mask.contiguous() - - criterionTV = TVLoss() - - if base_tex is not None: - base_tex = base_tex.cuda() - - for tex_resolution in [64, 128, 256, 512, 1024, 2048]: - tex = tex.detach() - tex = tex.permute(0, 3, 1, 2) - tex = F.interpolate(tex, (tex_resolution * 5 // 4, tex_resolution)) - tex = tex.permute(0, 2, 3, 1).contiguous() - - if base_tex is not None: - _base_tex = base_tex.permute(0, 3, 1, 2) - _base_tex = F.interpolate( - _base_tex, (tex_resolution * 5 // 4, tex_resolution)) - _base_tex = _base_tex.permute(0, 2, 3, 1).contiguous() - tex += _base_tex - - tex.requires_grad = True - optim = torch.optim.Adam([tex], lr=1e-2) - - texture_opt_iters = 100 - - if tex_resolution == 2048: - optim_mask = torch.optim.Adam([tex_mask], lr=1e-2) - - for i in range(int(texture_opt_iters)): - - if tex_resolution == 2048: - optim_mask.zero_grad() - rendered = dr.texture( - tex_mask, interp_out, filter_mode='linear') # , uv_da) - rendered = rendered * torch.clamp( - rast_out[..., -1:], 0, 1) # Mask out background. - tex_loss = torch.mean((target_img - rendered)**2) - - tex_loss.backward() - optim_mask.step() - - optim.zero_grad() - - img = dr.texture( - tex, interp_out, filter_mode='linear') # , uv_da) - img = img * torch.clamp(rast_out[..., -1:], 0, - 1) # Mask out background. - recon_loss = torch.mean((target_img - img)**2) - - if tex_resolution < 2048: - tv_loss = criterionTV(tex.permute(0, 3, 1, 2)) - - total_loss = recon_loss + tv_loss * 0.01 - else: - - total_loss = recon_loss - - total_loss.backward() - optim.step() - - tex_map = tex[0].detach().cpu().numpy()[..., ::-1] * 255.0 - - image = img.permute(0, 3, 1, 2) - - tex_mask = tex_mask[0].detach().cpu().numpy() * 255.0 - tex_mask = np.where(tex_mask[..., 1] > 250, 1.0, 0.0) * np.where( - tex_mask[..., 0] < 10, 1.0, 0) * np.where(tex_mask[..., 2] < 10, - 1.0, 0) - tex_mask = 1.0 - tex_mask - - return mask, depth, image, tex_map, tex_mask + return mask, depth, image diff --git a/modelscope/models/cv/face_reconstruction/models/opt.py b/modelscope/models/cv/face_reconstruction/models/opt.py index c979e64e..b8d97faf 100644 --- a/modelscope/models/cv/face_reconstruction/models/opt.py +++ b/modelscope/models/cv/face_reconstruction/models/opt.py @@ -11,3 +11,4 @@ use_ddp = False use_last_fc = False z_far = 15.0 z_near = 5.0 +lr = 0.001 diff --git a/modelscope/models/cv/face_reconstruction/models/pix2pix/__init__.py b/modelscope/models/cv/face_reconstruction/models/pix2pix/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/face_reconstruction/models/pix2pix/networks.py b/modelscope/models/cv/face_reconstruction/models/pix2pix/networks.py new file mode 100644 index 00000000..c18881ed --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/pix2pix/networks.py @@ -0,0 +1,848 @@ +# Part of the implementation is borrowed and modified from pix2pix, +# publicly available at https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix +import functools + +import torch +import torch.nn as nn +from torch.nn import init +from torch.optim import lr_scheduler + +############################################################################### +# Helper Functions +############################################################################### + + +class Identity(nn.Module): + + def forward(self, x): + return x + + +def get_norm_layer(norm_type='instance'): + """Return a normalization layer + + Parameters: + norm_type (str) -- the name of the normalization layer: batch | instance | none + + For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). + For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. + """ + if norm_type == 'batch': + norm_layer = functools.partial( + nn.BatchNorm2d, affine=True, track_running_stats=True) + elif norm_type == 'instance': + norm_layer = functools.partial( + nn.InstanceNorm2d, affine=False, track_running_stats=False) + elif norm_type == 'none': + + def norm_layer(x): + return Identity() + else: + raise NotImplementedError('normalization layer [%s] is not found' + % norm_type) + return norm_layer + + +def get_scheduler(optimizer, opt): + """Return a learning rate scheduler + + Parameters: + optimizer -- the optimizer of the network + opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　 + opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine + + For 'linear', we keep the same learning rate for the first epochs + and linearly decay the rate to zero over the next epochs. + For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. + See https://pytorch.org/docs/stable/optim.html for more details. + """ + if opt.lr_policy == 'linear': + + def lambda_rule(epoch): + lr_l = 1.0 - max(0, epoch + opt.epoch_count + - opt.n_epochs) / float(opt.n_epochs_decay + 1) + return lr_l + + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) + elif opt.lr_policy == 'step': + scheduler = lr_scheduler.StepLR( + optimizer, step_size=opt.lr_decay_iters, gamma=0.1) + elif opt.lr_policy == 'plateau': + scheduler = lr_scheduler.ReduceLROnPlateau( + optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) + elif opt.lr_policy == 'cosine': + scheduler = lr_scheduler.CosineAnnealingLR( + optimizer, T_max=opt.n_epochs, eta_min=0) + else: + return NotImplementedError( + 'learning rate policy [%s] is not implemented', opt.lr_policy) + return scheduler + + +def init_weights(net, init_type='normal', init_gain=0.02): + """Initialize network weights. + + Parameters: + net (network) -- network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + + We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might + work better for some applications. Feel free to try yourself. + """ + + def init_func(m): # define the initialization function + classname = m.__class__.__name__ + if hasattr(m, 'weight') and (classname.find('Conv') != -1 + or classname.find('Linear') != -1): + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, init_gain) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=init_gain) + elif init_type == 'kaiming': + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif init_type == 'orthogonal': + init.orthogonal_(m.weight.data, gain=init_gain) + else: + raise NotImplementedError( + 'initialization method [%s] is not implemented' + % init_type) + if hasattr(m, 'bias') and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif classname.find( + 'BatchNorm2d' + ) != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. + init.normal_(m.weight.data, 1.0, init_gain) + init.constant_(m.bias.data, 0.0) + + print('initialize network with %s' % init_type) + net.apply(init_func) # apply the initialization function + + +def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): + """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights + Parameters: + net (network) -- the network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Return an initialized network. + """ + if len(gpu_ids) > 0: + assert (torch.cuda.is_available()) + net.to(gpu_ids[0]) + net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs + init_weights(net, init_type, init_gain=init_gain) + return net + + +def define_G(input_nc, + output_nc, + ngf, + netG, + norm='batch', + use_dropout=False, + init_type='normal', + init_gain=0.02, + gpu_ids=[]): + """Create a generator + + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + ngf (int) -- the number of filters in the last conv layer + netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128 + norm (str) -- the name of normalization layers used in the network: batch | instance | none + use_dropout (bool) -- if use dropout layers. + init_type (str) -- the name of our initialization method. + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Returns a generator + + Our current implementation provides two types of generators: + U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images) + The original U-Net paper: https://arxiv.org/abs/1505.04597 + + Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks) + Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations. + We adapt Torch code from Justin Johnson's neural style transfer project + (https://github.com/jcjohnson/fast-neural-style). + + + The generator has been initialized by . It uses RELU for non-linearity. + """ + net = None + norm_layer = get_norm_layer(norm_type=norm) + + if netG == 'resnet_9blocks': + net = ResnetGenerator( + input_nc, + output_nc, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout, + n_blocks=9) + elif netG == 'resnet_6blocks': + net = ResnetGenerator( + input_nc, + output_nc, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout, + n_blocks=6) + elif netG == 'unet_32': + net = UnetGenerator( + input_nc, + output_nc, + 5, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout) + elif netG == 'unet_64': + net = UnetGenerator( + input_nc, + output_nc, + 6, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout) + elif netG == 'unet_128': + net = UnetGenerator( + input_nc, + output_nc, + 7, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout) + elif netG == 'unet_256': + net = UnetGenerator( + input_nc, + output_nc, + 8, + ngf, + norm_layer=norm_layer, + use_dropout=use_dropout) + else: + raise NotImplementedError('Generator model name [%s] is not recognized' + % netG) + return init_net(net, init_type, init_gain, gpu_ids) + + +def define_D(input_nc, + ndf, + netD, + n_layers_D=3, + norm='batch', + init_type='normal', + init_gain=0.02, + gpu_ids=[]): + """Create a discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the first conv layer + netD (str) -- the architecture's name: basic | n_layers | pixel + n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers' + norm (str) -- the type of normalization layers used in the network. + init_type (str) -- the name of the initialization method. + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Returns a discriminator + + Our current implementation provides three types of discriminators: + [basic]: 'PatchGAN' classifier described in the original pix2pix paper. + It can classify whether 70×70 overlapping patches are real or fake. + Such a patch-level discriminator architecture has fewer parameters + than a full-image discriminator and can work on arbitrarily-sized images + in a fully convolutional fashion. + + [n_layers]: With this mode, you can specify the number of conv layers in the discriminator + with the parameter (default=3 as used in [basic] (PatchGAN).) + + [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not. + It encourages greater color diversity but has no effect on spatial statistics. + + The discriminator has been initialized by . It uses Leakly RELU for non-linearity. + """ + net = None + norm_layer = get_norm_layer(norm_type=norm) + + if netD == 'basic': # default PatchGAN classifier + net = NLayerDiscriminator( + input_nc, ndf, n_layers=3, norm_layer=norm_layer) + elif netD == 'n_layers': # more options + net = NLayerDiscriminator( + input_nc, ndf, n_layers_D, norm_layer=norm_layer) + elif netD == 'pixel': # classify if each pixel is real or fake + net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer) + else: + raise NotImplementedError( + 'Discriminator model name [%s] is not recognized' % netD) + return init_net(net, init_type, init_gain, gpu_ids) + + +############################################################################## +# Classes +############################################################################## +class GANLoss(nn.Module): + """Define different GAN objectives. + + The GANLoss class abstracts away the need to create the target label tensor + that has the same size as the input. + """ + + def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0): + """ Initialize the GANLoss class. + + Parameters: + gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. + target_real_label (bool) - - label for a real image + target_fake_label (bool) - - label of a fake image + + Note: Do not use sigmoid as the last layer of Discriminator. + LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss. + """ + super(GANLoss, self).__init__() + self.register_buffer('real_label', torch.tensor(target_real_label)) + self.register_buffer('fake_label', torch.tensor(target_fake_label)) + self.gan_mode = gan_mode + if gan_mode == 'lsgan': + self.loss = nn.MSELoss() + elif gan_mode == 'vanilla': + self.loss = nn.BCEWithLogitsLoss() + elif gan_mode in ['wgangp']: + self.loss = None + else: + raise NotImplementedError('gan mode %s not implemented' % gan_mode) + + def get_target_tensor(self, prediction, target_is_real): + """Create label tensors with the same size as the input. + + Parameters: + prediction (tensor) - - tpyically the prediction from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + A label tensor filled with ground truth label, and with the size of the input + """ + + if target_is_real: + target_tensor = self.real_label + else: + target_tensor = self.fake_label + return target_tensor.expand_as(prediction) + + def __call__(self, prediction, target_is_real): + """Calculate loss given Discriminator's output and grount truth labels. + + Parameters: + prediction (tensor) - - tpyically the prediction output from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + the calculated loss. + """ + if self.gan_mode in ['lsgan', 'vanilla']: + target_tensor = self.get_target_tensor(prediction, target_is_real) + loss = self.loss(prediction, target_tensor) + elif self.gan_mode == 'wgangp': + if target_is_real: + loss = -prediction.mean() + else: + loss = prediction.mean() + return loss + + +def cal_gradient_penalty(netD, + real_data, + fake_data, + device, + type='mixed', + constant=1.0, + lambda_gp=10.0): + """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 + + Arguments: + netD (network) -- discriminator network + real_data (tensor array) -- real images + fake_data (tensor array) -- generated images from the generator + device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) + if self.gpu_ids else torch.device('cpu') + type (str) -- if we mix real and fake data or not [real | fake | mixed]. + constant (float) -- the constant used in formula ( ||gradient||_2 - constant)^2 + lambda_gp (float) -- weight for this loss + + Returns the gradient penalty loss + """ + if lambda_gp > 0.0: + if type == 'real': # either use real images, fake images, or a linear interpolation of two. + interpolatesv = real_data + elif type == 'fake': + interpolatesv = fake_data + elif type == 'mixed': + alpha = torch.rand(real_data.shape[0], 1, device=device) + alpha = alpha.expand( + real_data.shape[0], + real_data.nelement() + // real_data.shape[0]).contiguous().view(*real_data.shape) + interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) + else: + raise NotImplementedError('{} not implemented'.format(type)) + interpolatesv.requires_grad_(True) + disc_interpolates = netD(interpolatesv) + gradients = torch.autograd.grad( + outputs=disc_interpolates, + inputs=interpolatesv, + grad_outputs=torch.ones(disc_interpolates.size()).to(device), + create_graph=True, + retain_graph=True, + only_inputs=True) + gradients = gradients[0].view(real_data.size(0), -1) # flat the data + gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) + - constant)**2).mean() + gradient_penalty = gradient_penalty * lambda_gp + return gradient_penalty, gradients + else: + return 0.0, None + + +class ResnetGenerator(nn.Module): + """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. + + We adapt Torch code and idea from Justin Johnson's neural style transfer + project(https://github.com/jcjohnson/fast-neural-style) + """ + + def __init__(self, + input_nc, + output_nc, + ngf=64, + norm_layer=nn.BatchNorm2d, + use_dropout=False, + n_blocks=6, + padding_type='reflect'): + """Construct a Resnet-based generator + + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers + n_blocks (int) -- the number of ResNet blocks + padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero + """ + assert (n_blocks >= 0) + super(ResnetGenerator, self).__init__() + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + model = [ + nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), + norm_layer(ngf), + nn.ReLU(True) + ] + + n_downsampling = 2 + for i in range(n_downsampling): # add downsampling layers + mult = 2**i + model += [ + nn.Conv2d( + ngf * mult, + ngf * mult * 2, + kernel_size=3, + stride=2, + padding=1, + bias=use_bias), + norm_layer(ngf * mult * 2), + nn.ReLU(True) + ] + + mult = 2**n_downsampling + for i in range(n_blocks): # add ResNet blocks + + model += [ + ResnetBlock( + ngf * mult, + padding_type=padding_type, + norm_layer=norm_layer, + use_dropout=use_dropout, + use_bias=use_bias) + ] + + for i in range(n_downsampling): # add upsampling layers + mult = 2**(n_downsampling - i) + model += [ + nn.ConvTranspose2d( + ngf * mult, + int(ngf * mult / 2), + kernel_size=3, + stride=2, + padding=1, + output_padding=1, + bias=use_bias), + norm_layer(int(ngf * mult / 2)), + nn.ReLU(True) + ] + model += [nn.ReflectionPad2d(3)] + model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] + model += [nn.Tanh()] + + self.model = nn.Sequential(*model) + + def forward(self, input): + """Standard forward""" + return self.model(input) + + +class ResnetBlock(nn.Module): + """Define a Resnet block""" + + def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias): + """Initialize the Resnet block + + A resnet block is a conv block with skip connections + We construct a conv block with build_conv_block function, + and implement skip connections in function. + Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf + """ + super(ResnetBlock, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, + use_dropout, use_bias) + + def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, + use_bias): + """Construct a convolutional block. + + Parameters: + dim (int) -- the number of channels in the conv layer. + padding_type (str) -- the name of padding layer: reflect | replicate | zero + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers. + use_bias (bool) -- if the conv layer uses bias or not + + Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU)) + """ + conv_block = [] + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' + % padding_type) + + conv_block += [ + nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), + norm_layer(dim), + nn.ReLU(True) + ] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' + % padding_type) + conv_block += [ + nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), + norm_layer(dim) + ] + + return nn.Sequential(*conv_block) + + def forward(self, x): + """Forward function (with skip connections)""" + out = x + self.conv_block(x) # add skip connections + return out + + +class UnetGenerator(nn.Module): + """Create a Unet-based generator""" + + def __init__(self, + input_nc, + output_nc, + num_downs, + ngf=64, + norm_layer=nn.BatchNorm2d, + use_dropout=False): + """Construct a Unet generator + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7, + image of size 128x128 will become of size 1x1 # at the bottleneck + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + + We construct the U-Net from the innermost layer to the outermost layer. + It is a recursive process. + """ + super(UnetGenerator, self).__init__() + # construct unet structure + unet_block = UnetSkipConnectionBlock( + ngf * 8, + ngf * 8, + input_nc=None, + submodule=None, + norm_layer=norm_layer, + innermost=True) # add the innermost layer + for i in range(num_downs + - 5): # add intermediate layers with ngf * 8 filters + unet_block = UnetSkipConnectionBlock( + ngf * 8, + ngf * 8, + input_nc=None, + submodule=unet_block, + norm_layer=norm_layer, + use_dropout=use_dropout) + # gradually reduce the number of filters from ngf * 8 to ngf + unet_block = UnetSkipConnectionBlock( + ngf * 4, + ngf * 8, + input_nc=None, + submodule=unet_block, + norm_layer=norm_layer) + unet_block = UnetSkipConnectionBlock( + ngf * 2, + ngf * 4, + input_nc=None, + submodule=unet_block, + norm_layer=norm_layer) + unet_block = UnetSkipConnectionBlock( + ngf, + ngf * 2, + input_nc=None, + submodule=unet_block, + norm_layer=norm_layer) + self.model = UnetSkipConnectionBlock( + output_nc, + ngf, + input_nc=input_nc, + submodule=unet_block, + outermost=True, + norm_layer=norm_layer) # add the outermost layer + + def forward(self, input): + """Standard forward""" + return self.model(input) + + +class UnetSkipConnectionBlock(nn.Module): + """Defines the Unet submodule with skip connection. + X -------------------identity---------------------- + |-- downsampling -- |submodule| -- upsampling --| + """ + + def __init__(self, + outer_nc, + inner_nc, + input_nc=None, + submodule=None, + outermost=False, + innermost=False, + norm_layer=nn.BatchNorm2d, + use_dropout=False): + """Construct a Unet submodule with skip connections. + + Parameters: + outer_nc (int) -- the number of filters in the outer conv layer + inner_nc (int) -- the number of filters in the inner conv layer + input_nc (int) -- the number of channels in input images/features + submodule (UnetSkipConnectionBlock) -- previously defined submodules + outermost (bool) -- if this module is the outermost module + innermost (bool) -- if this module is the innermost module + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers. + """ + super(UnetSkipConnectionBlock, self).__init__() + self.outermost = outermost + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + if input_nc is None: + input_nc = outer_nc + downconv = nn.Conv2d( + input_nc, + inner_nc, + kernel_size=4, + stride=2, + padding=1, + bias=use_bias) + downrelu = nn.LeakyReLU(0.2, True) + downnorm = norm_layer(inner_nc) + uprelu = nn.ReLU(True) + upnorm = norm_layer(outer_nc) + + if outermost: + upconv = nn.ConvTranspose2d( + inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1) + down = [downconv] + up = [uprelu, upconv, nn.Tanh()] + model = down + [submodule] + up + elif innermost: + upconv = nn.ConvTranspose2d( + inner_nc, + outer_nc, + kernel_size=4, + stride=2, + padding=1, + bias=use_bias) + down = [downrelu, downconv] + up = [uprelu, upconv, upnorm] + model = down + up + else: + upconv = nn.ConvTranspose2d( + inner_nc * 2, + outer_nc, + kernel_size=4, + stride=2, + padding=1, + bias=use_bias) + down = [downrelu, downconv, downnorm] + up = [uprelu, upconv, upnorm] + + if use_dropout: + model = down + [submodule] + up + [nn.Dropout(0.5)] + else: + model = down + [submodule] + up + + self.model = nn.Sequential(*model) + + def forward(self, x): + if self.outermost: + return self.model(x) + else: # add skip connections + return torch.cat([x, self.model(x)], 1) + + +class NLayerDiscriminator(nn.Module): + """Defines a PatchGAN discriminator""" + + def __init__(self, + input_nc, + ndf=64, + n_layers=3, + norm_layer=nn.BatchNorm2d): + """Construct a PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + n_layers (int) -- the number of conv layers in the discriminator + norm_layer -- normalization layer + """ + super(NLayerDiscriminator, self).__init__() + if type( + norm_layer + ) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + kw = 4 + padw = 1 + sequence = [ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, True) + ] + nf_mult = 1 + nf_mult_prev = 1 + for n in range(1, + n_layers): # gradually increase the number of filters + nf_mult_prev = nf_mult + nf_mult = min(2**n, 8) + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=2, + padding=padw, + bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + nf_mult_prev = nf_mult + nf_mult = min(2**n_layers, 8) + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=1, + padding=padw, + bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + sequence += [ + nn.Conv2d( + ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw) + ] # output 1 channel prediction map + self.model = nn.Sequential(*sequence) + + def forward(self, input): + """Standard forward.""" + return self.model(input) + + +class PixelDiscriminator(nn.Module): + """Defines a 1x1 PatchGAN discriminator (pixelGAN)""" + + def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d): + """Construct a 1x1 PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + """ + super(PixelDiscriminator, self).__init__() + if type( + norm_layer + ) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + self.net = [ + nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0), + nn.LeakyReLU(0.2, True), + nn.Conv2d( + ndf, + ndf * 2, + kernel_size=1, + stride=1, + padding=0, + bias=use_bias), + norm_layer(ndf * 2), + nn.LeakyReLU(0.2, True), + nn.Conv2d( + ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias) + ] + + self.net = nn.Sequential(*self.net) + + def forward(self, input): + """Standard forward.""" + return self.net(input) diff --git a/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_model.py b/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_model.py new file mode 100644 index 00000000..54768fc1 --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_model.py @@ -0,0 +1,125 @@ +# Part of the implementation is borrowed and modified from pix2pix, +# publicly available at https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix +import torch +from torch import nn + +from . import networks + + +class Pix2PixModel(nn.Module): + """ This class implements the pix2pix model, for learning a mapping from input + images to output images given paired data. + + The model training requires '--dataset_mode aligned' dataset. + By default, it uses a '--netG unet256' U-Net generator, + a '--netD basic' discriminator (PatchGAN), + and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper). + + pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf + """ + + def __init__(self, opt): + """Initialize the pix2pix class. + + Parameters: + opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions + """ + super(Pix2PixModel, self).__init__() + # specify the training losses you want to print out. + # The training/test scripts will call + self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake'] + # specify the images you want to save/display. + # The training/test scripts will call + self.visual_names = ['real_A', 'fake_B', 'real_B'] + # specify the models you want to save to the disk. The training/test scripts will call + # and + if opt.isTrain: + self.model_names = ['G', 'D'] + else: # during test time, only load G + self.model_names = ['G'] + # define networks (both generator and discriminator) + self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, + opt.netG, opt.norm, not opt.no_dropout, + opt.init_type, opt.init_gain, + opt.gpu_ids) + # self.netG = UNet(opt.input_nc, opt.output_nc) + + # define a discriminator; conditional GANs need to take both input and output images; + # Therefore, #channels for D is input_nc + output_nc + if opt.isTrain: + self.netD = networks.define_D(opt.input_nc + opt.output_nc, + opt.ndf, opt.netD, opt.n_layers_D, + opt.norm, opt.init_type, + opt.init_gain, opt.gpu_ids) + + if opt.isTrain: + # define loss functions + self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) + self.criterionL1 = torch.nn.L1Loss() + # initialize optimizers; schedulers will be automatically created by function . + self.optimizer_G = torch.optim.Adam( + self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) + self.optimizer_D = torch.optim.Adam( + self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) + self.optimizers.append(self.optimizer_G) + self.optimizers.append(self.optimizer_D) + + def set_input(self, input): + """Unpack input data from the dataloader and perform necessary pre-processing steps. + + Parameters: + input (dict): include the data itself and its metadata information. + + The option 'direction' can be used to swap images in domain A and domain B. + """ + AtoB = self.opt.direction == 'AtoB' + self.real_A = input['A' if AtoB else 'B'].to(self.device) + self.real_B = input['B' if AtoB else 'A'].to(self.device) + self.image_paths = input['A_paths' if AtoB else 'B_paths'] + + def forward(self): + """Run forward pass; called by both functions and .""" + self.fake_B = self.netG(self.real_A) # G(A) + + def backward_D(self): + """Calculate GAN loss for the discriminator""" + # Fake; stop backprop to the generator by detaching fake_B + fake_AB = torch.cat( + (self.real_A, self.fake_B), 1 + ) # we use conditional GANs; we need to feed both input and output to the discriminator + pred_fake = self.netD(fake_AB.detach()) + self.loss_D_fake = self.criterionGAN(pred_fake, False) + # Real + real_AB = torch.cat((self.real_A, self.real_B), 1) + pred_real = self.netD(real_AB) + self.loss_D_real = self.criterionGAN(pred_real, True) + # combine loss and calculate gradients + self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5 + self.loss_D.backward() + + def backward_G(self): + """Calculate GAN and L1 loss for the generator""" + # First, G(A) should fake the discriminator + fake_AB = torch.cat((self.real_A, self.fake_B), 1) + pred_fake = self.netD(fake_AB) + self.loss_G_GAN = self.criterionGAN(pred_fake, True) + # Second, G(A) = B + self.loss_G_L1 = self.criterionL1(self.fake_B, + self.real_B) * self.opt.lambda_L1 + # combine loss and calculate gradients + self.loss_G = self.loss_G_GAN + self.loss_G_L1 + self.loss_G.backward() + + def optimize_parameters(self): + self.forward() # compute fake images: G(A) + # update D + self.set_requires_grad(self.netD, True) # enable backprop for D + self.optimizer_D.zero_grad() # set D's gradients to zero + self.backward_D() # calculate gradients for D + self.optimizer_D.step() # update D's weights + # update G + self.set_requires_grad( + self.netD, False) # D requires no gradients when optimizing G + self.optimizer_G.zero_grad() # set G's gradients to zero + self.backward_G() # calculate graidents for G + self.optimizer_G.step() # udpate G's weights diff --git a/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_options.py b/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_options.py new file mode 100644 index 00000000..27eb4c07 --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_options.py @@ -0,0 +1,26 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +class Pix2PixOptions(): + + def __init__(self): + self.gpu_ids = [] + self.input_nc = 3 + self.output_nc = 3 + self.ngf = 64 + self.ndf = 64 + self.netG = 'resnet_9blocks' + self.netD = 'basic' + self.norm = 'instance' + self.no_dropout = False + self.init_type = 'normal' + self.init_gain = 0.02 + self.n_layers_D = 3 + self.gan_mode = 'lsgan' + self.lr = 0.0002 + self.beta1 = 0.5 + self.isTrain = False + self.checkpoints_dir = './pix2pix_checkpoints' + self.name = 'mid_net' + self.lr_policy = 'linear' + self.direction = 'AtoB' + self.lambda_L1 = 100.0 + self.preprocess = 'resize_and_crop' diff --git a/modelscope/models/cv/face_reconstruction/models/renderer.py b/modelscope/models/cv/face_reconstruction/models/renderer.py new file mode 100755 index 00000000..d10fd560 --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/renderer.py @@ -0,0 +1,325 @@ +# Part of the implementation is borrowed and modified from pytorch3d, +# publicly available at https://github.com/facebookresearch/pytorch3d + +import imageio +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from skimage.io import imread + +from .. import utils +from ..utils import read_obj + + +def set_rasterizer(): + global Meshes, load_obj, rasterize_meshes + from pytorch3d.structures import Meshes + from pytorch3d.io import load_obj + from pytorch3d.renderer.mesh import rasterize_meshes + + +class Pytorch3dRasterizer(nn.Module): + # TODO: add support for rendering non-squared images, since pytorc3d supports this now + """ Borrowed from https://github.com/facebookresearch/pytorch3d + Notice: + x,y,z are in image space, normalized + can only render squared image now + """ + + def __init__(self, image_size=224): + """ + use fixed raster_settings for rendering faces + """ + super().__init__() + raster_settings = { + 'image_size': image_size, + 'blur_radius': 0.0, + 'faces_per_pixel': 1, + 'bin_size': None, + 'max_faces_per_bin': None, + 'perspective_correct': False, + } + raster_settings = utils.dict2obj(raster_settings) + self.raster_settings = raster_settings + + def forward(self, vertices, faces, attributes=None, h=None, w=None): + fixed_vertices = vertices.clone() + fixed_vertices[..., :2] = -fixed_vertices[..., :2] + raster_settings = self.raster_settings + if h is None and w is None: + image_size = raster_settings.image_size + else: + image_size = [h, w] + if h > w: + fixed_vertices[..., 1] = fixed_vertices[..., 1] * h / w + else: + fixed_vertices[..., 0] = fixed_vertices[..., 0] * w / h + + meshes_screen = Meshes( + verts=fixed_vertices.float(), faces=faces.long()) + pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( + meshes_screen, + image_size=image_size, + blur_radius=raster_settings.blur_radius, + faces_per_pixel=raster_settings.faces_per_pixel, + bin_size=raster_settings.bin_size, + max_faces_per_bin=raster_settings.max_faces_per_bin, + perspective_correct=raster_settings.perspective_correct, + ) + vismask = (pix_to_face > -1).float() + D = attributes.shape[-1] + attributes = attributes.clone() + attributes = attributes.view(attributes.shape[0] * attributes.shape[1], + 3, attributes.shape[-1]) + N, H, W, K, _ = bary_coords.shape + mask = pix_to_face == -1 + pix_to_face = pix_to_face.clone() + pix_to_face[mask] = 0 + idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D) + pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D) + pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2) + pixel_vals[mask] = 0 # Replace masked values in output. + pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) + pixel_vals = torch.cat( + [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) + return pixel_vals + + +class SRenderY(nn.Module): + + def __init__(self, image_size, obj_filename, uvcoords_path, uv_size=256): + super(SRenderY, self).__init__() + self.image_size = image_size + self.uv_size = uv_size + + self.rasterizer = Pytorch3dRasterizer(image_size) + self.uv_rasterizer = Pytorch3dRasterizer(uv_size) + + mesh = read_obj(obj_filename) + uvcoords = np.load(uvcoords_path)[None, ...] + uvcoords = torch.from_numpy(uvcoords) + verts = mesh['vertices'] + verts = torch.from_numpy(verts) + uvfaces = mesh['faces'][None, ...] - 1 + uvfaces = torch.from_numpy(uvfaces) + faces = mesh['faces'][None, ...] - 1 + faces = torch.from_numpy(faces) + + # faces + dense_triangles = utils.generate_triangles(uv_size, uv_size) + self.register_buffer( + 'dense_faces', + torch.from_numpy(dense_triangles).long()[None, :, :]) + self.register_buffer('faces', faces) + self.register_buffer('raw_uvcoords', uvcoords) + + # uv coords + uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.], + -1) # [bz, ntv, 3] + uvcoords = uvcoords * 2 - 1 + uvcoords[..., 1] = -uvcoords[..., 1] + face_uvcoords = utils.face_vertices(uvcoords, uvfaces) + self.register_buffer('uvcoords', uvcoords) + self.register_buffer('uvfaces', uvfaces) + self.register_buffer('face_uvcoords', face_uvcoords) + + # shape colors, for rendering shape overlay + colors = torch.tensor([180, 180, 180])[None, None, :].repeat( + 1, + faces.max() + 1, 1).float() / 255. + face_colors = utils.face_vertices(colors, faces) + self.register_buffer('face_colors', face_colors) + + # SH factors for lighting + pi = np.pi + value_1 = 1 / np.sqrt(4 * pi) + value_2 = ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))) + value_3 = ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))) + value_4 = ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))) + value_5 = (pi / 4) * 3 * (np.sqrt(5 / (12 * pi))) + value_6 = (pi / 4) * 3 * (np.sqrt(5 / (12 * pi))) + value_7 = (pi / 4) * 3 * (np.sqrt(5 / (12 * pi))) + value_8 = (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))) + value_9 = (pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi))) + constant_factor = torch.tensor([ + value_1, value_2, value_3, value_4, value_5, value_6, value_7, + value_8, value_9 + ]).float() + self.register_buffer('constant_factor', constant_factor) + + def forward(self, + vertices, + transformed_vertices, + albedos, + lights=None, + light_type='point'): + ''' + -- Texture Rendering + vertices: [batch_size, V, 3], vertices in world space, for calculating normals, then shading + transformed_vertices: [batch_size, V, 3], range:normalized to [-1,1], projected vertices in image space + (that is aligned to the iamge pixel), for rasterization + albedos: [batch_size, 3, h, w], uv map + lights: + spherical homarnic: [N, 9(shcoeff), 3(rgb)] + points/directional lighting: [N, n_lights, 6(xyzrgb)] + light_type: + point or directional + ''' + batch_size = vertices.shape[0] + # rasterizer near 0 far 100. move mesh so minz larger than 0 + transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] + 10 + # attributes + face_vertices = utils.face_vertices( + vertices, self.faces.expand(batch_size, -1, -1)) + normals = utils.vertex_normals(vertices, + self.faces.expand(batch_size, -1, -1)) + face_normals = utils.face_vertices( + normals, self.faces.expand(batch_size, -1, -1)) + transformed_normals = utils.vertex_normals( + transformed_vertices, self.faces.expand(batch_size, -1, -1)) + transformed_face_normals = utils.face_vertices( + transformed_normals, self.faces.expand(batch_size, -1, -1)) + + attributes = torch.cat([ + self.face_uvcoords.expand(batch_size, -1, -1, -1), + transformed_face_normals.detach(), + face_vertices.detach(), face_normals + ], -1) + # rasterize + rendering = self.rasterizer(transformed_vertices, + self.faces.expand(batch_size, -1, -1), + attributes) + + #### + # vis mask + alpha_images = rendering[:, -1, :, :][:, None, :, :].detach() + + # albedo + uvcoords_images = rendering[:, :3, :, :] + grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2] + albedo_images = F.grid_sample(albedos, grid, align_corners=False) + + # visible mask for pixels with positive normal direction + transformed_normal_map = rendering[:, 3:6, :, :].detach() + pos_mask = (transformed_normal_map[:, 2:, :, :] < -0.05).float() + + # shading + normal_images = rendering[:, 9:12, :, :] + if lights is not None: + if lights.shape[1] == 9: + shading_images = self.add_SHlight(normal_images, lights) + else: + if light_type == 'point': + vertice_images = rendering[:, 6:9, :, :].detach() + shading = self.add_pointlight( + vertice_images.permute(0, 2, 3, + 1).reshape([batch_size, -1, 3]), + normal_images.permute(0, 2, 3, + 1).reshape([batch_size, -1, 3]), + lights) + shading_images = shading.reshape([ + batch_size, albedo_images.shape[2], + albedo_images.shape[3], 3 + ]).permute(0, 3, 1, 2) + else: + shading = self.add_directionlight( + normal_images.permute(0, 2, 3, + 1).reshape([batch_size, -1, 3]), + lights) + shading_images = shading.reshape([ + batch_size, albedo_images.shape[2], + albedo_images.shape[3], 3 + ]).permute(0, 3, 1, 2) + images = albedo_images * shading_images + else: + images = albedo_images + shading_images = images.detach() * 0. + + outputs = { + 'images': images * alpha_images, + 'albedo_images': albedo_images * alpha_images, + 'alpha_images': alpha_images, + 'pos_mask': pos_mask, + 'shading_images': shading_images, + 'grid': grid, + 'normals': normals, + 'normal_images': normal_images * alpha_images, + 'transformed_normals': transformed_normals, + } + + return outputs + + def add_SHlight(self, normal_images, gamma, init_lit): + ''' + sh_coeff: [bz, 9, 3] + ''' + batch_size = gamma.shape[0] + gamma = gamma.reshape([batch_size, 3, 9]) + gamma = gamma + init_lit + sh_coeff = gamma.permute(0, 2, 1) + + N = normal_images + tmp_value = 3 * (N[:, 2]**2) - 1 + sh = torch.stack([ + N[:, 0] * 0. + 1., N[:, 0], N[:, 1], N[:, 2], N[:, 0] * N[:, 1], + N[:, 0] * N[:, 2], N[:, 1] * N[:, 2], N[:, 0]**2 - N[:, 1]**2, + tmp_value + ], 1) # [bz, 9, h, w] + sh = sh * self.constant_factor[None, :, None, None] + shading = torch.sum(sh_coeff[:, :, :, None, None] + * sh[:, :, None, :, :], 1) # [bz, 9, 3, h, w] + return shading + + def add_pointlight(self, vertices, normals, lights): + ''' + vertices: [bz, nv, 3] + lights: [bz, nlight, 6] + returns: + shading: [bz, nv, 3] + ''' + light_positions = lights[:, :, :3] + light_intensities = lights[:, :, 3:] + directions_to_lights = F.normalize( + light_positions[:, :, None, :] - vertices[:, None, :, :], dim=3) + # normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.) + normals_dot_lights = (normals[:, None, :, :] + * directions_to_lights).sum(dim=3) + shading = normals_dot_lights[:, :, :, + None] * light_intensities[:, :, None, :] + return shading.mean(1) + + def add_directionlight(self, normals, lights): + ''' + normals: [bz, nv, 3] + lights: [bz, nlight, 6] + returns: + shading: [bz, nv, 3] + ''' + light_direction = lights[:, :, :3] + light_intensities = lights[:, :, 3:] + directions_to_lights = F.normalize( + light_direction[:, :, None, :].expand(-1, -1, normals.shape[1], + -1), + dim=3) + # normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.) + # normals_dot_lights = (normals[:,None,:,:]*directions_to_lights).sum(dim=3) + normals_dot_lights = torch.clamp( + (normals[:, None, :, :] * directions_to_lights).sum(dim=3), 0., 1.) + shading = normals_dot_lights[:, :, :, + None] * light_intensities[:, :, None, :] + return shading.mean(1) + + def world2uv(self, vertices): + ''' + warp vertices from world space to uv space + vertices: [bz, V, 3] + uv_vertices: [bz, 3, h, w] + ''' + batch_size = vertices.shape[0] + face_vertices = utils.face_vertices( + vertices, self.faces.expand(batch_size, -1, -1)) + uv_vertices = self.uv_rasterizer( + self.uvcoords.expand(batch_size, -1, -1), + self.uvfaces.expand(batch_size, -1, -1), face_vertices)[:, :3] + return uv_vertices diff --git a/modelscope/models/cv/face_reconstruction/models/unet.py b/modelscope/models/cv/face_reconstruction/models/unet.py new file mode 100644 index 00000000..3ceed26d --- /dev/null +++ b/modelscope/models/cv/face_reconstruction/models/unet.py @@ -0,0 +1,173 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F + +warnings.filterwarnings(action='ignore') + + +def weights_init(init_type='kaiming', gain=0.02): + + def init_func(m): + classname = m.__class__.__name__ + if hasattr(m, 'weight') and (classname.find('Conv') != -1 + or classname.find('Linear') != -1): + + if init_type == 'normal': + nn.init.normal_(m.weight.data, 0.0, gain) + elif init_type == 'xavier': + nn.init.xavier_normal_(m.weight.data, gain=gain) + elif init_type == 'kaiming': + nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif init_type == 'orthogonal': + nn.init.orthogonal_(m.weight.data, gain=gain) + + if hasattr(m, 'bias') and m.bias is not None: + nn.init.constant_(m.bias.data, 0.0) + + elif classname.find('BatchNorm2d') != -1: + nn.init.normal_(m.weight.data, 1.0, gain) + nn.init.constant_(m.bias.data, 0.0) + + return init_func + + +class double_conv(nn.Module): + '''(conv => BN => ReLU) * 2''' + + def __init__(self, in_ch, out_ch): + super(double_conv, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_ch, out_ch, 3, padding=1), + nn.BatchNorm2d(out_ch), + # nn.InstanceNorm2d(out_ch), + nn.ReLU(inplace=True), + nn.Conv2d(out_ch, out_ch, 3, padding=1), + nn.BatchNorm2d(out_ch), + # nn.InstanceNorm2d(out_ch), + nn.ReLU(inplace=True)) + + def forward(self, x): + x = self.conv(x) + return x + + +class inconv(nn.Module): + + def __init__(self, in_ch, out_ch): + super(inconv, self).__init__() + self.conv = double_conv(in_ch, out_ch) + + def forward(self, x): + x = self.conv(x) + return x + + +class down(nn.Module): + + def __init__(self, in_ch, out_ch): + super(down, self).__init__() + self.mpconv = nn.Sequential( + nn.MaxPool2d(2), double_conv(in_ch, out_ch)) + + def forward(self, x): + x = self.mpconv(x) + return x + + +class up(nn.Module): + + def __init__(self, in_ch, out_ch, bilinear=True): + super(up, self).__init__() + + if bilinear: + self.up = nn.Upsample( + scale_factor=2, mode='bilinear', align_corners=True) + else: + self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2) + + self.conv = double_conv(in_ch, out_ch) + + def forward(self, x1, x2): + x1 = self.up(x1) + + diffY = x2.size()[2] - x1.size()[2] + diffX = x2.size()[3] - x1.size()[3] + + x1 = F.pad( + x1, + (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2)) + + x = torch.cat([x2, x1], dim=1) + x = self.conv(x) + return x + + +class outconv(nn.Module): + + def __init__(self, in_ch, out_ch): + super(outconv, self).__init__() + self.conv = nn.Conv2d(in_ch, out_ch, 1) + + def forward(self, x): + x = self.conv(x) + return x + + +class UNet(nn.Module): + + def __init__(self, + n_channels, + n_classes, + deep_supervision=False, + init_weights=True): + super(UNet, self).__init__() + self.deep_supervision = deep_supervision + self.inc = inconv(n_channels, 64) + self.down1 = down(64, 128) + self.down2 = down(128, 256) + self.down3 = down(256, 512) + self.down4 = down(512, 512) + self.up1 = up(1024, 256) + self.up2 = up(512, 128) + self.up3 = up(256, 64) + self.up4 = up(128, 64) + self.outc = outconv(64, n_classes) + + self.dsoutc4 = outconv(256, n_classes) + self.dsoutc3 = outconv(128, n_classes) + self.dsoutc2 = outconv(64, n_classes) + self.dsoutc1 = outconv(64, n_classes) + + self.sigmoid = nn.Sigmoid() + + if init_weights: + self.apply(weights_init()) + + def forward(self, x): + x1 = self.inc(x) + x2 = self.down1(x1) + x3 = self.down2(x2) + x4 = self.down3(x3) + x5 = self.down4(x4) + x44 = self.up1(x5, x4) + x33 = self.up2(x44, x3) + x22 = self.up3(x33, x2) + x11 = self.up4(x22, x1) + x0 = self.outc(x11) + x0 = self.sigmoid(x0) + if self.deep_supervision: + x11 = F.interpolate( + self.dsoutc1(x11), x0.shape[2:], mode='bilinear') + x22 = F.interpolate( + self.dsoutc2(x22), x0.shape[2:], mode='bilinear') + x33 = F.interpolate( + self.dsoutc3(x33), x0.shape[2:], mode='bilinear') + x44 = F.interpolate( + self.dsoutc4(x44), x0.shape[2:], mode='bilinear') + + return x0, x11, x22, x33, x44 + else: + return x0 diff --git a/modelscope/models/cv/face_reconstruction/utils.py b/modelscope/models/cv/face_reconstruction/utils.py index 9f2a25ed..655d8b2a 100644 --- a/modelscope/models/cv/face_reconstruction/utils.py +++ b/modelscope/models/cv/face_reconstruction/utils.py @@ -383,6 +383,52 @@ def load_lm3d(bfm_folder): return Lm3D +def mesh_to_string(mesh): + out_string = '' + out_string += '# Create by HRN\n' + + if 'colors' in mesh: + for i, v in enumerate(mesh['vertices']): + out_string += \ + 'v {:.6f} {:.6f} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format( + v[0], v[1], v[2], mesh['colors'][i][0], + mesh['colors'][i][1], mesh['colors'][i][2]) + else: + for v in mesh['vertices']: + out_string += 'v {:.6f} {:.6f} {:.6f}\n'.format(v[0], v[1], v[2]) + + if 'UVs' in mesh: + for uv in mesh['UVs']: + out_string += 'vt {:.6f} {:.6f}\n'.format(uv[0], uv[1]) + + if 'normals' in mesh: + for vn in mesh['normals']: + out_string += 'vn {:.6f} {:.6f} {:.6f}\n'.format( + vn[0], vn[1], vn[2]) + + if 'faces' in mesh: + for ind, face in enumerate(mesh['faces']): + if 'faces_uv' in mesh or 'faces_normal' in mesh or 'UVs' in mesh: + if 'faces_uv' in mesh: + face_uv = mesh['faces_uv'][ind] + else: + face_uv = face + if 'faces_normal' in mesh: + face_normal = mesh['faces_normal'][ind] + else: + face_normal = face + row = 'f ' + ' '.join([ + '{}/{}/{}'.format(face[i], face_uv[i], face_normal[i]) + for i in range(len(face)) + ]) + '\n' + else: + row = 'f ' + ' '.join( + ['{}'.format(face[i]) for i in range(len(face))]) + '\n' + out_string += row + + return out_string + + def write_obj(save_path, mesh): save_dir = os.path.dirname(save_path) save_name = os.path.splitext(os.path.basename(save_path))[0] @@ -392,7 +438,7 @@ def write_obj(save_path, mesh): os.path.join(save_dir, save_name + '.jpg'), mesh['texture_map']) with open(os.path.join(save_dir, save_name + '.mtl'), 'w') as wf: - wf.write('# Created by ModelScope\n') + wf.write('# Created by HRN\n') wf.write('newmtl material_0\n') wf.write('Ka 1.000000 0.000000 0.000000\n') wf.write('Kd 1.000000 1.000000 1.000000\n') @@ -404,29 +450,31 @@ def write_obj(save_path, mesh): with open(save_path, 'w') as wf: if 'texture_map' in mesh: - wf.write('# Create by ModelScope\n') + wf.write('# Create by HRN\n') wf.write('mtllib ./{}.mtl\n'.format(save_name)) if 'colors' in mesh: for i, v in enumerate(mesh['vertices']): - wf.write('v {} {} {} {} {} {}\n'.format( - v[0], v[1], v[2], mesh['colors'][i][0], - mesh['colors'][i][1], mesh['colors'][i][2])) + wf.write( + 'v {:.6f} {:.6f} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format( + v[0], v[1], v[2], mesh['colors'][i][0], + mesh['colors'][i][1], mesh['colors'][i][2])) else: for v in mesh['vertices']: - wf.write('v {} {} {}\n'.format(v[0], v[1], v[2])) + wf.write('v {:.6f} {:.6f} {:.6f}\n'.format(v[0], v[1], v[2])) if 'UVs' in mesh: for uv in mesh['UVs']: - wf.write('vt {} {}\n'.format(uv[0], uv[1])) + wf.write('vt {:.6f} {:.6f}\n'.format(uv[0], uv[1])) if 'normals' in mesh: for vn in mesh['normals']: - wf.write('vn {} {} {}\n'.format(vn[0], vn[1], vn[2])) + wf.write('vn {:.6f} {:.6f} {:.6f}\n'.format( + vn[0], vn[1], vn[2])) if 'faces' in mesh: for ind, face in enumerate(mesh['faces']): - if 'faces_uv' in mesh or 'faces_normal' in mesh: + if 'faces_uv' in mesh or 'faces_normal' in mesh or 'UVs' in mesh: if 'faces_uv' in mesh: face_uv = mesh['faces_uv'][ind] else: @@ -750,3 +798,203 @@ def estimate_normals(vertices, faces): norm[inds] = [0, 0, 1.0] result = normalize_v3(norm) return result + + +def draw_landmarks(img, landmark, color='r', step=2): + """ + Return: + img -- numpy.array, (B, H, W, 3) img with landmark, RGB order, range (0, 255) + + + Parameters: + img -- numpy.array, (B, H, W, 3), RGB order, range (0, 255) + landmark -- numpy.array, (B, 68, 2), y direction is opposite to v direction + color -- str, 'r' or 'b' (red or blue) + """ + if color == 'r': + c = np.array([255., 0, 0]) + else: + c = np.array([0, 0, 255.]) + + _, H, W, _ = img.shape + img, landmark = img.copy(), landmark.copy() + landmark[..., 1] = H - 1 - landmark[..., 1] + landmark = np.round(landmark).astype(np.int32) + for i in range(landmark.shape[1]): + x, y = landmark[:, i, 0], landmark[:, i, 1] + for j in range(-step, step): + for k in range(-step, step): + u = np.clip(x + j, 0, W - 1) + v = np.clip(y + k, 0, H - 1) + for m in range(landmark.shape[0]): + img[m, v[m], u[m]] = c + return img + + +def split_vis(img_path, target_dir=None): + img = cv2.imread(img_path) + h, w = img.shape[:2] + n_split = w // h + if target_dir is None: + target_dir = os.path.dirname(img_path) + base_name = os.path.splitext(os.path.basename(img_path))[0] + for i in range(n_split): + img_i = img[:, i * h:(i + 1) * h, :] + cv2.imwrite( + os.path.join(target_dir, '{}_{:0>2d}.jpg'.format(base_name, + i + 1)), img_i) + + +def write_video(image_list, save_path, fps=20.0): + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + # fourcc = cv2.VideoWriter_fourcc(*'MJPG') # avi格式 + + h, w = image_list[0].shape[:2] + + out = cv2.VideoWriter(save_path, fourcc, fps, (w, h), True) + + for frame in image_list: + out.write(frame) + + out.release() + + +# ---------------------------- process/generate vertices, normals, faces +def generate_triangles(h, w, margin_x=2, margin_y=5, mask=None): + # quad layout: + # 0 1 ... w-1 + # w w+1 + # . + # w*h + triangles = [] + for x in range(margin_x, w - 1 - margin_x): + for y in range(margin_y, h - 1 - margin_y): + triangle0 = [y * w + x, y * w + x + 1, (y + 1) * w + x] + triangle1 = [y * w + x + 1, (y + 1) * w + x + 1, (y + 1) * w + x] + triangles.append(triangle0) + triangles.append(triangle1) + triangles = np.array(triangles) + triangles = triangles[:, [0, 2, 1]] + return triangles + + +def face_vertices(vertices, faces): + """ + :param vertices: [batch size, number of vertices, 3] + :param faces: [batch size, number of faces, 3] + :return: [batch size, number of faces, 3, 3] + """ + assert (vertices.ndimension() == 3) + assert (faces.ndimension() == 3) + assert (vertices.shape[0] == faces.shape[0]) + assert (vertices.shape[2] == 3) + assert (faces.shape[2] == 3) + + bs, nv = vertices.shape[:2] + bs, nf = faces.shape[:2] + device = vertices.device + faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) + * nv)[:, None, None] + vertices = vertices.reshape((bs * nv, 3)) + # pytorch only supports long and byte tensors for indexing + return vertices[faces.long()] + + +def vertex_normals(vertices, faces): + """ + :param vertices: [batch size, number of vertices, 3] + :param faces: [batch size, number of faces, 3] + :return: [batch size, number of vertices, 3] + """ + assert (vertices.ndimension() == 3) + assert (faces.ndimension() == 3) + assert (vertices.shape[0] == faces.shape[0]) + assert (vertices.shape[2] == 3) + assert (faces.shape[2] == 3) + bs, nv = vertices.shape[:2] + bs, nf = faces.shape[:2] + device = vertices.device + normals = torch.zeros(bs * nv, 3).to(device) + + faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) + * nv)[:, None, None] # expanded faces + vertices_faces = vertices.reshape((bs * nv, 3))[faces.long()] + + faces = faces.reshape(-1, 3) + vertices_faces = vertices_faces.reshape(-1, 3, 3) + + normals.index_add_( + 0, faces[:, 1].long(), + torch.cross(vertices_faces[:, 2] - vertices_faces[:, 1], + vertices_faces[:, 0] - vertices_faces[:, 1])) + normals.index_add_( + 0, faces[:, 2].long(), + torch.cross(vertices_faces[:, 0] - vertices_faces[:, 2], + vertices_faces[:, 1] - vertices_faces[:, 2])) + normals.index_add_( + 0, faces[:, 0].long(), + torch.cross(vertices_faces[:, 1] - vertices_faces[:, 0], + vertices_faces[:, 2] - vertices_faces[:, 0])) + + normals = F.normalize(normals, eps=1e-6, dim=1) + normals = normals.reshape((bs, nv, 3)) + # pytorch only supports long and byte tensors for indexing + return normals + + +def dict2obj(d): + # if isinstance(d, list): + # d = [dict2obj(x) for x in d] + if not isinstance(d, dict): + return d + + class C(object): + pass + + o = C() + for k in d: + o.__dict__[k] = dict2obj(d[k]) + return o + + +def enlarged_bbox(bbox, img_width, img_height, enlarge_ratio=0.2): + ''' + :param bbox: [xmin,ymin,xmax,ymax] + :return: bbox: [xmin,ymin,xmax,ymax] + ''' + + left = bbox[0] + top = bbox[1] + + right = bbox[2] + bottom = bbox[3] + + roi_width = right - left + roi_height = bottom - top + + new_left = left - int(roi_width * enlarge_ratio) + new_left = 0 if new_left < 0 else new_left + + new_top = top - int(roi_height * enlarge_ratio) + new_top = 0 if new_top < 0 else new_top + + new_right = right + int(roi_width * enlarge_ratio) + new_right = img_width if new_right > img_width else new_right + + new_bottom = bottom + int(roi_height * enlarge_ratio) + new_bottom = img_height if new_bottom > img_height else new_bottom + + bbox = [new_left, new_top, new_right, new_bottom] + + bbox = [int(x) for x in bbox] + + return bbox + + +def draw_line(im, points, color, stroke_size=2, closed=False): + points = points.astype(np.int32) + for i in range(len(points) - 1): + cv2.line(im, tuple(points[i]), tuple(points[i + 1]), color, + stroke_size) + if closed: + cv2.line(im, tuple(points[0]), tuple(points[-1]), color, stroke_size) diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 28fb7439..d8b1cdc2 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -31,6 +31,8 @@ class OutputKeys(object): OUTPUT_PCM = 'output_pcm' OUTPUT_PCM_LIST = 'output_pcm_list' OUTPUT_WAV = 'output_wav' + OUTPUT_OBJ = 'output_obj' + OUTPUT_MESH = 'output_mesh' IMG_EMBEDDING = 'img_embedding' SPK_EMBEDDING = 'spk_embedding' SPO_LIST = 'spo_list' @@ -445,15 +447,19 @@ TASK_OUTPUTS = { # 3D face reconstruction result for single sample # { + # "output_obj": io.BytesIO, + # "output_img": np.array with shape(h, w, 3), # "output": { - # "vertices": np.array with shape(n, 3), - # "faces": np.array with shape(n, 3), - # "faces_uv": np.array with shape(n, 3), - # "faces_normal": np.array with shape(n, 3), - # "colors": np.array with shape(n, 3), - # "UVs": np.array with shape(n, 2), - # "normals": np.array with shape(n, 3), - # "texture_map": np.array with shape(h, w, 3), + # "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), + # }, + # "vis_image": np.array with shape(h, w, 3), + # "frame_list", [np.array with shape(h, w, 3), ...], # } # } Tasks.face_reconstruction: [OutputKeys.OUTPUT], diff --git a/modelscope/pipelines/cv/face_reconstruction_pipeline.py b/modelscope/pipelines/cv/face_reconstruction_pipeline.py index 64f8b3a9..f8240fc0 100644 --- a/modelscope/pipelines/cv/face_reconstruction_pipeline.py +++ b/modelscope/pipelines/cv/face_reconstruction_pipeline.py @@ -1,4 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import io import os import shutil from typing import Any, Dict @@ -12,14 +13,14 @@ import torch from scipy.io import loadmat, savemat from modelscope.metainfo import Pipelines -from modelscope.models import Model -from modelscope.models.cv.face_reconstruction.models.facelandmark.large_model_infer import \ - LargeModelInfer -from modelscope.models.cv.face_reconstruction.utils import (align_for_lm, - align_img, - load_lm3d, - read_obj, - write_obj) +from modelscope.models.cv.face_reconstruction.models.facelandmark.large_base_lmks_infer import \ + LargeBaseLmkInfer +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.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 @@ -54,7 +55,10 @@ class FaceReconstructionPipeline(Pipeline): >>> pipeline_faceRecon = pipeline('face-reconstruction', model='damo/cv_resnet50_face-reconstruction') >>> result = pipeline_faceRecon(test_image) - >>> write_obj('result_face_reconstruction.obj', result[OutputKeys.OUTPUT]) + >>> mesh = result[OutputKeys.OUTPUT]['mesh'] + >>> texture_map = result[OutputKeys.OUTPUT_IMG] + >>> mesh['texture_map'] = texture_map + >>> write_obj('hrn_mesh_mid.obj', mesh) """ super().__init__(model=model, device=device) @@ -62,16 +66,28 @@ class FaceReconstructionPipeline(Pipeline): bfm_folder = os.path.join(model_root, 'assets') checkpoint_path = os.path.join(model_root, ModelFile.TORCH_MODEL_FILE) - self.face_mark_model = LargeModelInfer( - os.path.join(model_root, 'large_base_net.pth'), device='cuda') + 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(0) - torch.cuda.set_device(device) + device = torch.device(self.device_name_) + self.model.set_device(device) self.model.setup(checkpoint_path) - self.model.device = device self.model.parallelize() self.model.eval() - self.model.set_render(image_res=1024) + self.model.set_render(image_res=512) save_ckpt_dir = os.path.join( os.path.expanduser('~'), '.cache/torch/hub/checkpoints') @@ -107,52 +123,6 @@ class FaceReconstructionPipeline(Pipeline): self.tex_size = 4096 - self.bald_tex_bg = cv2.imread( - '{}/assets/template_texture.jpg'.format(model_root)).astype( - np.float32) - - front_mask = cv2.imread( - '{}/assets/face_mask.jpg'.format(model_root)).astype( - np.float32) / 255 - front_mask = cv2.resize(front_mask, (1024, 1024)) - front_mask = cv2.resize(front_mask, (0, 0), fx=0.1, fy=0.1) - front_mask = cv2.erode(front_mask, - np.ones(shape=(7, 7), dtype=np.float32)) - front_mask = cv2.GaussianBlur(front_mask, (13, 13), 0) - self.front_mask = cv2.resize(front_mask, - (self.tex_size, self.tex_size)) - self.binary_front_mask = self.front_mask.copy() - self.binary_front_mask[(self.front_mask < 0.3) - + (self.front_mask > 0.7)] = 0 - self.binary_front_mask[self.binary_front_mask != 0] = 1.0 - self.binary_front_mask_ = self.binary_front_mask.copy() - self.binary_front_mask = np.zeros((4096 + 1024, 4096, 3), - dtype=np.float32) - self.binary_front_mask[:4096, :] = self.binary_front_mask_ - self.front_mask_ = self.front_mask.copy() - self.front_mask = np.zeros((4096 + 1024, 4096, 3), dtype=np.float32) - self.front_mask[:4096, :] = self.front_mask_ - - l_eye_mask = cv2.imread( - '{}/assets/l_eye_mask.png'.format(model_root))[:, :, :1] / 255.0 - l_eye_mask = cv2.erode(l_eye_mask, - np.ones(shape=(5, 5), dtype=np.float32)) - self.l_eye_mask = cv2.GaussianBlur(l_eye_mask, (7, 7), 0)[..., None] - self.l_eye_binary_mask = self.l_eye_mask.copy() - self.l_eye_binary_mask[(self.l_eye_mask < 0.3) - + (self.l_eye_mask > 0.7)] = 0 - self.l_eye_binary_mask[self.l_eye_binary_mask != 0] = 1.0 - - r_eye_mask = cv2.imread( - '{}/assets/r_eye_mask.png'.format(model_root))[:, :, :1] / 255.0 - r_eye_mask = cv2.dilate(r_eye_mask, - np.ones(shape=(7, 7), dtype=np.float32)) - self.r_eye_mask = cv2.GaussianBlur(r_eye_mask, (7, 7), 0)[..., None] - self.r_eye_binary_mask = self.r_eye_mask.copy() - self.r_eye_binary_mask[(self.r_eye_mask < 0.3) - + (self.r_eye_mask > 0.7)] = 0 - self.r_eye_binary_mask[self.r_eye_binary_mask != 0] = 1.0 - self.lm3d_std = load_lm3d(bfm_folder) self.align_params = loadmat( '{}/assets/BBRegressorParam_r.mat'.format(model_root)) @@ -168,47 +138,23 @@ class FaceReconstructionPipeline(Pipeline): result = {'img': img} return result - def read_data(self, - img, - lm, - lm3d_std, - to_tensor=True, - image_res=1024, - img_fat=None): + def read_data(self, img, lm, lm3d_std, to_tensor=True, image_res=1024): # to RGB im = PIL.Image.fromarray(img[..., ::-1]) W, H = im.size lm[:, -1] = H - 1 - lm[:, -1] - - im_lr_coeff, lm_lr_coeff = None, None - head_mask = None - - _, im_lr, lm_lr, mask_lr_head = align_img( - im, lm, lm3d_std, mask=head_mask) + _, im_lr, lm_lr, _ = align_img(im, lm, lm3d_std) _, im_hd, lm_hd, _ = align_img( im, lm, lm3d_std, target_size=image_res, - rescale_factor=102.0 * image_res / 224) + rescale_factor=102. * image_res / 224) mask_lr = self.face_sess.run( self.face_sess.graph.get_tensor_by_name('output_alpha:0'), feed_dict={'input_image:0': np.array(im_lr)}) - if img_fat is not None: - assert img_fat.shape == img.shape - im_fat = PIL.Image.fromarray(img_fat[..., ::-1]) - - _, im_hd, _, _ = align_img( - im_fat, - lm, - lm3d_std, - target_size=image_res, - rescale_factor=102.0 * image_res / 224) - - im_hd = np.array(im_hd).astype(np.float32) - if to_tensor: im_lr = torch.tensor( np.array(im_lr) / 255., @@ -219,12 +165,12 @@ class FaceReconstructionPipeline(Pipeline): mask_lr = torch.tensor( np.array(mask_lr) / 255., dtype=torch.float32)[None, None, :, :] - mask_lr_head = torch.tensor( - np.array(mask_lr_head) / 255., dtype=torch.float32)[ - None, None, :, :] if mask_lr_head is not None else None lm_lr = torch.tensor(lm_lr).unsqueeze(0) lm_hd = torch.tensor(lm_hd).unsqueeze(0) - return im_lr, lm_lr, im_hd, lm_hd, mask_lr, mask_lr_head, im_lr_coeff, lm_lr_coeff + return im_lr, lm_lr, im_hd, lm_hd, mask_lr + + def parse_label(self, label): + return torch.tensor(np.array(label).astype(np.float32)) def prepare_data(self, img, lm_sess, five_points=None): input_img, scale, bbox = align_for_lm( @@ -246,58 +192,276 @@ class FaceReconstructionPipeline(Pipeline): return landmark - def blend_eye_corner(self, tex_map, template_tex): - tex_map = tex_map.astype(np.float32) + def get_img_for_texture(self, input_img_tensor): + input_img = input_img_tensor.permute( + 0, 2, 3, 1).detach().cpu().numpy()[0] * 255. + input_img = input_img.astype(np.uint8) - x1 = int(288 * 4096 / 758) - y1 = int(235 * 4096 / 758) - w = int(90 * 4096 / 758) - h = int(50 * 4096 / 758) - template_tex_l = template_tex[y1:y1 + h, x1:x1 + w] - pred_tex_l = tex_map[y1:y1 + h, x1:x1 + w] - pred_tex_l_mean_rgb = np.sum( - pred_tex_l * self.l_eye_binary_mask, axis=(0, 1)) - template_tex_l_mean_rgb = np.sum( - template_tex_l * self.l_eye_binary_mask, axis=(0, 1)) - for ch in range(3): - template_tex_l[:, :, ch] *= pred_tex_l_mean_rgb[ - ch] / template_tex_l_mean_rgb[ch] - pred_tex_l = pred_tex_l * ( - 1 - self.l_eye_mask) + template_tex_l * self.l_eye_mask + input_img_for_texture = self.fat_face(input_img, degree=0.03) + input_img_for_texture_tensor = torch.tensor( + np.array(input_img_for_texture) / 255., + dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) + input_img_for_texture_tensor = input_img_for_texture_tensor.to( + self.model.device) + return input_img_for_texture_tensor - x2 = 4096 - x1 - w - y2 = y1 - template_tex_r = template_tex[y2:y2 + h, x2:x2 + w] - pred_tex_r = tex_map[y2:y2 + h, x2:x2 + w] - pred_tex_r_mean_rgb = np.sum( - pred_tex_r * self.r_eye_binary_mask, axis=(0, 1)) - template_tex_r_mean_rgb = np.sum( - template_tex_r * self.r_eye_binary_mask, axis=(0, 1)) - for ch in range(3): - template_tex_r[:, :, ch] *= pred_tex_r_mean_rgb[ - ch] / template_tex_r_mean_rgb[ch] - pred_tex_r = pred_tex_r * ( - 1 - self.r_eye_mask) + template_tex_r * self.r_eye_mask + def infer_lmks(self, img_bgr): + INPUT_SIZE = 224 + ENLARGE_RATIO = 1.35 - tex_map[y1:y1 + h, x1:x1 + w] = pred_tex_l - tex_map[y2:y2 + h, x2:x2 + w] = pred_tex_r + landmarks = [] - return tex_map + rgb_image = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + results = self.detector.predict_jsons(rgb_image) - def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: - rgb_image = input['img'].cpu().numpy().astype(np.uint8) + 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] + }) - bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR) + 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.1): + + _img, scale = resize_on_long_side(img, 800) + + contour_maps, boxes = self.find_face_contour(_img) + + contour_map = contour_maps[0] + + boxes = boxes[0] + + Flow = np.zeros( + shape=(contour_map.shape[0], contour_map.shape[1], 2), + dtype=np.float32) + + box_center = [(boxes['x1'] + boxes['x2']) / 2, + (boxes['y1'] + boxes['y2']) / 2] + + box_length = max( + abs(boxes['y1'] - boxes['y2']), abs(boxes['x1'] - boxes['x2'])) + + value_1 = 2 * (Flow.shape[0] - box_center[1] - 1) + value_2 = 2 * (Flow.shape[1] - box_center[0] - 1) + value_list = [ + box_length * 2, 2 * (box_center[0] - 1), 2 * (box_center[1] - 1), + value_1, value_2 + ] + flow_box_length = min(value_list) + flow_box_length = int(flow_box_length) + + sf = spread_flow(100, flow_box_length * degree) + sf = cv2.resize(sf, (flow_box_length, flow_box_length)) + + Flow[int(box_center[1] + - flow_box_length / 2):int(box_center[1] + + flow_box_length / 2), + int(box_center[0] + - flow_box_length / 2):int(box_center[0] + + flow_box_length / 2)] = sf + + Flow = Flow * np.dstack((contour_map, contour_map)) / 255.0 + + inter_face_maps = contour_maps[-1] + + Flow = Flow * (1.0 - np.dstack( + (inter_face_maps, inter_face_maps)) / 255.0) + + Flow = cv2.resize(Flow, (img.shape[1], img.shape[0])) + + Flow = Flow / scale + + pred, top_bound, bottom_bound, left_bound, right_bound = image_warp_grid1( + Flow[..., 0], Flow[..., 1], img, 1.0, [0, 0, 0, 0]) + + return pred + + def predict_base(self, img): + + if img.shape[0] > 2000 or img.shape[1] > 2000: + img, _ = resize_on_long_side(img, 1500) + + box, results = self.infer_lmks(img) - img = bgr_image - # preprocess - flag = 0 - box, results = self.face_mark_model.infer(img) if results is None or np.array(results).shape[0] == 0: - flag = 1 # no face - return flag, {} - - fatbgr = self.face_mark_model.fat_face(img, degree=0.02) + return {} landmarks = [] results = results[0] @@ -307,64 +471,79 @@ class FaceReconstructionPipeline(Pipeline): landmarks = self.prepare_data(img, self.lm_sess, five_points=landmarks) - im_tensor, lm_tensor, im_hd_tensor, lm_hd_tensor, mask, _, _, _ = self.read_data( - img, landmarks, self.lm3d_std, image_res=1024, img_fat=fatbgr) + im_tensor, lm_tensor, im_hd_tensor, lm_hd_tensor, mask = self.read_data( + img, landmarks, self.lm3d_std, image_res=512) data = { 'imgs': im_tensor, 'imgs_hd': im_hd_tensor, 'lms': lm_tensor, 'lms_hd': lm_hd_tensor, 'face_mask': mask, - 'img_name': 'temp', } - self.model.set_input(data) # unpack data from data loader + self.model.set_input_base(data) # unpack data from data loader - # reconstruct - out_dir = None - output = self.model(out_dir=out_dir) # run inference + output = self.model.predict_results_base() # run inference - # process texture map - tex_map = output['head_tex_map'].astype(np.float32) - tex_map = cv2.resize(tex_map, (self.tex_size, self.tex_size + 1024)) - bg_mean_rgb = np.sum( - self.bald_tex_bg * self.binary_front_mask, axis=(0, 1)) - pred_tex_mean_rgb = np.sum( - tex_map * self.binary_front_mask, axis=(0, 1)) * 1.05 - mid_mean_rgb = bg_mean_rgb * 0.8 + pred_tex_mean_rgb * 0.2 - tex_map += ( - (mid_mean_rgb - pred_tex_mean_rgb) - / np.sum(self.binary_front_mask, axis=(0, 1)))[None, None] * 0.5 - pred_tex_mean_rgb = np.sum( - tex_map * self.binary_front_mask, axis=(0, 1)) * 1.05 - _bald_tex_bg = self.bald_tex_bg.copy() - for ch in range(3): - _bald_tex_bg[:, :, ch] *= pred_tex_mean_rgb[ch] / bg_mean_rgb[ch] - tex_map = _bald_tex_bg * ( - 1. - self.front_mask) + tex_map * self.front_mask - tex_map = tex_map * 1.05 - tex_map = self.blend_eye_corner(tex_map, self.bald_tex_bg) + return output - # export mesh + 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 + + base_model_output = self.predict_base(img) + + input_img_for_tex = self.get_img_for_texture( + base_model_output['input_img']) + + hrn_input = { + 'input_img': base_model_output['input_img'], + 'input_img_for_tex': input_img_for_tex, + 'input_img_hd': base_model_output['input_img_hd'], + 'face_mask': base_model_output['face_mask'], + 'gt_lm': base_model_output['gt_lm'], + 'coeffs': base_model_output['coeffs'], + 'position_map': base_model_output['position_map'], + 'texture_map': base_model_output['texture_map'], + 'tex_valid_mask': base_model_output['tex_valid_mask'], + 'de_retouched_albedo_map': + base_model_output['de_retouched_albedo_map'] + } + + self.model.set_input_hrn(hrn_input) + self.model.get_edge_points_horizontal() + + self.model(visualize=True) + + results = self.model.save_results_hrn() + texture_map = results['texture_map'] results = { - 'vertices': output['head_vertices'], - 'faces': output['head_faces'], - 'UVs': output['head_UVs'], - 'faces_uv': output['head_faces_uv'], - 'normals': output['head_normals'], - 'texture_map': tex_map, + 'mesh': results['face_mesh'], + 'vis_image': results['vis_image'], + 'frame_list': results['frame_list'], } - if out_dir is not None: - face_mesh = { - 'vertices': output['face_vertices'], - 'faces': output['face_faces'], - 'colors': output['face_colors'], - } + return { + OutputKeys.OUTPUT_OBJ: None, + OutputKeys.OUTPUT_IMG: texture_map, + OutputKeys.OUTPUT: results + } - write_obj(os.path.join(out_dir, 'face.obj'), face_mesh) - write_obj(os.path.join(out_dir, 'head.obj'), 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] - return {OutputKeys.OUTPUT: results} + 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) - def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return inputs + result = { + OutputKeys.OUTPUT_OBJ: output_obj, + OutputKeys.OUTPUT_IMG: texture_map, + OutputKeys.OUTPUT: None if render else results, + } + return result diff --git a/tests/pipelines/test_face_reconstruction.py b/tests/pipelines/test_face_reconstruction.py index d4370da3..b35482fb 100644 --- a/tests/pipelines/test_face_reconstruction.py +++ b/tests/pipelines/test_face_reconstruction.py @@ -1,8 +1,13 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import io +import os import os.path as osp import sys import unittest +import cv2 +from moviepy.editor import ImageSequenceClip + from modelscope.hub.snapshot_download import snapshot_download from modelscope.models.cv.face_reconstruction.utils import write_obj from modelscope.outputs import OutputKeys @@ -22,17 +27,35 @@ class FaceReconstructionTest(unittest.TestCase, DemoCompatibilityCheck): self.model_id = 'damo/cv_resnet50_face-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, 'hrn_mesh_mid.obj'), mesh) + + # export rotation video + frame_list = result[OutputKeys.OUTPUT]['frame_list'] + video = ImageSequenceClip(sequence=frame_list, fps=30) + video.write_videofile( + os.path.join(save_root, 'rotate.mp4'), fps=30, audio=False) + del frame_list + + # save visualization image + vis_image = result[OutputKeys.OUTPUT]['vis_image'] + cv2.imwrite(os.path.join(save_root, 'vis_image.jpg'), vis_image) + + print(f'Output written to {osp.abspath(save_root)}') + def pipeline_inference(self, pipeline: Pipeline, input_location: str): result = pipeline(input_location) - mesh = result[OutputKeys.OUTPUT] - write_obj('result_face_reconstruction.obj', mesh) - print( - f'Output written to {osp.abspath("result_face_reconstruction.obj")}' - ) + self.save_results(result, './face_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) + model_dir = snapshot_download(self.model_id, revision='v2.0.0-HRN') face_reconstruction = pipeline( Tasks.face_reconstruction, model=model_dir) self.pipeline_inference(face_reconstruction, self.test_image) @@ -40,7 +63,9 @@ class FaceReconstructionTest(unittest.TestCase, DemoCompatibilityCheck): @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_modelhub(self): face_reconstruction = pipeline( - Tasks.face_reconstruction, model=self.model_id) + Tasks.face_reconstruction, + model=self.model_id, + model_revision='v2.0.0-HRN') self.pipeline_inference(face_reconstruction, self.test_image) @unittest.skip('demo compatibility test is only enabled on a needed-basis')