Function gridcloud3d

TEACh/utils/basic.py:130–139

(B, Z, Y, X, norm=False)

    return tensor

def gridcloud3d(B, Z, Y, X, norm=False):
    # we want to sample for each location in the grid
    grid_z, grid_y, grid_x = meshgrid3d(B, Z, Y, X, norm=norm)
    x = torch.reshape(grid_x, [B, -1])
    y = torch.reshape(grid_y, [B, -1])
    z = torch.reshape(grid_z, [B, -1])
    # these are B x N
    xyz = torch.stack([x, y, z], dim=2)
    # this is B x N x 3
    return xyz

def gridcloud2d(B, Y, X, norm=False):
    # we want to sample for each location in the grid