24 | 使用PyTorch完成医疗图像识别大项目：图像分割数据准备

import torchfrom torch import nnimport torch.nn.functional as Fclass UNet(nn.Module):#默认参数是原论文中使用的参数
    def __init__(self, in_channels=1, n_classes=2, depth=5, wf=6, padding=False,
                 batch_norm=False, up_mode='upconv'):
        """
    参数说明
        Args:
            in_channels (int): 输入通道数
            n_classes (int): 输出通道数
            depth (int): 网络深度
            wf (int): 过滤器数量，也就是用于控制内部的卷积块和上采样块的输入输出通道数，第一层过滤器数量是 2**wf
            padding (bool): 补充，是否使用padding
                            
            batch_norm (bool):是否使用batch norm
            up_mode (str):使用上卷积或上采样 'upconv' or 'upsample'.
                           'upconv' 使用转置卷积实现上采样
                           'upsample' 使用线性上采样
        """
        super(UNet, self).__init__()#参数处理
        assert up_mode in ('upconv', 'upsample')
        self.padding = padding
        self.depth = depth
        prev_channels = in_channels
        self.down_path = nn.ModuleList()#根据深度决定实现几个卷积块以及上采样
        for i in range(depth):
            self.down_path.append(UNetConvBlock(prev_channels, 2**(wf+i),
                                                padding, batch_norm))
            prev_channels = 2**(wf+i)

        self.up_path = nn.ModuleList()
        for i in reversed(range(depth - 1)):
            self.up_path.append(UNetUpBlock(prev_channels, 2**(wf+i), up_mode,
                                            padding, batch_norm))
            prev_channels = 2**(wf+i)#最后的卷积
        self.last = nn.Conv2d(prev_channels, n_classes, kernel_size=1)#前向传播
    def forward(self, x):
        blocks = []#首先执行卷积块
        for i, down in enumerate(self.down_path):
            x = down(x)
            if i != len(self.down_path)-1:
                blocks.append(x)#平均池化
                x = F.avg_pool2d(x, 2)#然后执行上采样块
        for i, up in enumerate(self.up_path):
            x = up(x, blocks[-i-1])

        return self.last(x)#卷积块的实现class UNetConvBlock(nn.Module):#在一个卷积块里面执行了两次卷积，两次激活函数变换，两次batch norm，对应架构图上每层向右的两个箭头
    def __init__(self, in_size, out_size, padding, batch_norm):
        super(UNetConvBlock, self).__init__()
        block = []

        block.append(nn.Conv2d(in_size, out_size, kernel_size=3,
                               padding=int(padding)))
        block.append(nn.ReLU())
        # block.append(nn.LeakyReLU()) 不同的激活函数
        if batch_norm:
            block.append(nn.BatchNorm2d(out_size))

        block.append(nn.Conv2d(out_size, out_size, kernel_size=3,
                               padding=int(padding)))
        block.append(nn.ReLU())
        # block.append(nn.LeakyReLU())
        if batch_norm:
            block.append(nn.BatchNorm2d(out_size))

        self.block = nn.Sequential(*block)#卷积块的前向传播
    def forward(self, x):
        out = self.block(x)
        return out#定义上采样块class UNetUpBlock(nn.Module):
    def __init__(self, in_size, out_size, up_mode, padding, batch_norm):
        super(UNetUpBlock, self).__init__()#使用转置卷积实现上采样
        if up_mode == 'upconv':
            self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2,
                                         stride=2)#使用双线性模型实现上采样
        elif up_mode == 'upsample':
            self.up = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2),
                                    nn.Conv2d(in_size, out_size, kernel_size=1))#卷积块
        self.conv_block = UNetConvBlock(in_size, out_size, padding, batch_norm)#中心裁剪
    def center_crop(self, layer, target_size):
        _, _, layer_height, layer_width = layer.size()
        diff_y = (layer_height - target_size[0]) // 2
        diff_x = (layer_width - target_size[1]) // 2
        return layer[:, :, diff_y:(diff_y + target_size[0]), diff_x:(diff_x + target_size[1])]#前向传播
    def forward(self, x, bridge):
        up = self.up(x)#在这里实现跳链以及裁剪
        crop1 = self.center_crop(bridge, up.shape[2:])
        out = torch.cat([up, crop1], 1)
        out = self.conv_block(out)

        return out

from util.unet import UNetclass UNetWrapper(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()

        self.input_batchnorm = nn.BatchNorm2d(kwargs['in_channels'])
        self.unet = UNet(**kwargs)
        self.final = nn.Sigmoid()#这里增加了一个权重初始化动作
        self._init_weights()

   def forward(self, input_batch):
       bn_output = self.input_batchnorm(input_batch)
       un_output = self.unet(bn_output)
       fn_output = self.final(un_output)
       return fn_output

#获取中心点信息，从xyz坐标转换到irc坐标
            center_irc = xyz2irc(
                candidateInfo_tup.center_xyz,
                self.origin_xyz,
                self.vxSize_xyz,
                self.direction_a,
            )
            ci = int(center_irc.index)
            cr = int(center_irc.row)
            cc = int(center_irc.col)#索引轴方向上的半径检索，下面还有行方向的检索和列方向的检索
            index_radius = 2
            try:
                while self.hu_a[ci + index_radius, cr, cc] > threshold_hu and \                        self.hu_a[ci - index_radius, cr, cc] > threshold_hu:
                    index_radius += 1
            except IndexError:  #异常捕获，如果超出阈值就返回1格
                index_radius -= 1

    def buildAnnotationMask(self, positiveInfo_list, threshold_hu = -700):
        boundingBox_a = np.zeros_like(self.hu_a, dtype=np.bool)

        for candidateInfo_tup in positiveInfo_list:#遍历所有的结节数据…………#这里省略的就是上面的检索代码
            boundingBox_a[#根据前面检索的半径生成整个框
                 ci - index_radius: ci + index_radius + 1,
                 cr - row_radius: cr + row_radius + 1,
                 cc - col_radius: cc + col_radius + 1] = True#然后对框内的数据进行判断，超过阈值的设为True
        mask_a = boundingBox_a & (self.hu_a > threshold_hu)

        return mask_a

class Ct:
    def __init__(self, series_uid):
     ……#省略了一些代码
        candidateInfo_list = getCandidateInfoDict()[self.series_uid] #获取候选信息

        self.positiveInfo_list = [  #正样本列表
            candidate_tup            for candidate_tup in candidateInfo_list            if candidate_tup.isNodule_bool        ]
        self.positive_mask = self.buildAnnotationMask(self.positiveInfo_list)  #正样本掩码构建#最后这行是把具有非0计数的掩码切片的索引存下来
        self.positive_indexes = (self.positive_mask.sum(axis=(1,2))
                                 .nonzero()[0].tolist())

    def getRawCandidate(self, center_xyz, width_irc):#获取中心位置
        center_irc = xyz2irc(center_xyz, self.origin_xyz, self.vxSize_xyz,
                             self.direction_a)#起始位置和终止位置索引
        slice_list = []
        for axis, center_val in enumerate(center_irc):
            start_ndx = int(round(center_val - width_irc[axis]/2))
            end_ndx = int(start_ndx + width_irc[axis])#断言 用来处理异常
            assert center_val >= 0 and center_val < self.hu_a.shape[axis], repr([self.series_uid, center_xyz, self.origin_xyz, self.vxSize_xyz, center_irc, axis])

            if start_ndx < 0:
                # log.warning("Crop outside of CT array: {} {}, center:{} shape:{} width:{}".format(
                #     self.series_uid, center_xyz, center_irc, self.hu_a.shape, width_irc))
                start_ndx = 0
                end_ndx = int(width_irc[axis])

            if end_ndx > self.hu_a.shape[axis]:
                # log.warning("Crop outside of CT array: {} {}, center:{} shape:{} width:{}".format(
                #     self.series_uid, center_xyz, center_irc, self.hu_a.shape, width_irc))
                end_ndx = self.hu_a.shape[axis]
                start_ndx = int(self.hu_a.shape[axis] - width_irc[axis])#加入切片信息
            slice_list.append(slice(start_ndx, end_ndx))

        ct_chunk = self.hu_a[tuple(slice_list)]#提取其中的正样本掩码数据
        pos_chunk = self.positive_mask[tuple(slice_list)]#返回数据        return ct_chunk, pos_chunk, center_irc

@raw_cache.memoize(typed=True)def getCtRawCandidate(series_uid, center_xyz, width_irc):
    ct = getCt(series_uid)
    ct_chunk, pos_chunk, center_irc = ct.getRawCandidate(center_xyz,
                                                         width_irc)
    ct_chunk.clip(-1000, 1000, ct_chunk)
    return ct_chunk, pos_chunk, center_irc

    candidateInfo_list = []
    with open('data/part2/luna/annotations_with_malignancy.csv', "r") as f:
        for row in list(csv.reader(f))[1:]:
            series_uid = row[0]
            annotationCenter_xyz = tuple([float(x) for x in row[1:4]])
            annotationDiameter_mm = float(row[4])
            isMal_bool = {'False': False, 'True': True}[row[5]]#从中取出我们需要的数据并合成结果
            candidateInfo_list.append(
                CandidateInfoTuple(
                    True,
                    True,
                    isMal_bool,
                    annotationDiameter_mm,
                    series_uid,
                    annotationCenter_xyz,
                )
            )#然后从candidates.csv取出候选信息
    with open('data/part2/luna/candidates.csv', "r") as f:
        for row in list(csv.reader(f))[1:]:
            series_uid = row[0]

            if series_uid not in presentOnDisk_set and requireOnDisk_bool:
                continue

            isNodule_bool = bool(int(row[4]))
            candidateCenter_xyz = tuple([float(x) for x in row[1:4]])#这里只使用非结节数据
            if not isNodule_bool:
                candidateInfo_list.append(
                    CandidateInfoTuple(
                        False,#是否结节
                        False,#是否恶性
                        False,#是否有标注
                        0.0,
                        series_uid,
                        candidateCenter_xyz,
                    )
                )

    candidateInfo_list.sort(reverse=True)
    return candidateInfo_list

class Luna2dSegmentationDataset(Dataset):
    def __init__(self,
                 val_stride=0, #验证集长度
                 isValSet_bool=None,  #是否验证集
                 series_uid=None, #序列uid
                 contextSlices_count=3, #上下文切片数，就是从中心点切片向上向下选取的切片数量，这里设置为3，就是选取7片切片作为我们的训练数据7个通道。
                 fullCt_bool=False,#这个参数用来标记是否使用所有切片，当评估的时候使用True，当训练的时候使用False，这样可以让模型更关注那些带有阳性掩码的切片，而评估的时候我们也不知道哪个有阳性，所以使用所有切片。
            ):#处理参数默认信息
        self.contextSlices_count = contextSlices_count        self.fullCt_bool = fullCt_bool        if series_uid:
            self.series_list = [series_uid]
        else:
            self.series_list = sorted(getCandidateInfoDict().keys())#在判断是否验证集的时候，如果不是验证集则需要按验证集长度剔除掉数据
        if isValSet_bool:
            assert val_stride > 0, val_stride            self.series_list = self.series_list[::val_stride]
            assert self.series_list
        elif val_stride > 0:
            del self.series_list[::val_stride]
            assert self.series_list#由于我们的CT数据切片数不一样，这里构建了一个缓存数据，用来存放CT数据的长度getCtSampleSize，等会看一下这个代码。
        self.sample_list = []
        for series_uid in self.series_list:
            index_count, positive_indexes = getCtSampleSize(series_uid)#这里处理fullct参数，可以看到如果True，则使用的index_count，如果是False使用的只有positive_indexes
            if self.fullCt_bool:
                self.sample_list += [(series_uid, slice_ndx)
                                     for slice_ndx in range(index_count)]
            else:
                self.sample_list += [(series_uid, slice_ndx)
                                     for slice_ndx in positive_indexes]#缓存数据
        self.candidateInfo_list = getCandidateInfoList()#建立一个set方便查找
        series_set = set(self.series_list)
        self.candidateInfo_list = [cit for cit in self.candidateInfo_list                                   if cit.series_uid in series_set]

        self.pos_list = [nt for nt in self.candidateInfo_list                            if nt.isNodule_bool]#记录日志信息
        log.info("{!r}: {} {} series, {} slices, {} nodules".format(
            self,
            len(self.series_list),
            {None: 'general', True: 'validation', False: 'training'}[isValSet_bool],
            len(self.sample_list),
            len(self.pos_list),
        ))

@raw_cache.memoize(typed=True)def getCtSampleSize(series_uid):
    ct = Ct(series_uid)
    return int(ct.hu_a.shape[0]), ct.positive_indexes

# getitem的功能都包装在getitemfullslice中了
    def __getitem__(self, ndx):
        series_uid, slice_ndx = self.sample_list[ndx % len(self.sample_list)]
        return self.getitem_fullSlice(series_uid, slice_ndx)#
    def getitem_fullSlice(self, series_uid, slice_ndx):
        ct = getCt(series_uid)#初始化tensor
        ct_t = torch.zeros((self.contextSlices_count * 2 + 1, 512, 512))#设置起始块和终止块位置
        start_ndx = slice_ndx - self.contextSlices_count
        end_ndx = slice_ndx + self.contextSlices_count + 1#把块数据取出来放进tensor中
        for i, context_ndx in enumerate(range(start_ndx, end_ndx)):
            context_ndx = max(context_ndx, 0)
            context_ndx = min(context_ndx, ct.hu_a.shape[0] - 1)
            ct_t[i] = torch.from_numpy(ct.hu_a[context_ndx].astype(np.float32))
        # 设置数据值的上下限
        ct_t.clamp_(-1000, 1000)#获取对应的掩码数据
        pos_t = torch.from_numpy(ct.positive_mask[slice_ndx]).unsqueeze(0)

        return ct_t, pos_t, ct.series_uid, slice_ndx

class TrainingLuna2dSegmentationDataset(Luna2dSegmentationDataset):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.ratio_int = 2

    def __len__(self):
        return 300000

    def shuffleSamples(self):
        random.shuffle(self.candidateInfo_list)
        random.shuffle(self.pos_list)

    def __getitem__(self, ndx):
        candidateInfo_tup = self.pos_list[ndx % len(self.pos_list)]
        return self.getitem_trainingCrop(candidateInfo_tup)

    def getitem_trainingCrop(self, candidateInfo_tup):
        ct_a, pos_a, center_irc = getCtRawCandidate(
            candidateInfo_tup.series_uid,
            candidateInfo_tup.center_xyz,
            (7, 96, 96),
        )
        pos_a = pos_a[3:4]#同时裁剪CT和掩码数据
        row_offset = random.randrange(0,32)
        col_offset = random.randrange(0,32)
        ct_t = torch.from_numpy(ct_a[:, row_offset:row_offset+64,
                                     col_offset:col_offset+64]).to(torch.float32)
        pos_t = torch.from_numpy(pos_a[:, row_offset:row_offset+64,
                                       col_offset:col_offset+64]).to(torch.long)

        slice_ndx = center_irc.index        return ct_t, pos_t, candidateInfo_tup.series_uid, slice_ndx