[人工智能]医学图像分割结果保存

由于医学图像的空间分辨率一般很高，一般训练集都是采用随机裁剪输入到网络进行训练，而对于验证集和测试集，则一般使用滑动窗口推理方法进行预测，滑动窗口方法一般分为在线切图和离线切图，离线切图比较劝退（大佬就可以忽略这句话），很那实现。在线切图在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上，则不需要这一步。