对XArray中的图像堆栈中的每个图像应用一个函数

Question

我有一堆图片，比如

import dask.array as da
import xarray as xr
import numpy as np
img_stack =  xr.DataArray(
da.from_array(np.random.random((5,100,100))), 
dims=('t','x','y'), 
coords={'t':range(5),'x':range(100),'y':range(100)}
)

和函数:

def filter(patch, N=3):
print(patch.shape)
return signal.convolve2d(
patch, (1 / (N * N)) * np.ones((N, N)), 
boundary="symm", mode="same")

我想对堆栈中的每个图像应用这个函数(沿着t维)。

我认为最好的方法是使用apply_ufunc，当dask未启用时，我可以这样做:

xr.apply_ufunc(
filter, img_stack.load(), 
input_core_dims=[['y','x']], 
output_core_dims=[['y','x']], vectorize=True)

这工作得很好，我得到了想要的结果。过滤器接收到一个(100,100)数组

但是如果我允许:

xr.apply_ufunc(
filter, 
img_stack, input_core_dims=[['y','x']], 
output_core_dims=[['y','x']], vectorize=True, dask='allowed')

我最终得到一个ValueError: convolve2d inputs must both be 2D arrays。

注意:在这种情况下，我没有load()数据集。如果您已经通过执行前一个代码块加载了数据集，那么通过运行第一个代码块重新初始化img_stack变量可能会更好。

发送给filter函数的数组形状为(5,100,100)，而不是前一种情况下的(100,100)，因此我遇到了一个错误(因为该方法期望m X n数组)。不知道为什么会发生这种情况，或者我可以做些什么来将filter函数应用于堆栈中的5个图像中的每一个。

我想知道最好的方法是什么，如果我正确地应用apply_ufunc。

Answer 1

问题是您指定了dask="allowed"，它告诉xarray您的函数可以本地处理dask数组-在您的情况下它不能。

文档说明:

dask ({"forbidden"， "allowed"， " paralleled"}， default: "forbidden")-如何处理应用到对象中包含惰性数据的形式任务数组:

' forbidden '(默认):如果遇到任务数组，则引发错误。

' allowed ':将任务数组直接传递给函数。如果函数本身支持任务数组。

' parallelized ':自动并行化任何输入的函数是通过使用dask.array.apply_gufunc创建的任务数组。多个输出支持参数。只有在func不使用时才使用此选项本机支持磁盘数组(例如，将它们转换为numpy数组)。

所以你需要设置dask="parallelized"，如果你这样做，你的例子将工作。

请注意，使用vectorize=True工作得很好，但它不是性能最好的解决方案-特别是当您的数据集变大时。如果您关心性能，请考虑使用numab.guvectorize将函数应用于3d数组的每个1d元素，就像这个冗长的xarray示例一样。