提取 python 中的 3D 图像补丁时出现大的分配错误

Question

I'm trying to extract small 3D patches (example patch size 20x20x4) from a 3D Image of size 250x250x250 with stride 1 for every axis. I'll be extracting all possible patches as I'll be running a function on each patch and returning the result in the form of a 3D image with the result of the current patch assigned to the center voxel of the patch. For extracting the patches I'll be using the code below:

import numpy as np
from numpy.lib import stride_tricks

def cutup(data, blck, strd):
    sh = np.array(data.shape)
    blck = np.asanyarray(blck)
    strd = np.asanyarray(strd)
    nbl = (sh - blck) // strd + 1
    strides = np.r_[data.strides * strd, data.strides]
    dims = np.r_[nbl, blck]
    data6 = stride_tricks.as_strided(data, strides=strides, shape=dims)
    return data6.reshape(-1, *blck)

#demo
x = np.zeros((250,250,250), int)
y = cutup(x, (20, 20, 4), (1, 1, 1))

I'm running this on google colab which has around 12gigs of ram. Since the result is large number of patches, I'm getting a large alloc error and then the kernel restarts. I think splitting the image in to parts would work, but If I do so how should I write the code in order for it to consider the neighbouring voxels? Is there a smart way to do this?

Answer 1

Don't reshape the newly strided array/view before returning.

def cutup(data, blck, strd):
    sh = np.array(data.shape)
    blck = np.asanyarray(blck)
    strd = np.asanyarray(strd)
    nbl = (sh - blck) // strd + 1
    strides = np.r_[data.strides * strd, data.strides]
    dims = np.r_[nbl, blck]
    data6 = stride_tricks.as_strided(data, strides=strides, shape=dims)
    return data6

Then iteratate over the patches.

p = np.zeros((250,250,250), int)
q = cutup(p, (20, 20, 4), (1, 1, 1))
print(f'windowed shape : {q.shape}')
print()
for i,x in enumerate(q):
    print(f'x.shape:{x.shape}')
    for j,y in enumerate(x):
        print(f'\ty.shape:{y.shape}')
        for k,z in enumerate(y):
            print(f'\t\tz.shape:{z.shape}')
            if k==5: break
        break
    break
>>>
windowed shape : (231, 231, 247, 20, 20, 4)

x.shape:(231, 247, 20, 20, 4)
        y.shape:(247, 20, 20, 4)
                z.shape:(20, 20, 4)
                z.shape:(20, 20, 4)
                z.shape:(20, 20, 4)
                z.shape:(20, 20, 4)
                z.shape:(20, 20, 4)
                z.shape:(20, 20, 4)

Your example will produce an array (or a view of the array) with a shape of (231,231, 247, 20, 20, 4) or thirteen million+ 3-d patches.

That will solve your memory allocation problem.

---

when I try to reshape it to (231,231,247,-1). I get large alloc error

If your operation requires the last three dimensions to be flattened, do that in your iteration.

for i,x in enumerate(q):
    for j,y in enumerate(x):
        for k,z in enumerate(y):
            z = z.reshape(-1)
            print(f'\t\tz.shape:{z.shape}')
            if k==5: break
        break
    break

Looks like you can do that reshape in the outermost loop - at least for a zeros array.

for i,x in enumerate(q):
    zero,one,*last = x.shape
    x = x.reshape(zero,one,-1)
    print(f'x.shape:{x.shape}')
    for j,y in enumerate(x):
        print(f'\ty.shape:{y.shape}')
        for k,z in enumerate(y):
            print(f'\t\tz.shape:{z.shape}')
            break
        break
    break
>>>
x.shape:(231, 247, 1600)
        y.shape:(247, 1600)
                z.shape:(1600,)

---

Is there a smart way to do this?

If you can figure out how to vectorize your operation so that you only need to iterate over the first dimension or the first and second dimensions you can speed up your processing. That should be a separate question if you encounter problems.