目录
用于 3D 超声标准平面定位的基于切线公式和解剖感知的智能体[MICCAI 2022] | 核心:强化学习(空间解剖奖励) 提升:重构动作空间、内容感知回归辅助任务(状态内容相似度预测(SCSP))、模仿学习的初始化(强化学习的预训练) 其他:CNN、特征点热图、SP 定位 |
摘要
1.将三维超声中的标准切面定位定义为强化学习中一个基于切线点的问题,以重组动作空间,显著减少搜索空间。
2.设计了一种辅助任务学习策略,以增强模型对平面搜索中非标准切面和标准切面之间的细微差异的识别能力。
3.提出了一种空间解剖学奖励,通过同时利用空间和解剖学信息来有效地指导学习轨迹。
探讨了该方法在子宫和胎儿大脑数据集上定位四个标准切面的有效性。
实验结果表明,该方法具有较高的定位精度和鲁棒性性能。
简介
三维空间定位标准平面:
基于监督学习方法:大多采用分类平面图像和回归平面参数的方式来得到标准平面,但三维搜索空间巨大,分类和回归没有充分利用环境的信息,只是学习了单一的映射关系,这种学习方式不够有效。同时网络难以有效学习到高维到低维图像的映射。
基于强化学习方法:临床中,医生手动定位标准切面时会根据经验对当前视图平面进行判断进而调整探头的角度和方向,以靠近目标平面。而强化学习的一个搜索过程,在该任务中就体现了天然的优势,智能体可以通过与环境进行交互,在交互的过程中,智能体可以获得当前时刻的状态,根据当前时刻的状态输出动作,并且获得奖励,进而到达下一个状态,整个过程可以整个过程的目标可以被建模成通过选择合适的动作来最大化奖励的过程,这种方式与医生手动定位标准切面的方式是非常近似的。
方法
agent智能体(神经网络)与环境(3D US Volume)交互,学习具有最大累积奖励的最优SP搜索策略。
此外,我们还为RL框架配备了一个辅助任务来预测 当前状态 和 目标状态 的相似性,从而提高了模型的识别能力。利用仿真学习模块初始化agent,加快优化速度。
图1 所提出的标准切面本地化框架的概述。SLR:空间位置奖励;ASR:解剖结构奖励。
法线式平面方程
xcosα+ycosβ+zcosγ=p [3] ,其中cosα、cosβ、cosγ是平面法矢量的方向余弦,p为原点到平面的距离。
通过调整方向余弦(α、β、γ)和到原点(d).的距离来控制平面运动
三维空间中的任何视平面都可以唯一地定义为其以原点为半径的球体上的切点(tx,ty,tz)
通过平移相切点的坐标(tx、ty、tz)来修改平面运动。
在不对齐的超声环境中,也能实现精确的SP定位
SP定位可以重新表述为切点搜索任务,其中动作空间只包含切点坐标的平移。
所提出的公式不受方向余弦耦合的限制,其动作空间更小(6<8)(坐标点3+切点平移3=6)(余弦(α、β、γ)3+原点(d)1+方向余弦移动3+原点(d)1移动=8),促进了智能体学习。我们在这里解释RL中元素的细节。
平面定位的强化学习
动作Action(8维->6维)
动作空间
给定步骤i中的一个动作,切线点坐标可以相应地修改,例如。
当相应的球面半径较小时,图像内容对步长很敏感。
为了解决这个问题,我们将agent-环境交互建模为一个多阶段的运动过程,当agent出现三步振荡的情况时,我们将步长从1.0逐步缩小为0.01。我们以60个步骤终止agent搜索。
状态State
根据当前平面参数从体中采样体素重建的二维超声图像。
由于重建后的图像大小可能会改变,我们将图像填充到一个正方形上,并将其大小调整为224×224。
此外,我们将前两次迭代得到的两幅图像与当前平面连接起来,丰富了类似于[31]的状态信息。
奖励Reward(SLR+ASR)
奖励函数指示agent采取正确行动的最优搜索策略。
最近的工作[17,18]根据相邻迭代中定义平面函数参数的差异计算奖励函数。
虽然有效,但我们认为这种设计可能会导致agent缺乏解剖感知和指导,从而影响agent对异常数据的表现。
在本研究中,我们设计了一种空间解剖奖励,包括
1)空间位置奖励(SLR)spatial location reward
2)解剖结构奖励(ASR)anatomical structure reward
SLR通过最小化当前平面与目标平面之间的平面参数的欧氏距离来激励agent接近目标位置
ASR则鼓励agent感知解剖信息。在解剖标志处用高斯核构建热图(见图1)来计算ASR。
奖励定义:
sgn(·)为符号函数,
Pt t时刻预测平面的平面参数(Pt,t时刻 plane)
Pg 目标真实平面(Pg,ground plane)
欧几里德距离
It和Ig分别表示预测平面和目标平面对应的热图值之和
(热图有助于了解一张图像的哪一部分让神经网络做出了最终的分类决策,热图的产生有两种方式,一是高斯热图,一个Grad-Cam产生的激活类热图
热图可以反映检测目标的位置
高斯热图,图中心高斯值接近1,边缘接近0
)
回放缓存Replay Buffer
内存容器,它存储agent的转换,以执行学习过程的经验重放。
元素转换通常用一个向量表示,
给定优先级的replay buffer存储agent-环境交互中每个步骤t的状态s、动作a、奖励r和下一步s0的状态的数据序列,这可以消除数据相关性,提高采样效率。
具有高误差的第i个序列元素将优先从buffer中采样,其采样概率可以计算为,
p = 0.6控制使用多少优先级
δ = 0.05被设置为调整误差ei。
采用了重要性抽样权值来纠正优先级重放中数据分布变化引起的偏差。
损失Loss
与[17,18]类似,用Dueling Q-Learning训练agent,
损失函数定义:
γ是平衡当前和未来奖励的重要性的折扣因素。
s、s0 a和a0是当前/下一步中的状态和动作。
M是回放缓存存储器,以避免频繁的数据采样。w和˜w为当前和目标网络参数。
在本研究中,当前Q-Learning和目标Q-Learning具有相同的网络架构(即ResNet[5]),
在训练过程中,目标Q-Learning每1800步复制一次当前Q-Learning的参数。
状态内容相似度预测的辅助任务
由于在搜索过程中标准切面集和非标准切面集之间的类间相似性较低,
SPs和非SPs之间的差异,甚至不同的SPs往往可以忽略不计,因为它们可以包含相同或不同的解剖结构,具有相似的外观。
而标准切面集的类内差异较高,因此在3D超声中定位标准切面集具有挑战性。
(类内差异(Intra-class variation):一类物体的个体之间的外形差异很大,比如椅子。这一类物体有许多不同的对象,每个都有自己的外形。)
大多数方法[10,17,18]缺乏使用图像级内容信息(如解剖先验)的适当策略,导致数据利用和agent学习效率低下。
RL[7,13]的辅助任务可以通过学习细粒度表示来提高学习效率,提高性能。为了方便agent学习内容的表示形式,设计了一个状态内容相似度预测(SCSP)的辅助任务。
如图1所示,我们在agent网络中使用了一个额外的回归分支来预测当前状态与目标状态的相似性。内容相似性通过归一化互相关(NCC)[20]来衡量。
辅助任务部分的损失函数定义为:
当前状态s和目标状态sgt之间的NCC
SCSP预测的NCC分数。
RL框架的总损失函数是,其中δ = 0.5是平衡Q-Learning损失和辅助任务损失的重要性的权重。
基于模仿学习的初始化
在与未对齐的超声环境交互时,agent很难获得有效的样本,因为在agent探索三维空间时,回放缓冲区会存储大量无效的数据。它可能会降低agent的学习效率,甚至会影响学习性能。
[6]指出,模仿学习可以通过对agent进行预训练,确保其在探索环境之前能够获得足够的知识,从而提高学习效率。因此,我们采用模仿学习作为agent的初始化。
具体来说,首先在每个训练数据中随机选择20个初始切点,然后根据到目标切点的距离采取最优动作来接近目标平面。
这可以模拟专家的操作,并获得许多有效的演示状态动作轨迹(例如,(s0、a0),(s1、a1),…,(sn、an))。
然后,基于随机采样的状态-动作对对agent进行交叉熵损失的监督学习。使用基于模仿学习初始化的训练良好的agent,可以简化和加速RL框架的学习。
实验
(详细结果见论文原文)
我们在两个数据集中的4个SPs上验证了我们提出的方法,包括子宫的冠状面和经脑室(TV)、经丘脑(TT)和经小脑(TC)平面。子宫数据集有363例正常患者和45例异常患者(先天性子宫异常,CUAs),平均体积大小为432×377×217,间距为0.3×0.3×0.3mm3;胎儿大脑数据集有432例患者,平均体积大小为270×207×235,间距为0.5×0.5×0.5mm3。6名经验丰富的超声医生在严格的质量控制下,使用配对注释软件包[11]手动对sp和地标进行了注释。我们随机分割每个数据集,分别进行训练、验证、测试子宫内290、20、53和胎儿大脑内330、30、72。为了验证我们的方法的可推广性,我们提出了一些建议我们只将健康受试者纳入我们的训练数据集,并独立测试45例CUAs患者。
NVIDIA RTX 2080Ti GPU的标准PC,
PyTorch
Adam优化器对模型进行训练,
学习率为5e-5,batch size 32 , 100 epochs
等式中的折扣系数γ2被设为0.85。有优先级的Replay Buffer的大小被设置为15000。
根据[17]建立了 − greedy 算法。我们计算了训练数据集中目标切点位置的平均值(µ)和标准差(σ),并在µ±2σ内随机初始化了训练的起点,以近似捕获95%的变异性。为了进行测试,该原点被设置为初始切线点。
标准切面;
预测平面,右下角为特征点热图;
标准切面(红色)和预测(绿色)之间的三维空间关系
图3 我们的方法的可视化示例结果
定量和定性分析
与五种SOTA方法进行了比较,
regression-based(即RGSingle,RGITN[10])、
registration-based(即Regist[4])
RL(即RLAVP[1],rlwsadt[17])。
使用了4个标准来评估平面定位
空间相似性:
1.两个平面之间的二面角(Ang)
np,ng表示预测平面和目标平面的法线
2.两个平面到原点的欧氏距离差(Dis)
dp,dg表示与体积的距离原点到预测平面,原点到真实平面
(Ang和Dis是基于平面采样函数,即cos(α)x + cos(β)y + cos(γ)z = d,
有效体素大小为0.5 mm3/voxel)
内容相似性:
3.峰值结构相似度(SSIM)。
并将迭代t中的ADI定义为距离起始平面的距离和角度的累积变化之和,如下
4.归一化互相关NCC
表1.不同方法对SP定位的定量比较。CN和CP分别为正常子宫和异常子宫的冠状面。(mean±std)
如表1所示,我们提出的方法在大多数指标上都优于所有其他的方法。
RL通过预配准获得了很大的性能提升,以确保方向的一致性(RLAVP vs。RLWSADT ).相比之下,我们的新公式可以使RL算法即使在没有预配准时也能获得优势。
先前的基于rl的方法在异常数据集(CP中的rlwsadt)中很容易失败。
相反,我们的方法在正常和异常数据集上都获得了一致的性能。
表2.分析SCSP和ASR的消融研究。
表2显示了研究每个设计模块的影响的消融研究的结果。
可以观察到,SCSP通过提高对SPs和非SPs的识别能力,提高了模型的通用性。
将ASR与基本的奖赏SLR组合起来,有利于促进主体对解剖结构的感知,使我们的模型能够归纳出与正常子宫数据集具有显著内容差异的外部异常子宫数据集。
图3还显示了与表1中报告的定量测量相关的SP定位性能的程度。
结论
我们提出了一种新的三维SP定位的RL框架。我们定义了一个基于切线点的平面公式来重组行动空间,并改进在不对准的环境中的代理优化。这种公式可以扩展到其他模式中的类似任务,如CT或MRI。
1.提出了一个内容感知的回归辅助任务,以提高代理对有噪声的环境的鲁棒性。
2.设计了一个空间解剖奖励,为代理提供空间和解剖知识。
3.通过模仿学习来初始化代理,以提高训练效率。
实验表明,该方法在包括异常情况在内的两个非对齐数据集的4个SP的定位具有较好的性能,表明该方法在随机初始化空间和异常情况下的定位具有很大的潜力
接下来的工作
连续动作的强化学习来提升效率和降低对初始化参数的依赖