def test_render_texture(self, uvs, faces, texture_maps,
face_vertices_camera, face_vertices_image,
face_camera_normals_z, height, width, dtype,
device):
batch_size = faces.shape[0]
# attributes with uvs
#uvs_with_mask = torch.nn.functional.pad(uvs, pad=(1, 0), value=1)
#attributes = uvs_with_mask
face_uvs = index_vertices_by_faces(uvs, faces)
face_attributes = [
torch.ones((*face_uvs.shape[:-1], 1), device=device, dtype=dtype),
face_uvs
]
(texmask, texcoord), improb, imfaceidx = dibr_rasterization(
height, width, face_vertices_camera[:, :, :, 2],
face_vertices_image, face_attributes, face_camera_normals_z)
texcolor = texture_mapping(texcoord, texture_maps, mode='bilinear')
image = texcolor * texmask
for bs in range(batch_size):
image_gt = torch.from_numpy(
np.array(
Image.open(os.path.join(SAMPLE_DIR, f'texture_{bs}.png'))))
image_gt = image_gt.to(device, dtype) / 255.
assert torch.allclose(image[bs], image_gt, atol=1. / 255.0)
def test_all_grads(self, height, width, face_vertices_camera,
face_vertices_image_zoom, face_vertex_colors,
face_camera_normals_z, dtype, target_all_grads):
_face_vertex_colors = face_vertex_colors.detach()
_face_vertex_colors.requires_grad = True
_face_vertices_image_zoom = face_vertices_image_zoom.detach()
_face_vertices_image_zoom.requires_grad = True
outputs = dibr_rasterization(
20, 30,
face_vertices_camera.to(dtype)[:, :, :, 2],
_face_vertices_image_zoom.to(dtype), _face_vertex_colors.to(dtype),
face_camera_normals_z.to(dtype))
# gradients will be provided through the two differentiable outputs
output = sum([torch.sum(output) for output in outputs])
output.backward()
if dtype == torch.float:
assert torch.allclose(target_all_grads[0],
_face_vertices_image_zoom.grad,
rtol=1e-5,
atol=1e-4)
else:
# TODO(cfujitsang): non-determinism?
assert torch.allclose(target_all_grads[0],
_face_vertices_image_zoom.grad)
assert torch.allclose(target_all_grads[1], _face_vertex_colors.grad)
def test_render_depths(self, face_vertices_camera, face_vertices_image,
face_camera_normals_z, height, width, dtype,
device):
batch_size = face_vertices_camera.shape[0]
# face_vertices_camera is of shape (num_batch, num_face, 9)
num_batch, num_faces = face_vertices_camera.shape[:2]
face_attributes = face_vertices_camera.reshape(num_batch, num_faces, 3,
3)[:, :, :, 2:3]
# imfeat is interpolated features
# improb is the soft mask
# imfaceidx is the face index map, which pixel is covered by which face
# it starts from 1, 0 is void.
imfeat, improb, imfaceidx = dibr_rasterization(
height, width, face_vertices_camera[:, :, :, 2],
face_vertices_image, face_attributes, face_camera_normals_z)
image = imfeat
image_valid_region = image < 0
image_valid_values = image[image_valid_region]
image_depth_norm = (image - image_valid_values.min()) / (
image_valid_values.max() - image_valid_values.min())
image = torch.where(image_valid_region, image_depth_norm, image)
for bs in range(batch_size):
image_gt = torch.from_numpy(
np.array(
Image.open(os.path.join(SAMPLE_DIR, f'depth_{bs}.png'))))
image_gt = image_gt.to(device, dtype).unsqueeze(2) / 255.
assert torch.allclose(image[bs], image_gt, atol=1 / 255.0)
def test_render_normal(self, face_vertices_camera, face_vertices_image,
face_camera_normals_z, height, width, dtype,
device):
batch_size = face_vertices_camera.shape[0]
face_normals_unit = face_normals(face_vertices_camera, unit=True)
face_attributes = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1)
# imfeat is interpolated features
# improb is the soft mask
# imfaceidx is the face index map, which pixel is covered by which face
# it starts from 1, 0 is void.
imfeat, improb, imfaceidx = dibr_rasterization(
height, width, face_vertices_camera[:, :, :, 2],
face_vertices_image, face_attributes, face_camera_normals_z)
images = (imfeat + 1) / 2
images_gt = [
torch.from_numpy(
np.array(
Image.open(
os.path.join(SAMPLE_DIR, f'vertex_normal_{bs}.png'))))
for bs in range(batch_size)
]
images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255.
if dtype == torch.double:
num_pix_diff_tol = 8
else:
num_pix_diff_tol = 0
num_pix_diff = torch.sum(
~torch.isclose(images, images_gt, atol=1. / 255.))
assert num_pix_diff <= num_pix_diff_tol
def test_render_vertex_colors(self, vertex_colors, faces,
face_vertices_camera, face_vertices_image,
face_camera_normals_z, height, width, dtype,
device):
batch_size = faces.shape[0]
# face_vertex_colors
attributes = vertex_colors
face_attributes_idx = faces
face_attributes = index_vertices_by_faces(attributes,
face_attributes_idx)
# imfeat is interpolated features
# improb is the soft mask
# imfaceidx is the face index map, which pixel is covered by which face
# it starts from 1, 0 is void.
imfeat, improb, imfaceidx = dibr_rasterization(
height, width, face_vertices_camera[:, :, :, 2],
face_vertices_image, face_attributes, face_camera_normals_z)
image = imfeat
images_gt = [
torch.from_numpy(
np.array(
Image.open(
os.path.join(SAMPLE_DIR, f'vertex_color_{bs}.png'))))
for bs in range(batch_size)
]
images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255.
# the rendered soft mask is only tested here
images_prob_gt = [
torch.from_numpy(
np.array(
Image.open(os.path.join(SAMPLE_DIR,
f"image_prob_{bs}.png"))))
for bs in range(batch_size)
]
images_prob_gt = torch.stack(images_prob_gt, dim=0).to(device,
dtype) / 255.
# the rendered face_idx is only tested here
images_face_idx_gt = [
torch.from_numpy(
np.array(
Image.open(
os.path.join(SAMPLE_DIR, f"image_face_idx_{bs}.png"))))
for bs in range(batch_size)
]
images_face_idx_gt = \
(torch.stack(images_face_idx_gt, dim=0).to(device, torch.long) // 100) - 1
assert torch.allclose(image, images_gt, atol=1. / 255.)
assert torch.allclose(improb, images_prob_gt, atol=1. / 255.0)
if dtype == torch.double:
num_pix_diff_tol = 4
else:
num_pix_diff_tol = 0
num_pix_diff = torch.sum(
~torch.isclose(imfaceidx, images_face_idx_gt, atol=1. / 255.))
assert num_pix_diff <= num_pix_diff_tol
def render(self, **attributes):
azimuths = attributes['azimuths']
elevations = attributes['elevations']
distances = attributes['distances']
batch_size = azimuths.shape[0]
device = azimuths.device
cam_proj = self.cam_proj.to(device)
vertices = attributes['vertices']
textures = attributes['textures']
lights = attributes['lights']
faces = self.faces.to(device)
face_uvs = self.face_uvs.to(device)
num_faces = faces.shape[0]
object_pos = torch.tensor([[0., 0., 0.]], dtype=torch.float, device=device).repeat(batch_size, 1)
camera_up = torch.tensor([[0., 1., 0.]], dtype=torch.float, device=device).repeat(batch_size, 1)
# camera_pos = torch.tensor([[0., 0., 4.]], dtype=torch.float, device=device).repeat(batch_size, 1)
camera_pos = camera_position_from_spherical_angles(distances, elevations, azimuths, degrees=True)
cam_transform = generate_transformation_matrix(camera_pos, object_pos, camera_up)
face_vertices_camera, face_vertices_image, face_normals = \
prepare_vertices(vertices=vertices,
faces=faces, camera_proj=cam_proj, camera_transform=cam_transform
)
face_normals_unit = kal.ops.mesh.face_normals(face_vertices_camera, unit=True)
face_normals_unit = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1)
face_attributes = [
torch.ones((batch_size, num_faces, 3, 1), device=device),
face_uvs.repeat(batch_size, 1, 1, 1),
face_normals_unit
]
image_features, soft_mask, face_idx = dibr_rasterization(
self.image_size, self.image_size, face_vertices_camera[:, :, :, -1],
face_vertices_image, face_attributes, face_normals[:, :, -1])
# image_features is a tuple in composed of the interpolated attributes of face_attributes
# texture_coords, mask = image_features
texmask, texcoord, imnormal = image_features
texcolor = texture_mapping(texcoord, textures, mode='bilinear')
coef = spherical_harmonic_lighting(imnormal, lights)
image = texcolor * texmask * coef.unsqueeze(-1) + torch.ones_like(texcolor) * (1 - texmask)
image = torch.clamp(image, 0, 1)
render_img = image
render_silhouttes = soft_mask[..., None]
rgbs = torch.cat([render_img, render_silhouttes], axis=-1).permute(0, 3, 1, 2)
attributes['face_normals'] = face_normals
attributes['faces_image'] = face_vertices_image.mean(dim=2)
attributes['visiable_faces'] = face_normals[:, :, -1] > 0.1
return rgbs, attributes
def target_face_vertex_colors_grad(self, height, width,
face_vertices_camera,
face_vertices_image, face_vertex_colors,
face_camera_normals_z):
_face_vertex_colors = face_vertex_colors.detach()
_face_vertex_colors.requires_grad = True
# This assume that test_vertex_colors_gradcheck pass
outputs = dibr_rasterization(20, 30, face_vertices_camera[:, :, :, 2],
face_vertices_image, _face_vertex_colors,
face_camera_normals_z)
# gradients will be provided through the two differentiable outputs
output = sum([torch.sum(output) for output in outputs])
output.backward()
return _face_vertex_colors.grad.clone()
def test_face_vertex_colors_grads(self, height, width,
face_vertices_camera,
face_vertices_image, face_vertex_colors,
face_camera_normals_z, dtype,
target_face_vertex_colors_grad):
_face_vertex_colors = face_vertex_colors.detach()
_face_vertex_colors.requires_grad = True
outputs = dibr_rasterization(
20, 30,
face_vertices_camera.to(dtype)[:, :, :, 2],
face_vertices_image.to(dtype), _face_vertex_colors.to(dtype),
face_camera_normals_z.to(dtype))
# gradients will be provided through the two differentiable outputs
output = sum([torch.sum(output) for output in outputs])
output.backward()
assert torch.allclose(_face_vertex_colors.grad,
target_face_vertex_colors_grad)
def test_render_texture_with_light(self, uvs, faces, texture_maps, lights,
face_vertices_camera,
face_vertices_image,
face_camera_normals_z, height, width,
dtype, device):
batch_size = faces.shape[0]
# Note: in this example uv face is the same as mesh face
# but they could be different
face_uvs = index_vertices_by_faces(uvs, faces)
# normal
face_normals_unit = face_normals(face_vertices_camera, unit=True)
face_normals_unit = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1)
# merge them together
face_attributes = [
torch.ones((*face_uvs.shape[:-1], 1), device=device, dtype=dtype),
face_uvs, face_normals_unit
]
(texmask, texcoord, imnormal), improb, imidx = dibr_rasterization(
height, width, face_vertices_camera[:, :, :, 2],
face_vertices_image, face_attributes, face_camera_normals_z)
texcolor = texture_mapping(texcoord, texture_maps, mode='nearest')
coef = spherical_harmonic_lighting(imnormal, lights)
images = torch.clamp(texmask * texcolor * coef.unsqueeze(-1), 0, 1)
if dtype == torch.double:
num_pix_diff_tol = 74 # (over 2 x 256 x 512 x 3 pixels)
else:
num_pix_diff_tol = 0
images_gt = [
torch.from_numpy(
np.array(
Image.open(
os.path.join(SAMPLE_DIR, f'texture_light_{bs}.png'))))
for bs in range(batch_size)
]
images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255.
num_pix_diff = torch.sum(
~torch.isclose(images, images_gt, atol=1. / 255.))
assert num_pix_diff <= num_pix_diff_tol
def test_optimize_vertex_position(self, vertices, faces, vertex_colors,
vertices_image, camera_rot, camera_trans,
camera_proj, height, width, dtype,
device):
batch_size = faces.shape[0]
# face_vertex_colors
camera_rot = camera_rot.to(device, dtype)
camera_rot.requires_grad = False
camera_trans = camera_trans.to(device, dtype)
camera_trans.requires_grad = False
camera_proj = camera_proj.to(device, dtype)
camera_proj.requires_grad = False
face_attributes = index_vertices_by_faces(
vertex_colors.to(device, dtype), faces)
vertices = vertices.to(device, dtype).clone().detach()
vertices.requires_grad = False
moved_vertices = vertices.to(device, dtype).clone()
moved_vertices[:, 0, :2] += 0.4
moved_vertices = moved_vertices.detach()
moved_vertices.requires_grad = True
images_gt = [
torch.from_numpy(
np.array(
Image.open(
os.path.join(SAMPLE_DIR, f'vertex_color_{bs}.png'))))
for bs in range(batch_size)
]
images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255.
with torch.no_grad():
moved_vertices_camera = rotate_translate_points(
moved_vertices, camera_rot, camera_trans)
moved_vertices_image = perspective_camera(moved_vertices_camera,
camera_proj)
# test that the vertex are far enough to fail the test.
assert not torch.allclose(
moved_vertices_image, vertices_image, atol=1e-2, rtol=1e-2)
with torch.no_grad():
face_moved_vertices_camera, face_moved_vertices_image, face_moved_normals = \
prepare_vertices(moved_vertices, faces, camera_proj, camera_rot, camera_trans)
face_moved_normals_z = face_moved_normals[:, :, 2]
imfeat, _, _ = dibr_rasterization(
height, width, face_moved_vertices_camera[:, :, :, 2],
face_moved_vertices_image, face_attributes,
face_moved_normals_z)
original_loss = torch.mean(torch.abs(imfeat - images_gt))
# test that the loss is high enough
assert original_loss > 0.01
optimizer = torch.optim.Adam([moved_vertices], lr=5e-3)
for i in range(100):
optimizer.zero_grad()
face_moved_vertices_camera, face_moved_vertices_image, face_moved_normals = \
prepare_vertices(moved_vertices, faces, camera_proj, camera_rot, camera_trans)
face_moved_normals_z = face_moved_normals[:, :, 2]
imfeat, _, _ = dibr_rasterization(
height, width, face_moved_vertices_camera[:, :, :, 2],
face_moved_vertices_image, face_attributes,
face_moved_normals_z)
loss = torch.mean(torch.abs(imfeat - images_gt))
loss.backward()
optimizer.step()
moved_vertices_camera = rotate_translate_points(
moved_vertices, camera_rot, camera_trans)
moved_vertices_image = perspective_camera(moved_vertices_camera,
camera_proj)
# test that the loss went down
assert loss < 0.001
# We only test on image plan since we don't change camera angle during training we don't expect depth to be correct.
# We could probably fine-tune the test to have a lower tolerance (TODO: cfujitsang)
assert torch.allclose(moved_vertices_image,
vertices_image,
atol=1e-2,
rtol=1e-2)
