Apache Beam 医疗影像预处理实战：Beam Summit 2025 黑客松一等奖 “PCollectors“ 项目解析-平芜编程栈

大数据
批处理
流处理
数据工程

【免费下载链接】beam

Apache Beam is a unified programming model for Batch and Streaming data processing.

项目地址：https://gitcode.com/gh_mirrors/beam4/beam

Apache Beam 在医疗 AI 场景下最典型的应用之一，是把大规模病理切片（Whole Slide Images, WSI）预处理成适合下游机器学习训练的格式。本文基于 Apache Beam 官方博客 beam-summit-2025-hackathon-pcollectors-blog.md（Beam Summit 2025 黑客松一等奖团队 PCollectors 的参赛记录），完整还原其"CSV 元数据 + .tif 病理切片 → 切块过滤 → EfficientNet 特征提取 → TFRecord 数据集"的可扩展预处理管线，并结合仓库中的 RunInference API、TFRecord I/O 等源码实现，深入解析其底层原理。读完本文，你将掌握如何用 Apache Beam 构建云原生、内存高效的医学影像批量预处理管线，并了解 Beam 在 ML 数据工程侧的完整能力边界。

一、背景：为什么医疗影像预处理需要 Beam

在医疗领域的机器学习中，大规模病理图像数据集的预处理长期是瓶颈所在。病理切片（WSI）是医学影像中体积最大的数据类型之一：单张高分辨率 .tif 切片可达数 GB 级别，传统 Python 图像工具（如 PIL）在内存压力下会直接失效，更遑论一次性处理成千上万张此类图像。对下游 ML 建模任务而言，这构成了标准工具无法跨越的障碍。

作者团队此前从事过目标检测领域的图像处理工作，深知"如何正确预处理并结构化图像数据"对下游任务的决定性影响——而这些挑战在医疗场景中更为严峻，因此天然适合 Apache Beam 这类可扩展数据处理框架来解决。在 Beam Summit 2025 Hackathon 中，他们以 "PCollectors" 队名参赛，目标是利用 Beam 处理大规模图像数据并转换为适合下游 ML 任务的数据格式，最终凭借该方案获得一等奖。

二、项目总览：可扩展的 WSI 预处理管线

2.1 管线目标

管线核心目标：处理患者数据（CSV）与 WSI 病理切片，提取图像 embedding，合并元数据，最终输出 TFRecord 格式的完整数据集，供大规模 ML 训练直接使用。

2.2 完整处理流程

管线依次处理以下环节：

读取患者元数据（CSV）
读取 WSI 文件（.tif）
将整张切片切分为 "tiles"（图像瓦片）
基于背景阈值过滤图像瓦片（剔除空白/背景区域）
使用 EfficientNet 为每位患者生成 max 与 avg 两种 embedding
将元数据 + embedding 合并写入 TFRecords

全部流程构建在 Apache Beam + Dataflow 之上，是一条可扩展、内存高效、云原生的流水线。

2.3 数据集结构（Mayo Clinic STRIP AI）

数据来源：Mayo Clinic STRIP AI Dataset（Kaggle）
元数据：每行 ={ image_id, center_id, patient_id, image_num, label }
每位患者对应多张切片图像
标签（label）仅存在于患者层级
图像：高分辨率 .tif 病理切片

关键设计点：由于标签只在患者层级，而 embedding 是按切片/瓦片计算的，因此管线的关键聚合操作是按patient_id分组，将同一患者的多个切片 embedding 聚合为 max/avg 后，再与患者级标签合并写盘。

2.4 技术栈

组件	职责
Apache Beam	编排引擎，负责整个数据处理 DAG
Google Cloud Dataflow	可扩展 runner，负责分布式执行与自动扩缩容
Google Cloud Storage	输入 TIFF 存放与输出 TFRecords 存储
TensorFlow	EfficientNet 特征提取与 TFRecord 序列化

三、核心技术拆解

3.1 高效图像加载：pyvips + openslide

管线集成 pyvips 与 openslide 实现高效图像加载。这两个库专为超大图像设计：openslide 面向病理切片的经典格式读取，支持对巨型 WSI 的随机区域访问；pyvips 则以惰性求值（lazy evaluation）和流式像素处理著称，可以在不把整张图像载入内存的前提下按需切取瓦片。对动辄数 GB 的 .tif 切片，这种"按需读取 + 区域裁剪"的能力正是管线内存高效的前提。

从实现角度，切块（tiling）逻辑对应 Beam 中的一个 ParDo/DoFn：输入一张 WSI 路径，输出一组经过裁剪与过滤的瓦片 PCollection，num_shards式的并行度由 Dataflow 按切片文件自动切分实现。

3.2 背景过滤：基于阈值的瓦片筛选

切块之后，管线基于背景阈值过滤瓦片。病理切片的显著特征是大部分区域是空白/背景（玻片、封片胶等），直接送入模型既浪费算力又引入噪声。通过计算瓦片的背景像素比例（例如基于亮度/强度阈值判断），保留组织区域瓦片、丢弃背景瓦片，既压缩了后续推理的数据量，又提升了 embedding 质量。

3.3 Beam RunInference API + TensorFlow

特征提取环节使用 Beam 的RunInference API。仓库中该 API 定义于 sdks/python/apache_beam/ml/inference/base.py（RunInference类定义于第 1373 行），是 Beam ML 的核心 PTransform：

输入：PCollection 的 examples（或 features）
输出：PCollection 的 PredictionResult，其中同时携带输入 example 与输出 inference（见 base.py 中PredictionResultNamedTuple 定义，含example、inference、model_id三字段）

RunInference 的底层实现做了三件对 ML 预处理至关重要的事：

自动批处理：内部通过beam.BatchElements对输入元素做分批（见 base.py），把推理请求聚合成 batch，最大化 GPU/CPU 利用率；
模型共享：_RunInferenceDoFn通过shared.Shared()句柄在多个线程/多个 bundle 之间复用已加载的模型实例（见 base.py），避免每个元素都重新加载权重；
指标收集：内置推理延迟、批大小、模型加载耗时等监控指标（metrics_namespace参数可自定义命名空间）。

针对 TensorFlow 模型，仓库提供专门的 ModelHandler 实现：

TFModelHandlerNumpy/TFModelHandlerTensor（定义于 sdks/python/apache_beam/ml/inference/tensorflow_inference.py）
支持两种加载方式：SAVED_MODEL（TF2 SavedModel 格式）与SAVED_WEIGHTS（仅权重 +create_model_fn重建结构），见 tensorflow_inference.py 的ModelType枚举
支持min_batch_size、max_batch_size、max_batch_duration_secs、model_copies、max_batch_weight等批处理与资源控制参数（tensorflow_inference.py）

典型用法示意：

from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerNumpy model_handler = TFModelHandlerNumpy( model_uri="gs://bucket/efficientnet_savedmodel", model_type=ModelType.SAVED_MODEL, max_batch_size=16, ) # tiles 为 PCollection[np.ndarray] results = tiles | beam.RunInference(model_handler)

需要说明：仓库源码中 RunInference 面向的是通用框架推理（TensorFlow、PyTorch、HuggingFace、Sklearn、ONNX 等均有对应 ModelHandler），EfficientNet 的 embedding 输出正是通过PredictionResult.inference字段在管线中流转。

3.4 患者级聚合：max/avg embedding

由于标签在患者层级，管线的核心聚合步骤是：

# embeddings: PCollection[(patient_id, tile_embedding)] patient_embeddings = ( embeddings | beam.GroupByKey() # 按 patient_id 聚合该患者所有切片的 embedding | beam.ParDo(ComputeMaxAvgEmbedding()) # 计算 max 与 avg 两种聚合 )

max-pooling 保留了每个患者最"显著"的组织特征，avg-pooling 则给出整体统计特征。两者都写入最终 TFRecord，让下游模型可以同时利用两种粒度的表征。

3.5 写出 TFRecord：WriteToTFRecord

最终结果通过 Beam 的 TFRecord I/O 落盘。仓库实现位于 sdks/python/apache_beam/io/tfrecordio.py，核心类是WriteToTFRecord（第 320 行起），其关键参数：

参数	含义
`file_path_prefix`	输出文件路径前缀，实际文件为"前缀 + 分片标识 + 后缀"
`coder`	每条记录的编码器，默认`BytesCoder()`，可换自定义 coder
`file_name_suffix`	输出文件扩展名，如`.tfrecord`
`num_shards`	输出分片数，未指定时由 runner 决定（Dataflow 下可充分发挥并行写吞吐）
`shard_name_template`	分片命名模板，默认`-SSSSS-of-NNNNN`（有界），流式默认`-W-SSSSS-of-NNNNN`
`compression_type`	压缩类型，默认`AUTO`（按扩展名自动检测）

实现层面，WriteToTFRecord将数据委托给_TFRecordSink，并保证 CRC32C 校验（依赖 snappy / google-crc32c / crcmod 三者之一，见 tfrecordio.py），确保 TFRecord 数据完整性。TensorFlow 侧则可用tf.data.TFRecordDataset直接消费这些文件用于训练。

典型写法：

result | beam.io.WriteToTFRecord( "gs://bucket/output/patient_embeddings", file_name_suffix=".tfrecord", num_shards=200, )

四、运行架构：Beam + Dataflow 的云原生执行

管线采用 Google Cloud Dataflow 作为 runner 运行。Dataflow 的自动扩缩容（autoscaling）机制使 Worker 数量随切片文件与瓦片数量动态伸缩，这正是"处理数千张数 GB 切片"的规模化保障。整体执行形态：

输入：GCS 上的 CSV 元数据 + .tif 切片（通过 Beam 的文件 I/O / 自定义 source 读取）
并行：Dataflow 将切片文件拆分为 bundle，分布式并行执行切块、过滤、推理
中间态：PCollection 流转，Beam 自动处理 shuffle（GroupByKey 的患者级聚合）
输出：GCS 上的 TFRecord 分片文件，供 TensorFlow 训练直接读取

仓库中 sdks/python/apache_beam/io/gcp 目录包含 GCS/Dataflow 相关的 I/O 组件，可作为输入输出接入的参考实现。

五、经验总结与后续演进

5.1 实战经验

Apache Beam 非常适合并行、云原生的 ML 预处理：切块、过滤、推理、聚合、写出全链路可以统一在一条 DAG 中声明式编排，天然获得分布式能力；
Dataflow 是处理大规模数据（如医学图像）的首选 runner：自动扩缩容 + 托管服务，让团队无需关心集群运维；
内存高效的关键在选对库：pyvips/openslide 的惰性读取避免了整图载入内存，与 Beam 的分布式切分配合，单机内存压力被彻底化解。

5.2 后续演进方向

原博客中提出的扩展方向，恰好对应 Beam 生态中已有的能力点：

语义化切块：从固定尺寸瓦片转向基于组织区域的图像分割，只聚焦相关区域生成更有意义的 patch，可提升模型性能；
视频数据适配：将帧视为带时间索引的图像切片，支撑运动分析、病程跟踪等时序建模任务；
多模态医疗 AI：结合组织病理图像与基因组数据、临床记录或影像学扫描，构建更全面的上下文感知模型。

这三个方向都可在现有 Beam 管线骨架内通过替换/新增 PTransform 实现，体现了该架构的可扩展性。

六、结论

PCollectors 团队在 Beam Summit 2025 黑客松中以这套"CSV + WSI → 切块过滤 → EfficientNet embedding → TFRecord"管线获得一等奖，验证了 Apache Beam 在大规模医疗影像预处理中的工程价值：通过 RunInference API 完成分布式模型推理、通过 GroupByKey 完成患者级聚合、通过 TFRecord I/O 完成训练数据落盘，全链路云原生、内存高效、按需扩缩容。对于任何需要将海量非结构化图像数据转化为 ML 就绪数据集的工程团队，本文的管线骨架与源码级原理都提供了可直接借鉴的范式。

相关仓库资源

关联博客原文：beam-summit-2025-hackathon-pcollectors-blog.md
RunInference 核心实现：sdks/python/apache_beam/ml/inference/base.py
TensorFlow ModelHandler：sdks/python/apache_beam/ml/inference/tensorflow_inference.py
TFRecord 读写实现：sdks/python/apache_beam/io/tfrecordio.py
推理示例管线：sdks/python/apache_beam/examples/inference/README.md
Beam ML 推理模块入口：sdks/python/apache_beam/ml/inference/init.py