def render_shape(self,
vertices,
transformed_vertices,
images=None,
lights=None):
batch_size = vertices.shape[0]
if lights is None:
light_positions = torch.tensor([[-0.1, -0.1, 0.2],
[0, 0, 1]])[None, :, :].expand(
batch_size, -1, -1).float()
light_intensities = torch.ones_like(light_positions).float()
lights = torch.cat((light_positions, light_intensities),
2).to(vertices.device)
## rasterizer near 0 far 100. move mesh so minz larger than 0
transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] + 10
# Attributes
face_vertices = util.face_vertices(
vertices, self.faces.expand(batch_size, -1, -1))
normals = util.vertex_normals(vertices,
self.faces.expand(batch_size, -1, -1))
face_normals = util.face_vertices(
normals, self.faces.expand(batch_size, -1, -1))
transformed_normals = util.vertex_normals(
transformed_vertices, self.faces.expand(batch_size, -1, -1))
transformed_face_normals = util.face_vertices(
transformed_normals, self.faces.expand(batch_size, -1, -1))
# render
attributes = torch.cat([
self.face_colors.expand(batch_size, -1, -1, -1),
transformed_face_normals.detach(),
face_vertices.detach(),
face_normals.detach()
], -1)
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes)
# albedo
albedo_images = rendering[:, :3, :, :]
# shading
normal_images = rendering[:, 9:12, :, :].detach()
if lights.shape[1] == 9:
shading_images = self.add_SHlight(normal_images, lights)
else:
print('directional')
shading = self.add_directionlight(
normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape([
batch_size, lights.shape[1], albedo_images.shape[2],
albedo_images.shape[3], 3
]).permute(0, 1, 4, 2, 3)
shading_images = shading_images.mean(1)
images = albedo_images * shading_images
return images
def __init__(self, image_size, obj_filename, uv_size=256, config=None):
super(Renderer, self).__init__()
self.image_size = image_size
self.uv_size = uv_size
verts, faces, aux = load_obj(obj_filename)
uvcoords = aux.verts_uvs[None, ...] # (N, V, 2)
uvfaces = faces.textures_idx[None, ...] # (N, F, 3)
faces = faces.verts_idx[None, ...]
if config is not None:
_tp = np.load(config.trim_path, allow_pickle=True)
tp = {}
for k in _tp.keys():
tp[k] = torch.LongTensor(_tp[k])
verts = verts[tp['idx_verts']]
#uvcoords = uvcoords[:,tp['idx_verts']]
uvfaces = uvfaces[:,tp['idx_faces']]
faces = tp['map_verts'][faces[:,tp['idx_faces']]]
self.rasterizer = Pytorch3dRasterizer(image_size)
self.uv_rasterizer = Pytorch3dRasterizer(uv_size)
# faces
self.register_buffer('faces', faces)
self.register_buffer('raw_uvcoords', uvcoords)
# uv coordsw
uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.], -1) # [bz, ntv, 3]
uvcoords = uvcoords * 2 - 1
uvcoords[..., 1] = -uvcoords[..., 1]
face_uvcoords = util.face_vertices(uvcoords, uvfaces)
self.register_buffer('uvcoords', uvcoords)
self.register_buffer('uvfaces', uvfaces)
self.register_buffer('face_uvcoords', face_uvcoords)
# shape colors
colors = torch.tensor([74, 120, 168])[None, None, :].repeat(1, faces.max() + 1, 1).float() / 255.
face_colors = util.face_vertices(colors, faces)
self.register_buffer('face_colors', face_colors)
## lighting
pi = np.pi
constant_factor = torch.tensor(
[1 / np.sqrt(4 * pi), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), \
((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), \
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi)))])
self.register_buffer('constant_factor', constant_factor)
def __init__(self, image_size, obj_filename, uv_size=256):
super(Renderer, self).__init__()
self.image_size = image_size
self.uv_size = uv_size
verts, faces, aux = load_obj(obj_filename)
uvcoords = aux.verts_uvs[None, ...] # (N, V, 2)
uvfaces = faces.textures_idx[None, ...] # (N, F, 3)
faces = faces.verts_idx[None, ...]
self.rasterizer = Pytorch3dRasterizer(image_size)
self.uv_rasterizer = Pytorch3dRasterizer(uv_size)
# faces
self.register_buffer('faces', faces)
self.register_buffer('raw_uvcoords', uvcoords)
# uv coordsw
uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.],
-1) # [bz, ntv, 3]
uvcoords = uvcoords * 2 - 1
uvcoords[..., 1] = -uvcoords[..., 1]
face_uvcoords = util.face_vertices(uvcoords, uvfaces)
self.register_buffer('uvcoords', uvcoords)
self.register_buffer('uvfaces', uvfaces)
self.register_buffer('face_uvcoords', face_uvcoords)
# shape colors
colors = torch.tensor([74, 120, 168])[None, None, :].repeat(
1,
faces.max() + 1, 1).float() / 255.
face_colors = util.face_vertices(colors, faces)
self.register_buffer('face_colors', face_colors)
## lighting
pi = np.pi
constant_factor = torch.tensor(
[1 / np.sqrt(4 * pi), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), \
((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), \
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi)))])
self.register_buffer('constant_factor', constant_factor)
def world2uv(self, vertices):
'''
sample vertices from world space to uv space
uv_vertices: [bz, 3, h, w]
'''
batch_size = vertices.shape[0]
face_vertices = util.face_vertices(vertices, self.faces.expand(batch_size, -1, -1)).clone().detach()
uv_vertices = self.uv_rasterizer(self.uvcoords.expand(batch_size, -1, -1),
self.uvfaces.expand(batch_size, -1, -1), face_vertices)[:, :3]
return uv_vertices
def render_normal(self, transformed_vertices, normals):
'''
-- rendering normal
'''
batch_size = normals.shape[0]
# Attributes
attributes = util.face_vertices(normals, self.faces.expand(batch_size, -1, -1))
# rasterize
rendering = self.rasterizer(transformed_vertices, self.faces.expand(batch_size, -1, -1), attributes)
####
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
normal_images = rendering[:, :3, :, :]
return normal_images
def forward(self, cam, vertices, images, cam_new):
full_vertices, N_bd = get_full_verts(vertices)
t_vertices = util.batch_orth_proj(full_vertices, cam)
t_vertices[..., 1:] = -t_vertices[..., 1:]
t_vertices[..., 2] = t_vertices[..., 2] + 10
t_vertices = image_meshing(t_vertices, N_bd)
t_vertices[..., :2] = torch.clamp(t_vertices[..., :2], -1, 1)
t_vertices[:, :, 2] = t_vertices[:, :, 2] - 9
batch_size = vertices.shape[0]
## rasterizer near 0 far 100. move mesh so minz larger than 0
uvcoords = t_vertices.clone()
# Attributes
uvcoords = torch.cat(
[uvcoords[:, :, :2], uvcoords[:, :, 0:1] * 0. + 1.],
-1) # [bz, ntv, 3]
face_vertices = util.face_vertices(
uvcoords, self.faces.expand(batch_size, -1, -1))
# render
attributes = face_vertices.detach()
full_vertices, N_bd = get_full_verts(vertices)
transformed_vertices = util.batch_orth_proj(full_vertices, cam_new)
transformed_vertices[..., 1:] = -transformed_vertices[..., 1:]
transformed_vertices[..., 2] = transformed_vertices[..., 2] + 10
transformed_vertices = image_meshing(transformed_vertices, N_bd)
transformed_vertices[..., :2] = torch.clamp(
transformed_vertices[..., :2], -1, 1)
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes)
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
# albedo
uvcoords_images = rendering[:, :3, :, :]
grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2]
results = F.grid_sample(images, grid, align_corners=False)
return {'rotate_images': results}
def forward(self, vertices, transformed_vertices, albedos, lights=None, light_type='point'):
'''
lihgts:
spherical homarnic: [N, 9(shcoeff), 3(rgb)]
vertices: [N, V, 3], vertices in work space, for calculating normals, then shading
transformed_vertices: [N, V, 3], range(-1, 1), projected vertices, for rendering
'''
batch_size = vertices.shape[0]
## rasterizer near 0 far 100. move mesh so minz larger than 0
transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] + 10
# Attributes
face_vertices = util.face_vertices(vertices, self.faces.expand(batch_size, -1, -1))
normals = util.vertex_normals(vertices, self.faces.expand(batch_size, -1, -1))
face_normals = util.face_vertices(normals, self.faces.expand(batch_size, -1, -1))
transformed_normals = util.vertex_normals(transformed_vertices, self.faces.expand(batch_size, -1, -1))
transformed_face_normals = util.face_vertices(transformed_normals, self.faces.expand(batch_size, -1, -1))
# render
attributes = torch.cat([self.face_uvcoords.expand(batch_size, -1, -1, -1), transformed_face_normals.detach(),
face_vertices.detach(), face_normals.detach()], -1)
# import ipdb;ipdb.set_trace()
rendering = self.rasterizer(transformed_vertices, self.faces.expand(batch_size, -1, -1), attributes)
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
# albedo
uvcoords_images = rendering[:, :3, :, :]
grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2]
albedo_images = F.grid_sample(albedos, grid, align_corners=False)
# remove inner mouth region
transformed_normal_map = rendering[:, 3:6, :, :].detach()
pos_mask = (transformed_normal_map[:, 2:, :, :] < -0.05).float()
# shading
if lights is not None:
normal_images = rendering[:, 9:12, :, :].detach()
if lights.shape[1] == 9:
shading_images = self.add_SHlight(normal_images, lights)
else:
if light_type == 'point':
vertice_images = rendering[:, 6:9, :, :].detach()
shading = self.add_pointlight(vertice_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]),
normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape(
[batch_size, lights.shape[1], albedo_images.shape[2], albedo_images.shape[3], 3]).permute(0, 1,
4, 2,
3)
shading_images = shading_images.mean(1)
else:
shading = self.add_directionlight(normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape(
[batch_size, lights.shape[1], albedo_images.shape[2], albedo_images.shape[3], 3]).permute(0, 1,
4, 2,
3)
shading_images = shading_images.mean(1)
images = albedo_images * shading_images
else:
images = albedo_images
shading_images = images.detach() * 0.
outputs = {
'images': images * alpha_images,
'albedo_images': albedo_images,
'alpha_images': alpha_images,
'pos_mask': pos_mask,
'shading_images': shading_images,
'grid': grid,
'normals': normals
}
return outputs