您现在的位置是:首页 >学无止境 >DRIVE视网膜血管分割——基于像素点分割(BP网络和CNN网络)网站首页学无止境

DRIVE视网膜血管分割——基于像素点分割(BP网络和CNN网络)

Chaoy6565 2024-06-17 10:43:26
简介DRIVE视网膜血管分割——基于像素点分割(BP网络和CNN网络)


前言

视网膜血管分割是一种计算机视觉任务,旨在从眼底图像中分割出视网膜图像的血管。这个任务对于糖尿病视网膜病变等眼部疾病的早期诊断和治疗至关重要。视网膜血管分割的方法主要有两种:基于规则的方法和基于深度学习的方法。基于规则的方法通常需要手工设计特征提取器和分类器,而基于深度学习的方法则利用深度神经网络来自动地学习特征并进行分类。近年来,基于深度学习的方法在视网膜血管分割中取得了许多优秀的成果,成为当前的主流方法。

像素点分割和语义分割

基于像素的分割是指将图像中的每一个像素都进行分类,通常会生成一个与原始图像大小相同的分割结果图,每一个像素点都被标记为属于哪一类别。这种方法适用于需要对图像进行详细分类的任务,例如边缘检测或者物体检测。而语义分割则是将图像中的每一个像素都赋予一个语义标签,即将图像中每一个像素点都标注为属于哪一个物体类别或者场景类别。这种方法适用于需要对图像进行高层次理解的任务,例如自动驾驶中的路面识别、物体分割等。
总体来说,基于像素的分割和语义分割都是图像分割技术的一种,但是它们的目标和应用场景有所不同。

DRIVE数据集

DRIVE(Digital Retinal Images for Vessel Extraction)是一个常用的眼底图像数据集,用于研究视网膜血管分割算法。该数据集包含了40个具有挑战性的视网膜图像,其中20个用于训练,另外20个用于测试。每张图像的大小为565×584像素,并且包含手动标注的视网膜血管。这些标注提供了血管的位置和形状信息,可以用于训练和评估血管分割算法。DRIVE数据集在医学图像处理领域被广泛使用,是许多血管分割算法的基准数据集之一。


一、BP网络分割

代码注释写的很详细,这里就不过多介绍了。

1、数据预处理

这里的工作主要是把DRIVE数据集的训练和测试数据读进来,然后把图片灰度化、对图片进行填充、根像素点切片(9*9)、归一化等。

import os
import cv2
import json
import torch
from PIL import Image
import numpy as np
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

class DriveDataset(Dataset):
    def __init__(self, root: str, train: bool, size=9):
        super(DriveDataset, self).__init__()
        self.size = size
        self.flag = "training" if train else "test"             # 选择训练还是测试
        data_root = os.path.join(root, "DRIVE", self.flag)      # 数据集路径拼接
        assert os.path.exists(data_root), f"path '{data_root}' does not exists."  # 判断路径是否存在

        img_names = [i for i in os.listdir(os.path.join(data_root, "images")) if i.endswith(".tif")] # 遍历images路径下的图片名
        self.img_list = [os.path.join(data_root, "images", i) for i in img_names]                    # images图片路径列表
        self.manual = [os.path.join(data_root, "1st_manual", i.split("_")[0] + "_manual1.gif")       # 1st_manual图片路径列表
                       for i in img_names]
        # check files
        for i in self.manual:     # 判断1st_manual下的图片是否存在
            if os.path.exists(i) is False:
                raise FileNotFoundError(f"file {i} does not exists.")

        data = []    # 总数据数组
        target = []  # 标签图片数组
        for image_filename, label_filename in zip(self.img_list, self.manual):  # 遍历每一张图片和对应的标签图片
            image = Image.open(image_filename).convert('L') # 读取图片、转灰度
            label = Image.open(label_filename).convert('L') # 读取标签图片

            image = np.expand_dims(np.expand_dims(image, axis=0), axis=0)   # 扩维
            label = np.expand_dims(label, axis=0) / 255                     # 扩维 转0 1数组
            data.append(image)      # 把所有图片存一个list
            target.append(label)    # 把所有标签图片存一个list

        self.data = torch.Tensor(np.concatenate(data, axis=0)).float()  # 转矩阵 转tensor 转float
        self.unfold = nn.Unfold(kernel_size=(self.size, self.size), dilation=1, padding=self.size // 2, stride=(1, 1)) #  定义切片函数
        self.data = self.unfold(self.data)     # 切片
        self.target = torch.Tensor(np.concatenate(target, axis=0)).long().view(-1,584*565)  # 转矩阵 转tensor 转float 拉平


    def __getitem__(self, idx):
        # 第 i 张
        i = idx//329960
        # 第 j 个索引
        j = idx%329960

        data = self.data[i,:,j]     # 取相应的片
        target = self.target[i,j]   # 取相应片的标签二值化值
        return data, target

    def __len__(self):
        return 584*565*20  # 索引总长度

    @staticmethod
    def collate_fn(batch):   # 打包函数
        images, targets = list(zip(*batch))
        batched_imgs = cat_list(images, fill_value=0)
        batched_targets = cat_list(targets, fill_value=255)
        return batched_imgs, batched_targets

if __name__ == "__main__":  # 测试dataset 显示函数
    dataset = DriveDataset(root=".", train = True)
    loader = DataLoader(dataset, batch_size=584, shuffle=False)
    print(len(dataset))
    for i, img in enumerate(loader):
        print(i, img[0].shape, img[1].shape)   # 584 81  ,584

2、构建BP网络

from torch import nn
import torch

class Net(nn.Module):
    def __init__(self, batch_size, in_chan, in_size, outs):   # 584 1 9 2
        super(Net, self).__init__()
        self.bs = batch_size
        self.size = in_size * in_size # 9*9
        self.layer1 = nn.Sequential(nn.Linear(in_chan * self.size, 64 * self.size),
                                    nn.BatchNorm1d(64 * self.size),
                                    nn.ReLU(True))
        self.layer2 = nn.Sequential(nn.Linear(64 * self.size, 64),
                                    nn.BatchNorm1d(64),
                                    nn.ReLU(True))
        self.layer3 = nn.Sequential(nn.Linear(64, 10),
                                    nn.ReLU(True))
        self.layer4 = nn.Sequential(nn.Linear(10, outs),
                                    nn.Softmax(dim=1))

    def forward(self, x):
        # 584*81 -> 584*5184
        x = self.layer1(x)
        # 584*5184 -> 584*64
        x = self.layer2(x)
        # 584*64 -> 584*10
        x = self.layer3(x)
        # 584*10 -> 584*2
        x = self.layer4(x)
        return x

3、训练代码

#-----------------------------------------------#
# 硬件
# 显卡:NVIDIA GeForce 940MX
# 软件
# CUDA Version: 11.7
# python3.8
# 其他安装包查看requirements.txt
#-----------------------------------------------#

import torch
import os
from dataset import DriveDataset
from torch.utils.data import DataLoader
from torch import nn, optim
from BpNet import Net
from tqdm import tqdm
#os.environ['CUDA_VISIBLE_DEVICES'] = '0'

batchsize = 584  # batchsize
size = 9         # 切片框的大小
eporch = 1       # eporch
# 数据集加载
train_dataset = DriveDataset(root=".", train = True, size=size)   # 数据集读取
train_loader = DataLoader(dataset=train_dataset, batch_size=batchsize, shuffle=True,drop_last = True)  # 打包

# 模型
model = Net(batchsize, 1, size, 2)

# 设备设定
device=torch.device("cuda:0"if torch.cuda.is_available()else"cpu")
model = model.to(device)

# 构造损失函数和优化函数
# 交叉熵损失 单标签多分类 标签是类别的序号
criterion = torch.nn.CrossEntropyLoss()
# 优化
optimizer = torch.optim.SGD(model.parameters(), lr=0.02, momentum=0.9)

for i in range(eporch):                   # eporch轮数
    print("第{}轮训练".format(i+1))
    model.train()  # 启用 batch normalization 和 dropout层
    for index, data in enumerate(tqdm(train_loader)):  # 遍历每一组打包好的数据
        img, label = data
        img, label = img.to(device), label.to(device)  #数据上传gpu
        out = model(img)              # 训练
        loss = criterion(out, label)  # 计算损失
        optimizer.zero_grad()         # 梯度清零
        loss.backward()               # 梯度反向传播
        optimizer.step()              # 更新网络参数

# 保存模型和权值
torch.save(model, 'weight/drive_bp.pth')

训练完后模型和权重文件保存在了weight文件夹中。

测试代码

# ---------------------------------------- #
# 测试
# 运行后测试图片在predict文件夹下
# ---------------------------------------- #
import torch
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from dataset import DriveDataset
import cv2
import numpy as np

eporch = 1
size = 9            # 切片框的大小
batchsize = 584     # batchsize

# 测试集读取
test_dataset = DriveDataset('.', train = False,size =size )
test_loader = DataLoader(dataset=test_dataset, batch_size=batchsize, shuffle=False,drop_last = True) # 打包

# 模型加载
model = torch.load("weight/drive_bp_10.pth")
# 设备选择
device=torch.device("cuda:0"if torch.cuda.is_available()else"cpu")#使用GPU进行计算
model.to(device) # 上传模型

# 损失
criterion = torch.nn.CrossEntropyLoss()

test_loss = 0  # 损失
acc = 0   # 准确率点计数值
model.eval()  # 保证 BN 层的均值和方差不变。
pre_img = torch.tensor([]).to(device)

with torch.no_grad(): # 不求梯度
    for index, data in enumerate(tqdm(test_loader)): # 遍历每一个打包的测试数据
        img, label = data
        img, label = img.to(device), label.to(device)   # 上传数据到gpu
        out = model(img)                                # 数据送入模型
        loss = criterion(out, label)                    # 计算损失
        test_loss += loss.item()                        # 损失累加
        acc += (out.argmax(1) == label).sum()           # 累加预测准确像素点的个数 ,out.argmax(1):返回out维度1里最大数的索引

        # 测试图片显示
        pre_img = torch.cat((pre_img, out.argmax(1)), 0) # 预测点收集
        if (index+1)%565==0:       # 判断是否到一张图片
            pre_img = (pre_img.reshape(584,565)*255).cpu().numpy()  # 转numpy
            # 存图片到本地
            cv2.imwrite("predict/test_e{eporch}_{size}_{num}.jpg".format(eporch=eporch,size=size,num=int((index+1)/565)), pre_img)  # 将图片保存为test.jpg
            pre_img = torch.tensor([]).to(device)  # 清空数组

print("test_loss:{}".format(test_loss / len(test_dataset))) # 打印loss
print("acc%: {:.6}".format(acc / len(test_dataset)))        # 打印准确率

运行test.py文件后,预测图片保存在了predict文件夹里。

二、CNN网络分割

cnn的代码与bp网络的类似,这里就不显示了。代码都在下面的连链接里。

四、遇到的Bug及学习小计

bug1——batchsize设置问题:

当batchsize设置不是584时运行会报错。这里主要是因为DataLoader送数据的原因。一张图片一共有329960个9*9的框,当batchsize为584时,DataLoader索引为565
当batchsize改变时有可能出现除不尽的情况 这里就要设置DataLoader设否舍弃最后一组数据,如果舍弃的话,运行是没有问题的但显示预测图片时就会乱码(这里主要是舍弃了图片的最后一组数据导致图片显示复原缺失)

bug2——设置多gpu运算

目前还没有搞定,网上说model = nn.DataParallel(model).cuda()可以将模型对象转变为多GPU并行运算的模型,但实际测试还是单卡在跑
数据集确实划分了

bug3

tensor 与 numpy格式转换的小问题 这里就不叙述了太多了 还有扩维 reshape 模型保存 图片保存 图片处理包 cv2 与Image等

bug4 分割训练完的结果进行预测,却只有背景

原因一:图像预处理时,压缩程度过大,图像变得很小,要识别的区域被挤压掉。 解决方法:修改resize的大小。
原因二:模型层数过多,网络架构过于庞大。 解决方法:降低层数,更换网络。
原因三:训练样本少。 解决方法:增加训练样本。
这里我删减了一些网络的线性层

tips

model.train():
在使用 pytorch 构建神经网络的时候,训练过程中会在程序上方添加一句model.train(),作用是 启用 batch normalization 和 dropout 。
如果模型中有BN层(Batch Normalization)和 Dropout ,需要在 训练时 添加 model.train()。
model.train() 是保证 BN 层能够用到 每一批数据 的均值和方差。对于 Dropout,model.train() 是 随机取一部分 网络连接来训练更新参数。
model.eval()
model.eval()的作用是 不启用 Batch Normalization 和 Dropout。
如果模型中有 BN 层(Batch Normalization)和 Dropout,在 测试时 添加 model.eval()。
model.eval() 是保证 BN 层能够用 全部训练数据 的均值和方差,即测试过程中要保证 BN 层的均值和方差不变。对于 Dropout,model.eval() 是利用到了 所有 网络连接,即不进行随机舍弃神经元。
nn.BatchNorm1d 是 PyTorch 中的一个模块,用于实现一维输入的 Batch Normalization(批归一化)操作。Batch Normalization 是一种通过规范化神经网络的激活函数输入以加速训练的技术。在每个 mini-batch 中,Batch Normalization 对每个输入特征进行标准化,使得每个特征的均值为 0,标准差为 1。
这样的标准化有助于缓解梯度消失和梯度爆炸问题,提高模型的训练速度和稳定性。nn.BatchNorm1d 可以在神经网络的任何一层中使用,通常放在激活函数之前。
nn.ReLU(True) : 表示对输入进行ReLU激活函数操作,并且inplace参数为True,表示直接对原始输入进行修改。ReLU是一种常用的激活函数,可以将负数输入变为0,而正数则保持不变。

Dropout层是一种常用的正则化技术,其作用为在神经网络中随机丢弃一部分神经元,从而减少过拟合的风险,提高模型的泛化能力。
通常,dropout层会随机将一定比例的神经元的输出设置为0,这些被丢弃的神经元不参与当前训练迭代的前向传播和反向传播。
具体地说,dropout层可以在每个训练迭代时随机丢弃一部分神经元,使得模型不会过度依赖于任何一个神经元的输出,从而减少过拟合风险。
在测试阶段,dropout层不起作用,所有神经元的输出都会被保留,以便对输入进行预测。因此,dropout层可以有效地提高神经网络的泛化能力,从而提高模型的性能。

nn.CrossEntropyLoss() 交叉熵损失函数 参考:http://t.csdn.cn/1mgdY

总代码和数据集

链接:https://download.csdn.net/download/weixin_45464524/87796822


风语者!平时喜欢研究各种技术,目前在从事后端开发工作,热爱生活、热爱工作。