forked from abdallahdib/NextFace
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pipeline.py
170 lines (141 loc) · 8.7 KB
/
pipeline.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
from sphericalharmonics import SphericalHarmonics
from morphablemodel import MorphableModel
from renderer import Renderer
from camera import Camera
from utils import *
class Pipeline:
def __init__(self, config):
'''
a pipeline can generate and render textured faces under different camera angles and lighting conditions
:param config: configuration file used to parameterize the pipeline
'''
self.config = config
self.device = config.device
self.camera = Camera(self.device)
self.sh = SphericalHarmonics(config.envMapRes, self.device)
if self.config.lamdmarksDetectorType == 'fan':
pathLandmarksAssociation = '/landmark_62.txt'
elif self.config.lamdmarksDetectorType == 'mediapipe':
pathLandmarksAssociation = '/landmark_62_mp.txt'
else:
raise ValueError(f'lamdmarksDetectorType must be one of [mediapipe, fan] but was {self.config.lamdmarksDetectorType}')
self.morphableModel = MorphableModel(path = config.path,
textureResolution= config.textureResolution,
trimPca= config.trimPca,
landmarksPathName=pathLandmarksAssociation,
device = self.device
)
self.renderer = Renderer(config.rtTrainingSamples, 1, self.device)
self.uvMap = self.morphableModel.uvMap.clone()
self.uvMap[:, 1] = 1.0 - self.uvMap[:, 1]
self.faces32 = self.morphableModel.faces.to(torch.int32).contiguous()
self.shBands = config.bands
self.sharedIdentity = False
def initSceneParameters(self, n, sharedIdentity = False):
'''
init pipeline parameters (face shape, albedo, exp coeffs, light and head pose (camera))
:param n: the the number of parameters (if negative than the pipeline variables are not allocated)
:param sharedIdentity: if true, the shape and albedo coeffs are equal to 1, as they belong to the same person identity
:return:
'''
if n <= 0:
return
self.sharedIdentity = sharedIdentity
nShape = 1 if sharedIdentity == True else n
self.vShapeCoeff = torch.zeros([nShape, self.morphableModel.shapeBasisSize], dtype = torch.float32, device = self.device)
self.vAlbedoCoeff = torch.zeros([nShape, self.morphableModel.albedoBasisSize], dtype=torch.float32, device=self.device)
self.vExpCoeff = torch.zeros([n, self.morphableModel.expBasisSize], dtype=torch.float32, device=self.device)
self.vRotation = torch.zeros([n, 3], dtype=torch.float32, device=self.device)
self.vTranslation = torch.zeros([n, 3], dtype=torch.float32, device=self.device)
self.vTranslation[:, 2] = 500.
self.vRotation[:, 0] = 3.14
self.vFocals = self.config.camFocalLength * torch.ones([n], dtype=torch.float32, device=self.device)
self.vShCoeffs = 0.0 * torch.ones([n, self.shBands * self.shBands, 3], dtype=torch.float32, device=self.device)
self.vShCoeffs[..., 0, 0] = 0.5
self.vShCoeffs[..., 2, 0] = -0.5
self.vShCoeffs[..., 1] = self.vShCoeffs[..., 0]
self.vShCoeffs[..., 2] = self.vShCoeffs[..., 0]
texRes = self.morphableModel.getTextureResolution()
self.vRoughness = 0.4 * torch.ones([nShape, texRes, texRes, 1], dtype=torch.float32, device=self.device)
def computeShape(self):
'''
compute shape vertices from the shape and expression coefficients
:return: tensor of 3d vertices [n, verticesNumber, 3]
'''
assert(self.vShapeCoeff is not None and self.vExpCoeff is not None)
vertices = self.morphableModel.computeShape(self.vShapeCoeff, self.vExpCoeff)
return vertices
def transformVertices(self, vertices = None):
'''
transform vertices to camera coordinate space
:param vertices: tensor of 3d vertices [n, verticesNumber, 3]
:return: transformed vertices [n, verticesNumber, 3]
'''
if vertices is None:
vertices = self.computeShape()
assert(vertices.dim() == 3 and vertices.shape[-1] == 3)
assert(self.vTranslation is not None and self.vRotation is not None)
assert(vertices.shape[0] == self.vTranslation.shape[0] == self.vRotation.shape[0])
transformedVertices = self.camera.transformVertices(vertices, self.vTranslation, self.vRotation)
return transformedVertices
def render(self, cameraVerts = None, diffuseTextures = None, specularTextures = None, roughnessTextures = None, renderAlbedo = False):
'''
ray trace an image given camera vertices and corresponding textures
:param cameraVerts: camera vertices tensor [n, verticesNumber, 3]
:param diffuseTextures: diffuse textures tensor [n, texRes, texRes, 3]
:param specularTextures: specular textures tensor [n, texRes, texRes, 3]
:param roughnessTextures: roughness textures tensor [n, texRes, texRes, 1]
:param renderAlbedo: if True render albedo else ray trace image
:return: ray traced images [n, resX, resY, 4]
'''
if cameraVerts is None:
vertices, diffAlbedo, specAlbedo = self.morphableModel.computeShapeAlbedo(self.vShapeCoeff, self.vExpCoeff, self.vAlbedoCoeff)
cameraVerts = self.camera.transformVertices(vertices, self.vTranslation, self.vRotation)
#compute normals
normals = self.morphableModel.meshNormals.computeNormals(cameraVerts)
if diffuseTextures is None:
diffuseTextures = self.morphableModel.generateTextureFromAlbedo(diffAlbedo)
if specularTextures is None:
specularTextures = self.morphableModel.generateTextureFromAlbedo(specAlbedo)
if roughnessTextures is None:
roughnessTextures = self.vRoughness
envMaps = self.sh.toEnvMap(self.vShCoeffs)
assert(envMaps.dim() == 4 and envMaps.shape[-1] == 3)
assert (cameraVerts.dim() == 3 and cameraVerts.shape[-1] == 3)
assert (diffuseTextures.dim() == 4 and diffuseTextures.shape[1] == diffuseTextures.shape[2] == self.morphableModel.getTextureResolution() and diffuseTextures.shape[-1] == 3)
assert (specularTextures.dim() == 4 and specularTextures.shape[1] == specularTextures.shape[2] == self.morphableModel.getTextureResolution() and specularTextures.shape[-1] == 3)
assert (roughnessTextures.dim() == 4 and roughnessTextures.shape[1] == roughnessTextures.shape[2] == self.morphableModel.getTextureResolution() and roughnessTextures.shape[-1] == 1)
assert(cameraVerts.shape[0] == envMaps.shape[0])
assert (diffuseTextures.shape[0] == specularTextures.shape[0] == roughnessTextures.shape[0])
scenes = self.renderer.buildScenes(cameraVerts, self.faces32, normals, self.uvMap, diffuseTextures,
specularTextures, torch.clamp(roughnessTextures, 1e-20, 10.0), self.vFocals, envMaps)
if renderAlbedo:
images = self.renderer.renderAlbedo(scenes)
else:
images = self.renderer.render(scenes)
return images
def landmarkLoss(self, cameraVertices, landmarks, focals, cameraCenters, debugDir = None):
'''
calculate scalar loss between vertices in camera space and 2d landmarks pixels
:param cameraVertices: 3d vertices [n, nVertices, 3]
:param landmarks: 2d corresponding pixels [n, nVertices, 2]
:param landmarks: camera focals [n]
:param cameraCenters: camera centers [n, 2
:param debugDir: if not none save landmarks and vertices to an image file
:return: scalar loss (float)
'''
assert (cameraVertices.dim() == 3 and cameraVertices.shape[-1] == 3)
assert (focals.dim() == 1)
assert(cameraCenters.dim() == 2 and cameraCenters.shape[-1] == 2)
assert (landmarks.dim() == 3 and landmarks.shape[-1] == 2)
assert cameraVertices.shape[0] == landmarks.shape[0] == focals.shape[0] == cameraCenters.shape[0]
headPoints = cameraVertices[:, self.morphableModel.landmarksAssociation]
assert (landmarks.shape[-2] == headPoints.shape[-2])
projPoints = focals.view(-1, 1, 1) * headPoints[..., :2] / headPoints[..., 2:]
projPoints += cameraCenters.unsqueeze(1)
loss = torch.norm(projPoints - landmarks, 2, dim=-1).pow(2).mean()
if debugDir:
for i in range(projPoints.shape[0]):
image = saveLandmarksVerticesProjections(self.inputImage.tensor[i], projPoints[i], self.landmarks[i])
cv2.imwrite(debugDir + '/lp' + str(i) +'.png', cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
return loss