In Numpy, given a stack of large images A of size(N,hl,wl), and coordinates x of size(N) and y of size(N) I want to get smaller images of size (N,16,16)
In a for loop it would look like this:
B=numpy.zeros((N,16,16)) for i in range(0,N): B[i,:,:]=A[i,y[i]:y[i]+16,x[i]:x[i]+16]
But can I do this just with indexing?
Bonus question: Will this indexing also work in pytorch? If not how can I implement this there?
Advertisement
Answer
Pretty simple really with view_as_windows from scikit-image, to get those sliding windowed views as a 6D array with the fourth axis being singleton. Then, use advanced-indexing to select the ones we want based off the y and x indices for indexing into the second and third axes of the windowed array to get our B.
Hence, the implementation would be –
from skimage.util.shape import view_as_windows BSZ = 16, 16 # Blocksize A6D = view_as_windows(A,(1,BSZ[0],BSZ[1])) B_out = A6D[np.arange(N),y,x,0]
Explanation
To explain to other readers on what’s really going on with the problem, here’s a sample run on a smaller dataset and with a blocksize of (2,2) –
1) Input array (3D) :
In [78]: A
Out[78]:
array([[[ 5, 5, 3, 5, 3, 8],
[ 5, *2, 6, 2, 2, 4],
[ 4, 3, 4, 9, 3, 8],
[ 6, 3, 3, 10, 4, 5],
[10, 2, 5, 7, 6, 7],
[ 5, 4, 2, 5, 2, 10]],
[[ 4, 9, 8, 4, 9, 8],
[ 7, 10, 8, 2, 10, 9],
[10, *9, 3, 2, 4, 7],
[ 5, 10, 8, 3, 5, 4],
[ 6, 8, 2, 4, 10, 4],
[ 2, 8, 6, 2, 7, 5]],
[[ *4, 8, 7, 2, 9, 9],
[ 2, 10, 2, 3, 8, 8],
[10, 7, 5, 8, 2, 10],
[ 7, 4, 10, 9, 6, 9],
[ 3, 4, 9, 9, 10, 3],
[ 6, 4, 10, 2, 6, 3]]])
2) y and x indices to index into the second and third axes :
In [79]: y Out[79]: array([1, 2, 0]) In [80]: x Out[80]: array([1, 1, 0])
3) Finally the desired output, which is a block each from each of the 2D slice along the first axis and whose starting point (top left corner point) is (y,x) on that 2D slice. Refer to the asterisks in A for those –
In [81]: B
Out[81]:
array([[[ 2, 6],
[ 3, 4]],
[[ 9, 3],
[10, 8]],
[[ 4, 8],
[ 2, 10]]])