First non zero value in an image dim

I have a stack of MRI images that look like

say in shape [train_size, h,w,d]

There is alot of black space making my convolution very heavy so I want to trim the images from the minimum index of the first non-zero pixel value from all the 4 corners. i.e identify the maximum distance I can trim my images from the four corners without losing any non-zero values pixel from any of the images in whole training size. This is only for the height and width. How Could I go about doing this?

This code should work for a [batch, channel, height, width]-shaped tensor:

x = torch.zeros(1, 3, 24, 24)
x[0, :, 5:10, 6:11] = torch.randn(1, 3, 5, 5)

non_zero = (x!=0).nonzero()  # Contains non-zero indices for all 4 dims
h_min = non_zero[:, 2].min()
h_max = non_zero[:, 2].max()
w_min = non_zero[:, 3].min()
w_max = non_zero[:, 3].max()

@ptrblck this solution gives the min and max nonzero cell across all images in the batch, which is unlikely to be the desired behavior.

I’m currently trying to solve the same problem, and so far haven’t found anything that can be done without a loop.

That’s right and I think it would be the only way to keep the batch intact, wouldn’t it?
If you crop each image separately, you could end up with variable spatial sizes, which would disallow creating a batch afterwards.

Let me know, if I’m missing something.

Yup @ptrblck you are missing something :wink:

I’ve figured out how to do this now so I can show you! Here’s how I did the batchwise nonzero thing:

And here’s how I then use that to create cropped zoomed sections from a batch:

As you see, the trick is to use grid_sample to resize them all to the same size after the crop. This is a nice way to do things like RandomResizeCrop on the GPU.

1 Like

Thanks for sharing! :slight_smile:
Did you know what speedup to expect, as I assume it should be way faster than the loop.

Yeah my preprocessing overall is 150x faster on CUDA, and I think that’s nearly all in these two methods.

1 Like

Updated github link to this solution:

1 Like