pytorch医学图片分割10层

import os,sys,cv2,torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image class EyeSegDataset(Dataset): def init(self, img_dir, mask_dir, transform=None): self.img_dir = img_dir self.mask_dir = mask_dir self.transform = transform self.images = os.listdir(img_dir) def __len__(self): return len(self.images) def __getitem__(self, idx): # 读取原始图片和标记图片 img_path = os.path.join(self.img_dir, self.images[idx]) mask_path = os.path.join(self.mask_dir, self.images[idx].split('.')[0] + '_mask.png') image = Image.open(img_path).convert('RGB') mask = Image.open(mask_path).convert('L') if self.transform: # 对原始图片和标记图片进行相同的变换 image = self.transform(image) mask = self.transform(mask) # 将标记图片转化为0-9的整数值 mask = torch.squeeze(mask) mask[mask == 255] = 0 return image, maskclass SegmentationModel(nn.Module): def init(self, num_classes=10): super(SegmentationModel, self).init() # 定义特征提取器 self.feature_extractor = nn.Sequential( nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2) )

#############

这里定义了一个特征提取器，它是由一些卷积层、批标准化层和ReLU非线性激活函数组成的序列#（nn.Sequential），用于从输入图像中提取重要的特征。具体来说，它包括四个卷积层和两个最大池化层。 #第一个卷积层(nn.Conv2d)的输入通道数为3（因为我们的图像是RGB三通道），输出通道数为64，kernel_size为3x3，stride为1，padding为1（保证输出和输入尺寸相同）。后接BatchNorm2d层进行批标准化，ReLU层作为激活函数。接着是一个与之类似的卷积层，其输入通道数为64，输出通道数为128。然后是一个最大池化层(nn.MaxPool2d)，kernel_size为2x2，stride为2，将特征图的尺寸缩小为原来的一半，同时增加了抗噪声能力。接下来是另外两个卷积层，其前两者输入输出通道数分别为128和256，而后两者输入输出通道数分别为256和512，kernel_size、stride和padding大小均相同。每个卷积层后都跟着一个BatchNorm2d层和ReLU激活函数。 ############################### # 定义分类器 self.classifier = nn.Sequential( nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(1024), nn.ReLU(inplace=True), nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.Conv2d(512, num_classes, kernel_size=1, stride=1), nn.ReLU(inplace=True) ) ###########################

这里定义了一个分类器，它将特征图转换为对每个像素点进行分类的概率分布。具体来说，它包括三个卷积层和ReLU非线性激活函数。最后一层卷积输出的通道数与我们要分类的类别数相同，kernel_size为1x1，stride也为1，不使用padding操作。 首先是一个Conv2d层(nn.Conv2d)，其输入通道数为512，输出通道数为1024，kernel_size为3x3，stride为1，padding为1。接着又跟着一个BatchNorm2d层和ReLU层作为激活函数。然后是另外一个类似的组合，其输入通道数为1024，输出通道数为512。最后是另一个只有一个输出通道的卷积层，通过ReLU激活函数得到最终的预测结果。 这个分类器采用了全卷积结构，没有采用全连接层，因为全连接层需要固定大小的输入，而卷积层则可以处理任意尺寸的输入，这样我们就可以使用任何大小的图像进行预测。同时，由于我们使用的是1x1的卷积核，因此在这一层中，每个卷积核所学习的信息都只与单个像素点相关联。这使得我们能够在每个像素点上计算类别的概率分布，从而实现像素级别的语义分割。 总之，这个分类器将把我们的特征图转化成一个10 通道的特征图，其中每个像素点代表了该像素属于哪个类别的概率。

########################### def forward(self, x): # 前向传播 out = self.feature_extractor(x) out = self.classifier(out) return out

#数据增强 transform = transforms.Compose([ transforms.Resize((256,256)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) 初始化模型 model = SegmentationModel(num_classes=10) 检查可用的GPU数量，并使用所有可用的GPU加速训练 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device.type == 'cuda': model = nn.DataParallel(model) model.to(device) 加载数据集并定义DataLoader img_dir = 'image_dir' mask_dir = 'mask_dir' dataset = EyeSegDataset(img_dir, mask_dir, transform=transform) dataloader = DataLoader(dataset, batch_size=16, shuffle=True) 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) 训练模型 epochs = 10 for epoch in range(epochs): for i, (images, labels) in enumerate(dataloader): images =images.to(device) labels = labels.long().to(device) # 将标签转化为长整型并送到GPU上 # 前向传播 outputs = model(images) loss = criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() # 每100个batch打印一次loss信息 if (i+1) % 100 == 0:  print('训练轮次 [{}/{}], 步骤 [{}/{}], 损失: {:.4f}'.format(epoch+1, epochs, i+1, len(dataloader), loss.item()))#保存模型权重 torch.save(model.state_dict(), 'model.ckpt')

在代码中，我们定义了一个名为EyeSegDataset的数据集类，用于读取原始图片和对应的标注图片，并将它们转换为PyTorch张量。然后定义了一个SegmentationModel类来实现语义分割模型。我们使用了nn.DataParallel函数来进行多GPU加速，以便更快地训练模型。接着使用transforms模块定义了数据增强操作，包括调整尺寸、转换为张量和标准化等。最后，我们使用DataLoader创建了一个数据加载器，它会从文件夹中读取样本批次，并自动将其转换为Tensor。 在训练过程中，我们使用了Adam优化器和交叉熵损失函数，同时使用GPU加速训练。在每个epoch结束后，我们使用torch.save()函数保存了模型权重，以备日后使用。

以下是一个简单的测试方法

import torch from torch.utils.data import DataLoader from torchvision.datasets import MNIST from torchvision.transforms import ToTensor from model import MyModel # 导入您自己定义的模型 # 加载MNIST数据集 test_dataset = MNIST(root='data/', train=False, transform=ToTensor()) test_dataloader = DataLoader(dataset=test_dataset, batch_size=64) # 加载并加载已训练好的模型参数 model = MyModel() model.load_state_dict(torch.load('my_model_params.pth')) # 在测试集上评估模型 model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_dataloader: outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('测试集准确率: {:.2f}%'.format(100 * correct / total))

以上内容为本人原创内容，转载或引用使需要标注来自转载的页面。