基于深度学习的医学图像分割学习笔记（八）HF U-Net

HF-UNet: Learning Hierarchically Inter-Task Relevance in Multi-Task U-Net for Accurate Prostate Segmentation in CT Images

IEEE Transactions on Medical Imaging(TMI)
原文链接
dataset：339张CT图像

快速定位
精确分割
将多任务学习和深度神经网络相结合，提出了多任务卷积网络框架（主任务分割前列腺，辅助任务划定前列腺边界）

与整个盆腔CT图像相比，前列腺位于一个相对较小的区域，因此大的背景噪声会削弱网路。为了解决上述问题，提出了一个two-stage分割框架：第一阶段进行快速定位；第二阶段进行精确分割，网络框架如下图所示：

A. 快速定位Automatic Prostate Region Detection

首先采用区域定位网络裁剪出包含整个前列腺区域的子图像。
通过一个U-Net架构自动定位前列腺区域，然后根据粗略预测的分割图确定前列腺中心。根据获得的中心裁剪得到一个128×128×128的完全覆盖前列腺器官的区域。为了进行更精确的分割，从识别区域裁剪出的patchs代替原始CT图像用于训练后续的网络（HF U-Net）。

B.精确分割 HF-UNet

为HF-Unet设置了两个任务以实现精确分割：Branch a: 分割前列腺 。Branch b: 划分前列腺边界(轮廓感知子网络，为网络提供关键的形态学指导)。

1.器官边界的划分

将划分任务看作一个回归问题，轮廓附近的像素点作为sub-groundtruths。制定每个轮廓点和其周围的像素形成高斯分布$\sigma$=5.
将损失层的输入表示为($\xi$$x$,$y$),其中$\xi$$x$表示网络的输出特征，$y$表示前列腺轮廓的ground truth标签。使用均方误差计算两个任务之间的差值，轮廓感知的目标函数：

其中$\varphi$($X$,$\theta$)是轮廓感知分支的输出，$S$($y_{ij}$)高斯滤波器的集合，可表示为：($K,L$)表示与($i,j$)距离小于$\sigma$的邻集

2.任务一致性学习TCL模块

两个分支的信息通过TCL模块组合和学习，从全局角度对两个分支学习到的知识进行建模。本文构建了两种TCL模块：一种是加权残差TCL块来评估共享信息和任务特定信息的有效性，另一种是基于注意力的TCL块来自适应学习上述两种信息的比率。

一个网络总共有$M$层TCL块，分割分支，轮廓感知和TCL模块在第$m$层的特征输出分别为$\xi$${}_{seg.,m}$,$\xi$${}_{cont.,m}$,$\xi$${}_{tcl.,m}$，TCL的最终损失是将两个任务对的计算损失与TCL块的输出的融合特征进行聚合：（平均平方误差可以建立融合特征与分割子网络和轮廓划分子网络特征的一致性）

3.多任务学习
通过多任务学习共同解决分割，回归和任务一致性学习三个任务，上述三个任务提供不同类型的信息指导HF-UNet的训练，最终的目标函数：

其中$L$${}_{cls}$表示前列腺分割时分类任务的损失，$L$${}_{reg}$ 表示前列腺轮廓划分时回归任务的损失，权重$\Theta$$=${$\theta$${}_{cls}$,$\theta$${}_{reg}$,$\theta$${}_{tcl}$}由$R$(·)正则化。$L$${}_{cls}$ 使用交叉熵损失进行分割，$\hat{p}$表示预测的概率，$y$表示ground truth标签。