(self, vertices, faces, attributes=None, h=None, w=None)
| 60 | self.w = w = width |
| 61 | |
| 62 | def forward(self, vertices, faces, attributes=None, h=None, w=None): |
| 63 | device = vertices.device |
| 64 | if h is None: |
| 65 | h = self.h |
| 66 | if w is None: |
| 67 | w = self.h |
| 68 | bz = vertices.shape[0] |
| 69 | depth_buffer = torch.zeros([bz, h, w]).float().to(device) + 1e6 |
| 70 | triangle_buffer = torch.zeros([bz, h, w]).int().to(device) - 1 |
| 71 | baryw_buffer = torch.zeros([bz, h, w, 3]).float().to(device) |
| 72 | vert_vis = torch.zeros([bz, vertices.shape[1]]).float().to(device) |
| 73 | |
| 74 | vertices = vertices.clone().float() |
| 75 | vertices[..., 0] = vertices[..., 0]*w/2 + w/2 |
| 76 | vertices[..., 1] = vertices[..., 1]*h/2 + h/2 |
| 77 | vertices[..., 2] = vertices[..., 2]*w/2 |
| 78 | f_vs = util.face_vertices(vertices, faces) |
| 79 | |
| 80 | standard_rasterize(f_vs, depth_buffer, |
| 81 | triangle_buffer, baryw_buffer, h, w) |
| 82 | pix_to_face = triangle_buffer[:, :, :, None].long() |
| 83 | bary_coords = baryw_buffer[:, :, :, None, :] |
| 84 | vismask = (pix_to_face > -1).float() |
| 85 | D = attributes.shape[-1] |
| 86 | attributes = attributes.clone() |
| 87 | attributes = attributes.view( |
| 88 | attributes.shape[0]*attributes.shape[1], 3, attributes.shape[-1]) |
| 89 | N, H, W, K, _ = bary_coords.shape |
| 90 | mask = pix_to_face == -1 |
| 91 | pix_to_face = pix_to_face.clone() |
| 92 | pix_to_face[mask] = 0 |
| 93 | idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D) |
| 94 | pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D) |
| 95 | pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2) |
| 96 | pixel_vals[mask] = 0 # Replace masked values in output. |
| 97 | pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) |
| 98 | pixel_vals = torch.cat( |
| 99 | [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) |
| 100 | return pixel_vals |
| 101 | |
| 102 | |
| 103 | class Pytorch3dRasterizer(nn.Module): |
nothing calls this directly
no test coverage detected