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

add head_reconstruction and text_to_head model

Browse Source 
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14099746

* add head_reconstruction and text_to_head model

* change savedir
This commit is contained in:
biwen.lbw
2023-09-21 10:41:25 +08:00
committed by wenmeng.zwm
parent 8c80b0c3f5
commit 10d1fbe86f
20 changed files with 3875 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
modelscope/metainfo.py
View File

@@ -366,6 +366,8 @@ class Pipelines(object):
hand_detection = 'yolox-pai_hand-detection'
skin_retouching = 'unet-skin-retouching'
face_reconstruction = 'resnet50-face-reconstruction'
head_reconstruction = 'HRN-head-reconstruction'
text_to_head = 'HRN-text-to-head'
tinynas_classification = 'tinynas-classification'
easyrobust_classification = 'easyrobust-classification'
tinynas_detection = 'tinynas-detection'

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

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

673
modelscope/models/cv/head_reconstruction/models/bfm.py Normal file
View File

@@ -0,0 +1,673 @@
# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
import os
import numpy as np
import torch
import torch.nn.functional as F
from scipy.io import loadmat
from modelscope.models.cv.face_reconstruction.utils import read_obj
def perspective_projection(focal, center):
# return p.T (N, 3) @ (3, 3)
return np.array([focal, 0, center, 0, focal, center, 0, 0,
1]).reshape([3, 3]).astype(np.float32).transpose()
class SH:
def __init__(self):
self.a = [np.pi, 2 * np.pi / np.sqrt(3.), 2 * np.pi / np.sqrt(8.)]
self.c = [
1 / np.sqrt(4 * np.pi),
np.sqrt(3.) / np.sqrt(4 * np.pi),
3 * np.sqrt(5.) / np.sqrt(12 * np.pi)
]
class ParametricFaceModel:
def __init__(self,
assets_root='assets',
recenter=True,
camera_distance=10.,
init_lit=np.array([0.8, 0, 0, 0, 0, 0, 0, 0, 0]),
focal=1015.,
center=112.,
is_train=True,
default_name='BFM_model_front.mat'):
model = loadmat(os.path.join(assets_root, '3dmm/BFM', default_name))
model_bfm_front = loadmat(
os.path.join(assets_root, '3dmm/BFM/BFM_model_front.mat'))
self.mean_shape_ori = model_bfm_front['meanshape'].astype(np.float32)
# mean face shape. [3*N,1]
self.mean_shape = model['meanshape'].astype(np.float32) # (1, 107127)
# identity basis. [3*N,80]
self.id_base = model['idBase'].astype(np.float32) # (107127, 80)
# expression basis. [3*N,64]
self.exp_base = model['exBase'].astype(np.float32) # (107127, 64)
# mean face texture. [3*N,1] (0-255)
self.mean_tex = model['meantex'].astype(np.float32) # (1, 107127)
# texture basis. [3*N,80]
self.tex_base = model['texBase'].astype(np.float32) # (107127, 80)
self.bfm_keep_inds = np.load(
os.path.join(assets_root, '3dmm/inds/bfm_keep_inds.npy'))
self.ours_hair_area_inds = np.load(
os.path.join(assets_root, '3dmm/inds/ours_hair_area_inds.npy'))
if default_name == 'ourRefineFull_model.mat':
self.mean_tex = self.mean_tex.reshape(1, -1, 3)
mean_tex_keep = self.mean_tex[:, self.bfm_keep_inds]
self.mean_tex[:, :len(self.bfm_keep_inds)] = mean_tex_keep
self.mean_tex[:,
len(self.bfm_keep_inds):] = np.array([200, 146,
118])[None,
None]
self.mean_tex[:, self.ours_hair_area_inds] = 40.0
self.mean_tex = self.mean_tex.reshape(1, -1)
self.mean_tex = np.ascontiguousarray(self.mean_tex)
self.tex_base = self.tex_base.reshape(-1, 3, 80)
tex_base_keep = self.tex_base[self.bfm_keep_inds]
self.tex_base[:len(self.bfm_keep_inds)] = tex_base_keep
self.tex_base[len(self.bfm_keep_inds):] = 0.0
self.tex_base = self.tex_base.reshape(-1, 80)
self.tex_base = np.ascontiguousarray(self.tex_base)
# face indices for each vertex that lies in. starts from 0. [N,8]
self.point_buf = model['point_buf'].astype(np.int64) - 1 # (35709, 8)
# vertex indices for each face. starts from 0. [F,3]
self.face_buf = model['tri'].astype(np.int64) - 1 # (70789, 3)
# vertex indices for 68 landmarks. starts from 0. [68,1]
self.keypoints = np.squeeze(model['keypoints']).astype(np.int64) - 1
if default_name == 'ourRefineFull_model.mat':
self.keypoints = np.load(
os.path.join(
assets_root,
'3dmm/inds/our_refine0223_basis_withoutEyes_withUV_keypoints_inds.npy'
)).astype(np.int64)
self.point_buf = self.point_buf[:, :8] + 1
if is_train:
# vertex indices for small face region to compute photometric error. starts from 0.
self.front_mask = np.squeeze(model['frontmask2_idx']).astype(
np.int64) - 1
# vertex indices for each face from small face region. starts from 0. [f,3]
self.front_face_buf = model['tri_mask2'].astype(np.int64) - 1
# vertex indices for pre-defined skin region to compute reflectance loss
self.skin_mask = np.squeeze(model['skinmask'])
if default_name == 'ourRefineFull_model.mat':
nose_reduced_mesh = read_obj(
os.path.join(assets_root,
'3dmm/adjust_part/our_full/145_nose.obj'))
self.nose_reduced_part = nose_reduced_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
neck_mesh = read_obj(
os.path.join(assets_root,
'3dmm/adjust_part/our_full/154_neck.obj'))
self.neck_adjust_part = neck_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
eyes_mesh = read_obj(
os.path.join(
assets_root,
'3dmm/adjust_part/our_full/our_mean_adjust_eyes.obj'))
self.eyes_adjust_part = eyes_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
self.neck_slim_part = None
self.neck_stretch_part = None
elif default_name == 'ourRefineBFMEye0504_model.mat':
nose_reduced_mesh = read_obj(
os.path.join(assets_root,
'3dmm/adjust_part/our_full_bfmEyes/145_nose.obj'))
self.nose_reduced_part = nose_reduced_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
neck_mesh = read_obj(
os.path.join(assets_root,
'3dmm/adjust_part/our_full_bfmEyes/146_neck.obj'))
self.neck_adjust_part = neck_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
self.eyes_adjust_part = None
neck_slim_mesh = read_obj(
os.path.join(
assets_root,
'3dmm/adjust_part/our_full_bfmEyes/147_neckSlim2.obj'))
self.neck_slim_part = neck_slim_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
neck_stretch_mesh = read_obj(
os.path.join(
assets_root,
'3dmm/adjust_part/our_full_bfmEyes/148_neckLength.obj'))
self.neck_stretch_part = neck_stretch_mesh['vertices'].reshape(
(1, -1)) - self.mean_shape
else:
self.nose_reduced_part = None
self.neck_adjust_part = None
self.eyes_adjust_part = None
self.neck_slim_part = None
self.neck_stretch_part = None
if recenter:
mean_shape = self.mean_shape.reshape([-1, 3])
mean_shape_ori = self.mean_shape_ori.reshape([-1, 3])
mean_shape = mean_shape - np.mean(
mean_shape_ori[:35709, ...], axis=0, keepdims=True)
self.mean_shape = mean_shape.reshape([-1, 1])
eye_corner_inds = np.load(
os.path.join(assets_root, '3dmm/inds/eye_corner_inds.npy'))
self.eye_corner_inds = torch.from_numpy(eye_corner_inds).long()
eye_lines = np.load(
os.path.join(assets_root, '3dmm/inds/eye_corner_lines.npy'))
self.eye_lines = torch.from_numpy(eye_lines).long()
self.center = center
self.persc_proj = perspective_projection(focal, self.center)
self.camera_distance = camera_distance
self.SH = SH()
self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32)
def to(self, device):
self.device = device
for key, value in self.__dict__.items():
if type(value).__module__ == np.__name__:
setattr(self, key, torch.tensor(value).to(device))
def compute_shape(self,
id_coeff,
exp_coeff,
nose_coeff=0.0,
neck_coeff=0.0,
eyes_coeff=0.0,
neckSlim_coeff=0.0,
neckStretch_coeff=0.0):
"""
Return:
face_shape -- torch.tensor, size (B, N, 3)
Parameters:
id_coeff -- torch.tensor, size (B, 80), identity coeffs
exp_coeff -- torch.tensor, size (B, 64), expression coeffs
"""
batch_size = id_coeff.shape[0]
id_part = torch.einsum('ij,aj->ai', self.id_base, id_coeff)
exp_part = torch.einsum('ij,aj->ai', self.exp_base, exp_coeff)
face_shape = id_part + exp_part + self.mean_shape.reshape([1, -1])
if nose_coeff != 0:
face_shape = face_shape + nose_coeff * self.nose_reduced_part
if neck_coeff != 0:
face_shape = face_shape + neck_coeff * self.neck_adjust_part
if eyes_coeff != 0 and self.eyes_adjust_part is not None:
face_shape = face_shape + eyes_coeff * self.eyes_adjust_part
if neckSlim_coeff != 0 and self.neck_slim_part is not None:
face_shape = face_shape + neckSlim_coeff * self.neck_slim_part
if neckStretch_coeff != 0 and self.neck_stretch_part is not None:
neck_stretch_part = self.neck_stretch_part.reshape(1, -1, 3)
neck_stretch_part_top = neck_stretch_part[0, 37476, 1]
neck_stretch_part_bottom = neck_stretch_part[0, 37357, 1]
neck_stretch_height = neck_stretch_part_top - neck_stretch_part_bottom
face_shape_ = face_shape.reshape(1, -1, 3)
face_shape_top = face_shape_[0, 37476, 1]
face_shape_bottom = face_shape_[0, 37357, 1]
face_shape_height = face_shape_top - face_shape_bottom
target_neck_height = 0.72 # top ind 37476, bottom ind 37357
neckStretch_coeff = (target_neck_height
- face_shape_height) / neck_stretch_height
face_shape = face_shape + neckStretch_coeff * self.neck_stretch_part
return face_shape.reshape([batch_size, -1, 3])
def compute_texture(self, tex_coeff, normalize=True):
"""
Return:
face_texture -- torch.tensor, size (B, N, 3), in RGB order, range (0, 1.)
Parameters:
tex_coeff -- torch.tensor, size (B, 80)
"""
batch_size = tex_coeff.shape[0]
face_texture = torch.einsum('ij,aj->ai', self.tex_base,
tex_coeff) + self.mean_tex
if normalize:
face_texture = face_texture / 255.
return face_texture.reshape([batch_size, -1, 3])
def compute_norm(self, face_shape):
"""
Return:
vertex_norm -- torch.tensor, size (B, N, 3)
Parameters:
face_shape -- torch.tensor, size (B, N, 3)
"""
v1 = face_shape[:, self.face_buf[:, 0]]
v2 = face_shape[:, self.face_buf[:, 1]]
v3 = face_shape[:, self.face_buf[:, 2]]
e1 = v1 - v2
e2 = v2 - v3
face_norm = torch.cross(e1, e2, dim=-1)
face_norm = F.normalize(face_norm, dim=-1, p=2)
face_norm = torch.cat(
[face_norm,
torch.zeros(face_norm.shape[0], 1, 3).to(self.device)],
dim=1)
vertex_norm = torch.sum(face_norm[:, self.point_buf], dim=2)
vertex_norm = F.normalize(vertex_norm, dim=-1, p=2)
return vertex_norm
def compute_color(self, face_texture, face_norm, gamma):
"""
Return:
face_color -- torch.tensor, size (B, N, 3), range (0, 1.)
Parameters:
face_texture -- torch.tensor, size (B, N, 3), from texture model, range (0, 1.)
face_norm -- torch.tensor, size (B, N, 3), rotated face normal
gamma -- torch.tensor, size (B, 27), SH coeffs
"""
batch_size = gamma.shape[0]
a, c = self.SH.a, self.SH.c
gamma = gamma.reshape([batch_size, 3, 9])
gamma = gamma + self.init_lit
gamma = gamma.permute(0, 2, 1)
y1 = a[0] * c[0] * torch.ones_like(face_norm[..., :1]).to(self.device)
y2 = -a[1] * c[1] * face_norm[..., 1:2]
y3 = a[1] * c[1] * face_norm[..., 2:]
y4 = -a[1] * c[1] * face_norm[..., :1]
y5 = a[2] * c[2] * face_norm[..., :1] * face_norm[..., 1:2]
y6 = -a[2] * c[2] * face_norm[..., 1:2] * face_norm[..., 2:]
y7 = 0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm[..., 2:]**2 - 1)
y8 = -a[2] * c[2] * face_norm[..., :1] * face_norm[..., 2:]
y9 = 0.5 * a[2] * c[2] * (
face_norm[..., :1]**2 - face_norm[..., 1:2]**2)
Y = torch.cat([y1, y2, y3, y4, y5, y6, y7, y8, y9], dim=-1)
r = Y @ gamma[..., :1]
g = Y @ gamma[..., 1:2]
b = Y @ gamma[..., 2:]
face_color = torch.cat([r, g, b], dim=-1) * face_texture
return face_color
def compute_rotation(self, angles):
"""
Return:
rot -- torch.tensor, size (B, 3, 3) pts @ trans_mat
Parameters:
angles -- torch.tensor, size (B, 3), radian
"""
batch_size = angles.shape[0]
ones = torch.ones([batch_size, 1]).to(self.device)
zeros = torch.zeros([batch_size, 1]).to(self.device)
x, y, z = angles[:, :1], angles[:, 1:2], angles[:, 2:],
value_list = [
ones, zeros, zeros, zeros,
torch.cos(x), -torch.sin(x), zeros,
torch.sin(x),
torch.cos(x)
]
rot_x = torch.cat(value_list, dim=1).reshape([batch_size, 3, 3])
value_list = [
torch.cos(y), zeros,
torch.sin(y), zeros, ones, zeros, -torch.sin(y), zeros,
torch.cos(y)
]
rot_y = torch.cat(value_list, dim=1).reshape([batch_size, 3, 3])
value_list = [
torch.cos(z), -torch.sin(z), zeros,
torch.sin(z),
torch.cos(z), zeros, zeros, zeros, ones
]
rot_z = torch.cat(value_list, dim=1).reshape([batch_size, 3, 3])
rot = rot_z @ rot_y @ rot_x
return rot.permute(0, 2, 1)
def to_camera(self, face_shape):
face_shape[..., -1] = self.camera_distance - face_shape[..., -1]
return face_shape
def to_image(self, face_shape):
"""
Return:
face_proj -- torch.tensor, size (B, N, 2), y direction is opposite to v direction
Parameters:
face_shape -- torch.tensor, size (B, N, 3)
"""
# to image_plane
face_proj = face_shape @ self.persc_proj
face_proj = face_proj[..., :2] / face_proj[..., 2:]
return face_proj
def transform(self, face_shape, rot, trans):
"""
Return:
face_shape -- torch.tensor, size (B, N, 3) pts @ rot + trans
Parameters:
face_shape -- torch.tensor, size (B, N, 3)
rot -- torch.tensor, size (B, 3, 3)
trans -- torch.tensor, size (B, 3)
"""
return face_shape @ rot + trans.unsqueeze(1)
def get_landmarks(self, face_proj):
"""
Return:
face_lms -- torch.tensor, size (B, 68, 2)
Parameters:
face_proj -- torch.tensor, size (B, N, 2)
"""
return face_proj[:, self.keypoints]
def split_coeff(self, coeffs):
"""
Return:
coeffs_dict -- a dict of torch.tensors
Parameters:
coeffs -- torch.tensor, size (B, 256)
"""
if type(coeffs) == dict and 'id' in coeffs:
return coeffs
id_coeffs = coeffs[:, :80]
exp_coeffs = coeffs[:, 80:144]
tex_coeffs = coeffs[:, 144:224]
angles = coeffs[:, 224:227]
gammas = coeffs[:, 227:254]
translations = coeffs[:, 254:]
return {
'id': id_coeffs,
'exp': exp_coeffs,
'tex': tex_coeffs,
'angle': angles,
'gamma': gammas,
'trans': translations
}
def merge_coeff(self, coeffs):
"""
Return:
coeffs_dict -- a dict of torch.tensors
Parameters:
coeffs -- torch.tensor, size (B, 256)
"""
names = ['id', 'exp', 'tex', 'angle', 'gamma', 'trans']
coeffs_merge = []
for name in names:
coeffs_merge.append(coeffs[name].detach())
coeffs_merge = torch.cat(coeffs_merge, dim=1)
return coeffs_merge
def reverse_recenter(self, face_shape):
batch_size = face_shape.shape[0]
face_shape = face_shape.reshape([-1, 3])
mean_shape_ori = self.mean_shape_ori.reshape([-1, 3])
face_shape = face_shape + torch.mean(
mean_shape_ori[:35709, ...], dim=0, keepdim=True)
face_shape = face_shape.reshape([batch_size, -1, 3])
return face_shape
def add_nonlinear_offset_eyes(self, face_shape, shape_offset):
assert face_shape.shape[0] == 1 and shape_offset.shape[0] == 1
face_shape = face_shape[0]
shape_offset = shape_offset[0]
corner_shape = face_shape[-625:, :]
corner_offset = shape_offset[self.eye_corner_inds]
for i in range(len(self.eye_lines)):
corner_shape[self.eye_lines[i]] += corner_offset[i][None, ...]
face_shape[-625:, :] = corner_shape
l_eye_landmarks = [11540, 11541]
r_eye_landmarks = [4271, 4272]
l_eye_offset = torch.mean(
shape_offset[l_eye_landmarks], dim=0, keepdim=True)
face_shape[37082:37082 + 609] += l_eye_offset
r_eye_offset = torch.mean(
shape_offset[r_eye_landmarks], dim=0, keepdim=True)
face_shape[37082 + 609:37082 + 609 + 608] += r_eye_offset
face_shape = face_shape[None, ...]
return face_shape
def add_nonlinear_offset(self, face_shape, shape_offset_uv, UVs):
"""
Args:
face_shape: torch.tensor, size (1, N, 3)
shape_offset_uv: torch.tensor, size (1, h, w, 3)
UVs: torch.tensor, size (N, 2)
Returns:
"""
assert face_shape.shape[0] == 1 and shape_offset_uv.shape[0] == 1
face_shape = face_shape[0]
shape_offset_uv = shape_offset_uv[0]
h, w = shape_offset_uv.shape[:2]
UVs_coords = UVs.clone()
UVs_coords[:, 0] *= w
UVs_coords[:, 1] *= h
UVs_coords_int = torch.floor(UVs_coords)
UVs_coords_float = UVs_coords - UVs_coords_int
UVs_coords_int = UVs_coords_int.long()
shape_lt = shape_offset_uv[(h - 1 - UVs_coords_int[:, 1]).clamp(
0, h - 1), UVs_coords_int[:, 0].clamp(0, w - 1)] # (N, 3)
shape_lb = shape_offset_uv[(h - UVs_coords_int[:, 1]).clamp(0, h - 1),
UVs_coords_int[:, 0].clamp(0, w - 1)]
shape_rt = shape_offset_uv[(h - 1
- UVs_coords_int[:, 1]).clamp(0, h - 1),
(UVs_coords_int[:, 0] + 1).clamp(0, w - 1)]
shape_rb = shape_offset_uv[(h - UVs_coords_int[:, 1]).clamp(0, h - 1),
(UVs_coords_int[:, 0] + 1).clamp(0, w - 1)]
value_1 = shape_lt * (
1 - UVs_coords_float[:, :1]) * UVs_coords_float[:, 1:]
value_2 = shape_lb * (1 - UVs_coords_float[:, :1]) * (
1 - UVs_coords_float[:, 1:])
value_3 = shape_rt * UVs_coords_float[:, :1] * UVs_coords_float[:, 1:]
value_4 = shape_rb * UVs_coords_float[:, :1] * (
1 - UVs_coords_float[:, 1:])
offset_shape = value_1 + value_2 + value_3 + value_4 # (B, N, 3)
face_shape = (face_shape + offset_shape)[None, ...]
return face_shape, offset_shape[None, ...]
def compute_for_render_head_fitting(self,
coeffs,
shape_offset_uv,
texture_offset_uv,
shape_offset_uv_head,
texture_offset_uv_head,
UVs,
reverse_recenter=True,
get_eyes=False,
get_neck=False,
nose_coeff=0.0,
neck_coeff=0.0,
eyes_coeff=0.0):
if type(coeffs) == dict:
coef_dict = coeffs
elif type(coeffs) == torch.Tensor:
coef_dict = self.split_coeff(coeffs)
face_shape = self.compute_shape(
coef_dict['id'],
coef_dict['exp'],
nose_coeff=nose_coeff,
neck_coeff=neck_coeff,
eyes_coeff=eyes_coeff) # (1, n, 3)
if reverse_recenter:
face_shape_ori_noRecenter = self.reverse_recenter(
face_shape.clone())
else:
face_shape_ori_noRecenter = face_shape.clone()
face_vertex_ori = self.to_camera(face_shape_ori_noRecenter)
face_shape[:, :35241, :], shape_offset = self.add_nonlinear_offset(
face_shape[:, :35241, :], shape_offset_uv,
UVs[:35709, ...][self.bfm_keep_inds]) # (1, n, 3)
if get_eyes:
face_shape = self.add_nonlinear_offset_eyes(
face_shape, shape_offset)
if get_neck:
face_shape[:, 35241:37082, ...], _ = self.add_nonlinear_offset(
face_shape[:, 35241:37082, ...], shape_offset_uv_head,
UVs[35709:, ...]) # (1, n, 3)
else:
face_shape[:, self.ours_hair_area_inds,
...], _ = self.add_nonlinear_offset(
face_shape[:, self.ours_hair_area_inds,
...], shape_offset_uv_head,
UVs[self.ours_hair_area_inds + (35709 - 35241),
...]) # (1, n, 3)
if reverse_recenter:
face_shape_offset_noRecenter = self.reverse_recenter(
face_shape.clone())
else:
face_shape_offset_noRecenter = face_shape.clone()
face_vertex_offset = self.to_camera(face_shape_offset_noRecenter)
rotation = self.compute_rotation(coef_dict['angle'])
face_shape_transformed = self.transform(face_shape, rotation,
coef_dict['trans'])
face_vertex = self.to_camera(face_shape_transformed)
face_proj = self.to_image(face_vertex)
landmark = self.get_landmarks(face_proj)
face_texture = self.compute_texture(coef_dict['tex']) # (1, n, 3)
face_texture[:, :35241, :], texture_offset = self.add_nonlinear_offset(
face_texture[:, :35241, :], texture_offset_uv,
UVs[:35709, ...][self.bfm_keep_inds])
face_texture[:, 35241:37082, :], _ = self.add_nonlinear_offset(
face_texture[:, 35241:37082, :], texture_offset_uv_head,
UVs[35709:, ...])
face_norm = self.compute_norm(face_shape)
face_norm_roted = face_norm @ rotation
face_color = self.compute_color(face_texture, face_norm_roted,
coef_dict['gamma'])
return face_vertex, face_texture, face_color, landmark, face_vertex_ori, face_vertex_offset, face_proj
def compute_for_render_head(self,
coeffs,
shape_offset_uv,
texture_offset_uv,
shape_offset_uv_head,
texture_offset_uv_head,
UVs,
reverse_recenter=True,
nose_coeff=0.0,
neck_coeff=0.0,
eyes_coeff=0.0,
neckSlim_coeff=0.0,
neckStretch_coeff=0.0):
if type(coeffs) == dict:
coef_dict = coeffs
elif type(coeffs) == torch.Tensor:
coef_dict = self.split_coeff(coeffs)
face_shape = self.compute_shape(
coef_dict['id'],
coef_dict['exp'],
nose_coeff=nose_coeff,
neck_coeff=neck_coeff,
eyes_coeff=eyes_coeff,
neckSlim_coeff=neckSlim_coeff,
neckStretch_coeff=neckStretch_coeff) # (1, n, 3)
if reverse_recenter:
face_shape_ori_noRecenter = self.reverse_recenter(
face_shape.clone())
else:
face_shape_ori_noRecenter = face_shape.clone()
face_vertex_ori = self.to_camera(face_shape_ori_noRecenter)
face_shape[:, :35709, :], shape_offset = self.add_nonlinear_offset(
face_shape[:, :35709, :], shape_offset_uv, UVs[:35709,
...]) # (1, n, 3)
face_shape[:, 35709:,
...], _ = self.add_nonlinear_offset(face_shape[:, 35709:,
...],
shape_offset_uv_head,
UVs[35709:,
...]) # (1, n, 3)
if reverse_recenter:
face_shape_offset_noRecenter = self.reverse_recenter(
face_shape.clone())
else:
face_shape_offset_noRecenter = face_shape.clone()
face_vertex_offset = self.to_camera(face_shape_offset_noRecenter)
rotation = self.compute_rotation(coef_dict['angle'])
face_shape_transformed = self.transform(face_shape, rotation,
coef_dict['trans'])
face_vertex = self.to_camera(face_shape_transformed)
face_proj = self.to_image(face_vertex)
landmark = self.get_landmarks(face_proj)
face_texture = self.compute_texture(coef_dict['tex']) # (1, n, 3)
face_texture[:, :35709, :], texture_offset = self.add_nonlinear_offset(
face_texture[:, :35709, :], texture_offset_uv, UVs[:35709, ...])
face_texture[:, 35709:, :], _ = self.add_nonlinear_offset(
face_texture[:, 35709:, :], texture_offset_uv_head, UVs[35709:,
...])
face_norm = self.compute_norm(face_shape)
face_norm_roted = face_norm @ rotation
face_color = self.compute_color(face_texture, face_norm_roted,
coef_dict['gamma'])
return face_vertex, face_texture, face_color, landmark, face_vertex_ori, face_vertex_offset, face_proj

196
modelscope/models/cv/head_reconstruction/models/head_segmentation.py Normal file
View File

@@ -0,0 +1,196 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import cv2
import json
import numpy as np
import tensorflow as tf
if tf.__version__ >= '2.0':
tf = tf.compat.v1
tf.disable_eager_execution()
class HeadSegmentor():
def __init__(self, model_root):
"""The HeadSegmentor is implemented based on https://arxiv.org/abs/2004.04955
Args:
model_root: the root directory of the model files
"""
self.sess = self.load_sess(
os.path.join(model_root, 'head_segmentation',
'Matting_headparser_6_18.pb'))
self.sess_detect = self.load_sess(
os.path.join(model_root, 'head_segmentation', 'face_detect.pb'))
self.sess_face = self.load_sess(
os.path.join(model_root, 'head_segmentation', 'segment_face.pb'))
def load_sess(self, model_path):
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
with tf.gfile.FastGFile(model_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
sess.run(tf.global_variables_initializer())
return sess
def process(self, image):
""" image: bgr
"""
h, w, c = image.shape
faceRects = self.detect_face(image)
face_num = len(faceRects)
all_head_alpha = []
all_face_mask = []
for i in range(face_num):
y1 = faceRects[i][0]
y2 = faceRects[i][1]
x1 = faceRects[i][2]
x2 = faceRects[i][3]
pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box(
y1, y2, x1, x2, 0.15, 0.15, 0.15, 0.15, h, w)
temp_img = image.copy()
roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2]
output_alpha = self.sess_face.run(
self.sess_face.graph.get_tensor_by_name('output_alpha_face:0'),
feed_dict={'input_image_face:0': roi_img[:, :, ::-1]})
face_mask = np.zeros((h, w, 3))
face_mask[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha
all_face_mask.append(face_mask)
cv2.imwrite(str(i) + 'face.jpg', face_mask)
cv2.imwrite(str(i) + 'face_roi.jpg', roi_img)
for i in range(face_num):
y1 = faceRects[i][0]
y2 = faceRects[i][1]
x1 = faceRects[i][2]
x2 = faceRects[i][3]
pad_y1, pad_y2, pad_x1, pad_x2 = self.pad_box(
y1, y2, x1, x2, 1.47, 1.47, 1.3, 2.0, h, w)
temp_img = image.copy()
for j in range(face_num):
y1 = faceRects[j][0]
y2 = faceRects[j][1]
x1 = faceRects[j][2]
x2 = faceRects[j][3]
small_y1, small_y2, small_x1, small_x2 = self.pad_box(
y1, y2, x1, x2, -0.1, -0.1, -0.1, -0.1, h, w)
small_width = small_x2 - small_x1
small_height = small_y2 - small_y1
if (small_x1 < 0 or small_y1 < 0 or small_width < 3
or small_height < 3 or small_x2 > w or small_y2 > h):
continue
# if(i!=j):
# temp_img[small_y1:small_y2,small_x1:small_x2]=0
if (i != j):
temp_img = temp_img * (1.0 - all_face_mask[j] / 255.0)
roi_img = temp_img[pad_y1:pad_y2, pad_x1:pad_x2]
output_alpha = self.sess.run(
self.sess.graph.get_tensor_by_name('output_alpha:0'),
feed_dict={'input_image:0': roi_img[:, :, ::-1]})
head_alpha = np.zeros((h, w))
head_alpha[pad_y1:pad_y2, pad_x1:pad_x2] = output_alpha[:, :, 0]
if np.sum(head_alpha) > 255 * w * h * 0.01 * 0.01:
all_head_alpha.append(head_alpha)
head_num = len(all_head_alpha)
head_elements = []
if head_num == 0:
return head_elements
for i in range(head_num):
head_alpha = all_head_alpha[i]
head_elements.append(head_alpha)
return head_elements
def pad_box(self, y1, y2, x1, x2, left_ratio, right_ratio, top_ratio,
bottom_ratio, h, w):
box_w = x2 - x1
box_h = y2 - y1
pad_y1 = np.maximum(np.int32(y1 - top_ratio * box_h), 0)
pad_y2 = np.minimum(np.int32(y2 + bottom_ratio * box_h), h - 1)
pad_x1 = np.maximum(np.int32(x1 - left_ratio * box_w), 0)
pad_x2 = np.minimum(np.int32(x2 + right_ratio * box_w), w - 1)
return pad_y1, pad_y2, pad_x1, pad_x2
def detect_face(self, img):
h, w, c = img.shape
input_img = cv2.resize(img[:, :, ::-1], (512, 512))
boxes, scores, num_detections = self.sess_detect.run(
[
self.sess_detect.graph.get_tensor_by_name('tower_0/boxes:0'),
self.sess_detect.graph.get_tensor_by_name('tower_0/scores:0'),
self.sess_detect.graph.get_tensor_by_name(
'tower_0/num_detections:0')
],
feed_dict={
'tower_0/images:0': input_img[np.newaxis],
'training_flag:0': False
})
faceRects = []
for i in range(num_detections[0]):
if scores[0, i] < 0.5:
continue
y1 = np.int32(boxes[0, i, 0] * h)
x1 = np.int32(boxes[0, i, 1] * w)
y2 = np.int32(boxes[0, i, 2] * h)
x2 = np.int32(boxes[0, i, 3] * w)
if x2 <= x1 + 3 or y2 <= y1 + 3:
continue
faceRects.append((y1, y2, x1, x2, y2 - y1, x2 - x1))
sorted(faceRects, key=lambda x: x[4] * x[5], reverse=True)
return faceRects
def generate_json(self, status_code, status_msg, ori_url, result_element,
track_id):
data = {}
data['originUri'] = ori_url
data['elements'] = result_element
data['statusCode'] = status_code
data['statusMessage'] = status_msg
data['requestId'] = track_id
return json.dumps(data)
def get_box(self, alpha):
h, w = alpha.shape
start_h = 0
end_h = 0
start_w = 0
end_w = 0
for i in range(0, h, 3):
line = alpha[i, :]
if np.max(line) >= 1:
start_h = i
break
for i in range(0, w, 3):
line = alpha[:, i]
if np.max(line) >= 1:
start_w = i
break
for i in range(0, h, 3):
i = h - 1 - i
line = alpha[i, :]
if np.max(line) >= 1:
end_h = i
if end_h < h - 1:
end_h = end_h + 1
break
for i in range(0, w, 3):
i = w - 1 - i
line = alpha[:, i]
if np.max(line) >= 1:
end_w = i
if end_w < w - 1:
end_w = end_w + 1
break
return start_h, start_w, end_h, end_w

564
modelscope/models/cv/head_reconstruction/models/headrecon_model.py Normal file
View File

@@ -0,0 +1,564 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from collections import OrderedDict
import cv2
import numpy as np
import torch
from modelscope.models import MODELS, TorchModel
from modelscope.models.cv.face_reconstruction.utils import (estimate_normals,
read_obj)
from . import networks, opt
from .bfm import ParametricFaceModel
from .losses import (BinaryDiceLoss, TVLoss, TVLoss_std, landmark_loss,
perceptual_loss, photo_loss, points_loss_horizontal,
reflectance_loss, reg_loss)
from .nv_diffrast import MeshRenderer
@MODELS.register_module('head-reconstruction', 'head_reconstruction')
class HeadReconModel(TorchModel):
def __init__(self, model_dir, *args, **kwargs):
"""The HeadReconModel is implemented based on HRN, publicly available at
https://github.com/youngLBW/HRN
Args:
model_dir: the root directory of the model files
"""
super().__init__(model_dir, *args, **kwargs)
self.model_dir = model_dir
opt.bfm_folder = os.path.join(model_dir, 'assets')
self.opt = opt
self.isTrain = opt.isTrain
self.visual_names = ['output_vis']
self.model_names = ['net_recon']
self.parallel_names = self.model_names + [
'renderer', 'renderer_fitting'
]
# networks
self.net_recon = networks.define_net_recon(
net_recon=opt.net_recon,
use_last_fc=opt.use_last_fc,
init_path=None)
# assets
self.headmodel = ParametricFaceModel(
assets_root=opt.bfm_folder,
camera_distance=opt.camera_d,
focal=opt.focal,
center=opt.center,
is_train=self.isTrain,
default_name='ourRefineBFMEye0504_model.mat')
self.headmodel_for_fitting = ParametricFaceModel(
assets_root=opt.bfm_folder,
camera_distance=opt.camera_d,
focal=opt.focal,
center=opt.center,
is_train=self.isTrain,
default_name='ourRefineFull_model.mat')
# renderer
fov = 2 * np.arctan(opt.center / opt.focal) * 180 / np.pi
self.renderer = MeshRenderer(
rasterize_fov=fov,
znear=opt.z_near,
zfar=opt.z_far,
rasterize_size=int(2 * opt.center))
self.renderer_fitting = MeshRenderer(
rasterize_fov=fov,
znear=opt.z_near,
zfar=opt.z_far,
rasterize_size=int(2 * opt.center))
template_obj_path = os.path.join(
model_dir,
'assets/3dmm/template_mesh/template_ourFull_bfmEyes.obj')
self.template_output_mesh = read_obj(template_obj_path)
self.nonlinear_UVs = self.template_output_mesh['uvs']
self.nonlinear_UVs = torch.from_numpy(self.nonlinear_UVs)
self.jaw_edge_mask = cv2.imread(
os.path.join(model_dir,
'assets/texture/jaw_edge_mask2.png'))[..., 0].astype(
np.float32) / 255.0
self.jaw_edge_mask = cv2.resize(self.jaw_edge_mask, (300, 300))[...,
None]
self.input_imgs = []
self.input_img_hds = []
self.input_fat_img_hds = []
self.atten_masks = []
self.gt_lms = []
self.gt_lm_hds = []
self.trans_ms = []
self.img_names = []
self.face_masks = []
self.head_masks = []
self.input_imgs_coeff = []
self.gt_lms_coeff = []
self.loss_names = [
'all', 'feat', 'color', 'lm', 'reg', 'gamma', 'reflc'
]
self.compute_feat_loss = perceptual_loss
self.comupte_color_loss = photo_loss
self.compute_lm_loss = landmark_loss
self.compute_reg_loss = reg_loss
self.compute_reflc_loss = reflectance_loss
if opt.isTrain:
self.optimizer = torch.optim.Adam(
self.net_recon.parameters(), lr=opt.lr)
self.optimizers = [self.optimizer]
self.parallel_names += ['net_recog']
def set_device(self, device):
self.device = device
self.net_recon = self.net_recon.to(self.device)
self.headmodel.to(self.device)
self.headmodel_for_fitting.to(self.device)
self.nonlinear_UVs = self.nonlinear_UVs.to(self.device)
def load_networks(self, load_path):
state_dict = torch.load(load_path, map_location=self.device)
print('loading the model from %s' % load_path)
for name in self.model_names:
if isinstance(name, str):
net = getattr(self, name)
if isinstance(net, torch.nn.DataParallel):
net = net.module
net.load_state_dict(state_dict[name], strict=False)
def setup(self, checkpoint_path):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
self.load_networks(checkpoint_path)
def parallelize(self, convert_sync_batchnorm=True):
if not self.opt.use_ddp:
for name in self.parallel_names:
if isinstance(name, str):
module = getattr(self, name)
setattr(self, name, module.to(self.device))
else:
for name in self.model_names:
if isinstance(name, str):
module = getattr(self, name)
if convert_sync_batchnorm:
module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
module)
setattr(
self, name,
torch.nn.parallel.DistributedDataParallel(
module.to(self.device),
device_ids=[self.device.index],
find_unused_parameters=True,
broadcast_buffers=True))
# DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient.
for name in self.parallel_names:
if isinstance(name, str) and name not in self.model_names:
module = getattr(self, name)
setattr(self, name, module.to(self.device))
# put state_dict of optimizer to gpu device
if self.opt.phase != 'test':
if self.opt.continue_train:
for optim in self.optimizers:
for state in optim.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(self.device)
def eval(self):
"""Make models eval mode"""
for name in self.model_names:
if isinstance(name, str):
net = getattr(self, name)
net.eval()
def set_render(self, image_res=1024):
fov = 2 * np.arctan(self.opt.center / self.opt.focal) * 180 / np.pi
if image_res is None:
image_res = int(2 * self.opt.center)
self.renderer = MeshRenderer(
rasterize_fov=fov,
znear=self.opt.z_near,
zfar=self.opt.z_far,
rasterize_size=image_res)
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input: a dictionary that contains the data itself and its metadata information.
"""
self.input_img = input['imgs'].to(self.device)
self.input_img_hd = input['imgs_hd'].to(
self.device) if 'imgs_hd' in input else None
if 'imgs_fat_hd' not in input or input['imgs_fat_hd'] is None:
self.input_fat_img_hd = self.input_img_hd
else:
self.input_fat_img_hd = input['imgs_fat_hd'].to(self.device)
self.atten_mask = input['msks'].to(
self.device) if 'msks' in input else None
self.gt_lm = input['lms'].to(self.device) if 'lms' in input else None
self.gt_lm_hd = input['lms_hd'].to(
self.device) if 'lms_hd' in input else None
self.trans_m = input['M'].to(self.device) if 'M' in input else None
self.image_paths = input['im_paths'] if 'im_paths' in input else None
self.img_name = input['img_name'] if 'img_name' in input else None
self.face_mask = input['face_mask'].to(
self.device) if 'face_mask' in input else None
self.head_mask = input['head_mask'].to(
self.device) if 'head_mask' in input else None
self.gt_normals = input['normals'].to(
self.device) if 'normals' in input else None
self.input_img_coeff = input['imgs_coeff'].to(
self.device) if 'imgs_coeff' in input else None
self.gt_lm_coeff = input['lms_coeff'].to(
self.device) if 'lms_coeff' in input else None
def check_head_pose(self, coeffs):
pi = 3.14
if coeffs[0, 225] > pi / 6 or coeffs[0, 225] < -pi / 6:
return False
elif coeffs[0, 224] > pi / 6 or coeffs[0, 224] < -pi / 6:
return False
elif coeffs[0, 226] > pi / 6 or coeffs[0, 226] < -pi / 6:
return False
else:
return True
def get_fusion_mask(self, keep_forehead=True):
self.without_forehead_inds = torch.from_numpy(
np.load(
os.path.join(self.model_dir,
'assets/3dmm/inds/bfm_withou_forehead_inds.npy'))
).long().to(self.device)
h, w = self.shape_offset_uv.shape[1:3]
self.fusion_mask = torch.zeros((h, w)).to(self.device).float()
if keep_forehead:
UVs_coords = self.nonlinear_UVs.clone()[:35709][
self.without_forehead_inds]
else:
UVs_coords = self.nonlinear_UVs.clone()[:35709]
UVs_coords[:, 0] *= w
UVs_coords[:, 1] *= h
UVs_coords_int = torch.floor(UVs_coords)
UVs_coords_int = UVs_coords_int.long()
self.fusion_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:,
0]] = 1
# blur mask
self.fusion_mask = self.fusion_mask.cpu().numpy()
new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
new_kernel2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8))
self.fusion_mask = cv2.dilate(self.fusion_mask, new_kernel1, 1)
self.fusion_mask = cv2.erode(self.fusion_mask, new_kernel2, 1)
self.fusion_mask = cv2.blur(self.fusion_mask, (17, 17))
self.fusion_mask = torch.from_numpy(self.fusion_mask).float().to(
self.device)
def get_edge_mask(self):
h, w = self.shape_offset_uv.shape[1:3]
self.edge_mask = torch.zeros((h, w)).to(self.device).float()
UVs_coords = self.nonlinear_UVs.clone()[self.edge_points_inds]
UVs_coords[:, 0] *= w
UVs_coords[:, 1] *= h
UVs_coords_int = torch.floor(UVs_coords)
UVs_coords_int = UVs_coords_int.long()
self.edge_mask[h - 1 - UVs_coords_int[:, 1], UVs_coords_int[:, 0]] = 1
# blur mask
self.edge_mask = self.edge_mask.cpu().numpy()
new_kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8))
self.edge_mask = cv2.dilate(self.edge_mask, new_kernel1, 1)
self.edge_mask = cv2.blur(self.edge_mask, (5, 5))
self.edge_mask = torch.from_numpy(self.edge_mask).float().to(
self.device)
def blur_shape_offset_uv(self, global_blur=False, blur_size=3):
if self.edge_mask is not None:
shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu(
).numpy()
shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur, (15, 15))
shape_offset_uv_blur = torch.from_numpy(
shape_offset_uv_blur).float().to(self.device)[None, ...]
self.shape_offset_uv = shape_offset_uv_blur * self.edge_mask[
None, ..., None] + self.shape_offset_uv * (
1 - self.edge_mask[None, ..., None])
self.shape_offset_uv = self.shape_offset_uv * self.fusion_mask[None,
...,
None]
if global_blur and blur_size > 0:
shape_offset_uv_blur = self.shape_offset_uv[0].detach().cpu(
).numpy()
shape_offset_uv_blur = cv2.blur(shape_offset_uv_blur,
(blur_size, blur_size))
shape_offset_uv_blur = torch.from_numpy(
shape_offset_uv_blur).float().to(self.device)[None, ...]
self.shape_offset_uv = shape_offset_uv_blur
def blur_offset_edge(self):
shape_offset_uv = self.shape_offset_uv[0].detach().cpu().numpy()
shape_offset_uv_head = self.shape_offset_uv_head[0].detach().cpu(
).numpy()
shape_offset_uv_head = cv2.resize(shape_offset_uv_head, (300, 300))
shape_offset_uv_head = shape_offset_uv_head * (
1 - self.jaw_edge_mask) + shape_offset_uv * self.jaw_edge_mask
shape_offset_uv_head = cv2.resize(shape_offset_uv_head, (100, 100))
self.shape_offset_uv_head = torch.from_numpy(
shape_offset_uv_head).float().to(self.device)[None, ...]
def fitting_nonlinear(self, coeff, n_iters=250):
output_coeff = coeff.detach().clone()
output_coeff = self.headmodel_for_fitting.split_coeff(output_coeff)
output_coeff['id'].requires_grad = True
output_coeff['exp'].requires_grad = True
output_coeff['tex'].requires_grad = True
output_coeff['angle'].requires_grad = True
output_coeff['gamma'].requires_grad = True
output_coeff['trans'].requires_grad = True
self.shape_offset_uv = torch.zeros((1, 300, 300, 3),
dtype=torch.float32).to(self.device)
self.shape_offset_uv.requires_grad = True
self.texture_offset_uv = torch.zeros(
(1, 300, 300, 3), dtype=torch.float32).to(self.device)
self.texture_offset_uv.requires_grad = True
self.shape_offset_uv_head = torch.zeros(
(1, 100, 100, 3), dtype=torch.float32).to(self.device)
self.shape_offset_uv_head.requires_grad = True
self.texture_offset_uv_head = torch.zeros(
(1, 100, 100, 3), dtype=torch.float32).to(self.device)
self.texture_offset_uv_head.requires_grad = True
head_face_inds = np.load(
os.path.join(self.model_dir,
'assets/3dmm/inds/ours_head_face_inds.npy'))
head_face_inds = torch.from_numpy(head_face_inds).to(self.device)
head_faces = self.headmodel_for_fitting.face_buf[head_face_inds]
# print('before fitting', output_coeff)
opt_parameters = [
self.shape_offset_uv, self.texture_offset_uv,
self.shape_offset_uv_head, self.texture_offset_uv_head,
output_coeff['id'], output_coeff['exp'], output_coeff['tex'],
output_coeff['gamma']
]
optim = torch.optim.Adam(opt_parameters, lr=1e-3)
optim_pose = torch.optim.Adam([output_coeff['trans']], lr=1e-1)
self.get_edge_points_horizontal()
for i in range(n_iters): # 500
self.pred_vertex_head, self.pred_tex, self.pred_color_head, self.pred_lm, face_shape, \
face_shape_offset, self.verts_proj_head = \
self.headmodel_for_fitting.compute_for_render_head_fitting(output_coeff, self.shape_offset_uv,
self.texture_offset_uv,
self.shape_offset_uv_head,
self.texture_offset_uv_head,
self.nonlinear_UVs)
self.pred_vertex = self.pred_vertex_head[:, :35241]
self.pred_color = self.pred_color_head[:, :35241]
self.verts_proj = self.verts_proj_head[:, :35241]
self.pred_mask_head, _, self.pred_head, self.occ_head = self.renderer_fitting(
self.pred_vertex_head, head_faces, feat=self.pred_color_head)
self.pred_mask, _, self.pred_face, self.occ_face = self.renderer_fitting(
self.pred_vertex,
self.headmodel_for_fitting.face_buf[:69732],
feat=self.pred_color)
self.pred_coeffs_dict = self.headmodel_for_fitting.split_coeff(
output_coeff)
self.compute_losses_fitting()
if i < 150:
optim_pose.zero_grad()
(self.loss_lm + self.loss_color * 0.1).backward()
optim_pose.step()
else:
optim.zero_grad()
self.loss_all.backward()
optim.step()
output_coeff = self.headmodel_for_fitting.merge_coeff(output_coeff)
self.get_edge_mask()
self.get_fusion_mask(keep_forehead=False)
self.blur_shape_offset_uv(global_blur=True)
self.blur_offset_edge()
return output_coeff
def forward(self):
with torch.no_grad():
output_coeff = self.net_recon(self.input_img_coeff)
if not self.check_head_pose(output_coeff):
return None
with torch.enable_grad():
output_coeff = self.fitting_nonlinear(output_coeff)
output_coeff = self.headmodel.split_coeff(output_coeff)
eye_coeffs = output_coeff['exp'][0, 16] + output_coeff['exp'][
0, 17] + output_coeff['exp'][0, 19]
if eye_coeffs > 1.0:
degree = 0.5
else:
degree = 1.0
# degree = 0.5
output_coeff['exp'][0, 16] += 1 * degree
output_coeff['exp'][0, 17] += 1 * degree
output_coeff['exp'][0, 19] += 1.5 * degree
output_coeff = self.headmodel.merge_coeff(output_coeff)
self.pred_vertex, _, _, _, face_shape_ori, face_shape, _ = \
self.headmodel.compute_for_render_head(output_coeff,
self.shape_offset_uv.detach(),
self.texture_offset_uv.detach(),
self.shape_offset_uv_head.detach() * 0,
self.texture_offset_uv_head.detach(),
self.nonlinear_UVs,
nose_coeff=0.1,
neck_coeff=0.3,
neckSlim_coeff=0.5,
neckStretch_coeff=0.5)
UVs = np.array(self.template_output_mesh['uvs'])
UVs_tensor = torch.tensor(UVs, dtype=torch.float32)
UVs_tensor = torch.unsqueeze(UVs_tensor, 0).to(self.pred_vertex.device)
target_img = self.input_fat_img_hd
target_img = target_img.permute(0, 2, 3, 1)
face_buf = self.headmodel.face_buf
# get texture map
with torch.enable_grad():
pred_mask, _, pred_face, texture_map, texture_mask = self.renderer.pred_shape_and_texture(
self.pred_vertex, face_buf, UVs_tensor, target_img, None)
self.pred_coeffs_dict = self.headmodel.split_coeff(output_coeff)
recon_shape = face_shape # get reconstructed shape, [1, 35709, 3]
recon_shape[
...,
-1] = 10 - recon_shape[..., -1] # from camera space to world space
recon_shape = recon_shape.cpu().numpy()[0]
tri = self.headmodel.face_buf.cpu().numpy()
output = {}
output['flag'] = 0
output['tex_map'] = texture_map
output['tex_mask'] = texture_mask * 255.0
'''
coeffs
{
'id': id_coeffs,
'exp': exp_coeffs,
'tex': tex_coeffs,
'angle': angles,
'gamma': gammas,
'trans': translations
}
'''
output['coeffs'] = self.pred_coeffs_dict
normals = estimate_normals(recon_shape, tri)
output['vertices'] = recon_shape
output['triangles'] = tri
output['uvs'] = UVs
output['faces_uv'] = self.template_output_mesh['faces_uv']
output['normals'] = normals
return output
def get_edge_points_horizontal(self):
left_points = []
right_points = []
for i in range(self.face_mask.shape[2]):
inds = torch.where(self.face_mask[0, 0, i, :] > 0.5) # 0.9
if len(inds[0]) > 0: # i > 112 and len(inds[0]) > 0
left_points.append(int(inds[0][0]) + 1)
right_points.append(int(inds[0][-1]))
else:
left_points.append(0)
right_points.append(self.face_mask.shape[3] - 1)
self.left_points = torch.tensor(left_points).long().to(self.device)
self.right_points = torch.tensor(right_points).long().to(self.device)
def compute_losses_fitting(self):
face_mask = self.pred_mask
face_mask = face_mask.detach()
self.loss_color = self.opt.w_color * self.comupte_color_loss(
self.pred_face, self.input_img, face_mask) # 1.0
loss_reg, loss_gamma = self.compute_reg_loss(
self.pred_coeffs_dict,
w_id=self.opt.w_id,
w_exp=self.opt.w_exp,
w_tex=self.opt.w_tex)
self.loss_reg = self.opt.w_reg * loss_reg # 1.0
self.loss_gamma = self.opt.w_gamma * loss_gamma # 1.0
self.loss_lm = self.opt.w_lm * self.compute_lm_loss(
self.pred_lm, self.gt_lm) * 0.1 # 0.1
self.loss_smooth_offset = TVLoss()(self.shape_offset_uv.permute(
0, 3, 1, 2)) * 10000 # 10000
self.loss_reg_textureOff = torch.mean(
torch.abs(self.texture_offset_uv)) * 10 # 10
self.loss_smooth_offset_std = TVLoss_std()(
self.shape_offset_uv.permute(0, 3, 1, 2)) * 50000 # 50000
self.loss_points_horizontal, self.edge_points_inds = points_loss_horizontal(
self.verts_proj, self.left_points, self.right_points) # 20
self.loss_points_horizontal *= 20
self.loss_all = self.loss_color + self.loss_lm + self.loss_reg + self.loss_gamma
self.loss_all += self.loss_smooth_offset + self.loss_smooth_offset_std + self.loss_reg_textureOff
self.loss_all += self.loss_points_horizontal
head_mask = self.pred_mask_head
head_mask = head_mask.detach()
self.loss_color_head = self.opt.w_color * self.comupte_color_loss(
self.pred_head, self.input_img, head_mask) # 1.0
self.loss_smooth_offset_head = TVLoss()(
self.shape_offset_uv_head.permute(0, 3, 1, 2)) * 100 # 10000
self.loss_smooth_offset_std_head = TVLoss_std()(
self.shape_offset_uv_head.permute(0, 3, 1, 2)) * 500 # 50000
self.loss_mask = BinaryDiceLoss()(self.occ_head, self.head_mask) * 20
self.loss_all += self.loss_mask + self.loss_color_head
self.loss_all += self.loss_smooth_offset_head + self.loss_smooth_offset_std_head

367
modelscope/models/cv/head_reconstruction/models/losses.py Normal file
View File

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

577
modelscope/models/cv/head_reconstruction/models/networks.py Normal file
View File

@@ -0,0 +1,577 @@
# Part of the implementation is borrowed and modified from Deep3DFaceRecon_pytorch,
# publicly available at https://github.com/sicxu/Deep3DFaceRecon_pytorch
import os
from typing import Any, Callable, List, Optional, Type, Union
import torch
import torch.nn as nn
from kornia.geometry import warp_affine
from torch import Tensor
from torch.optim import lr_scheduler
try:
from torch.hub import load_state_dict_from_url
except ImportError:
from torch.utils.model_zoo import load_url as load_state_dict_from_url
def resize_n_crop(image, M, dsize=112):
# image: (b, c, h, w)
# M : (b, 2, 3)
return warp_affine(image, M, dsize=(dsize, dsize))
def filter_state_dict(state_dict, remove_name='fc'):
new_state_dict = {}
for key in state_dict:
if remove_name in key:
continue
new_state_dict[key] = state_dict[key]
return new_state_dict
def define_net_recon(net_recon, use_last_fc=False, init_path=None):
return ReconNetWrapper(
net_recon, use_last_fc=use_last_fc, init_path=init_path)
def define_net_recon2(net_recon, use_last_fc=False, init_path=None):
return ReconNetWrapper2(
net_recon, use_last_fc=use_last_fc, init_path=init_path)
class ReconNetWrapper(nn.Module):
fc_dim = 257
def __init__(self, net_recon, use_last_fc=False, init_path=None):
super(ReconNetWrapper, self).__init__()
self.use_last_fc = use_last_fc
if net_recon not in func_dict:
return NotImplementedError('network [%s] is not implemented',
net_recon)
func, last_dim = func_dict[net_recon]
backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_dim)
if init_path and os.path.isfile(init_path):
state_dict = filter_state_dict(
torch.load(init_path, map_location='cpu'))
backbone.load_state_dict(state_dict)
print('loading init net_recon %s from %s' % (net_recon, init_path))
self.backbone = backbone
if not use_last_fc:
self.final_layers = nn.ModuleList([
conv1x1(last_dim, 80, bias=True), # id layer
conv1x1(last_dim, 64, bias=True), # exp layer
conv1x1(last_dim, 80, bias=True), # tex layer
conv1x1(last_dim, 3, bias=True), # angle layer
conv1x1(last_dim, 27, bias=True), # gamma layer
conv1x1(last_dim, 2, bias=True), # tx, ty
conv1x1(last_dim, 1, bias=True) # tz
])
for m in self.final_layers:
nn.init.constant_(m.weight, 0.)
nn.init.constant_(m.bias, 0.)
def forward(self, x):
x = self.backbone(x)
if not self.use_last_fc:
output = []
for layer in self.final_layers:
output.append(layer(x))
x = torch.flatten(torch.cat(output, dim=1), 1)
return x
class ReconNetWrapper2(nn.Module):
fc_dim = 264
def __init__(self, net_recon, use_last_fc=False, init_path=None):
super(ReconNetWrapper2, self).__init__()
self.use_last_fc = use_last_fc
if net_recon not in func_dict:
return NotImplementedError('network [%s] is not implemented',
net_recon)
func, last_dim = func_dict[net_recon]
backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_dim)
if init_path and os.path.isfile(init_path):
state_dict = filter_state_dict(
torch.load(init_path, map_location='cpu'))
backbone.load_state_dict(state_dict)
print('loading init net_recon %s from %s' % (net_recon, init_path))
self.backbone = backbone
if not use_last_fc:
self.final_layers2 = nn.ModuleList([
conv1x1(last_dim, 80, bias=True), # id layer
conv1x1(last_dim, 51, bias=True), # exp layer
conv1x1(last_dim, 100, bias=True), # tex layer
conv1x1(last_dim, 3, bias=True), # angle layer
conv1x1(last_dim, 27, bias=True), # gamma layer
conv1x1(last_dim, 2, bias=True), # tx, ty
conv1x1(last_dim, 1, bias=True) # tz
])
for m in self.final_layers2:
nn.init.constant_(m.weight, 0.)
nn.init.constant_(m.bias, 0.)
def forward(self, x):
x = self.backbone(x)
if not self.use_last_fc:
output = []
for layer in self.final_layers2:
output.append(layer(x))
x = torch.flatten(torch.cat(output, dim=1), 1)
return x
# adapted from https://github.com/pytorch/vision/edit/master/torchvision/models/resnet.py
__all__ = [
'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',
'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2',
'wide_resnet101_2'
]
model_urls = {
'resnet18':
'https://download.pytorch.org/models/resnet18-f37072fd.pth',
'resnet34':
'https://download.pytorch.org/models/resnet34-b627a593.pth',
'resnet50':
'https://download.pytorch.org/models/resnet50-0676ba61.pth',
'resnet101':
'https://download.pytorch.org/models/resnet101-63fe2227.pth',
'resnet152':
'https://download.pytorch.org/models/resnet152-394f9c45.pth',
'resnext50_32x4d':
'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
'resnext101_32x8d':
'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
'wide_resnet50_2':
'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
'wide_resnet101_2':
'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
}
def conv3x3(in_planes: int,
out_planes: int,
stride: int = 1,
groups: int = 1,
dilation: int = 1) -> nn.Conv2d:
"""3x3 convolution with padding"""
return nn.Conv2d(
in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
groups=groups,
bias=False,
dilation=dilation)
def conv1x1(in_planes: int,
out_planes: int,
stride: int = 1,
bias: bool = False) -> nn.Conv2d:
"""1x1 convolution"""
return nn.Conv2d(
in_planes, out_planes, kernel_size=1, stride=stride, bias=bias)
class BasicBlock(nn.Module):
expansion: int = 1
def __init__(
self,
inplanes: int,
planes: int,
stride: int = 1,
downsample: Optional[nn.Module] = None,
groups: int = 1,
base_width: int = 64,
dilation: int = 1,
norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
super(BasicBlock, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
if groups != 1 or base_width != 64:
raise ValueError(
'BasicBlock only supports groups=1 and base_width=64')
if dilation > 1:
raise NotImplementedError(
'Dilation > 1 not supported in BasicBlock')
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = norm_layer(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = norm_layer(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x: Tensor) -> Tensor:
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class Bottleneck(nn.Module):
# Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
# while original implementation places the stride at the first 1x1 convolution(self.conv1)
# according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
# This variant is also known as ResNet V1.5 and improves accuracy according to
# https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
expansion: int = 4
def __init__(
self,
inplanes: int,
planes: int,
stride: int = 1,
downsample: Optional[nn.Module] = None,
groups: int = 1,
base_width: int = 64,
dilation: int = 1,
norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
super(Bottleneck, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
width = int(planes * (base_width / 64.)) * groups
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
self.conv1 = conv1x1(inplanes, width)
self.bn1 = norm_layer(width)
self.conv2 = conv3x3(width, width, stride, groups, dilation)
self.bn2 = norm_layer(width)
self.conv3 = conv1x1(width, planes * self.expansion)
self.bn3 = norm_layer(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x: Tensor) -> Tensor:
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(
self,
block: Type[Union[BasicBlock, Bottleneck]],
layers: List[int],
num_classes: int = 1000,
zero_init_residual: bool = False,
use_last_fc: bool = False,
groups: int = 1,
width_per_group: int = 64,
replace_stride_with_dilation: Optional[List[bool]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
super(ResNet, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
self._norm_layer = norm_layer
self.inplanes = 64
self.dilation = 1
if replace_stride_with_dilation is None:
# each element in the tuple indicates if we should replace
# the 2x2 stride with a dilated convolution instead
replace_stride_with_dilation = [False, False, False]
if len(replace_stride_with_dilation) != 3:
raise ValueError('replace_stride_with_dilation should be None '
'or a 3-element tuple, got {}'.format(
replace_stride_with_dilation))
self.use_last_fc = use_last_fc
self.groups = groups
self.base_width = width_per_group
self.conv1 = nn.Conv2d(
3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = norm_layer(self.inplanes)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(
block,
128,
layers[1],
stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(
block,
256,
layers[2],
stride=2,
dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(
block,
512,
layers[3],
stride=2,
dilate=replace_stride_with_dilation[2])
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
if self.use_last_fc:
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(
m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
# Zero-initialize the last BN in each residual branch,
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
nn.init.constant_(m.bn3.weight,
0) # type: ignore[arg-type]
elif isinstance(m, BasicBlock):
nn.init.constant_(m.bn2.weight,
0) # type: ignore[arg-type]
def _make_layer(self,
block: Type[Union[BasicBlock, Bottleneck]],
planes: int,
blocks: int,
stride: int = 1,
dilate: bool = False) -> nn.Sequential:
norm_layer = self._norm_layer
downsample = None
previous_dilation = self.dilation
if dilate:
self.dilation *= stride
stride = 1
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv1x1(self.inplanes, planes * block.expansion, stride),
norm_layer(planes * block.expansion),
)
layers = []
layers.append(
block(self.inplanes, planes, stride, downsample, self.groups,
self.base_width, previous_dilation, norm_layer))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(
block(
self.inplanes,
planes,
groups=self.groups,
base_width=self.base_width,
dilation=self.dilation,
norm_layer=norm_layer))
return nn.Sequential(*layers)
def _forward_impl(self, x: Tensor) -> Tensor:
# See note [TorchScript super()]
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
if self.use_last_fc:
x = torch.flatten(x, 1)
x = self.fc(x)
return x
def forward(self, x: Tensor) -> Tensor:
return self._forward_impl(x)
def _resnet(arch: str, block: Type[Union[BasicBlock,
Bottleneck]], layers: List[int],
pretrained: bool, progress: bool, **kwargs: Any) -> ResNet:
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
model_urls[arch], progress=progress)
model.load_state_dict(state_dict)
return model
def resnet18(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNet-18 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
**kwargs)
def resnet34(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNet-34 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
**kwargs)
def resnet50(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNet-50 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
**kwargs)
def resnet101(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNet-101 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained,
progress, **kwargs)
def resnet152(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNet-152 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained,
progress, **kwargs)
def resnext50_32x4d(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNeXt-50 32x4d model from
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['groups'] = 32
kwargs['width_per_group'] = 4
return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained,
progress, **kwargs)
def resnext101_32x8d(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""ResNeXt-101 32x8d model from
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['groups'] = 32
kwargs['width_per_group'] = 8
return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained,
progress, **kwargs)
def wide_resnet50_2(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""Wide ResNet-50-2 model from
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
The model is the same as ResNet except for the bottleneck number of channels
which is twice larger in every block. The number of channels in outer 1x1
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['width_per_group'] = 64 * 2
return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], pretrained,
progress, **kwargs)
def wide_resnet101_2(pretrained: bool = False,
progress: bool = True,
**kwargs: Any) -> ResNet:
r"""Wide ResNet-101-2 model from
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
The model is the same as ResNet except for the bottleneck number of channels
which is twice larger in every block. The number of channels in outer 1x1
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
kwargs['width_per_group'] = 64 * 2
return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], pretrained,
progress, **kwargs)
func_dict = {'resnet18': (resnet18, 512), 'resnet50': (resnet50, 2048)}

414
modelscope/models/cv/head_reconstruction/models/nv_diffrast.py Normal file
View File

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

21
modelscope/models/cv/head_reconstruction/models/opt.py Normal file
View File

@@ -0,0 +1,21 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
bfm_folder = ''
bfm_model = 'head_model_for_maas.mat'
camera_d = 10.0
center = 112.0
focal = 1015.0
isTrain = False
net_recon = 'resnet50'
phase = 'test'
use_ddp = False
use_last_fc = False
z_far = 15.0
z_near = 5.0
lr = 0.001
w_color = 1.92
w_reg = 3.0e-4
w_gamma = 10.0
w_lm = 1.6e-3
w_id = 1.0
w_exp = 0.8
w_tex = 1.7e-2

145
modelscope/models/cv/head_reconstruction/models/tex_processor.py Normal file
View File

@@ -0,0 +1,145 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import cv2
import numpy as np
def get_fade_out_mask(length, start_value, end_value, fade_start_ratio,
fade_end_ratio):
fade_start_ind = int(length * fade_start_ratio)
fade_end_ind = int(length * fade_end_ratio)
left_part = np.array([start_value] * fade_start_ind)
fade_part = np.linspace(start_value, end_value,
fade_end_ind - fade_start_ind)
len_right = length - len(left_part) - len(fade_part)
right_part = np.array([end_value] * len_right)
fade_out_mask = np.concatenate([left_part, fade_part, right_part], axis=0)
return fade_out_mask
class TexProcesser():
def __init__(self, model_root):
self.tex_size = 4096
self.bald_tex_bg = cv2.imread(
os.path.join(model_root,
'assets/texture/template_bald_tex_2.jpg')).astype(
np.float32)
self.hair_tex_bg = cv2.imread(
os.path.join(model_root,
'assets/texture/template_withHair_tex.jpg')).astype(
np.float32)
self.hair_mask = cv2.imread(
os.path.join(model_root,
'assets/texture/hair_mask_male.png'))[..., 0].astype(
np.float32) / 255.0
self.hair_mask = cv2.resize(self.hair_mask, (4096, 4096 + 1024))
front_mask = cv2.imread(
os.path.join(model_root,
'assets/texture/face_mask_singleview.jpg')).astype(
np.float32) / 255
front_mask = cv2.resize(front_mask, (1024, 1024))
front_mask = cv2.resize(front_mask, (0, 0), fx=0.1, fy=0.1)
front_mask = cv2.erode(front_mask,
np.ones(shape=(7, 7), dtype=np.float32))
front_mask = cv2.GaussianBlur(front_mask, (13, 13), 0)
self.front_mask = cv2.resize(front_mask,
(self.tex_size, self.tex_size))
self.binary_front_mask = self.front_mask.copy()
self.binary_front_mask[(self.front_mask < 0.3)
+ (self.front_mask > 0.7)] = 0
self.binary_front_mask[self.binary_front_mask != 0] = 1.0
self.binary_front_mask_ = self.binary_front_mask.copy()
self.binary_front_mask_[:int(4096 * 375 / 950)] = 0
self.binary_front_mask_[int(4096 * 600 / 950):] = 0
self.binary_front_mask = np.zeros((4096 + 1024, 4096, 3),
dtype=np.float32)
self.binary_front_mask[:4096, :] = self.binary_front_mask_
self.front_mask_ = self.front_mask.copy()
self.front_mask = np.zeros((4096 + 1024, 4096, 3), dtype=np.float32)
self.front_mask[:4096, :] = self.front_mask_
self.fg_mask = cv2.imread(
os.path.join(model_root,
'assets/texture/fg_mask.png'))[..., 0].astype(
np.float32) / 255.0
self.fg_mask = cv2.resize(self.fg_mask, (256, 256))
self.fg_mask = cv2.dilate(self.fg_mask,
np.ones(shape=(13, 13), dtype=np.float32))
self.fg_mask = cv2.blur(self.fg_mask, (27, 27), 0)
self.fg_mask = cv2.resize(self.fg_mask, (4096, 4096 + 1024))
self.fg_mask = self.fg_mask[..., None]
self.cheek_mask = cv2.imread(
os.path.join(model_root,
'assets/texture/cheek_area_mask.png'))[..., 0].astype(
np.float32) / 255.0
self.cheek_mask = cv2.resize(self.cheek_mask, (4096, 4096 + 1024))
self.cheek_mask = self.cheek_mask[..., None]
self.bald_tex_bg = self.bald_tex_bg[:4096]
self.hair_tex_bg = self.hair_tex_bg[:4096]
self.fg_mask = self.fg_mask[:4096]
self.hair_mask = self.hair_mask[:4096]
self.front_mask = self.front_mask[:4096]
self.binary_front_mask = self.binary_front_mask[:4096]
self.front_mask_ = self.front_mask_[:4096]
self.cheek_mask_left = self.cheek_mask[:4096]
self.cheek_mask_right = self.cheek_mask[:4096].copy()[:, ::-1]
def post_process_texture(self, tex_map, hair_tex=True):
tex_map = cv2.resize(tex_map, (self.tex_size, self.tex_size))
# if hair_tex is true and there is a dark side, use the mirror texture
if hair_tex:
left_cheek_light_mean = np.mean(
tex_map[self.cheek_mask_left[..., 0] == 1.0])
right_cheek_light_mean = np.mean(
tex_map[self.cheek_mask_right[..., 0] == 1.0])
tex_map_flip = tex_map[:, ::-1, :]
w = tex_map.shape[1]
half_w = w // 2
if left_cheek_light_mean > right_cheek_light_mean * 1.5:
tex_map[:, half_w:, :] = tex_map_flip[:, half_w:, :]
elif right_cheek_light_mean > left_cheek_light_mean * 2:
tex_map[:, :half_w, :] = tex_map_flip[:, :half_w, :]
# change the color of template texture
bg_mean_rgb = np.mean(
self.bald_tex_bg[self.binary_front_mask[..., 0] == 1.0],
axis=0)[None, None]
pred_tex_mean_rgb = np.mean(
tex_map[self.binary_front_mask[..., 0] == 1.0], axis=0)[None,
None] * 1.1
_bald_tex_bg = self.bald_tex_bg.copy()
_bald_tex_bg = self.bald_tex_bg + (pred_tex_mean_rgb - bg_mean_rgb)
if hair_tex:
# inpaint hair
tex_gray = cv2.cvtColor(
tex_map.astype(np.uint8),
cv2.COLOR_BGR2GRAY).astype(np.float32)
hair_mask = (self.hair_mask == 1.0) * (tex_gray < 120)
hair_bgr = np.mean(tex_map[hair_mask, :], axis=0) * 0.5
if hair_bgr is None:
hair_bgr = 20.0
_bald_tex_bg[self.hair_mask == 1.0] = hair_bgr
# fuse
tex_map = _bald_tex_bg * (1.
- self.fg_mask) + tex_map * self.fg_mask
else:
# fuse
tex_map = _bald_tex_bg * (
1. - self.front_mask) + tex_map * self.front_mask
return tex_map

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

55
modelscope/models/cv/text_to_head/text_to_head_model.py Normal file
View File

@@ -0,0 +1,55 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from collections import OrderedDict
import cv2
import numpy as np
import torch
from diffusers import (ControlNetModel, DDIMScheduler,
StableDiffusionControlNetPipeline)
from diffusers.utils import load_image
from modelscope.models import MODELS, TorchModel
@MODELS.register_module('text-to-head', 'text_to_head')
class TextToHeadModel(TorchModel):
def __init__(self, model_dir, *args, **kwargs):
"""The HeadReconModel is implemented based on HRN, publicly available at
https://github.com/youngLBW/HRN
Args:
model_dir: the root directory of the model files
"""
super().__init__(model_dir, *args, **kwargs)
self.model_dir = model_dir
base_model_path = os.path.join(model_dir, 'base_model')
controlnet_path = os.path.join(model_dir, 'control_net')
controlnet = ControlNetModel.from_pretrained(
controlnet_path, torch_dtype=torch.float16)
self.face_gen_pipeline = StableDiffusionControlNetPipeline.from_pretrained(
base_model_path, controlnet=controlnet, torch_dtype=torch.float16)
self.face_gen_pipeline.scheduler = DDIMScheduler.from_config(
self.face_gen_pipeline.scheduler.config)
self.face_gen_pipeline.enable_model_cpu_offload()
self.add_prompt = ', 4K, good looking face, epic realistic, Sony a7, sharp, ' \
'skin detail pores, soft light, uniform illumination'
self.neg_prompt = 'ugly, cross eye, bangs, teeth, glasses, hat, dark, shadow'
control_pose_path = os.path.join(self.model_dir, 'control_pose.jpg')
self.control_pose = load_image(control_pose_path)
def forward(self, input):
prompt = input['text'] + self.add_prompt
image = self.face_gen_pipeline(
prompt,
negative_prompt=self.neg_prompt,
image=self.control_pose,
num_inference_steps=20).images[0] # PIL.Image
return image

35
modelscope/outputs/outputs.py
View File

@@ -861,6 +861,41 @@ TASK_OUTPUTS = {
# }
Tasks.face_reconstruction: [OutputKeys.OUTPUT],
# 3D head reconstruction result for single sample
# {
# "output_obj": io.BytesIO,
# "output_img": np.array with shape(h, w, 3),
# "output": {
# "mesh": {
# "vertices": np.array with shape(n, 3),
# "faces": np.array with shape(n, 3),
# "faces_uv": np.array with shape(n, 3),
# "faces_normal": np.array with shape(n, 3),
# "UVs": np.array with shape(n, 2),
# "normals": np.array with shape(n, 3),
# },
# }
# }
Tasks.head_reconstruction: [OutputKeys.OUTPUT],
# text to head result for text input
# {
# "output_obj": io.BytesIO,
# "output_img": np.array with shape(h, w, 3),
# "output": {
# "mesh": {
# "vertices": np.array with shape(n, 3),
# "faces": np.array with shape(n, 3),
# "faces_uv": np.array with shape(n, 3),
# "faces_normal": np.array with shape(n, 3),
# "UVs": np.array with shape(n, 2),
# "normals": np.array with shape(n, 3),
# },
# },
# "image": np.array with shape(h, w, 3),
# }
Tasks.text_to_head: [OutputKeys.OUTPUT],
# 3D human reconstruction result for single sample
# {
# "output": {

4
modelscope/pipeline_inputs.py
View File

@@ -126,6 +126,10 @@ TASK_INPUTS = {
InputType.IMAGE,
Tasks.face_reconstruction:
InputType.IMAGE,
Tasks.head_reconstruction:
InputType.IMAGE,
Tasks.text_to_head:
InputType.TEXT,
Tasks.human_detection:
InputType.IMAGE,
Tasks.face_image_generation:

607
modelscope/pipelines/cv/head_reconstruction_pipeline.py Normal file
View File

@@ -0,0 +1,607 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import io
import os
import shutil
from typing import Any, Dict
import cv2
import face_alignment
import numpy as np
import PIL.Image
import tensorflow as tf
import torch
from scipy.io import loadmat, savemat
from modelscope.metainfo import Pipelines
from modelscope.models.cv.face_reconstruction.models.facelandmark.large_base_lmks_infer import \
LargeBaseLmkInfer
from modelscope.models.cv.face_reconstruction.utils import (
POS, align_for_lm, draw_line, enlarged_bbox, extract_5p, image_warp_grid1,
load_lm3d, mesh_to_string, read_obj, resize_n_crop_img,
resize_on_long_side, spread_flow, write_obj)
from modelscope.models.cv.head_reconstruction.models.head_segmentation import \
HeadSegmentor
from modelscope.models.cv.head_reconstruction.models.tex_processor import \
TexProcesser
from modelscope.models.cv.skin_retouching.retinaface.predict_single import \
Model
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.pipelines.base import Input, Pipeline
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors import LoadImage
from modelscope.utils.config import Config
from modelscope.utils.constant import ModelFile, Tasks
from modelscope.utils.device import create_device, device_placement
from modelscope.utils.logger import get_logger
try:
from torch.hub import get_dir
except BaseException:
from torch.hub import _get_torch_home as get_dir
if tf.__version__ >= '2.0':
tf = tf.compat.v1
tf.disable_eager_execution()
logger = get_logger()
@PIPELINES.register_module(
Tasks.head_reconstruction, module_name=Pipelines.head_reconstruction)
class HeadReconstructionPipeline(Pipeline):
def __init__(self, model: str, device: str, hair_tex=False):
"""The inference pipeline for head reconstruction task.
Args:
model (`str` or `Model` or module instance): A model instance or a model local dir
or a model id in the model hub.
device ('str'): device str, should be either cpu, cuda, gpu, gpu:X or cuda:X.
Example:
>>> from modelscope.pipelines import pipeline
>>> test_image = 'data/test/images/face_reconstruction.jpg'
>>> pipeline_headRecon = pipeline('head-reconstruction',
model='damo/cv_HRN_head-reconstruction')
>>> result = pipeline_headRecon(test_image)
>>> mesh = result[OutputKeys.OUTPUT]['mesh']
>>> texture_map = result[OutputKeys.OUTPUT_IMG]
>>> mesh['texture_map'] = texture_map
>>> write_obj('head_reconstruction.obj', mesh)
"""
super().__init__(model=model, device=device)
model_root = model
bfm_folder = os.path.join(model_root, 'assets')
checkpoint_path = os.path.join(model_root, ModelFile.TORCH_MODEL_FILE)
config_path = os.path.join(model_root, ModelFile.CONFIGURATION)
logger.info(f'loading config from {config_path}')
self.cfg = Config.from_file(config_path)
self.hair_tex = hair_tex
if 'gpu' in device:
self.device_name_ = 'cuda'
else:
self.device_name_ = device
self.device_name_ = self.device_name_.lower()
lmks_cpkt_path = os.path.join(model_root, 'large_base_net.pth')
self.large_base_lmks_model = LargeBaseLmkInfer.model_preload(
lmks_cpkt_path, self.device_name_ == 'cuda')
self.detector = Model(max_size=512, device=self.device_name_)
detector_ckpt_name = 'retinaface_resnet50_2020-07-20_old_torch.pth'
state_dict = torch.load(
os.path.join(os.path.dirname(lmks_cpkt_path), detector_ckpt_name),
map_location='cpu')
self.detector.load_state_dict(state_dict)
self.detector.eval()
device = torch.device(self.device_name_)
self.model.set_device(device)
self.model.setup(checkpoint_path)
self.model.parallelize()
self.model.eval()
self.model.set_render()
hub_dir = get_dir()
save_ckpt_dir = os.path.join(hub_dir, 'checkpoints')
if not os.path.exists(save_ckpt_dir):
os.makedirs(save_ckpt_dir)
shutil.copy(
os.path.join(model_root, 'face_alignment', 's3fd-619a316812.pth'),
save_ckpt_dir)
shutil.copy(
os.path.join(model_root, 'face_alignment',
'3DFAN4-4a694010b9.zip'), save_ckpt_dir)
shutil.copy(
os.path.join(model_root, 'face_alignment', 'depth-6c4283c0e0.zip'),
save_ckpt_dir)
self.lm_sess = face_alignment.FaceAlignment(
face_alignment.LandmarksType.THREE_D,
flip_input=False) # face_alignment.LandmarksType._3D
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = 0.2
config.gpu_options.allow_growth = True
g1 = tf.Graph()
self.face_sess = tf.Session(graph=g1, config=config)
with self.face_sess.as_default():
with g1.as_default():
with tf.gfile.FastGFile(
os.path.join(model_root, 'segment_face.pb'),
'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
self.face_sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
self.face_sess.run(tf.global_variables_initializer())
self.head_segmentor = HeadSegmentor(model_root=model_root)
self.tex_processor = TexProcesser(model_root=model_root)
self.lm3d_std = load_lm3d(bfm_folder)
self.align_params = loadmat(
'{}/assets/BBRegressorParam_r.mat'.format(model_root))
device = create_device(self.device_name)
self.device = device
def preprocess(self, input: Input) -> Dict[str, Any]:
if isinstance(input, str):
img = LoadImage.convert_to_ndarray(input)
if len(img.shape) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
img = img.astype(float)
else:
img = input.astype(float)
result = {'img': img}
return result
def align_img(self,
img,
lm,
lm3D,
mask=None,
target_size=224.,
rescale_factor=102.):
"""
Return:
transparams --numpy.array (raw_W, raw_H, scale, tx, ty)
img_new --PIL.Image (target_size, target_size, 3)
lm_new --numpy.array (68, 2), y direction is opposite to v direction
mask_new --PIL.Image (target_size, target_size)
Parameters:
img --PIL.Image (raw_H, raw_W, 3)
lm --numpy.array (68, 2), y direction is opposite to v direction
lm3D --numpy.array (5, 3)
mask --PIL.Image (raw_H, raw_W, 3)
"""
w0, h0 = img.size
if lm.shape[0] != 5:
lm5p = extract_5p(lm)
else:
lm5p = lm
# calculate translation and scale factors using 5 facial landmarks and standard landmarks of a 3D face
t, s = POS(lm5p.transpose(), lm3D.transpose())
s = rescale_factor / s
# processing the image
img_new, lm_new, mask_new = resize_n_crop_img(
img, lm, t, s, target_size=target_size, mask=mask)
trans_params = np.array([w0, h0, s, t[0][0], t[1][0]])
return trans_params, img_new, lm_new, mask_new
def read_data(self,
img,
lm,
lm3d_std,
to_tensor=True,
image_res=1024,
img_fat=None,
head_mask=None,
rescale_factor=75.0):
# to RGB
im = PIL.Image.fromarray(img[..., ::-1])
W, H = im.size
lm[:, -1] = H - 1 - lm[:, -1]
head_mask = PIL.Image.fromarray(head_mask)
im_fat = PIL.Image.fromarray(img_fat[..., ::-1])
_, im_lr_coeff, lm_lr_coeff, _ = self.align_img(im, lm, lm3d_std)
_, im_lr, lm_lr, mask_lr_head = self.align_img(
im, lm, lm3d_std, mask=head_mask, rescale_factor=rescale_factor)
_, im_hd, lm_hd, _ = self.align_img(
im_fat,
lm,
lm3d_std,
target_size=image_res,
rescale_factor=rescale_factor * image_res / 224)
mask_lr = self.face_sess.run(
self.face_sess.graph.get_tensor_by_name('output_alpha:0'),
feed_dict={'input_image:0': np.array(im_lr)})
if to_tensor:
im_lr = torch.tensor(
np.array(im_lr) / 255.,
dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
im_hd = torch.tensor(
np.array(im_hd) / 255.,
dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
mask_lr = torch.tensor(
np.array(mask_lr) / 255., dtype=torch.float32)[None,
None, :, :]
mask_lr_head = torch.tensor(
np.array(mask_lr_head) / 255., dtype=torch.float32)[
None, None, :, :] if mask_lr_head is not None else None
lm_lr = torch.tensor(lm_lr).unsqueeze(0)
lm_hd = torch.tensor(lm_hd).unsqueeze(0)
im_lr_coeff = torch.tensor(
np.array(im_lr_coeff) / 255.,
dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
lm_lr_coeff = torch.tensor(lm_lr_coeff).unsqueeze(0)
return im_lr, lm_lr, im_hd, lm_hd, mask_lr, mask_lr_head, im_lr_coeff, lm_lr_coeff
def prepare_data(self, img, lm_sess, five_points=None):
input_img, scale, bbox = align_for_lm(
img, five_points,
self.align_params) # align for 68 landmark detection
if scale == 0:
return None
# detect landmarks
input_img = np.reshape(input_img, [1, 224, 224, 3]).astype(np.float32)
input_img = input_img[0, :, :, ::-1]
landmark = lm_sess.get_landmarks_from_image(input_img)[0]
landmark = landmark[:, :2] / scale
landmark[:, 0] = landmark[:, 0] + bbox[0]
landmark[:, 1] = landmark[:, 1] + bbox[1]
return landmark
def infer_lmks(self, img_bgr):
INPUT_SIZE = 224
ENLARGE_RATIO = 1.35
landmarks = []
rgb_image = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
results = self.detector.predict_jsons(rgb_image)
boxes = []
for anno in results:
if anno['score'] == -1:
break
boxes.append({
'x1': anno['bbox'][0],
'y1': anno['bbox'][1],
'x2': anno['bbox'][2],
'y2': anno['bbox'][3]
})
for detect_result in boxes:
x1 = detect_result['x1']
y1 = detect_result['y1']
x2 = detect_result['x2']
y2 = detect_result['y2']
w = x2 - x1 + 1
h = y2 - y1 + 1
cx = (x2 + x1) / 2
cy = (y2 + y1) / 2
sz = max(h, w) * ENLARGE_RATIO
x1 = cx - sz / 2
y1 = cy - sz / 2
trans_x1 = x1
trans_y1 = y1
x2 = x1 + sz
y2 = y1 + sz
height, width, _ = rgb_image.shape
dx = max(0, -x1)
dy = max(0, -y1)
x1 = max(0, x1)
y1 = max(0, y1)
edx = max(0, x2 - width)
edy = max(0, y2 - height)
x2 = min(width, x2)
y2 = min(height, y2)
crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)]
if dx > 0 or dy > 0 or edx > 0 or edy > 0:
crop_img = cv2.copyMakeBorder(
crop_img,
int(dy),
int(edy),
int(dx),
int(edx),
cv2.BORDER_CONSTANT,
value=(103.94, 116.78, 123.68))
crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE))
base_lmks = LargeBaseLmkInfer.infer_img(
crop_img, self.large_base_lmks_model,
self.device_name_ == 'cuda')
inv_scale = sz / INPUT_SIZE
affine_base_lmks = np.zeros((106, 2))
for idx in range(106):
affine_base_lmks[idx][
0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1
affine_base_lmks[idx][
1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1
x1 = np.min(affine_base_lmks[:, 0])
y1 = np.min(affine_base_lmks[:, 1])
x2 = np.max(affine_base_lmks[:, 0])
y2 = np.max(affine_base_lmks[:, 1])
w = x2 - x1 + 1
h = y2 - y1 + 1
cx = (x2 + x1) / 2
cy = (y2 + y1) / 2
sz = max(h, w) * ENLARGE_RATIO
x1 = cx - sz / 2
y1 = cy - sz / 2
trans_x1 = x1
trans_y1 = y1
x2 = x1 + sz
y2 = y1 + sz
height, width, _ = rgb_image.shape
dx = max(0, -x1)
dy = max(0, -y1)
x1 = max(0, x1)
y1 = max(0, y1)
edx = max(0, x2 - width)
edy = max(0, y2 - height)
x2 = min(width, x2)
y2 = min(height, y2)
crop_img = rgb_image[int(y1):int(y2), int(x1):int(x2)]
if dx > 0 or dy > 0 or edx > 0 or edy > 0:
crop_img = cv2.copyMakeBorder(
crop_img,
int(dy),
int(edy),
int(dx),
int(edx),
cv2.BORDER_CONSTANT,
value=(103.94, 116.78, 123.68))
crop_img = cv2.resize(crop_img, (INPUT_SIZE, INPUT_SIZE))
base_lmks = LargeBaseLmkInfer.infer_img(
crop_img, self.large_base_lmks_model,
self.device_name_.lower() == 'cuda')
inv_scale = sz / INPUT_SIZE
affine_base_lmks = np.zeros((106, 2))
for idx in range(106):
affine_base_lmks[idx][
0] = base_lmks[0][idx * 2 + 0] * inv_scale + trans_x1
affine_base_lmks[idx][
1] = base_lmks[0][idx * 2 + 1] * inv_scale + trans_y1
landmarks.append(affine_base_lmks)
return boxes, landmarks
def find_face_contour(self, image):
boxes, landmarks = self.infer_lmks(image)
landmarks = np.array(landmarks)
args = [[0, 33, False], [33, 38, False], [42, 47, False],
[51, 55, False], [57, 64, False], [66, 74, True],
[75, 83, True], [84, 96, True]]
roi_bboxs = []
for i in range(len(boxes)):
roi_bbox = enlarged_bbox([
boxes[i]['x1'], boxes[i]['y1'], boxes[i]['x2'], boxes[i]['y2']
], image.shape[1], image.shape[0], 0.5)
roi_bbox = [int(x) for x in roi_bbox]
roi_bboxs.append(roi_bbox)
people_maps = []
for i in range(landmarks.shape[0]):
landmark = landmarks[i, :, :]
maps = []
whole_mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
roi_box = roi_bboxs[i]
roi_box_width = roi_box[2] - roi_box[0]
roi_box_height = roi_box[3] - roi_box[1]
short_side_length = roi_box_width if roi_box_width < roi_box_height else roi_box_height
line_width = short_side_length // 10
if line_width == 0:
line_width = 1
kernel_size = line_width * 2
gaussian_kernel = kernel_size if kernel_size % 2 == 1 else kernel_size + 1
for t, arg in enumerate(args):
mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
draw_line(mask, landmark[arg[0]:arg[1]], (255, 255, 255),
line_width, arg[2])
mask = cv2.GaussianBlur(mask,
(gaussian_kernel, gaussian_kernel), 0)
if t >= 1:
draw_line(whole_mask, landmark[arg[0]:arg[1]],
(255, 255, 255), line_width * 2, arg[2])
maps.append(mask)
whole_mask = cv2.GaussianBlur(whole_mask,
(gaussian_kernel, gaussian_kernel),
0)
maps.append(whole_mask)
people_maps.append(maps)
return people_maps[0], boxes
def fat_face(self, img, degree=0.04):
_img, scale = resize_on_long_side(img, 800)
contour_maps, boxes = self.find_face_contour(_img)
contour_map = contour_maps[0]
boxes = boxes[0]
Flow = np.zeros(
shape=(contour_map.shape[0], contour_map.shape[1], 2),
dtype=np.float32)
box_center = [(boxes['x1'] + boxes['x2']) / 2,
(boxes['y1'] + boxes['y2']) / 2]
box_length = max(
abs(boxes['y1'] - boxes['y2']), abs(boxes['x1'] - boxes['x2']))
value_1 = 2 * (Flow.shape[0] - box_center[1] - 1)
value_2 = 2 * (Flow.shape[1] - box_center[0] - 1)
value_list = [
box_length * 2, 2 * (box_center[0] - 1), 2 * (box_center[1] - 1),
value_1, value_2
]
flow_box_length = min(value_list)
flow_box_length = int(flow_box_length)
sf = spread_flow(100, flow_box_length * degree)
sf = cv2.resize(sf, (flow_box_length, flow_box_length))
Flow[int(box_center[1]
- flow_box_length / 2):int(box_center[1]
+ flow_box_length / 2),
int(box_center[0]
- flow_box_length / 2):int(box_center[0]
+ flow_box_length / 2)] = sf
Flow = Flow * np.dstack((contour_map, contour_map)) / 255.0
inter_face_maps = contour_maps[-1]
Flow = Flow * (1.0 - np.dstack(
(inter_face_maps, inter_face_maps)) / 255.0)
Flow = cv2.resize(Flow, (img.shape[1], img.shape[0]))
Flow = Flow / scale
pred, top_bound, bottom_bound, left_bound, right_bound = image_warp_grid1(
Flow[..., 0], Flow[..., 1], img, 1.0, [0, 0, 0, 0])
return pred
def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
rgb_image = input['img'].cpu().numpy().astype(np.uint8)
bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
img = bgr_image
if img.shape[0] > 2000 or img.shape[1] > 2000:
img, _ = resize_on_long_side(img, 1500)
box, results = self.infer_lmks(img)
if results is None or np.array(results).shape[0] == 0:
return {}
fatbgr = self.fat_face(img)
landmarks = []
results = results[0]
for idx in [74, 83, 54, 84, 90]:
landmarks.append([results[idx][0], results[idx][1]])
landmarks = np.array(landmarks)
landmarks = self.prepare_data(img, self.lm_sess, five_points=landmarks)
head_mask = self.head_segmentor.process(img)[0]
im_tensor, lm_tensor, im_hd_tensor, lm_hd_tensor, mask, head_mask, im_co, lm_co = self.read_data(
img, landmarks, self.lm3d_std, img_fat=fatbgr, head_mask=head_mask)
data = {
'imgs': im_tensor,
'imgs_hd': im_hd_tensor,
'lms': lm_tensor,
'lms_hd': lm_hd_tensor,
'face_mask': mask,
'head_mask': head_mask,
'imgs_coeff': im_co,
'lms_coeff': lm_co,
}
self.model.set_input(data) # unpack data from data loader
output = self.model() # run inference
assert output is not None
tex_map = output['tex_map'].astype(np.float32)
# post-process texture map
tex_map = self.tex_processor.post_process_texture(
tex_map, hair_tex=self.hair_tex)
head_mesh = {
'vertices': output['vertices'],
'faces': output['triangles'] + 1,
'UVs': output['uvs'],
'faces_uv': output['faces_uv'],
'normals': output['normals'],
'texture_map': tex_map
}
results = {
'mesh': head_mesh,
}
return {
OutputKeys.OUTPUT_OBJ: None,
OutputKeys.OUTPUT_IMG: tex_map,
OutputKeys.OUTPUT: results
}
def postprocess(self, inputs, **kwargs) -> Dict[str, Any]:
render = kwargs.get('render', False)
output_obj = inputs[OutputKeys.OUTPUT_OBJ]
texture_map = inputs[OutputKeys.OUTPUT_IMG]
results = inputs[OutputKeys.OUTPUT]
if render:
output_obj = io.BytesIO()
mesh_str = mesh_to_string(results['mesh'])
mesh_bytes = mesh_str.encode(encoding='utf-8')
output_obj.write(mesh_bytes)
result = {
OutputKeys.OUTPUT_OBJ: output_obj,
OutputKeys.OUTPUT_IMG: texture_map,
OutputKeys.OUTPUT: None if render else results,
}
return result

91
modelscope/pipelines/cv/text_to_head_pipeline.py Normal file
View File

@@ -0,0 +1,91 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import io
import os
import shutil
from typing import Any, Dict
import numpy as np
from modelscope.metainfo import Pipelines
from modelscope.models.cv.face_reconstruction.utils import (
align_for_lm, align_img, draw_line, enlarged_bbox, image_warp_grid1,
load_lm3d, mesh_to_string, read_obj, resize_on_long_side, spread_flow,
write_obj)
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.pipelines.base import Input, Pipeline
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors import LoadImage
from modelscope.utils.constant import ModelFile, Tasks
from modelscope.utils.device import create_device, device_placement
from modelscope.utils.logger import get_logger
logger = get_logger()
@PIPELINES.register_module(
Tasks.text_to_head, module_name=Pipelines.text_to_head)
class TextToHeadPipeline(Pipeline):
def __init__(self, model: str, device: str, hair_tex=True):
"""The inference pipeline for text-to-head task.
Args:
model (`str` or `Model` or module instance): A model instance or a model local dir
or a model id in the model hub.
device ('str'): device str, should be either cpu, cuda, gpu, gpu:X or cuda:X.
Example:
>>> from modelscope.pipelines import pipeline
>>> from modelscope.models.cv.face_reconstruction.utils import write_obj
>>> test_prompt = "a clown with red nose"
>>> pipeline_textToHead = pipeline('text-to-head',
model='damo/cv_HRN_text-to-head')
>>> result = pipeline_textToHead(test_prompt)
>>> mesh = result[OutputKeys.OUTPUT]['mesh']
>>> texture_map = result[OutputKeys.OUTPUT_IMG]
>>> mesh['texture_map'] = texture_map
>>> write_obj('text_to_head.obj', mesh)
"""
super().__init__(model=model, device=device)
self.hair_tex = hair_tex
head_recon_model_id = 'damo/cv_HRN_head-reconstruction'
self.head_reconstructor = pipeline(
Tasks.head_reconstruction,
model=head_recon_model_id,
model_revision='v0.1',
hair_tex=hair_tex)
def preprocess(self, input: Input) -> Dict[str, Any]:
result = {'text': input}
return result
def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
image = self.model(input)
image = np.array(image)
results = self.head_reconstructor(image)
results['image'] = image
return results
def postprocess(self, inputs, **kwargs) -> Dict[str, Any]:
render = kwargs.get('render', False)
output_obj = inputs[OutputKeys.OUTPUT_OBJ]
texture_map = inputs[OutputKeys.OUTPUT_IMG]
results = inputs[OutputKeys.OUTPUT]
if render:
output_obj = io.BytesIO()
mesh_str = mesh_to_string(results['mesh'])
mesh_bytes = mesh_str.encode(encoding='utf-8')
output_obj.write(mesh_bytes)
result = {
OutputKeys.OUTPUT_OBJ: output_obj,
OutputKeys.OUTPUT_IMG: texture_map,
OutputKeys.OUTPUT: None if render else results,
'image': inputs['image']
}
return result

2
modelscope/utils/constant.py
View File

@@ -148,6 +148,8 @@ class CVTasks(object):
# 3d face reconstruction
face_reconstruction = 'face-reconstruction'
head_reconstruction = 'head-reconstruction'
text_to_head = 'text-to-head'
# 3d human reconstruction
human_reconstruction = 'human-reconstruction'

60
tests/pipelines/test_head_reconstruction.py Normal file
View File

@@ -0,0 +1,60 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import io
import os
import os.path as osp
import sys
import unittest
import cv2
from modelscope.hub.snapshot_download import snapshot_download
from modelscope.models.cv.face_reconstruction.utils import write_obj
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.pipelines.base import Pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.test_utils import test_level
sys.path.append('.')
class HeadReconstructionTest(unittest.TestCase):
def setUp(self) -> None:
self.task = Tasks.head_reconstruction
self.model_id = 'damo/cv_HRN_head-reconstruction'
self.test_image = 'data/test/images/face_reconstruction.jpg'
def save_results(self, result, save_root):
os.makedirs(save_root, exist_ok=True)
# export obj and texture
mesh = result[OutputKeys.OUTPUT]['mesh']
texture_map = result[OutputKeys.OUTPUT_IMG]
mesh['texture_map'] = texture_map
write_obj(os.path.join(save_root, 'head_recon_result.obj'), mesh)
print(f'Output written to {osp.abspath(save_root)}')
def pipeline_inference(self, pipeline: Pipeline, input_location: str):
result = pipeline(input_location)
self.save_results(result, './head_reconstruction_results')
@unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
def test_run_by_direct_model_download(self):
model_dir = snapshot_download(self.model_id, revision='v0.2')
head_reconstruction = pipeline(
Tasks.head_reconstruction, model=model_dir)
self.pipeline_inference(head_reconstruction, self.test_image)
@unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
def test_run_modelhub(self):
head_reconstruction = pipeline(
Tasks.head_reconstruction,
model=self.model_id,
model_revision='v0.2')
self.pipeline_inference(head_reconstruction, self.test_image)
if __name__ == '__main__':
unittest.main()

62
tests/pipelines/test_text_to_head.py Normal file
View File

@@ -0,0 +1,62 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import io
import os
import os.path as osp
import sys
import unittest
import cv2
from modelscope.hub.snapshot_download import snapshot_download
from modelscope.models.cv.face_reconstruction.utils import write_obj
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.pipelines.base import Pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.test_utils import test_level
sys.path.append('.')
class TextToHeadTest(unittest.TestCase):
def setUp(self) -> None:
self.task = Tasks.text_to_head
self.model_id = 'damo/cv_HRN_text-to-head'
self.test_prompt = 'a clown with red nose'
def save_results(self, result, save_root):
os.makedirs(save_root, exist_ok=True)
# export obj and texture
mesh = result[OutputKeys.OUTPUT]['mesh']
texture_map = result[OutputKeys.OUTPUT_IMG]
mesh['texture_map'] = texture_map
write_obj(os.path.join(save_root, 'text_to_head_result.obj'), mesh)
image = result['image']
image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite(
os.path.join(save_root, 'text_to_head_image.jpg'), image_bgr)
print(f'Output written to {osp.abspath(save_root)}')
def pipeline_inference(self, pipeline: Pipeline, prompt: str):
result = pipeline(prompt)
self.save_results(result, './text_to_head_results')
@unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
def test_run_by_direct_model_download(self):
model_dir = snapshot_download(self.model_id, revision='v0.1')
text_to_head = pipeline(Tasks.text_to_head, model=model_dir)
self.pipeline_inference(text_to_head, self.test_prompt)
@unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
def test_run_modelhub(self):
face_reconstruction = pipeline(
Tasks.text_to_head, model=self.model_id, model_revision='v0.1')
self.pipeline_inference(face_reconstruction, self.test_prompt)
if __name__ == '__main__':
unittest.main()