医学图像分割结果保存

由于医学图像的空间分辨率一般很高,一般训练集都是采用随机裁剪输入到网络进行训练,而对于验证集和测试集,则一般使用滑动窗口推理方法进行预测,滑动窗口方法一般分为在线切图和离线切图,离线切图比较劝退(大佬就可以忽略这句话),很那实现。在线切图在monai中已经被实现了,调用代码如下:

from monai.inferers import sliding_window_inference

通过滑动窗口推理方法推理出来的分割结果在空间尺寸上(h,W,D三个维度)和原始图像是一致的。然后再使用monai.data中的Nitisaver类对分割结果进行保存,完整的代码如下:

from monai.inferers import sliding_window_inference
from monai.data import NiftiSaver

def validation_step(self, batch, batch_idx):
        self.eval()
        cfg = self.cfg
        x = batch['image']
        y = batch['label']
        y_hat = sliding_window_inference(x, roi_size=cfg.FinalShape, sw_batch_size=cfg.BatchSize,
                                         predictor=self.net,
                                         overlap=cfg.slid_window_overlap)
        loss, mean_dice = self.shared_step(y_hat=y_hat, y=y)
        # if mean_dice > 0.65:
        meta_dict = batch['image_meta_dict']  # 将meta_dict中的值转成cpu()向量,原来位于GPU上
        for k, v in meta_dict.items():
            if isinstance(v, torch.Tensor):
                meta_dict[k] = v.detach().cpu()
        saver = NiftiSaver(output_dir=self.cfg.ValidSegDir, mode="nearest", print_log=False)
        y_hat = torch.argmax(y_hat, dim=1, keepdim=True)
        saver.save_batch(y_hat, meta_dict)  # fixme 检查此处用法是否正确
        self.log('valid_mean_dice', mean_dice, prog_bar=True)
        self.log('valid_loss', loss, prog_bar=True)
        return {'loss': loss}

注意上述中将meta_dict转换成cpu张量的时候,由于我使用的是pytorch-lightning框架,这个框架在训练过程中Trainer会自动把所有张量加载到GPU上去,所以在此处我需要将其转放到CPU上。如果batch本身就在CPU上,则不需要这一步。

  • 4
    点赞
  • 13
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
以下是使用PyTorch实现U-Net医学图像分割的示例代码: ```python import torch import torch.nn as nn # 定义U-Net模型 class UNet(nn.Module): def __init__(self): super(UNet, self).__init__() # 定义卷积模块 self.conv1 = nn.Conv2d(1, 64, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1) self.conv5 = nn.Conv2d(128, 256, kernel_size=3, padding=1) self.conv6 = nn.Conv2d(256, 256, kernel_size=3, padding=1) self.conv7 = nn.Conv2d(256, 512, kernel_size=3, padding=1) self.conv8 = nn.Conv2d(512, 512, kernel_size=3, padding=1) self.conv9 = nn.Conv2d(512, 1024, kernel_size=3, padding=1) self.conv10 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1) # 定义反卷积模块 self.upconv1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) self.conv11 = nn.Conv2d(1024, 512, kernel_size=3, padding=1) self.conv12 = nn.Conv2d(512, 512, kernel_size=3, padding=1) self.upconv2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.conv13 = nn.Conv2d(512, 256, kernel_size=3, padding=1) self.conv14 = nn.Conv2d(256, 256, kernel_size=3, padding=1) self.upconv3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.conv15 = nn.Conv2d(256, 128, kernel_size=3, padding=1) self.conv16 = nn.Conv2d(128, 128, kernel_size=3, padding=1) self.upconv4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.conv17 = nn.Conv2d(128, 64, kernel_size=3, padding=1) self.conv18 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.conv19 = nn.Conv2d(64, 2, kernel_size=1) # 定义前向传播函数 def forward(self, x): # 编码器部分 x1 = nn.functional.relu(self.conv1(x)) x2 = nn.functional.relu(self.conv2(x1)) x3 = nn.functional.max_pool2d(x2, kernel_size=2, stride=2) x4 = nn.functional.relu(self.conv3(x3)) x5 = nn.functional.relu(self.conv4(x4)) x6 = nn.functional.max_pool2d(x5, kernel_size=2, stride=2) x7 = nn.functional.relu(self.conv5(x6)) x8 = nn.functional.relu(self.conv6(x7)) x9 = nn.functional.max_pool2d(x8, kernel_size=2, stride=2) x10 = nn.functional.relu(self.conv7(x9)) x11 = nn.functional.relu(self.conv8(x10)) x12 = nn.functional.max_pool2d(x11, kernel_size=2, stride=2) x13 = nn.functional.relu(self.conv9(x12)) x14 = nn.functional.relu(self.conv10(x13)) # 解码器部分 x15 = nn.functional.relu(self.upconv1(x14)) x15 = torch.cat((x15, x11), dim=1) x16 = nn.functional.relu(self.conv11(x15)) x17 = nn.functional.relu(self.conv12(x16)) x18 = nn.functional.relu(self.upconv2(x17)) x18 = torch.cat((x18, x8), dim=1) x19 = nn.functional.relu(self.conv13(x18)) x20 = nn.functional.relu(self.conv14(x19)) x21 = nn.functional.relu(self.upconv3(x20)) x21 = torch.cat((x21, x5), dim=1) x22 = nn.functional.relu(self.conv15(x21)) x23 = nn.functional.relu(self.conv16(x22)) x24 = nn.functional.relu(self.upconv4(x23)) x24 = torch.cat((x24, x2), dim=1) x25 = nn.functional.relu(self.conv17(x24)) x26 = nn.functional.relu(self.conv18(x25)) x27 = self.conv19(x26) return x27 # 定义数据加载器 class Dataset(torch.utils.data.Dataset): def __init__(self, images, labels): self.images = images self.labels = labels def __getitem__(self, index): image = self.images[index] label = self.labels[index] return image, label def __len__(self): return len(self.images) # 定义训练函数 def train(model, train_loader, criterion, optimizer, device): model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) epoch_loss = running_loss / len(train_loader.dataset) return epoch_loss # 定义测试函数 def test(model, test_loader, criterion, device): model.eval() running_loss = 0.0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) epoch_loss = running_loss / len(test_loader.dataset) return epoch_loss # 加载数据集 images_train = # 包含训练图像的numpy数组 labels_train = # 包含训练标签的numpy数组 images_test = # 包含测试图像的numpy数组 labels_test = # 包含测试标签的numpy数组 # 定义超参数 batch_size = 4 learning_rate = 0.001 num_epochs = 10 # 将数据转换为PyTorch张量 images_train = torch.from_numpy(images_train).float() labels_train = torch.from_numpy(labels_train).long() images_test = torch.from_numpy(images_test).float() labels_test = torch.from_numpy(labels_test).long() # 创建数据集 train_dataset = Dataset(images_train, labels_train) test_dataset = Dataset(images_test, labels_test) # 创建数据加载器 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 创建模型和优化器 model = UNet() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 将模型移动到GPU上 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) # 定义损失函数 criterion = nn.CrossEntropyLoss() # 训练模型 for epoch in range(num_epochs): train_loss = train(model, train_loader, criterion, optimizer, device) test_loss = test(model, test_loader, criterion, device) print('Epoch [{}/{}], Train Loss: {:.4f}, Test Loss: {:.4f}'.format(epoch+1, num_epochs, train_loss, test_loss)) # 保存模型 torch.save(model.state_dict(), 'unet.pth') ``` 请注意,上述示例代码仅包含U-Net模型的实现和训练代码,并且需要自己准备数据和标签。在实际应用中,还需要进行数据预处理、数据增强和模型评估等操作。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值