执行摘要

本文探讨了使用 Dask 数组加载大型图像数据栈的简单工作流程。

特别地，我们从一个包含 TIFF 文件的目录开始，其中的图像如下所示

$ $ ls raw/ | head
ex6-2_CamA_ch1_CAM1_stack0000_560nm_0000000msec_0001291795msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0001_560nm_0043748msec_0001335543msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0002_560nm_0087497msec_0001379292msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0003_560nm_0131245msec_0001423040msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0004_560nm_0174993msec_0001466788msecAbs_000x_000y_000z_0000t.tif

并展示如何使用 dask-image 库将这些文件拼接成大型延迟数组

>>> import dask_image
>>> x = dask_image.imread.imread('raw/*.tif')

或通过编写自己的 Dask 延迟图像读取函数来实现。

Array Chunk Bytes 3.16 GB 316.15 MB Shape (2010, 1024, 768) (201, 1024, 768) Count 30 Tasks 10 Chunks Type uint16 numpy.ndarray 76810242010

将来有一天，我们将能够对这个 Dask 数组执行复杂的计算。

免责声明：我们不会在本文中生成如上所示的渲染图像。这些图像是使用 NVidiaIndeX 创建的，这是一个与此处讨论的内容完全独立的工具链。本文仅介绍图像加载的第一步。

系列概述

在获取大量成像数据的领域中，常见的情况是将较小的采集数据写入许多小文件。这些文件可以平铺更大的空间，从更长的时间段中进行子采样，并且可能包含多个通道。采集技术本身通常是最先进的，并且在可获取的视场范围、分辨率和质量方面不断突破极限。

一旦获取到数据，就会出现许多挑战。通常设计和测试用于处理数据中非常小部分的算法需要扩展以处理整个数据集。一开始可能并不清楚实际哪些方法有效，因此探索仍然是整个过程中的重要组成部分。

从历史上看，这种分析过程涉及大量自定义代码。分析过程通常由一系列可能使用不同语言编写的脚本串联而成，这些脚本将各种中间结果写入磁盘。得益于现代工具的进步，这些过程可以得到显著改进。在本系列博客文章中，我们将概述图像科学家如何利用不同的工具，构建一个高级、友好、连贯、交互式的分析管道。

文章概述

本文特别关注如何从 Python 中并行加载和管理大型图像数据栈。

加载大型图像数据可能是一个复杂且通常独特的问题。不同的团队可能会选择将其存储在磁盘上的许多文件中、使用商业或定制数据库解决方案，或者选择将其存储在云端。由于各种原因，同一团队内的所有数据集可能不会被同等对待。简而言之，这意味着数据加载是一个困难且代价高昂的问题。

尽管数据存储方式多种多样，但团队通常希望将相同的分析管道重新应用于这些数据集。然而，如果数据管道与特定数据加载方式紧密耦合以用于后续分析步骤，那么重用现有管道可能会非常困难，甚至不可能。换句话说，加载和分析步骤之间存在摩擦，这阻碍了实现可重用性的努力。

拥有模块化和通用的数据加载方式可以轻松地以标准方式呈现不同存储方式的数据。此外，拥有将数据呈现给分析管道的标准方式，可以让分析管道专注于其最擅长的部分——分析！总的来说，这应该能够以改善参与管道所有部分的用户体验的方式解耦这些组件。

我们将使用图像数据，这些数据由加州大学伯克利分校高级生物成像中心的Gokul Upadhyayula慷慨提供，并在本文（预印本）中讨论过，尽管此处介绍的工作负载适用于任何类型的成像数据或一般的数组数据。

使用 Dask 加载图像数据

让我们从文章开头处的图像数据再次开始

$ $ ls /path/to/files/raw/ | head
ex6-2_CamA_ch1_CAM1_stack0000_560nm_0000000msec_0001291795msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0001_560nm_0043748msec_0001335543msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0002_560nm_0087497msec_0001379292msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0003_560nm_0131245msec_0001423040msecAbs_000x_000y_000z_0000t.tif
ex6-2_CamA_ch1_CAM1_stack0004_560nm_0174993msec_0001466788msecAbs_000x_000y_000z_0000t.tif

使用 Scikit-Image 加载单个样本图像

要加载单个图像，我们使用 Scikit-Image

>>> import glob
>>> filenames = glob.glob("/path/to/files/raw/*.tif")
>>> len(filenames)
597

>>> import imageio
>>> sample = imageio.imread(filenames[0])
>>> sample.shape
(201, 1024, 768)

每个文件名对应于更大图像的某个 3D 块。我们可以查看这个单个 3D 块的一些 2D 切片，以获取一些背景信息。

import matplotlib.pyplot as plt
import skimage.io
plt.figure(figsize=(10, 10))
skimage.io.imshow(sample[:, :, 0])

plt.figure(figsize=(10, 10))
skimage.io.imshow(sample[:, 0, :])

plt.figure(figsize=(10, 10))
skimage.io.imshow(sample[0, :, :])

研究文件名结构

这些切片仅来自更大聚合图像的一个块。我们在这里的兴趣是将这些片段组合成一个大型图像栈。文件名中常见命名结构。每个文件名可能指示通道、时间步长和空间位置，其中 <i> 是某个数值（可能带单位）。单个文件名可能包含更多或更少的信息，并且其标记方式可能与我们有所不同。

mydata_ch<i>_<j>t_<k>x_<l>y_<m>z.tif

原则上，我们可以使用 NumPy 分配一个巨大的数组，然后迭代加载图像并将其放入巨大的数组中。

full_array = np.empty((..., ..., ..., ..., ...), dtype=sample.dtype)

for fn in filenames
img = imageio.imread(fn)
index = get_location_from_filename(fn) # 我们需要编写这个函数
full_array[index, :, :, :] = img

然而，如果我们的数据很大，我们就不能像这样一次性将所有数据加载到内存中形成一个单一的 NumPy 数组，而是需要更巧妙地高效处理它。一种方法是使用 Dask，它可以轻松处理超出内存的工作负载。

使用 Dask 数组延迟加载图像

现在我们学习如何使用 Dask 数组延迟加载和拼接图像数据。我们将先从简单的示例开始，然后再介绍这个更复杂数据集的完整示例。

我们可以使用 DaskDelayed 延迟 imageio.imread 调用。

import dask
import dask.array as da

lazy_arrays = [dask.delayed(imageio.imread)(fn) for fn in filenames]
lazy_arrays = [da.from_delayed(x, shape=sample.shape, dtype=sample.dtype)
for x in lazy_arrays]

注意：这里我们假设所有图像都具有与上面加载的样本文件相同的形状和 dtype。情况并非总是如此。有关替代方案，请参阅下面“未来工作”部分中关于 dask_image 的说明。

我们尚未将它们拼接在一起。我们有数百个单块 Dask 数组，每个数组都延迟加载磁盘上的一个 3D 数据块。让我们看一个数组。

>>> lazy_arrays[0]

Array Chunk Bytes 316.15 MB 316.15 MB Shape (201, 1024, 768) (201, 1024, 768) Count 2 Tasks 1 Chunks Type uint16 numpy.ndarray 7681024201

这是一个由一个 单个 300MB 数据块组成的延迟 3 维 Dask 数组。该块是通过加载特定的 TIFF 文件创建的。通常 Dask 数组由 许多 块组成。我们可以使用诸如 da.concatenate 和 da.stack 之类的函数将许多这些单块 Dask 数组连接成一个多块 Dask 数组。

这里我们将前十个 Dask 数组沿几个轴连接起来，以便更容易理解其外观。请注意在左侧表格和右侧图像中，随着我们更改 axis= 参数，形状如何变化。

da.concatenate(lazy_arrays[:10], axis=0)

Array Chunk Bytes 3.16 GB 316.15 MB Shape (2010, 1024, 768) (201, 1024, 768) Count 30 Tasks 10 Chunks Type uint16 numpy.ndarray 76810242010

da.concatenate(lazy_arrays[:10], axis=1)

Array Chunk Bytes 3.16 GB 316.15 MB Shape (201, 10240, 768) (201, 1024, 768) Count 30 Tasks 10 Chunks Type uint16 numpy.ndarray 76810240201

da.concatenate(lazy_arrays[:10], axis=2)

Array Chunk Bytes 3.16 GB 316.15 MB Shape (201, 1024, 7680) (201, 1024, 768) Count 30 Tasks 10 Chunks Type uint16 numpy.ndarray 76801024201

或者，如果我们想创建一个新维度，我们将使用 da.stack。在这种情况下，请注意我们已经用完了易于查看的维度，因此您应该更多地关注左侧表格输入中列出的形状，而不是右侧的图片。注意我们将这些 3D 图像堆叠成了 4D 图像。

da.stack(lazy_arrays[:10])

Array Chunk Bytes 3.16 GB 316.15 MB Shape (10, 201, 1024, 768) (1, 201, 1024, 768) Count 30 Tasks 10 Chunks Type uint16 numpy.ndarray 101 7681024201

这些是常见情况，即您希望沿着单个轴将图像拼接在一起。

完整示例

这沿着单个轴组合数据效果很好。但是，如果我们需要跨多个轴组合，则需要执行多个连接步骤。幸运的是，有一个更简单的选项 da.block，如果您给它一个嵌套的 dask 数组列表，它可以一次沿多个轴进行连接。

a = da.block([[laxy_array_00, lazy_array_01],
[lazy_array_10, lazy_array_11]])

现在我们执行以下步骤

• 解析每个文件名，了解它在更大数组中的位置
• 查看每个相关维度中有多少文件
• 分配一个适当大小的 NumPy object-dtype 数组，该数组的每个元素将容纳一个单块 Dask 数组
• 遍历文件名，并将正确的 Dask 数组插入到正确的位置
• 对结果调用 da.block

这段代码有点复杂，但展示了在实际场景中它是如何工作的

# 获取各种维度

fn_comp_sets = dict()
for fn in filenames
for i, comp in enumerate(os.path.splitext(fn)[0].split("_"))
fn_comp_sets.setdefault(i, set())
fn_comp_sets[i].add(comp)
fn_comp_sets = list(map(sorted, fn_comp_sets.values()))

remap_comps = [
dict(map(reversed, enumerate(fn_comp_sets[2]))),
dict(map(reversed, enumerate(fn_comp_sets[4])))
]

# 创建一个空的 object 数组来组织加载 TIFF 的每个块
a = np.empty(tuple(map(len, remap_comps)) + (1, 1, 1), dtype=object)

for fn, x in zip(filenames, lazy_arrays)
channel = int(fn[fn.index("_ch") + 3:].split("_")[0])
stack = int(fn[fn.index("_stack") + 6:].split("_")[0])

a[channel, stack, 0, 0, 0] = x

# 将许多块拼接成一个单一数组
a = da.block(a.tolist())

Array Chunk Bytes 188.74 GB 316.15 MB Shape (3, 199, 201, 1024, 768) (1, 1, 201, 1024, 768) Count 2985 Tasks 597 Chunks Type uint16 numpy.ndarray 1993 7681024201

这是一个 180 GB 的逻辑数组，由大约 600 个块组成，每个块大小为 300MB。现在我们可以使用 DaskArray 对这个数组进行类似 NumPy 的常规计算，但这将在未来的文章中介绍。

>>> # 数组计算会正常工作，并且会在低内存环境下运行
>>> # 但我们将把实际计算留待未来的文章介绍
>>> a.sum().compute()

保存数据

为了简化将来的数据加载，我们使用 to_zarr 方法将其存储为像 Zarr 这样的大型分块数组格式。

a.to_zarr("mydata.zarr")

我们可以将有关图像数据的附加信息添加为属性。这既简化了未来用户的使用（他们可以使用 da.from_zarr 通过一行代码读取完整数据集），也显著提高了性能，因为 Zarr 是一种 分析就绪格式，经过高效编码以用于计算。

Zarr 默认使用 Blosc 库进行压缩。对于科学成像数据，我们可以选择传递压缩选项，这些选项可以在压缩率和速度之间取得良好的平衡，并优化压缩性能。

from numcodecs import Blosc
a.to_zarr("mydata.zarr", compressor=Blosc(cname='zstd', clevel=3, shuffle=Blosc.BITSHUFFLE))

未来工作

上述工作负载通用且直接。它在简单情况下工作良好，并且也可以很好地扩展到更复杂的情况，前提是您愿意围绕自定义逻辑编写一些循环和解析代码。它可以在小型笔记本电脑上运行，也可以在大型 HPC 或云集群上运行。如果您有一个将文件名转换为 NumPy 数组的函数，您可以使用该函数以及 DaskDelayed 和 DaskArray 生成大型延迟 Dask 数组。

Dask Image

然而，如果我们进行一些特化，我们可以让用户使用起来更方便。例如，Dask Image 库提供了一个并行图像读取函数，它在简单情况下自动化了我们上面大部分工作。

>>> import dask_image
>>> x = dask_image.imread.imread('raw/*.tif')

类似地，像 Xarray 这样的库也提供了对其他文件格式（如 GeoTIFF）的读取器。

随着各领域越来越多地进行类似我们上面所做的工作，他们倾向于将常见模式编写成领域特定库，从而提高这些工具的可访问性和用户群。

GPU

如果我们有一些特殊的硬件，比如几个 GPU，我们可以将数据转移到上面，并使用像 CuPy 这样的库进行计算，CuPy 非常接近 NumPy。因此，可以受益于上面列出的相同操作，但同时享受 GPU 带来的额外性能。

import cupy as cp
a_gpu = a.map_blocks(cp.asarray)

计算

最后，在未来的博客文章中，我们计划讨论如何使用常见的图像处理工作负载（如重叠模板函数、分割和反卷积）以及与其他库（如 ITK）集成，对大型 Dask 数组进行计算。