How to obtain 3D grid from example of 2D grid?

I have a code to generate 2D grid

        B, C, H, W = x.size()
        # mesh grid 
        xx = torch.arange(0, W).view(1,-1).repeat(H,1)
        yy = torch.arange(0, H).view(-1,1).repeat(1,W)
        xx = xx.view(1,1,H,W).repeat(B,1,1,1)
        yy = yy.view(1,1,H,W).repeat(B,1,1,1)
        grid = torch.cat((xx,yy),1).float()

How to obtain 3D grid by adding zz in the code? This is what I tried

       B,C,D,H,W=input.size()
        xx = torch.arange(0, W).view(1, 1,-1).repeat(D, H, 1)
        yy = torch.arange(0, H).view(1, -1,1).repeat(D, 1, W)
        zz = torch.arange(0, D).view(-1, 1,1).repeat(1,H,W)
        xx = xx.view(1,1,D,H,W).repeat(B,1,1,1,1)
        yy = yy.view(1,1,D,H,W).repeat(B,1,1,1,1)
        zz = zz.view(1,1,D,H,W).repeat(B,1,1,1,1)
        grid = torch.cat((xx,yy,zz),1).float().to('cuda')

But i feel it wrong in somewhere. Could you verify it?

Was this suggestion not working?

@ptrblck: Yes. Because it is care about the order. I can use it to make grid but when I feed to grid_sample () function. The output looks rotated.

For example xx = torch.arange(0, W).view(1, 1,-1).repeat(D, H, 1) or xx = torch.arange(0, D).view(-1, 1,1).repeat(1, H, W)

The easiest way to obtain a grid for use with the grid_sample() function is probably to use torch.nn.functional.affine_grid().

You give it an affine matrix, and it returns a grid that you can then pass to grid_sample(). If you want an identity grid (no transformation), you can just pass it the identity affine matrix, which for 3D data is

aff = torch.FloatTensor([[[1, 0, 0, 0],[0, 1, 0, 0],[0, 0, 1, 0]]])

For example,

aff = aff.expand(B, 3, 4)  # expand to the number of batches you need
grid = torch.nn.functional.affine_grid(aff, size=(B,C,D,H,W))
torch.nn.functional.grid_sample(3d_data, grid)

So yeah, you could probably get the same effect using your code sample if you do it just right, but this should be much easier.
I hope this helps.

1 Like