使用MedNIST数据集进行医学图像分类教程

在本教程中，我们将基于MedNIST数据集介绍一个end-to-end训练和评估示例。

我们将按照以下步骤进行操作：

为训练和测试创建数据集
使用MONAI转换预处理数据
使用MONAI中的DenseNet进行分类
使用PyTorch程序训练模型
在测试数据集上进行评估

Setup environment

!python -c "import monai" || pip install -q "monai-weekly[pillow, tqdm]"
!python -c "import matplotlib" || pip install -q matplotlib
%matplotlib inline

Setup imports

import os
import shutil
import tempfile
import matplotlib.pyplot as plt
import PIL
import torch
import numpy as np
from sklearn.metrics import classification_report

from monai.apps import download_and_extract
from monai.config import print_config
from monai.data import decollate_batch, DataLoader
from monai.metrics import ROCAUCMetric
from monai.networks.nets import DenseNet121
from monai.transforms import (
    Activations,
    EnsureChannelFirst,
    AsDiscrete,
    Compose,
    LoadImage,
    RandFlip,
    RandRotate,
    RandZoom,
    ScaleIntensity,
)
from monai.utils import set_determinism

print_config()

设置数据目录Setup data directory

您可以使用MONAI_DATA_DIRECTORY环境变量指定一个目录。
这样可以保存结果并重用下载的数据。
如果未指定，则会使用临时目录。

directory = os.environ.get("MONAI_DATA_DIRECTORY")
root_dir = tempfile.mkdtemp() if directory is None else directory
print(root_dir)

下载数据集Download dataset

MedNIST数据集从TCIA、RSNA骨龄挑战和NIH胸部X射线数据集中收集而来。

该数据集由Bradley J. Erickson博士（Mayo Clinic放射学部门）慷慨提供，根据Creative Commons CC BY-SA 4.0许可证发布。

如果您使用MedNIST数据集，请注明数据来源。

resource = "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/MedNIST.tar.gz"
md5 = "0bc7306e7427e00ad1c5526a6677552d"

compressed_file = os.path.join(root_dir, "MedNIST.tar.gz")
data_dir = os.path.join(root_dir, "MedNIST")
if not os.path.exists(data_dir):
    download_and_extract(resource, compressed_file, root_dir, md5)

Set deterministic training for reproducibility

set_determinism(seed=0)

从数据集文件夹中读取图像文件名

首先，检查数据集文件并显示一些统计信息。
数据集中有6个文件夹：Hand、AbdomenCT、CXR、ChestCT、BreastMRI、HeadCT，
这些文件夹应该作为标签用于训练我们的分类模型。

class_names = sorted(x for x in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, x)))  #这行代码使用列表推导式获取数据集文件夹中的所有文件夹名称，并按字母顺序进行排序。这些文件夹名称将作为训练分类模型的标签。
num_class = len(class_names) #这行代码获取标签的数量，即数据集中的文件夹数量。
image_files = [
    [os.path.join(data_dir, class_names[i], x) for x in os.listdir(os.path.join(data_dir, class_names[i]))]
    for i in range(num_class)
]
num_each = [len(image_files[i]) for i in range(num_class)] #获取每个标签下的图像文件数量，并创建一个列表来保存这些数量。
image_files_list = []
image_class = []
for i in range(num_class):
    image_files_list.extend(image_files[i])
    image_class.extend([i] * num_each[i])
num_total = len(image_class)
image_width, image_height = PIL.Image.open(image_files_list[0]).size

print(f"Total image count: {num_total}")
print(f"Image dimensions: {image_width} x {image_height}")
print(f"Label names: {class_names}")
print(f"Label counts: {num_each}")

在数据集中随机选择一些图像进行可视化和检查。

plt.subplots(3, 3, figsize=(8, 8))
for i, k in enumerate(np.random.randint(num_total, size=9)):
    im = PIL.Image.open(image_files_list[k])
    arr = np.array(im)
    plt.subplot(3, 3, i + 1)
    plt.xlabel(class_names[image_class[k]])
    plt.imshow(arr, cmap="gray", vmin=0, vmax=255)
plt.tight_layout()
plt.show()

准备训练、验证和测试数据列表

从数据集中随机选择10%作为验证集，10%作为测试集。
以下这段代码将整个数据集划分为三个部分：训练数据、验证数据和测试数据

# 分别指定验证集和测试集的比例，例如这里都是10%。
val_frac = 0.1
test_frac = 0.1

length = len(image_files_list)
indices = np.arange(length) #为整个数据集生成一个索引数组。
np.random.shuffle(indices) #随机打乱索引数组。

test_split = int(test_frac * length) #计算测试集的数量
val_split = int(val_frac * length) + test_split #计算验证集的数量
test_indices = indices[:test_split] #从打乱后的索引数组中选取前test_split个作为测试集下标。
val_indices = indices[test_split:val_split] #从打乱后的索引数组中选取[test_split, val_split)区间作为验证集下标。
train_indices = indices[val_split:] #从打乱后的索引数组中选取后面的所有元素作为训练集下标。

train_x = [image_files_list[i] for i in train_indices] #根据划分好的下标，将所有的训练集图像的路径存储在一个列表中
train_y = [image_class[i] for i in train_indices] #根据划分好的下标，将所有的训练集图像的类别标签保存到一个列表中
val_x = [image_files_list[i] for i in val_indices] # 根据划分好的下标，将所有的验证集图像的路径存储在一个列表中
val_y = [image_class[i] for i in val_indices] #根据划分好的下标，将所有的验证集图像的类别标签保存到一个列表中
test_x = [image_files_list[i] for i in test_indices] #根据划分好的下标，将所有的测试集图像的路径存储在一个列表中
test_y = [image_class[i] for i in test_indices] #根据划分好的下标，将所有的测试集图像的类别标签保存到一个列表中

print(f"Training count: {len(train_x)}, Validation count: " f"{len(val_x)}, Test count: {len(test_x)}")

定义MONAI转换、数据集和数据加载器，以预处理数据。

这段代码定义了三个MONAI转换函数：train_transforms、val_transforms和y_pred_trans与y_trans。分别用于训练、验证和测试时对数据进行预处理，并将标签转换为模型可接受的格式。

train_transforms = Compose( #用于训练时对数据进行增强。具体来说，它包括以下的转换：
    [
        LoadImage(image_only=True), #加载图像，返回一个变换后的图像。
        EnsureChannelFirst(), #将每个样本的通道数放在第一维。
        ScaleIntensity(), #对每个样本的强度进行线性缩放，使其范围在[0,1]之间。
        RandRotate(range_x=np.pi / 12, prob=0.5, keep_size=True), #随机旋转每个样本，旋转的角度在[-π/12,π/12]之间，概率为0.5，并保持图像尺寸不变。
        RandFlip(spatial_axis=0, prob=0.5), #随机翻转每个样本，空间轴（x/y/z）为0，概率为0.5。
        RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5), #随机缩放每个样本，缩放范围在[0.9,1.1]之间，概率为0.5。
    ]
)

#val_transforms用于验证和测试时对数据进行预处理。这里只包括了以下的转换：
val_transforms =  Compose([LoadImage(image_only=True), EnsureChannelFirst(), ScaleIntensity()])
    #LoadImage(image_only=True)：加载图像，返回一个变换后的图像。
    #EnsureChannelFirst()：将每个样本的通道数放在第一维。
    #ScaleIntensity()：对每个样本的强度进行线性缩放，使其范围在[0,1]之间。
y_pred_trans = Compose([Activations(softmax=True)]) #转换输出的预测结果，用来返回模型的输出概率。这里包括了一个转换：Activations(softmax=True)：将输出值在softmax函数中进行标准化，返回模型的输出概率
y_trans = Compose([AsDiscrete(to_onehot=num_class)]) #转换模型标签。这里包括一个转换：将预测值转换为独热编码格式。

以下这段代码定义了三个MONAI数据集和数据加载器，用于训练、验证和测试

class MedNISTDataset(torch.utils.data.Dataset):
    def __init__(self, image_files, labels, transforms):
        self.image_files = image_files
        self.labels = labels
        self.transforms = transforms

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

    def __getitem__(self, index):
        return self.transforms(self.image_files[index]), self.labels[index]

train_ds = MedNISTDataset(train_x, train_y, train_transforms)
train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=10)

val_ds = MedNISTDataset(val_x, val_y, val_transforms)
val_loader = DataLoader(val_ds, batch_size=300, num_workers=10)

test_ds = MedNISTDataset(test_x, test_y, val_transforms)
test_loader = DataLoader(test_ds, batch_size=300, num_workers=10)

定义神经网络模型以及优化器

1. 设置学习率，以确定每个批次更新模型权重的速度。
2. 设置总的epoch数，由于我们使用了打乱和随机转换，因此每个epoch的训练数据都是不同的。

由于这仅是一个入门教程，所以我们只训练4个epochs。如果训练10个epochs，则模型在测试数据集上可以达到100％的准确率。

1. 使用MONAI的DenseNet模型，并将其移动到GPU设备上。该模型可以支持2D和3D分类任务。
2. 使用Adam优化器。

以下这段代码的作用是定义了设备类型，模型，损失函数，优化器，最大训练轮数以及一个用于计算AUC指标的ROCAUCMetric。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #根据是否可用GPU，选择性使用cuda设备或cpu设备
model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=num_class).to(device) #接下来，创建一个DenseNet121模型，该模型适用于2D图像数据，输入通道数为1，输出通道数为num_class（类别的数量）。将模型推送到设备上。

loss_function = torch.nn.CrossEntropyLoss() #定义了一个交叉熵损失函数，它在多类别分类任务中常用于衡量预测和真实标签之间的差异。

optimizer = torch.optim.Adam(model.parameters(), 1e-5) #使用Adam优化器来更新模型的参数，学习率为1e-5。

max_epochs = 4 #设置最大的训练轮数为4，意味着模型将从训练数据中遍历4次。

val_interval = 1 #表示每1个epoch结束后就会在验证集上进行一次模型评估。

auc_metric = ROCAUCMetric() #创建了一个ROCAUCMetric实例，用于计算AUC指标，该指标常用于二分类任务中评估模型的性能。

Model training

执行一个典型的PyTorch训练，它运行epoch循环和step循环，并在每个epoch结束时进行验证。如果在验证集上得到最佳准确度，则将保存模型的权重文件。

best_metric = -1 #保存最好的性能指标，初始化为-1
best_metric_epoch = -1 #保存达到最佳性能的epoch数，初始化为-1
epoch_loss_values = [] #保存每个epoch的平均损失值
metric_values = [] #保存每个epoch的模型性能指标

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}") #这是一个用来迭代每个训练轮次的循环。它打印出当前的训练轮次。
    model.train()  #这一行将模型设置为训练模式，这对于包含批量标准化和dropout等层很重要，以确保在训练过程中正确运行。
    epoch_loss = 0
    step = 0 #这里初始化了每个epoch的总损失和步数。
    for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data[0].to(device), batch_data[1].to(device) #这部分循环迭代训练数据集中的每个批次。inputs是输入数据，labels是对应的标签。
        optimizer.zero_grad() #这里将优化器中的梯度清零，为了进行下一步的反向传播。
        outputs = model(inputs) #将输入数据传递给模型，得到模型的输出。
        loss = loss_function(outputs, labels) #计算预测输出与真实标签之间的损失
        loss.backward() #计算损失对模型参数的梯度。
        optimizer.step() #通过按照已计算的梯度更新模型参数进行优化。
        epoch_loss += loss.item()
        print(f"{step}/{len(train_ds) // train_loader.batch_size}, " f"train_loss: {loss.item():.4f}") #将每个批次的损失加到 epoch 的总损失上，并打印每个批次的训练损失。
        epoch_len = len(train_ds) // train_loader.batch_size
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}") #计算并记录每个epoch的平均损失，并打印每个epoch的平均损失值。最后一个print语句会输出当前epoch的平均损失。
    #以上这段代码是一个简单的训练循环，用于在每个epoch中训练模型并计算损失。每个epoch结束后，会记录该epoch的平均损失值，并将其添加到epoch_loss_values列表中，以便后续绘制损失值的折线图。

#针对每个验证间隔（val_interval）对模型进行验证，并计算并打印出当前的AUC、准确率以及最佳AUC的信息。
    if (epoch + 1) % val_interval == 0:
        model.eval() #设置模型为评估模式，即不进行梯度计算和参数更新。
        with torch.no_grad(): #在此上下文中，不计算梯度，以节省内存和计算资源。
            y_pred = torch.tensor([], dtype=torch.float32, device=device)
            y = torch.tensor([], dtype=torch.long, device=device)
            #创建空的y_pred和y张量，用于保存模型的预测值和验证集的真实标签。
            for val_data in val_loader: #遍历验证数据加载器val_loader
                val_images, val_labels = (
                    val_data[0].to(device),
                    val_data[1].to(device),
                ) #使用val_data获取验证集中的图像和标签，使用to方法将它们移动到指定的计算设备上，通常是GPU
                y_pred = torch.cat([y_pred, model(val_images)], dim=0)
                y = torch.cat([y, val_labels], dim=0) 
                #使用模型对验证图像进行预测，将预测结果拼接到y_pred张量上，验证集的真实标签拼接到y张量上。
            y_onehot = [y_trans(i) for i in decollate_batch(y, detach=False)]
            y_pred_act = [y_pred_trans(i) for i in decollate_batch(y_pred)]
            #将存储真实标签的张量y进行解码和处理，将其转换为one-hot编码格式的y_onehot。将存储模型预测结果的张量y_pred进行处理和转换，将其转化为激活值y_pred_act
            auc_metric(y_pred_act, y_onehot)
            result = auc_metric.aggregate()
            auc_metric.reset()
            #使用用于计算AUC的auc_metric评价指标计算y_pred_act和y_onehot的AUC值，保存结果到result变量中，在计算前需要进行重置（reset）。
            del y_pred_act, y_onehot #释放之前创建的中间变量y_pred_act和y_onehot的内存空间。
            metric_values.append(result) #将当前的AUC值添加到metric_values列表中保存。
            acc_value = torch.eq(y_pred.argmax(dim=1), y)
            acc_metric = acc_value.sum().item() / len(acc_value)
            #计算预测准确率，即使用模型预测结果与真实标签进行比较，计算预测正确的样本数除以总样本数。
            if result > best_metric:
                best_metric = result
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(root_dir, "best_metric_model.pth"))
                print("saved new best metric model")
                #如果当前AUC值超过了历史最佳AUC值，则更新最佳AUC值和最佳AUC对应的epoch，保存当前模型状态到文件，并打印出保存模型的信息。
            print(
                f"current epoch: {epoch + 1} current AUC: {result:.4f}"
                f" current accuracy: {acc_metric:.4f}"
                f" best AUC: {best_metric:.4f}"
                f" at epoch: {best_metric_epoch}"
            )
            #打印当前epoch的信息，包括当前AUC、当前准确率、历史最佳AUC以及历史最佳AUC对应的epoch。

print(f"train completed, best_metric: {best_metric:.4f} " f"at epoch: {best_metric_epoch}")
#这句代码是在训练完成后打印训练结果的信息。它打印了训练完成后的最佳指标值best_metric和最佳指标值对应的epochbest_metric_epoch。

Plot the loss and metric绘制损失和指标的线图

plt.figure("train", (12, 6)) #创建一个名为train的图片窗口，并设定大小为(12, 6)。
plt.subplot(1, 2, 1) #使用plt.subplot方法创建一个1行2列的图表，选择第一个子图
plt.title("Epoch Average Loss")
x = [i + 1 for i in range(len(epoch_loss_values))] #x存储了1到epoch数目的列表
y = epoch_loss_values #y存储了每个epoch的平均损失值
plt.xlabel("epoch") #使用plt.xlabel为横坐标添加标签
plt.plot(x, y) #使用plt.plot方法绘制曲线图
plt.subplot(1, 2, 2)
plt.title("Val AUC") #在第二个子图中绘制每个epoch的AUC值。
x = [val_interval * (i + 1) for i in range(len(metric_values))] #x存储了验证间隔乘以1到epoch数目的列表
y = metric_values #y存储了每个epoch的AUC值
plt.xlabel("epoch") #使用plt.xlabel为横坐标添加标签
plt.plot(x, y) #使用plt.plot方法绘制曲线图
plt.show() #显示图表。这个调用要在所有绘图命令完成之后，且只需调用一次，才能将图表显示出来。

Evaluate the model on test dataset

在训练和验证完成后，我们已经得到了在验证集上表现最好的模型。现在我们需要在测试集上评估这个模型，以检查它是否具有鲁棒性并且没有过拟合。我们将使用这些预测结果生成分类报告。

model.load_state_dict(torch.load(os.path.join(root_dir, "best_metric_model.pth")))#该行代码加载通过在验证集上表现最好的模型的参数，也就是上文中通过torch.save()函数保存在文件名为best_metric_model.pth的文件中。具体来说，torch.load()函数被用于从文件中读取模型的状态字典，并通过load_state_dict()方法将其加载到模型中
model.eval()#该行代码将模型设置为评估模式，用于在测试集上对模型进行评估。
y_true = []
y_pred = []
#定义空列表用于存储测试集中的真实标签和预测标签，用于后续生成分类报告。
with torch.no_grad():
    for test_data in test_loader: #遍历测试集数据加载器test_loader中的测试数据。
        test_images, test_labels = (
            test_data[0].to(device),
            test_data[1].to(device),
        )
        #获取当前测试数据对应的图像数据和标签，并将它们放在GPU上（如果有），以便快速计算模型预测结果。
        pred = model(test_images).argmax(dim=1) #用当前测试数据传入模型中，得到模型预测结果，将结果中概率最大的类别作为预测标签pred
        for i in range(len(pred)):
            y_true.append(test_labels[i].item())
            y_pred.append(pred[i].item())
            #用于遍历当前测试数据得到的所有预测标签，将它们添加到存储真实标签和预测标签的列表y_true和y_pred中去
            #.item()方法用于将PyTorch张量中的值取出，转换为Python中的标量。

通过上面这段代码我们可以得到测试集的所有预测结果和对应的真实标签，以便后续对模型进行评估和生成分类报告。

print(classification_report(y_true, y_pred, target_names=class_names, digits=4)) 

classification_report是Scikit-learn库中的一个函数，用于生成分类报告，提供有关模型性能的详细分类指标。

y_true和y_pred是存储测试集中的真实标签和预测标签的列表。

target_names是目标类别的名称列表，用于标记分类指标的行和列。

digits表示显示的小数位数。

清理数据目录

如果使用了临时目录，则删除该目录。
在机器学习项目中，我们通常需要从原始数据中构建模型所需的数据集。在上下文中，“数据目录”是指保存原始数据和构建模型所需的派生数据的文件夹。

如果在载入数据时使用了一个临时目录，并希望在模型训练结束之后删除该目录，则可以使用这句话完成目录的清理工作。此举旨在防止不必要的磁盘空间占用，确保系统的稳定性和安全性。

if directory is None:
    shutil.rmtree(root_dir)