U-Net：大脑MRI海马体语义分割

相关技术：Python、Pytorch、U-Net

Step 01 数据集

​ 大脑MRI数据集源于Kaggle平台，链接如下：https://www.kaggle.com/sabermalek/mrihs

原数据集中包含100个病人的大脑MRI图像，存储在文件夹original中，每张MRI图像对应的海马体分割结果则存储在文件夹label中。

​ 从原数据集中随机抽取150张大脑MRI图像及其对应的分割结果图像，即为此次实验所采用的数据集。此次实验数据集获取方式如下：链接: https://pan.baidu.com/s/1FVsVG6XgHqF201fPlRkCYg 密码: 655b

Step 02 自定义Dataset

​ 在Pytorch中，一般通过DataLoader来完成数据的批量导入，而如果想将数据载入DataLoader，则需要继承Dataset类并实现自己的数据集。

​ 数据集的实现和关键代码的注释如下：

class MRIDataset(Dataset):
    def __init__(self, url, transform=None):
        self.imgs = glob.glob(os.path.join(url, 'img/*.jpg'))
        self.transform = transform

    def augment(self, img, flip):
        # flip = 1: 水平翻转
        # flip = 0: 垂直翻转
        # flip = -1: 同时进行水平翻转和垂直翻转
        img_flipped = cv2.flip(img, flip)
        return img_flipped

    def __getitem__(self, index):
        # 通过索引获取原图和标签的URL
        img_url = self.imgs[index]
        label_url = img_url.replace("img", "label").replace("ACPC", "L")
        # 通过cv2库读取图像
        image = cv2.imread(img_url)
        label = cv2.imread(label_url)
        # 转换为灰度图，3通道转换为1通道
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
        # 修正标签，将海马体部分像素设置为1，非海马体部分设置为0
        if label.max() > 1:
            label = label / 255
        label[label >= 0.5] = 1
        label[label < 0.5] = 0
        # 进行随机的数据增强
        flip = random.choice([-1, 0, 1, 2])
        if flip != 2:
            image = self.augment(image, flip)
            label = self.augment(label, flip)
        # 转换数据
        if self.transform:
            image = self.transform(image)
            label = self.transform(label)
        return image, label

    def __len__(self):
        return len(self.imgs)

Step 03 构建U-Net网络

​ U-Net网络在医疗领域进行语义分割时效果十分优异，其网络结构如下：

​ 根据图示网络结构，可以知道U-Net由4个下采样层、4个上采样层以及若干卷积层构成。除此之外，还可以发现卷积往往会连续进行两次。因此，可以考虑将两个卷积层封装为一个卷积单元。

• 卷积单元ConvUnit：

class ConvUnit(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvUnit, self).__init__()
        self.unit = nn.Sequential(
            # 保持图像大小
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            # 保持图像大小
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.unit(x)

• 下采样层DownSampling：

class DownSampling(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownSampling, self).__init__()
        self.layer = nn.Sequential(
            nn.MaxPool2d(2), # 最大池化，将图像大小变为原来的1/2
            ConvUnit(in_channels, out_channels) # 卷积单元
        )

    def forward(self, x):
        return self.layer(x)

• 上采样层UpSampling：

class UpSampling(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UpSampling, self).__init__()
        # 上采样层，将通道数变为1/2是为了在concat后保持通道数不变
        self.layer = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        # 卷积单元
        self.conv = ConvUnit(in_channels, out_channels)

    def forward(self, x, r):
        # 对x进行上采样，同时通道数减半
        x = self.layer(x)
        # 将x与r在通道维度连接，恢复原本通道数
        x = torch.cat((x, r), dim=1)
        return self.conv(x)

​ 定义好了卷积单元、下采样层以及上采样层，便可以搭建一个U-Net网络，其具体实现和关键代码注释如下：

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        # 输入层
        self.conv = ConvUnit(in_channels, 64)
        # 定义四个下采样层和四个上采样层
        self.D1 = DownSampling(64, 128)
        self.D2 = DownSampling(128, 256)
        self.D3 = DownSampling(256, 512)
        self.D4 = DownSampling(512, 1024)
        self.U1 = UpSampling(1024, 512)
        self.U2 = UpSampling(512, 256)
        self.U3 = UpSampling(256, 128)
        self.U4 = UpSampling(128, 64)
        # 输出层，输出图像像素保持在0～1以进行二分类
        self.out = nn.Sequential(
            nn.Conv2d(64, out_channels, kernel_size=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        # 输入和UNet的左半部分
        L1 = self.conv(x)
        L2 = self.D1(L1)
        L3 = self.D2(L2)
        L4 = self.D3(L3)
        Bottom = self.D4(L4)
        # UNet的右半部分和输出
        R1 = self.U1(Bottom, L4)
        R2 = self.U2(R1, L3)
        R3 = self.U3(R2, L2)
        R4 = self.U4(R3, L1)
        return self.out(R4)

Step 04 训练与预测

​ 首先定义数据加载函数如下：

def load_data(url="data", batch_size=4, shuffle=True, split=None):
    # 实例化自定义Dataset，加载数据集
    mri = MRIDataset(url, transform=transforms.Compose([
        transforms.ToTensor()
    ]))
    # 分割数据集
    if split is None:
        split = [0.7, 0.3]
    n_total = len(mri)
    n_train = int(split[0] * n_total)
    n_test = n_total - n_train
    train_set, test_set = random_split(mri, [n_train, n_test])
    # 将数据集封装入DataLoader
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=shuffle)
    return train_loader, test_loader

​ 定义训练函数：

def train(loader, lr=1e-3, epochs=10, model="model/model.pth", record="model/record.txt"):
    assert(isinstance(loader, DataLoader))
    net = UNet(1, 1) # 实例化网络
    optimizer = optim.Adam(net.parameters(), lr=lr) # 优化器
    criterion = nn.BCELoss() # 损失函数
    record = open(record, "w")
    # 开始训练
    for epoch in range(epochs):
        net.train()
        batch_id = 1
        total_batch = len(loader)
        for image, label in loader:
            logit = net(image)
            loss = criterion(logit, label.float()) # 计算损失
            # 更新模型参数
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if batch_id % 4 == 0:
                r = "[Epoch: {}/{}][Batch: {}/{}][Loss: {}]"\
                    .format(epoch + 1, epochs, batch_id, total_batch, loss.item())
                print(r)
                record.write(r + "\n")
            batch_id += 1
    record.write("[Training Finished]")
    # 存储网络参数
    torch.save(net.state_dict(), model)

​ 定义预测函数：

def predict(loader, model="model/model.pth", pred_dir="data/pred"):
    assert (isinstance(loader, DataLoader))
    # 初始化网络并加载权重
    net = UNet(1, 1)
    net.load_state_dict(torch.load(model))
    net.eval()
    # 预测
    order = 1
    for image, label in loader:
        pred = net(image)
        for i, p in zip(image.detach().numpy(), pred.detach().numpy()):
            # i: [1, w, h], p: [1, w, h]
            i, p = i[0], p[0]
            p[p >= 0.5] = 255
            p[p < 0.5] = 0
            # 存储预测结果
            cv2.imwrite("{}/{}.jpg".format(pred_dir, order), i * 255)
            cv2.imwrite("{}/{}_L.jpg".format(pred_dir, order), p * 255)
            order += 1

​ 训练30个epochs后效果如下：