Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-09-01 19:49:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

update face reconstruction to HRN(CVPR2023)

Browse Source 
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
This commit is contained in:
biwen.lbw
2023-04-06 14:03:27 +08:00
committed by wenmeng.zwm
parent 8d8de93779
commit 9ec8f14f6f
17 changed files with 3135 additions and 1497 deletions
Download Patch File Download Diff File Expand all files Collapse all files

663
modelscope/models/cv/face_reconstruction/models/bfm.py
View File

@@ -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

47
modelscope/models/cv/face_reconstruction/models/de_retouching_module.py Normal file
View File

@@ -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

430
modelscope/models/cv/face_reconstruction/models/facelandmark/large_model_infer.py
View File

@@ -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

894
modelscope/models/cv/face_reconstruction/models/facerecon_model.py
View File

File diff suppressed because it is too large Load Diff

46
modelscope/models/cv/face_reconstruction/models/losses.py
View File

@@ -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):

186
modelscope/models/cv/face_reconstruction/models/nv_diffrast.py
View File

@@ -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

1
modelscope/models/cv/face_reconstruction/models/opt.py
View File

@@ -11,3 +11,4 @@ use_ddp = False
use_last_fc = False
z_far = 15.0
z_near = 5.0
lr = 0.001

0
modelscope/models/cv/face_reconstruction/models/pix2pix/__init__.py Normal file
View File

848
modelscope/models/cv/face_reconstruction/models/pix2pix/networks.py Normal file
View File

@@ -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 <opt.n_epochs> epochs
and linearly decay the rate to zero over the next <opt.n_epochs_decay> 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 <init_func>
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 <init_net>. 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 <n_layers_D> (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 <init_net>. 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 <forward> 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)

125
modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_model.py Normal file
View File

@@ -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 <BaseModel.get_current_losses>
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 <BaseModel.get_current_visuals>
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
# <BaseModel.save_networks> and <BaseModel.load_networks>
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 <BaseModel.setup>.
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 <optimize_parameters> and <test>."""
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

26
modelscope/models/cv/face_reconstruction/models/pix2pix/pix2pix_options.py Normal file
View File

@@ -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'

325
modelscope/models/cv/face_reconstruction/models/renderer.py Executable file
View File

@@ -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

173
modelscope/models/cv/face_reconstruction/models/unet.py Normal file
View File

@@ -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

266
modelscope/models/cv/face_reconstruction/utils.py
View File

@@ -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)

22
modelscope/outputs/outputs.py
View File

@@ -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],

541
modelscope/pipelines/cv/face_reconstruction_pipeline.py
View File

@@ -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

39
tests/pipelines/test_face_reconstruction.py
View File

@@ -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')