Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from pytorch3d.renderer import ( | |
| PerspectiveCameras, | |
| PointsRasterizationSettings, | |
| PointsRasterizer, | |
| AlphaCompositor, | |
| ) | |
| def homogenize_pt(coord): | |
| return torch.cat([coord, torch.ones_like(coord[..., :1])], dim=-1) | |
| def unproject_pts_pt(intrinsics, coords, depth): | |
| if coords.shape[-1] == 2: | |
| coords = homogenize_pt(coords) | |
| intrinsics = intrinsics.squeeze()[:3, :3] | |
| coords = torch.inverse(intrinsics).mm(coords.T) * depth.reshape(1, -1) | |
| return coords.T # [n, 3] | |
| def get_coord_grids_pt(h, w, device, homogeneous=False): | |
| """ | |
| create pxiel coordinate grid | |
| :param h: height | |
| :param w: weight | |
| :param device: device | |
| :param homogeneous: if homogeneous coordinate | |
| :return: coordinates [h, w, 2] | |
| """ | |
| y = torch.arange(0, h).to(device) | |
| x = torch.arange(0, w).to(device) | |
| grid_y, grid_x = torch.meshgrid(y, x) | |
| if homogeneous: | |
| return torch.stack([grid_x, grid_y, torch.ones_like(grid_x)], dim=-1) | |
| return torch.stack([grid_x, grid_y], dim=-1) # [h, w, 2] | |
| class PointsRenderer(nn.Module): | |
| """ | |
| A class for rendering a batch of points. The class should | |
| be initialized with a rasterizer and compositor class which each have a forward | |
| function. | |
| """ | |
| def __init__(self, rasterizer, compositor) -> None: | |
| super().__init__() | |
| self.rasterizer = rasterizer | |
| self.compositor = compositor | |
| def to(self, device): | |
| self.rasterizer = self.rasterizer.to(device) | |
| self.compositor = self.compositor.to(device) | |
| return self | |
| def forward(self, point_clouds, **kwargs) -> torch.Tensor: | |
| fragments = self.rasterizer(point_clouds, **kwargs) | |
| r = self.rasterizer.raster_settings.radius | |
| if type(r) == torch.Tensor: | |
| if r.shape[-1] > 1: | |
| idx = fragments.idx.clone() | |
| idx[idx == -1] = 0 | |
| r = r[:, idx.squeeze().long()] | |
| r = r.permute(0, 3, 1, 2) | |
| dists2 = fragments.dists.permute(0, 3, 1, 2) | |
| weights = 1 - dists2 / (r * r) | |
| images = self.compositor( | |
| fragments.idx.long().permute(0, 3, 1, 2), | |
| weights, | |
| point_clouds.features_packed().permute(1, 0), | |
| **kwargs, | |
| ) | |
| # permute so image comes at the end | |
| images = images.permute(0, 2, 3, 1) | |
| return images | |
| def create_pcd_renderer(h, w, intrinsics, R=None, T=None, radius=None, device="cuda"): | |
| fx = intrinsics[0, 0] | |
| fy = intrinsics[1, 1] | |
| if R is None: | |
| R = torch.eye(3)[None] # (1, 3, 3) | |
| if T is None: | |
| T = torch.zeros(1, 3) # (1, 3) | |
| cameras = PerspectiveCameras(R=R, T=T, | |
| device=device, | |
| focal_length=((-fx, -fy),), | |
| principal_point=(tuple(intrinsics[:2, -1]),), | |
| image_size=((h, w),), | |
| in_ndc=False, | |
| ) | |
| if radius is None: | |
| radius = 1.5 / min(h, w) * 2.0 | |
| raster_settings = PointsRasterizationSettings( | |
| image_size=(h, w), | |
| radius=radius, | |
| points_per_pixel=8, | |
| ) | |
| rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings) | |
| renderer = PointsRenderer( | |
| rasterizer=rasterizer, | |
| compositor=AlphaCompositor(background_color=(1, 1, 1)) | |
| ) | |
| return renderer | |