UNet 眼底血管分割实战教程

最新推荐文章于 2025-04-18 16:31:42 发布

白乐天_ξ( ✿＞◡❛)

最新推荐文章于 2025-04-18 16:31:42 发布

阅读量1.7k

点赞数 26

分类专栏：活动文章文章标签：人工智能交互人机交互

本文链接：https://blog.youkuaiyun.com/lusanjiu/article/details/142457825

版权

活动文章专栏收录该内容

9 篇文章

订阅专栏

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

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

一、云实例配置与启动

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

二、云存储与数据集处理

文件存储的优势
1. 可在不同实例间共享。
2. 支持多点读写。
3. 不受实例释放影响。
4. 存储后端有多冗余副本，数据可靠性高。
文件存储的不足
IO 性能一般。
使用建议
1. 将重要数据或代码存放于文件存储中，以便所有实例共享，保障数据可靠性。
2. 在训练时，对于需要高 IO 性能的数据（如训练数据），先将其拷贝到实例本地数据盘，从本地盘读取数据以获得更好的 IO 性能。
上传训练数据的方法
使用 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 案例背景

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

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

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

3.3 网络搭建

结合信息提高准确性
U-Net 通过编解码器架构，有效结合局部和全局信息，提高分割准确性。
保留细节边缘信息
跳跃连接结构有助于保留和恢复图像细节及边缘信息。
小样本表现好
在小样本情况下表现优异，能充分利用有限数据有效训练。
广泛用于医学分割
广泛应用于医学图像分割任务。
网络架构如下：

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 的神经网络训练流程可以分为以下步骤(不考虑前期数据准备和模型结构)：

定义损失函数根据任务类型选择合适的损失函数（loss function），如分类任务常用的交叉熵损失（Cross-Entropy Loss）或回归任务中的均方误差（Mean Square Error）。
选择优化器选择合适的优化器（optimizer），如随机梯度下降（SGD）、Adam 或 RMSprop，并设置初始学习率及其它优化参数。
训练模型在训练过程中，通过迭代训练数据集来调整模型参数。每个迭代周期称为一个epoch。对于每个 epoch，数据会被分成多个 batch，每个 batch 被输入到模型中进行前向传播、计算损失、反向传播更新梯度，并最终优化模型参数。
保存模型当满足需求时，可以将训练好的模型保存下来，以便后续部署和使用。

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)