psa-for-prostate代码梳理
psa-for-prostate代码梳理
将psa(Pixel-level Semantic Affinity)这篇文章应用到prostate图像当中,用于对prostate图像进行弱监督语义分割。代码主要分为四个部分:
- 训练分类网络模型
- 生成不同形式的CAMs以供使用
- 训练affinity网络模型
- 对原CAMs进行fine-tuning
下面对代码的整个结构以及其中的一些细节进行记录和整理。
1.训练分类网络模型
2.生成不同形式的CAMs以供使用
3.训练affinity网络模型
3.1 数据预处理
这一部分的训练数据在data.py
文件中由AffDataset
类定义,该类继承ImageDataset
类。训练数据包括原图片和affinity标签(关于affinity标签的定义参照论文,此处不祥说),其中原图片易取(包括一些裁剪、翻转、归一化等预处理操作)而affinity标签难得。因为affinity标签需要根据之前生成的各种CAMs来生成,具体生成过程见data.py
中ExtractAffinityLabelInRadius
类。生成affinity标签所需要的数据有降采样后尺寸cropsize=cropsize//8
,范围半径radius=radius
和阈值化后的CAMs标签(label_la,label_ha)。其中,
cropsize是降采样后的尺寸,注意这里的操作都是在8倍降采样(平均池化,步长为8)后的尺寸里操作的;
范围半径radius是指在选取affinity标签对的时候所考虑的点对之间的距离是在半径radius内的;
阈值化后的CAMs标签包括确定的前景,确定的背景和非确定区域。label_la是将背景区域值偏大估计的二维mask,label_ha是将背景区域偏小估计的二维mask,通过这两个mask可以形成前面所提到的阈值化后的CAMs标签,这个阈值化后的mask有以下几个特点:1.是降采样后的尺寸,2.包括确定的背景和确定的前景,3.是单通道的,不同类别用不同像素值表示。然后对affinity标签的提取过程详细说明如下:
首先在给定半径范围选择符合要求的搜索点对,这里假设半径为5,则符合要求的点对有34个;
然后根据将采样后的尺寸结合半径确定之后标签提取过程在(crop_height,cop_width)的尺寸上进行;
接着确定起始点和目的点的坐标,分别用labels_from和labels_to来表示;
最后根据affinity标签的定义用逻辑与或非等操作确定背景正样本、前景正样本和负样本。具体代码如下:
bg_pos_affinity_label = np.logical_and(pos_affinity_label, np.equal(bc_labels_from, 0)).astype(np.float32)
fg_pos_affinity_label = np.logical_and(np.logical_and(pos_affinity_label, np.not_equal(bc_labels_from, 0)), concat_valid_pair).astype(np.float32)
neg_affinity_label = np.logical_and(np.logical_not(pos_affinity_label), concat_valid_pair).astype(np.float32)
3.2 网络模型
使用的backbone网络还是vgg16结构,然后对后三个卷积层的输出进行连接,再进行一个1x1卷积操作得到最后的多通道的affinity特征图,生成的特征图用于后续的affinity以及loss的计算。
3.3 模型训练
训练的前向传播过程先生成上一部分提到的多通道(448)的特征图,然后利用3.1中确定的搜索点对坐标确定特征图层面的起始点和目的点矩阵ff,ft。
ff = torch.index_select(x, dim=2, index=ind_from.cuda(non_blocking=True))
ft = torch.index_select(x, dim=2, index=ind_to.cuda(non_blocking=True))
之后利用公式
aff = torch.exp(-torch.mean(torch.abs(ft-ff), dim=1))
得到34通道的affinity矩阵aff
,每个通道表示到不同距离(指定半径范围内)的点的相似度情况。利用这个aff
,shape为((1,34,1/8原图尺寸))结合3.1中得到的affinity标签,即可通过以下公式计算出loss,随后便可进行梯度下降训练。
bg_loss = torch.sum(- bg_label * torch.log(aff + 1e-5)) / bg_count
fg_loss = torch.sum(- fg_label * torch.log(aff + 1e-5)) / fg_count
neg_loss = torch.sum(- neg_label * torch.log(1. + 1e-5 - aff)) / neg_count
loss = bg_loss / 4 + fg_loss / 4 + neg_loss / 2
3.4 模型推理
利用训练得到的affinity模型,对相同的训练图片进行一个infer的过程。整个前向传播的过程和训练时的过程完全一样,就是在返回aff
的时候需要对该矩阵进行稀疏化,变成(1/8原图高1/8原图宽,1/8原图高1/8原图宽)的正方形矩阵,对角线元素全为1,表示自己和自己肯定相似度最高,相关代码如下:
aff_mat = sparse.FloatTensor(torch.cat([indices, indices_id, indices_tp], dim=1),
torch.cat([aff, torch.ones([area]), aff])).to_dense().cuda()
4.对原CAMs进行fine-tuning
利用第3部分最后推理得到的affinity矩阵aff_mat
,将其转换成概率矩阵(归一化),然后对8倍降采样后的CAMs(第2部分得到)进行random walk处理(简单的矩阵相乘多次),变得到最后的fine-tuned的结果。
aff_mat = torch.pow(model.forward(img.cuda(), True), args.beta)
trans_mat = aff_mat / torch.sum(aff_mat, dim=0, keepdim=True)
for _ in range(args.logt):
trans_mat = torch.matmul(trans_mat, trans_mat)
cam_full_arr = torch.from_numpy(cam_full_arr)
cam_full_arr = F.avg_pool2d(cam_full_arr, 8, 8)
cam_vec = cam_full_arr.view(2, -1)
cam_rw = torch.matmul(cam_vec.cuda(), trans_mat)
cam_rw = cam_rw.view(1, 2, dheight, dwidth)
cam_rw = torch.nn.Upsample((img.shape[2], img.shape[3]), mode='bilinear')(cam_rw)
_, cam_rw_pred = torch.max(cam_rw, 1)
res = np.uint8(cam_rw_pred.cpu().data[0])[:orig_shape[2], :orig_shape[3]]