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)