使用U-Net进行乳腺癌图像分割

使用U-Net进行乳腺癌图像分割

1. 数据准备

首先，我们需要准备数据。在这个例子中，我们的数据集包含两类图像：原始图像和相应的mask图像。原始图像是乳腺癌组织的图片，mask图像是与原始图像对应的分割标签。我们需要将这些图像预处理，以便能够输入到我们的神经网络模型中。

1.1 数据预处理

我们使用了Albumentations库进行数据增强，对图像进行旋转、翻转、亮度对比度调整等操作。同时，我们将数据集划分为训练集、验证集和测试集。注意，Albumentations库需要通过pip进行安装，具体安装为：

pip install Albumentations

代码如下：

import albumentations as A
from albumentations.pytorch import ToTensorV2
from torch.utils.data import DataLoader

train_transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Rotate(limit=30, p=0.3),
    A.RandomResizedCrop(height=256, width=256, scale=(0.8, 1.0), p=0.2),
    A.Normalize(),
    ToTensorV2()
])

val_transform = A.Compose([
    A.Normalize(),
    ToTensorV2()
])

test_transform = A.Compose([
    A.Normalize(),
    ToTensorV2()
])

1.2 自定义数据集

我们创建了一个自定义数据集类，该类负责从磁盘读取图像，应用预处理操作并返回图像和对应的mask。

代码如下：

from torchvision import transforms

class BreastCancerSegmentationDataset(Dataset):
    """
    乳腺癌分割数据集
    """
    def __init__(self, img_dir, mask_dir, transform=None, one_hot_encode=True, target_size=(256, 256)):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.one_hot_encode = one_hot_encode
        self.target_size = target_size
        self.img_filenames = os.listdir(img_dir)

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

    def __getitem__(self, index):
        # 根据文件存放方式设置os，便于建立原始图像与mask图像的联系
        img_name = self.img_filenames[index]
        img_path = os.path.join(self.img_dir, img_name)
        mask_path = os.path.join(self.mask_dir, img_name[:-4] + '_mask'+img_name[-4:])
        # Skip .ipynb_checkpoints files
        if img_path.endswith(".ipynb_checkpoints") or mask_path.endswith(".ipynb_checkpoints"):
            return self.__getitem__((index + 1) % len(self))
        image = cv2.imread(img_path, cv2.IMREAD_COLOR)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        # 适应度处理，检查是否将图像添加入os
        if image is None:
            raise FileNotFoundError(f"Image not found at {img_path}")
        if mask is None:
            raise FileNotFoundError(f"Mask not found at {mask_path}")

        # 将 image 和 mask 的图像进行强制转化，转化成同样大小
        image = cv2.resize(image, self.target_size, interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, self.target_size, interpolation=cv2.INTER_NEAREST)

        # 将mask进行独热处理
        if self.one_hot_encode:
            mask = one_hot_encode(mask, num_classes=3)

        # 定义transform构架，为数据增强做准备
        if self.transform:
            augmented = self.transform(image=image, mask=mask)
            image = augmented['image']
            mask = augmented['mask']

        # 将mask的属性值转化为float类型
        mask = np.asarray(mask)   # 转换为NumPy数组
        mask = mask.astype(np.float32)   # 变换dtype
        mask = torch.from_numpy(mask) # 转换为Tensor
        
        return image, mask

2. 构建U-Net模型

我们选择使用预训练的FCN-ResNet50作为我们的基本模型。我们创建了一个名为UNetTrainer的类，用于处理训练、评估和测试过程。

代码如下：

class UNetTrainer:
    """实际上我们使用的是一个全卷积网络（FCN）的ResNet50实现，而不是U-Net"""
    def __init__(self, num_classes=3, lr=1e-4):
        # 我们使用的fcn_resnet50是U-Net模型的变体,其中编码器部分初始化了ResNet50的权重。这可以加速模型训练和提高最终性能。但解码器部分仍需要我们从零训练
        self.model = fcn_resnet50(pretrained=False, num_classes=num_classes)
        # 损失函数使用交叉熵损失函数
        self.criterion = nn.CrossEntropyLoss()
        # 使用adam优化器
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
        # 设置GPU为训练设备
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

3. 模型训练和评估

我们使用交叉熵损失作为损失函数，并使用Adam优化器进行模型训练。在训练过程中，我们使用训练集进行训练，并在每个epoch结束时使用验证集对模型进行评估。我们使用IoU（Intersection over Union）和F1分数作为性能指标。注意，train、evaluate与test都是UNetTrainer的成员函数。

代码如下：

def train(self, num_epochs, train_loader, val_loader):
    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}/{num_epochs}")
        print("-" * 10)

        start_time = time.time()

        self.evaluate(epoch, train_loader, "train")
        self.evaluate(epoch, val_loader, "val")

        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"Epoch time: {elapsed_time:.4f}s")

    print("Training complete")

def evaluate(self, epoch, dataloader, phase):
    if phase == "train":
        self.model.train()
    else:
        self.model.eval()

    running_loss = 0.0
    running_iou = 0.0
    running_f1_score = 0.0

    for images, masks in dataloader:
        images = images.to(self.device)
        masks = masks.to(self.device)

        masks = torch.mean(masks, dim=3, keepdim=False).long()

        self.optimizer.zero_grad()

        with torch.set_grad_enabled(phase == "train"):
            outputs = self.model(images)['out']
            preds = torch.argmax(outputs, dim=1)
            loss = self.criterion(outputs, masks)

            if phase == "train":
                loss.backward()
                self.optimizer.step()

        running_loss += loss.item() * images.size(0)
        running_iou += self.calculate_iou(preds, masks.data)
        running_f1_score += f1_score(masks.cpu().numpy().ravel(), preds.cpu().numpy().ravel(), average="macro")

    epoch_loss = running_loss / len(dataloader.dataset)
    epoch_iou = running_iou / len(dataloader)
    epoch_f1_score = running_f1_score / len(dataloader)

    print(f"{phase} Loss: {epoch_loss:.4f} IoU: {epoch_iou:.4f} F1: {epoch_f1_score:.4f}")

4. 模型测试

在训练完成后，我们将训练集和验证集合并，并在这个组合数据集上重新训练模型。然后，我们使用测试集对最终模型进行评估。

代码如下：

def test(self, train_loader, val_loader, test_loader):
    # Combine train and val datasets
    combined_dataset = torch.utils.data.ConcatDataset([train_loader.dataset, val_loader.dataset])
    combined_loader = torch.utils.data.DataLoader(combined_dataset, batch_size=train_loader.batch_size, shuffle=True)

    # Retrain the model on the combined dataset
    print("Retraining the model on the combined dataset")
    self.train(20, combined_loader, test_loader)

    # Evaluate the model on the test dataset
    print("Evaluating the model on the test dataset")
    self.evaluate(0, test_loader, "test")

5. 改进模型性能的建议

为了进一步提高模型性能和训练效率，我提供了一些建议：

1. 早停法（Early Stopping）：在训练过程中，如果验证集的损失长时间没有明显改善，可以提前停止训练。这可以防止模型过拟合，并节省训练时间。
2. 学习率调整策略：使用学习率调整策略，如学习率衰减或者周期性学习率，可以帮助模型更快地收敛，并在一定程度上提高最终性能。
3. 数据增强策略：尝试使用更多的数据增强方法，如随机裁剪、弹性变形等，以提高模型的泛化能力。
4. 正则化方法：添加正则化项（如L1、L2正则化或Dropout）可以帮助防止模型过拟合，并提高模型的泛化能力。
5. 模型结构调整：尝试使用其他更先进的网络结构，如Attention U-Net或DeepLabv3等，来改进模型性能。
6. 交叉验证：使用k折交叉验证来评估模型性能。这可以提供一个更稳定的性能评估指标，并有助于避免过拟合。
7. 模型集成：将多个模型的预测结果进行组合，可以进一步提高模型性能。这可以通过简单的平均法、投票法或更复杂的集成策略来实现。

将这些方法应用于模型训练过程中，可以帮助我们进一步提高乳腺癌图像分割任务的性能。