Render depth maps for batched meshes. vertices: [batch_size, num_vertices, 3] faces: [batch_size, num_faces, 3]
(
vertices: torch.Tensor,
faces: torch.Tensor,
image_size: Tuple[int, int],
focal_length: Optional[float] = None,
R: Optional[torch.Tensor] = None,
T: Optional[torch.Tensor] = None,
reverse_axis: bool = False,
)
| 154 | |
| 155 | |
| 156 | def render_depth_maps( |
| 157 | vertices: torch.Tensor, |
| 158 | faces: torch.Tensor, |
| 159 | image_size: Tuple[int, int], |
| 160 | focal_length: Optional[float] = None, |
| 161 | R: Optional[torch.Tensor] = None, |
| 162 | T: Optional[torch.Tensor] = None, |
| 163 | reverse_axis: bool = False, |
| 164 | ) -> torch.Tensor: |
| 165 | """ |
| 166 | Render depth maps for batched meshes. |
| 167 | vertices: [batch_size, num_vertices, 3] |
| 168 | faces: [batch_size, num_faces, 3] |
| 169 | """ |
| 170 | device = vertices.device |
| 171 | fov = 2 * np.arctan(min(image_size) / (2 * focal_length)) |
| 172 | camera_kwargs = {"fov": fov, "znear": 0.005, "zfar": 1000, "device": device, "degrees": False} |
| 173 | if R is not None: |
| 174 | vertices = torch.matmul(vertices, R) |
| 175 | if T is not None: |
| 176 | vertices = vertices + T[:, None] |
| 177 | cameras = FoVPerspectiveCameras(**camera_kwargs) |
| 178 | projected_vertices = cameras.transform_points(vertices) |
| 179 | projected_vertices[..., -1] = vertices[..., -1] |
| 180 | depth = pytorch3d_rasterize(projected_vertices, faces, image_size=image_size, reverse_axis=reverse_axis) |
| 181 | depth_maps = depth.unsqueeze(1).detach() |
| 182 | return depth_maps |
| 183 | |
| 184 | |
| 185 | def rendering_batches( |
no test coverage detected