自学内容网 自学内容网

UNet 眼底血管分割实战教程


Blog’s 主页: 白乐天_ξ( ✿>◡❛)
🌈 个人Motto:他强任他强,清风拂山冈!
💫 欢迎来到我的学习笔记!

在医学影像分析领域,准确地分割眼底血管对于眼科疾病的诊断和治疗至关重要。本教程将详细介绍如何利用 UNet 进行眼底血管分割,包括云实例配置、数据集处理以及模型训练和测试。

一、云实例配置与启动

  1. 登录注册
    打开丹摩平台,进入登录界面进行注册并登录账号,为后续操作奠定基础。
  2. 配置 SSH 密钥对
    SSH 密钥对的配置能让远程登录服务器更加便捷。首先在本地.ssh 目录下输入ssh-keygen -o命令创建本地公钥,可自行设置文件名。生成的两个文件中,id_dsa.pub为所需的公钥文件。接着进入密钥对配置,将公钥文件内容复制到此处完成配置。
  3. 创建实例
    进入GPU云实例页面,根据需求选择合适的 GPU 型号和镜像。在配置过程中,务必记得选择之前创建的密钥对。确认所有选项无误后,点击立即创建,等待实例创建完成。
  4. 登录云实例
    实例创建完成后,复制访问链接。然后在任意一个 SSH 连接终端,如 VSCode 中进行云实例登录。登录成功后,可通过 nvidia-smitorch.cuda.is_available () 等命令简单验证功能是否正常。

二、云存储与数据集处理

  1. 文件存储的优势
    1. 可在不同实例间共享。
    2. 支持多点读写。
    3. 不受实例释放影响。
    4. 存储后端有多冗余副本,数据可靠性高。
  2. 文件存储的不足
    IO 性能一般。
  3. 使用建议
    1. 将重要数据或代码存放于文件存储中,以便所有实例共享,保障数据可靠性。
    2. 在训练时,对于需要高 IO 性能的数据(如训练数据),先将其拷贝到实例本地数据盘,从本地盘读取数据以获得更好的 IO 性能。
  4. 上传训练数据的方法
    使用 scp 工具,命令如scp -rP 35740./DRIVE-SEG-DATA root@cn-north-b.ssh.damodel.com:/root/workspace,其中35740为端口号,cn-north-b.ssh.damodel.com为远程地址,./DRIVE-SEG-DATA是本地数据集路径,/root/workspace是远程实例数据集路径,需根据实际情况替换这些参数。数据下载的命令与之类似。

三、云开发之眼底血管分割案例

3.1 案例背景

  1. 眼底结构:包含黄斑、视网膜和视网膜中央动静脉等。
  2. 眼底图像作用:在眼科医生诊断中重要。
  3. 深度学习对医学影像分割的影响:
    • 卷积神经网络能学习高级特征表示,实现更精确分割。
    • 训练后的模型泛化能力好,对未见过数据预测准确。
    • 支持端到端学习,简化开发流程,提高效率和准确性。
    • 能处理多模态数据,融合信息,提高分割准确性和全面性。

3.2 UNet 在眼底血管分割中的优势

  1. 编码器 - 解码器结构和跳跃连接:有效捕获不同尺度特征信息,效果好。
  2. 推理阶段全卷积网络结构:快速分割新眼底图像,提供实时性支持。

3.3 网络搭建

  1. 结合信息提高准确性
    U-Net 通过编解码器架构,有效结合局部和全局信息,提高分割准确性。
  2. 保留细节边缘信息
    跳跃连接结构有助于保留和恢复图像细节及边缘信息。
  3. 小样本表现好
    在小样本情况下表现优异,能充分利用有限数据有效训练。
  4. 广泛用于医学分割
    广泛应用于医学图像分割任务。
  5. 网络架构如下:
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

3.4 网络训练

基于 PyTorch 的神经网络训练流程可以分为以下步骤(不考虑前期数据准备和模型结构):

  1. 定义损失函数 根据任务类型选择合适的损失函数(loss function),如分类任务常用的交叉熵损失(Cross-Entropy Loss)或回归任务中的均方误差(Mean Square Error)。
  2. 选择优化器 选择合适的优化器(optimizer),如随机梯度下降(SGD)、Adam 或 RMSprop,并设置初始学习率及其它优化参数。
  3. 训练模型 在训练过程中,通过迭代训练数据集来调整模型参数。每个迭代周期称为一个epoch。对于每个 epoch,数据会被分成多个 batch,每个 batch 被输入到模型中进行前向传播、计算损失、反向传播更新梯度,并最终优化模型参数。
  4. 保存模型 当满足需求时,可以将训练好的模型保存下来,以便后续部署和使用。
def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    dataset = Dateset_Loader(data_path)
    per_epoch_num = len(dataset) / batch_size
    train_loader = torch.utils.data.DataLoader(dataset=dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    optimizer = optim.Adam(net.parameters(),lr=lr,betas=(0.9, 0.999),eps=1e-08, weight_decay=1e-08,amsgrad=False)
    criterion = nn.BCEWithLogitsLoss()
    best_loss = float('inf')
    loss_record = []
    with tqdm(total=epochs*per_epoch_num) as pbar:
        for epoch in range(epochs):
            net.train()
            for image, label in train_loader:
                optimizer.zero_grad()
                image = image.to(device=device, dtype=torch.float32)
                label = label.to(device=device, dtype=torch.float32)
                pred = net(image)
                loss = criterion(pred, label)
                pbar.set_description("Processing Epoch: {} Loss: {}".format(epoch+1, loss))
                if loss < best_loss:
                    best_loss = loss
                    torch.save(net.state_dict(), 'best_model.pth')
                loss.backward()
                optimizer.step()
                pbar.update(1)
            loss_record.append(loss.item())

    plt.figure()
    plt.plot([i+1 for i in range(0, len(loss_record))], loss_record)
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig('/root/shared-storage/results/training_loss.png')

按照这个脚本进行运行,可以看见进度:

训练损失函数如下,可以看出来已经收敛了:

3.5 模型测试

测试逻辑如下所示,主要是计算 IoU 指标

def cal_miou(test_dir="/root/workspace/DRIVE-SEG-DATA/Test_Images",
             pred_dir="/root/workspace/DRIVE-SEG-DATA/results", gt_dir="/root/workspace/DRIVE-SEG-DATA/Test_Labels",
             model_path='best_model_drive.pth'):
    name_classes = ["background", "vein"]
    num_classes = len(name_classes)
    if not os.path.exists(pred_dir):
        os.makedirs(pred_dir)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    net = UNet(n_channels=1, n_classes=1)
    net.to(device=device)
    net.load_state_dict(torch.load(model_path, map_location=device))
    net.eval()

    img_names = os.listdir(test_dir)
    image_ids = [image_name.split(".")[0] for image_name in img_names]

    time.sleep(1)
    for image_id in tqdm(image_ids):
        image_path = os.path.join(test_dir, image_id + ".png")
        img = cv2.imread(image_path)
        origin_shape = img.shape
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = cv2.resize(img, (512, 512))
        img = img.reshape(1, 1, img.shape[0], img.shape[1])
        img_tensor = torch.from_numpy(img)
        img_tensor = img_tensor.to(device=device, dtype=torch.float32)
        pred = net(img_tensor)
        pred = np.array(pred.data.cpu()[0])[0]
        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0
        pred = cv2.resize(pred, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_NEAREST)
        cv2.imwrite(os.path.join(pred_dir, image_id + ".png"), pred)

    hist, IoUs, PA_Recall, Precision = compute_mIoU_gray(gt_dir, pred_dir, image_ids, num_classes, name_classes)
    miou_out_path = "/root/shared-storage/results/"
    show_results(miou_out_path, hist, IoUs, PA_Recall, Precision, name_classes)

模型保存的时候保存到共享存储路径/root/shared-storage,其他实例可以直接从共享存储中获取训练后的模型:


原文地址:https://blog.csdn.net/lusanjiu/article/details/142457825

免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!