OTLesMix：基于Wasserstein重心的医学影像病灶数据增强

医学影像 AI 里病灶标注是最贵的那一环。一个肺结节、一个肝肿瘤从影像上被医生发现到专家逐层核对并给出掩码成本远高于自然图像上的框选。很多团队手里只有几十个病例、上百个病灶却要训练一个能在真实场景里不漏检、不误报的模型。数据增强在这种情况下不是“锦上添花”而是决定项目能不能启动的关键环节。这一篇想深入聊一种更接近病灶生成本质的数据增强思路OTLesMix。先说我的判断。OTLesMix 的目标不是替代分类、检测或分割模型而是解决训练数据里真实病灶数量不足、形态单一的问题。它的核心贡献在于把“少数真实病灶”变成“大量合理的新病灶”而不是简单复制粘贴。理解它本质上是在理解两件事Wasserstein 重心到底在求什么最优传输映射能如何把一张病灶“搬”到另一张图上。如果只读论文标题很多人会觉得这是一个偏理论的研究。但它实际解决的是非常工程化的问题合成病灶的形状多样性和位置合理性。这篇文章会从医学影像小样本场景的痛点切入把 Wasserstein 距离、重心、最优传输映射三个概念讲清楚再给出一个可运行的参考实现最后列出实际项目中容易踩的坑。无论你是做分类、分割还是检测只要数据里病灶样本稀缺这篇文章都有参考价值。1. 这篇文章真正要解决的问题先想一个场景如果手上只有 50 个带肺结节标注的 CT 切面你打算怎么把模型练好常规答案无非是旋转、翻转、缩放、加噪、弹性形变。这些操作对整张图像做全局变换病灶本身没有本质变化。继续用这些方法模型很快会把“结节长什么样”记住但很难学会“结节还可以长成什么样”。当测试集里出现一个形态和训练集差异较大的病灶时模型大概率会漏检因为数据增强没有给模型提供足够的形态先验。第二种思路是 Mixup、CutMix 这类图像级混合。把两张图的像素按权重叠加或者把 A 图的一个区域贴到 B 图上。它们能产生新样本但对医学影像有个明显问题邻域上下文会被破坏。医学影像不是自然照片组织之间有解剖连续性随便贴一块上去模型学到的是“异常强度块出现在任意位置”而不是“符合局部组织环境的病灶”。如果只是在分类任务上追求涨点这类方法还能用一旦涉及分割或检测标签和像素的不对齐会让模型训练变得很不稳定。第三种思路是从病灶本身出发把真实病灶“搬”到健康图像上。这就是 lesion mix、Copy-Paste、ROI transplant 一类的做法。它比全局变换更进一步但仍有两个核心痛点。第一个痛点是形状多样性不足。复制粘贴使用的病灶形状全部来自原始标注。100 个病灶就是 100 种固定形状再多也就 100 种。模型反复看见同一个病灶即便每次放置位置不同对形状的过拟合依然存在。第二个痛点是位置不合理。很多实现把病灶随机放到一个坐标完全不考虑这个位置有没有对应的解剖组织。一个肾肿瘤被贴到胸腔里在像素数值上可行在医学上却毫无意义甚至会给训练集引入错误标签。OTLesMix 试图同时处理这两个问题。从名字拆开看LesMix 说明是病灶级混合Wasserstein Barycenter 负责生成介于多个真实病灶之间的新形状Optimal Transport Map 负责把病灶的强度分布映射到目标位置让合成区域在边界上更连续、在内容上更接近真实病灶。换句话说它把病灶混合从“粘贴像素”升级成了“搬运分布”。所以这篇文章适合三类人正在处理小样本医学影像分类和分割任务的工程师做数据增强方向的研究者想看看最优传输如何被用到医学图像以及对 CutMix 系列增强比较熟悉、想进一步拓展方法论的人。如果你只想找一个现成 pip 包立刻开始跑那可能还需要等一等但原理部分对任何做医学影像预处理的人来说都有参考价值。2. 从 Mixup、CutMix 到 OTLesMix数据增强的演进路线要把 OTLesMix 说清楚得先把它放在数据增强的发展脉络里观察。这个脉络可以从两个维度拆操作粒度和空间约束。操作粒度上最早是像素级。加噪声、改变对比度、调整亮度都是为了增强模型对成像参数变化的鲁棒性。接着是全局几何变换旋转、翻转、缩放、弹性形变模型开始学习形状和位置的不变性。这些方法实现简单但在样本极度稀缺时它们只是在同一批真实样本附近反复扰动无法创造真正新的结构。然后是图像级混合。Mixup 在整张图上做线性插值标签也做线性插值。CutMix 从一张图中截取一块贴到另一张图标签按面积比例混合。这类方法在自然图像分类任务上表现不错但它们生成的图像常常“看起来很奇怪”。原因是线性插值和直接拼接都忽略了图像内容之间的语义边界对医学影像来说这种语义混乱尤其致命。再往下就是对象级混合。先检测出目标对象再把对象区域粘贴到其他图上。医学影像里的 Copy-Paste、ROI transplant、Lesion Mix 都属于这一类。它们保留了病灶内部纹理也保留了标签边界但正如前面所说形状来源固定放置位置通常随机。OTLesMix 在这个链条里的位置可以理解为“对象级混合 分布约束”。它不直接粘贴原始病灶块而是先在病灶这一层做分布层面的合成再通过最优传输把合成结果贴合到目标背景上。相比传统 Copy-Paste变化不在于多了一个模块而在于把“复制像素”变成了“复制分布”。这个转变带来两个直接收益一是新形状可以指数级增多因为组合方式来自多个病灶的加权融合二是位置不再完全随机而是通过传输映射与目标区域建立联系。空间约束这个维度也很关键。普通图像分类允许对象出现在任何位置所以随机位置增强是合理的。但医学影像不同器官的位置相对固定病灶的出现位置与该器官、血管、周边组织有强相关性。一个建模得当的增强方法应当让合成病灶落在符合解剖语义的位置而不是一个均匀随机的坐标点。OTLesMix 对位置的处理从方法设计上更倾向于“利用现有病灶与周围组织的关系”。更稳妥的判断是它不是用一个手工规则去限制位置而是在传输映射过程中让位置分布参与计算通过重心和映射关系决定病灶在不同目标上的落位方式。这样生成的样本位置多样性来自真实病灶的空间分布统计而不是拍脑袋的随机数。理解这条演进路线之后再看 Wasserstein 相关概念动机就清楚了。3. 核心概念Wasserstein 距离、Wasserstein 重心与最优传输映射这里我先做一个概念澄清很多人把 Wasserstein 距离当成 WGAN 的附属品实际上它是最优传输理论的一小部分而最优传输是一个成熟的分析工具用途远不止生成对抗网络。3.1 为什么是 Wasserstein 而不是 L2先看 L2 距离在病灶比较上的尴尬。两个形状几乎相同的病灶如果其中一个整体平移了一个像素L2 距离会很大因为每个像素位置上都有差异。这不符合人类直觉形状基本一样只是错位了一点点。Wasserstein 距离的直觉是“搬土”。把 A 分布的土搬到 B 分布搬了多少体积、走了多远距离加在一起就是搬土成本。错位一个像素的两个相似形状搬土距离很短Wasserstein 距离也就很小。这个性质让它在衡量分布相似性时天然对轻微形变更友好。在病灶合成场景里“病灶强度分布”和“病灶空间位置分布”都可以当作概率分布来处理Wasserstein 给了一种度量差异的方式而同样的框架也能用来生成新分布。3.2 Wasserstein 重心多个病灶的共同“平均形状”平均值的概念大家都熟。几个数字加一起除以个数就是算术平均值。但一群形状的平均是什么逐像素平均当然可以得到一张图但结果会非常模糊边界被抹平强度被稀释这样的“平均病灶”反而失去真实性直接拿来训练会让模型学到模糊的边界。Wasserstein 重心解决的是“在分布空间里做平均”的问题。给定多个概率分布找到一个分布使得到这些分布的最优传输成本之和最小。这个结果不是像素的简单平均而是在分布意义上最折中的那个形状。它的价值在于新的病灶形状不完全来自任何一个真实样本而是多个真实样本在传输意义上的融合。因为一个重心对应一个合成形状通过随机选取不同真实病灶子集、调整每个参与样本的权重OTLesMix 能在一个由真实样本张成的分布空间中采样出大量新病灶。形状多样性由此而来而且新形状在结构上与真实病灶保持了较高的解剖合理性。日常项目里这种加权融合的思想也可以用在别的数据模态上并不局限于医学影像。3.3 最优传输映射把一张病灶“搬”成另一张病灶Wasserstein 距离负责“度量差异”最优传输映射负责“找到怎么搬”。给定源分布和目标分布最优传输映射回答源分布的每个单位质量应该被搬运到哪里能使总搬运成本最小。在离散图像上这会变成一个线性规划问题。Kantorovich 形式允许质量切分很适合图像这类连续密度。实际实现中常用熵正则化的 Sinkhorn 迭代来逼近最优传输计划计算效率高也容易结合 GPU。POT 库提供了相当完整的实现这也是后面代码部分能快速跑通的原因。映射在病灶合成中的意义有两层。一是在形状层面把源病灶的掩码或概率图映射到目标区域的网格上使得贴上去的病灶形状自然过渡。二是在强度层面把源病灶的灰度直方图映射到背景区域的强度分布上避免出现亮度跳跃。这两层映射共同作用合成的病灶才不像是“抠图贴图”。一旦理解了这三个概念OTLesMix 的整体框架可以描述成一个式子选取一组真实病灶求出它们的 Wasserstein 重心得到新形状再通过最优传输映射把这个新形状渲染到一个目标背景的候选位置。4. OTLesMix 的方法拆解再次强调以下拆解基于论文标题、方法命名和最优传输在图像任务中的常规用法不等同于论文官方实现细节。如果论文已经开源代码请以官方实现为准。这里的目的是帮读者建立可迁移的理解框架。整体流程可以分成五个阶段病灶区域提取、分布建模与重心计算、传输映射生成、样本融合、标签处理。4.1 病灶区域提取与掩码表示第一步是从训练集中已有的病灶标注中提取病灶块。对于二维切片病灶是掩码上的连通区域对于三维 CT 或 MRI病灶是一组体素。实际操作中通常以病灶中心为基准裁剪固定尺寸的小块同时保留病灶掩码和距离变换图。距离变换能描述“到病灶边界的距离”是后续形状重心的常用输入。提取时要注意统一尺寸。不同病例中病灶大小差异很大直接随机裁剪会导致形状重心计算不稳定。更常见的做法是 Resize 到统一尺寸但不要丢掉原始空间分辨率信息后续映射回原图时要乘回缩放系数。如果病灶贴近图像边缘裁剪区域可能超出边界需要做 padding 或采集镜像像素。这一阶段的输出质量直接决定后面的合成质量值得多花时间做可视化检查。4.2 病灶强度与空间位置的分布建模病灶块本质上是强度场加掩码。在最优传输框架里可以把病灶看成定义在像素网格上的分布。空间位置分布用掩码或概率图表示强度分布用灰度直方图表示。这两个分布是 OTLesMix 做重心和传输的对象。把强度直方图当作分布传输映射就能控制合成病灶的灰度统计把掩码当作分布重心就能控制形状。两者独立计算再结合比直接对原始像素做传输更稳定也更可控。这也解释了很多医学影像合成方法采用“先形状、后纹理”的设计。实际实现里强度分布可以只统计病灶内部像素也可以把周边一圈背景也纳入统计后者会让边界过渡更自然。4.3 Wasserstein 重心生成新形状给定若干病灶掩码计算其 Wasserstein 重心得到一张概率图。概率图中每个像素的值表示“这个像素属于病灶的可能性”。这个概率图本身可以当作软掩码使用也可以阈值化后生成硬掩码。合成时不一定每次都使用全部病灶。更合理的做法是随机选 2 到 4 个病灶并为每个病灶分配一个权重。权重越大重心形状越接近该病灶。通过控制权重分布和参与样本能生成连续过渡的病灶形状。这一步回答的是“病灶可以长成什么样”。4.4 最优传输映射生成位置与强度新形状确定后面临的是放到哪里的问题。一般会先在目标图像上选出病灶候选区域集合。候选区域可以来自训练样本的位置先验也可以是对应器官的解剖区域。接着通过最优传输映射把新形状的掩码概率分配给候选位置。从工程角度看这一步等价于在候选位置中寻找一个与生成形状在传输成本上最匹配的位置分布。随后是强度合成。把源病灶的强度直方图通过传输映射匹配到目标区域的局部强度分布或者在生成形状边缘用 Sinkhorn 平滑地混合边界。这样合成的病灶在边界上会有过渡而不是一条生硬的分割线。顺序上建议先做位置映射再做强度匹配因为强度匹配依赖目标区域的具体像素统计位置没定之前无从谈起。4.5 标签处理与损失函数病灶级增强必然会改变标签。如果任务是分割新样本的掩码就是重心生成的软掩码。如果任务是分类标签可以用软标签或按病灶覆盖面积加权。如果任务是检测需要记录新病灶的边界框。还有一个容易忽略的环节合成的病灶不应该让模型把背景误判为病灶。因此可以在增强后做一个置信度检查或人工抽检确保新样本的标签没有明显错误。具体做法包括跨专家抽检、与原始病灶形态相似度对比、以及在验证集上单独观察合成样本的错误分布。标签处理策略往往决定了增强方法能否真正落地到业务模型里。5. 基于论文思路的 Python 参考实现本节给出一个基于论文思路整理的最小 Python 流程目的是跑通概念不是复现论文的全部算法。使用的主要库是 NumPy 和 Python Optimal TransportPOT。如果只是想理解原理这个流程已经足够如果要投入到真实项目还需要根据数据规模做内存和性能优化。5.1 环境准备建议环境如下pip install numpy pot opencv-python-headless如果处理三维医学影像还需要 SimpleITK 或 nibabel具体版本请以项目实际情况为准。以下代码不依赖 GPUCPU 即可运行非常适合在本地先验证概念。POT 是计算最优传输的常用 Python 库提供了 Sinkhorn、emd、重心计算等接口能省去大量底层实现工作。5.2 病灶块提取下面函数从二维图像中提取以病灶中心为中心的正方形小块。实际项目中如果病灶是三维体素可以按切片先做二维预筛选再扩展到三维块。# 文件路径lesion_utils.py import numpy as np def extract_lesion_patches(image, mask, patch_size64): 从二维图像中提取包含病灶的正方形小块。 image: (H, W) 强度图 mask: (H, W) 二值掩码 patch_size: 裁剪边长 ys, xs np.where(mask 0) if len(ys) 0: return [] cy, cx int(