def rasterize():
rasterize_meshes(
meshes_batch,
image_size,
blur_radius,
8,
bin_size,
max_faces_per_bin,
)
torch.cuda.synchronize()
def test_mesh_image_size_arg(self):
meshes = Meshes(verts=[verts0], faces=[faces0])
with self.assertRaises(ValueError) as cm:
rasterize_meshes(
meshes,
(100, 200, 3),
0.0001,
faces_per_pixel=1,
)
self.assertTrue("tuple/list of (H, W)" in cm.msg)
with self.assertRaises(ValueError) as cm:
rasterize_meshes(
meshes,
(0, 10),
0.0001,
faces_per_pixel=1,
)
self.assertTrue("sizes must be positive" in cm.msg)
with self.assertRaises(ValueError) as cm:
rasterize_meshes(
meshes,
(100.5, 120.5),
0.0001,
faces_per_pixel=1,
)
self.assertTrue("sizes must be integers" in cm.msg)
def _rasterize(self, meshes, image_size, bin_size, blur):
"""
Simple wrapper around the rasterize function to return
the fragment data.
"""
face_idxs, zbuf, bary_coords, pix_dists = rasterize_meshes(
meshes,
image_size,
blur,
faces_per_pixel=1,
bin_size=bin_size,
)
return Fragments(
pix_to_face=face_idxs,
zbuf=zbuf,
bary_coords=bary_coords,
dists=pix_dists,
)
def compute_mesh_visibility(vert, tri, bfm_torch, H_img, W_img):
N, nver, _ = vert.shape
ntri = tri.shape[0]
tri = torch.from_numpy(tri).to(vert.device).unsqueeze(0).expand(N, ntri, 3)
# Transform to NDC
vert_t = vert + torch.tensor((0.5, 0.5, 0), dtype=torch.float, device=vert.device).view(1, 1, 3)
vert_t = vert_t * torch.tensor((-1, 1, -1), dtype=torch.float, device=vert.device).view(1, 1, 3)
vert_t = vert_t * torch.tensor((2. / float(W_img), 2. / float(H_img), 1),
dtype=torch.float, device=vert.device).view(1, 1, 3)
# vert_t = vert * torch.tensor((-1, 1, -1), dtype=torch.float, device=vert.device).view(1, 1, 3)
# vert_t = vert_t * torch.tensor((2. / (float(W_img) - 1.), 2. / (float(H_img) - 1.), 1),
# dtype=torch.float, device=vert.device).view(1, 1, 3)
z_offset = -vert_t[:, :, 2].min() + 10.
vert_t = vert_t + torch.tensor((0, 0, z_offset.item()), dtype=torch.float, device=vert.device).view(1, 1, 3)
# Render
mesh_torch = Meshes(verts=vert_t, faces=tri)
pix_to_face, zbuf, _, _ = rasterize_meshes(mesh_torch, image_size=H_img, faces_per_pixel=1, cull_backfaces=True)
assert H_img == W_img
# Compute visibility
zbuf = (zbuf.view(N, 1, H_img, W_img) - z_offset) * -1 - 1.
zbuf = torch.where(torch.eq(pix_to_face.view(N, 1, H_img, W_img), -1), torch.zeros_like(zbuf) - 2e2, zbuf)
vert_zbuf = sample_per_vert_feat(zbuf, vert, bfm_torch, H_img, W_img).squeeze(1)
mask = torch.gt(vert[:, :, 2], vert_zbuf)
return mask, zbuf
def test_cpp_vs_cuda_naive_vs_cuda_binned(self):
# Make sure that the backward pass runs for all pathways
image_size = 64 # test is too slow for very large images.
N = 1
radius = 0.1**2
faces_per_pixel = 3
grad_zbuf = torch.randn(N, image_size, image_size, faces_per_pixel)
grad_dist = torch.randn(N, image_size, image_size, faces_per_pixel)
grad_bary = torch.randn(N, image_size, image_size, faces_per_pixel, 3)
device = torch.device("cpu")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
# Option I: CPU, naive
args = (meshes, image_size, radius, faces_per_pixel)
idx1, zbuf1, bary1, dist1 = rasterize_meshes(*args)
loss = ((zbuf1 * grad_zbuf).sum() + (dist1 * grad_dist).sum() +
(bary1 * grad_bary).sum())
loss.backward()
idx1 = idx1.data.cpu().clone()
zbuf1 = zbuf1.data.cpu().clone()
dist1 = dist1.data.cpu().clone()
grad1 = verts.grad.data.cpu().clone()
# Option II: CUDA, naive
device = torch.device("cuda:0")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
args = (meshes, image_size, radius, faces_per_pixel, 0, 0)
idx2, zbuf2, bary2, dist2 = rasterize_meshes(*args)
grad_zbuf = grad_zbuf.cuda()
grad_dist = grad_dist.cuda()
grad_bary = grad_bary.cuda()
loss = ((zbuf2 * grad_zbuf).sum() + (dist2 * grad_dist).sum() +
(bary2 * grad_bary).sum())
loss.backward()
idx2 = idx2.data.cpu().clone()
zbuf2 = zbuf2.data.cpu().clone()
dist2 = dist2.data.cpu().clone()
grad2 = verts.grad.data.cpu().clone()
# Option III: CUDA, binned
device = torch.device("cuda:0")
meshes = ico_sphere(0, device)
verts, faces = meshes.get_mesh_verts_faces(0)
verts.requires_grad = True
meshes = Meshes(verts=[verts], faces=[faces])
args = (meshes, image_size, radius, faces_per_pixel, 32, 500)
idx3, zbuf3, bary3, dist3 = rasterize_meshes(*args)
loss = ((zbuf3 * grad_zbuf).sum() + (dist3 * grad_dist).sum() +
(bary3 * grad_bary).sum())
loss.backward()
idx3 = idx3.data.cpu().clone()
zbuf3 = zbuf3.data.cpu().clone()
dist3 = dist3.data.cpu().clone()
grad3 = verts.grad.data.cpu().clone()
# Make sure everything was the same
self.assertTrue((idx1 == idx2).all().item())
self.assertTrue((idx1 == idx3).all().item())
self.assertTrue(torch.allclose(zbuf1, zbuf2, atol=1e-6))
self.assertTrue(torch.allclose(zbuf1, zbuf3, atol=1e-6))
self.assertTrue(torch.allclose(dist1, dist2, atol=1e-6))
self.assertTrue(torch.allclose(dist1, dist3, atol=1e-6))
self.assertTrue(torch.allclose(grad1, grad2, rtol=5e-3)) # flaky test
self.assertTrue(torch.allclose(grad1, grad3, rtol=5e-3))
self.assertTrue(torch.allclose(grad2, grad3, rtol=5e-3))
def rasterize():
rasterize_meshes(meshes_batch, image_size, blur_radius, bin_size=0)
