基于强化学习的图像配准 - Image Registration: Reinforcement Learning Approaches

1. 配准定义

给定参考图像 I_f 和浮动图像 I_m ，所谓的配准就是寻找一个图像变换T，将浮动图像I_m变换到和 I_f 相同的坐标空间下，使得两个图像中对应的点处于同一坐标下，从而达到信息聚合的目的。在医学图像配准中，最常见的例子便是将同一个部位的CT图像变换移动到与MRI图像的位置相同，如下图：

​

而配准又分为刚性配准（rigid registration）和非刚性配准（non-rigid registration），而非刚性配准也叫做柔性配准（deformable registration）。其区别是：刚性配准的变化通常都是仿射变换，一般通过全图的平移、旋转、放缩等几个操作就可以对齐图像；而柔性配准则是对每一个局部区域（甚至每个像素点）都会有一个单独的变换，有学者们使用向量场（vector field）或者密集形变场（dense displacement field）来形容这种变换是比较合理的。换一种说法的话柔性变换也可以认为是光流（optical flow）对齐的一种医学特例。一般来说柔性配准更具有挑战性一些。

图像刚性配准

图像柔性配准

2. 基于深度学习的图像配准（DL-based registration）

此类方法是最常见的，最主流的做法。由于本篇着重介绍的是基于强化学习的配准，这里我们就简单介绍一下传统经典的Deep Learning的方法，如VoxelMorph （DL 配准的开山之作）：

VoxelMorph的结构图

其余的DL方法差不多都是这样，一般就是改改结构、改改Loss什么的。

3. 基于强化学习的图像配准（RL-based registration）

这里说到重点了。但在开始之前，我们先简单聊一下强化学习。在图像配准中，所谓的强化学习其实就是把DL方法中一步到位估计的图像变换T分解为多个小的变换。比如在刚性配准中，DL的方法是直接预测一个最终的仿射变换（2x3的矩阵），而强化学习则会学习一个策略，让浮动图像先平移一点、再旋转一点、然后又放缩一下什么的，在N个步骤之后才会获得比较大的配准结果。通俗一点，你可以设想有一个智能体（机器人）使用双手和眼睛去不断修正自己配准的结果。而如何学习呢，则是通过一个校验器去判断每次配准的得分，得分高则给予高奖励，得分低则只给少许奖励（Reward）或者给一定惩罚。

​

具体来说，根据当前状态，智能体会做出一个动作（平移旋转等），然后环境对浮动图像进行变换，并给予反馈。这里的环境是我们人为设置的，它的作用就是：1、提供每次训练所需的数据；2、可以衡量配准结果的好坏（如Dice等评价指标）并反馈奖励给智能体；3、根据智能体的动作对浮动图像进行操作。

那为什么要用强化学习来做配准呢？大多数给出的答案是当遇到多模态时，一般的DL方法没法从形态差异较大的图像间找出最优解（关键是不好设计Loss）。且图像变换比较复杂或数据较少时，由于DL的loss通常是针对全局的，它们会更倾向于学习粗略对齐而忽略掉一些细节。

• 刚性配准中的强化学习

刚性配准其实是比较适合用强化学习来解决的，因为它就只需要处理全局的变换，一般就上下左右移动、左右旋转、放大或缩小8个离散动作。而奖励的设置则是通过计算两图的相应特征点（SIFT点）之间的距离，再用当前本的距离减去上一步的距离得到Reward。或者就直接距离越远则奖励越小，距离越近则奖励越高。其代表是Kai Ma 、Rui Liao 以及Shanhui Sun 等人的工作。

在强化学习算法的选用上，Kai Ma用的是Dueling DQN ，Rui Liao用的是DQN 加上一个对动作的强监督，而Shanhui Sun则是用的A3C 。从RL算法上来说，A3C属于异步学习框架，是三者中最高效的，而且目前大部分工业RL算法也都是异步学习的，从论文看来其在配准中效果也最好。基于A3C，后面有一系列扩展的配准相关工作，如改进奖励函数、增加更强的时序和空间信息（ConvLSTM/ConvGRU）等发表在医学顶刊Medical Image Analysis上 或者ACCESS上 。这系列基于A3C的方法整体框架其实是差不多的，如下图：

​

基于A3C异步强化学习的图像配准框架

• 柔性配准中的强化学习

前面我们提到过，相比于刚性配准，柔性配准要难得多。这是由于它的变换比较复杂，且输出的是全图每个像素点的变换（动不动就是128x128x128的向量场）。这对强化学习来说无疑是巨大的挑战。

目前绝大部分的强化学习算法都是针对低维度的、离散的动作，也就是说根本无法直接将强化学习搬运过来搞柔性配准任务。那怎么办呢，目前已知的有两种方法。一是取巧的方法：我搞不定高维动作的RL算法那又怎么样呢，我选择直接将高维连续动作（向量场）转化为低维离散的动作 。

其具体做法是将图像网格化，再利用B曲线（B-splines）来控制某点带动周围网格的形变。你可以想象你用手去提一个平铺在地上的渔网，你用手去拉渔网上的一个点（网线的交界处），则该点周围全都会被拉起来，且靠的越近，受到的影响越大。而渔网上有数个这样的点，你每次拉一点点，拉的点越多，则渔网的形变就越大。Krebs等人就是这么做的 ，但他们还遇到一个问题，那就是图像上的点实在是太多了，而智能体受到RL的算法框架限制，每次又只能拉一点，多了就拉不了了。于是他们使用PCA对这些点进行降维（3D图像从16x16x8个点降到25个点，2D图像从16x16个点降到15个点），由于每个点都可以向上或向下拉，所以最终需要预测50个动作（for 3D）或30个动作（for 2D）。

​

强制问题匹配算法倾向：解决不了RL算法，就解决问题本身。

但这个方法主要有两个大问题。其一：该算法需要Ground Truth Deformation Field。就算你的动作已经从128x128x128(x2)降到了50，但这对于DQN此类RL算法来说还是太难了。因此作者采用GT形变场来计算Reward，但这也就意味着只能自己合成数据了，且合成的方式只能是基于B-Spline的，且只能是单模态。其二：整个算法流程实在是太复杂了= =，包含了各种传统机器学习和图像算法：DDF（Free Form Deformations）、PCA等等，再加上强化学习，基本就是来一个劝退一个。而且这篇paper没有公开代码，我复现的时候都快哭了好吧！

当然，作为第一篇处理柔性配准的RL算法它已经很屌了，毕竟算是开了一个先河。那么不改变问题本身，难道就没有RL算法可以直接处理这种高维连续动作的问题了吗？答案是：非常难！就算在深度学习飞速发展的近几年，RL+柔性配准依然劝退了一波又一波人。直到今年（2022年）才出现了第二篇paper使用强化学习来处理柔性配准任务 （发表在AAAI）。

这篇paper（这里）的主要思想其实很简单：高维动作我解决不了，我就学习一个低维的representation vector，再学一个Decoder把这个vector还原成高维动作，此外我再增加一个监督或无监督loss来加速智能体对连续动作的学习，确保万无一失。

​如图所示，除了Actor和Critic以外，这篇paper还多加了一个Planner用来学习高维动作的低维表示，而Actor负责还原高维的动作。这里的动作其实就是完整的Deformation Field，整个过程就类似于coarse-to-fine进行配准的过程。这篇文章还使用了K-means来得到每个图像自己粗略的分割label，并利用这个这些label来作为reward function，因此是无监督的，不需要合成数据，且可以用到任意的公开数据集中。此外，这篇paper的RL算法看得出来是基于SAC（Soft Actor-Critic）的，也是目前最流行的连续动作算法之一吧。并且代码还开源了！

paper：Stochastic Planner-Actor-Critic for Unsupervised Deformable Image Registration

文章地址：https://arxiv.org/abs/2112.07415​

Github 地址：https://github.com/Algolzw/SPAC-Deformable-Registration