Save the results with nii.gz in 3D segmentation：SaveImage Error

[jjhhan]

code:

from init import Options
from Network.networks import build_net, update_learning_rate, build_UNETR, build_unet, build_vnet, build_swinUNet, build_HRNet
import logging
import sys
from glob import glob
import numpy as np
import torch
import nibabel as nib
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import monai
from monai.data import list_data_collate, decollate_batch, CacheDataset
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.transforms import (EnsureType, Compose, LoadImaged, AddChanneld, AsDiscrete, RandGaussianSmoothd,
CropForegroundd, SpatialPadd,
ScaleIntensityd, ToTensord, RandSpatialCropd, Rand3DElasticd, RandAffined, Spacingd,
RandShiftIntensityd,
RandGaussianNoised, RandAdjustContrastd, NormalizeIntensityd, RandFlipd, SaveImage,
EnsureChannelFirstd)
from monai.visualize import plot_2d_or_3d_image
from tqdm import tqdm
import glob
from time import time
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

def main():
opt = Options().parse()
# monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
device = torch.device("cuda")
train_images = sorted(
glob.glob(os.path.join(opt.images_folder, ".nii.gz")))
train_labels = sorted(
glob.glob(os.path.join(opt.labels_folder, ".nii.gz")))
data_dicts = [
{"image": image_name, "label": label_name}
for image_name, label_name in zip(train_images, train_labels)
]
val_dicts = data_dicts[-10:]

print('Number of validating images per epoch:', len(val_dicts))


# Creation of data directories for data_loader
# Transforms for training and validation

val_transforms = [
        LoadImaged(keys=['image', 'label']),
        #AddChanneld(keys=['image', 'label']),
        EnsureChannelFirstd(keys=['image', 'label']),
        CropForegroundd(keys=['image', 'label'], source_key='image'),                   # crop CropForeground
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.5, 1.5, 2.0),
            mode=("bilinear", "nearest"),
        ),
        SpatialPadd(keys=['image', 'label'], spatial_size=opt.patch_size, method= 'end'),  # pad if the image is smaller than patch
        ToTensord(keys=['image', 'label'])
    ]

val_transforms = Compose(val_transforms)


# Define CacheDataset and DataLoader for training and validation
# create a training data loader

val_ds = CacheDataset(
    data=val_dicts, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=opt.workers, collate_fn=list_data_collate, pin_memory=True)
post_pred = Compose([EnsureType(),  AsDiscrete(argmax=True, to_onehot=6)])
post_label = Compose([EnsureType(),  AsDiscrete(to_onehot=6)])
saver_ori = SaveImage(output_dir=opt.output_folder, output_ext=".nii.gz", output_postfix="ori",print_log=True)
saver_gt = SaveImage(output_dir=opt.output_folder, output_ext=".nii.gz", output_postfix="gt",print_log=True)
saver_seg = SaveImage(output_dir=opt.output_folder, output_ext=".nii.gz", output_postfix="seg",print_log=True)


# Create Model, Loss, Optimizer
# build the network
if opt.network == 'nnunet':
    net = build_net()  # nn build_net
    print('net == nnunet')
elif opt.network == 'unetr':
    net = build_UNETR() # UneTR
    print('net == unetr')
elif opt.network == 'unet':
    net = build_unet() # unet3D
elif opt.network == 'vnet':
    net = build_vnet()
elif opt.network == 'HRNet':
    net = build_HRNet()

net.to(device)

torch.backends.cudnn.benchmark = opt.benchmark  # for accelerating Convs
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)  # Mean Dice caculate
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")


best_metric = -1
metric_values = []
metric_values_tb = []
metric_values_tc = []
metric_values_fb = []
metric_values_lfc = []
metric_values_rfc = []

val_time = time()


net.load_state_dict(torch.load('best_metric_model_test_0824_hrnet_C32.pth'))
print("Use pretrained model 824 HRNet ! Attention")

net.eval()
saver_flag = 1
with torch.no_grad():
        pbar = tqdm(val_loader, ncols=50, dynamic_ncols=True)
        for val_data in pbar:
            # for val_data in val_loader:
            pbar.set_description(desc='Valid_Epoch ')
            val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
            )

            roi_size = opt.patch_size
            sw_batch_size = 4
            val_outputs = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, net)
            val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
            val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                # compute metric for current iteration
            dice_metric(y_pred=val_outputs, y=val_labels)
            dice_metric_batch(y_pred=val_outputs, y=val_labels)

            if (saver_flag > 0):
                    meta_data = decollate_batch(val_data["image_meta_dict"])
                    saver_flag = 0
                    print(val_inputs[0].shape, val_outputs[0].shape, val_labels[0].shape)
                    saver_ori(val_inputs[0], meta_data[0])
                    saver_seg(torch.argmax(val_outputs[0], dim=0, keepdim=True), meta_data[0])
                    saver_gt(torch.argmax(val_labels[0], dim=0, keepdim=True), meta_data[0])

            val_end = time.time()
            print('val takes:{:.4f}'.format(val_end-val_time))
            # aggregate the final mean dice result
            dice = dice_metric.aggregate().item()
            metric_values.append(dice)
            metric_batch = dice_metric_batch.aggregate()
            # print('metric_batch.shape', metric_batch.shape)
            metric_tb = metric_batch[1].item()
            metric_values_tb.append(metric_tb)
            metric_tc = metric_batch[2].item()
            metric_values_tc.append(metric_tc)
            metric_fb = metric_batch[3].item()
            metric_values_fb.append(metric_fb)
            metric_lfc = metric_batch[4].item()
            metric_values_lfc.append(metric_lfc)
            metric_rfc = metric_batch[5].item()
            metric_values_rfc.append(metric_rfc)

            print(
                f"average dice:({metric_tb:.4f}, {metric_tc:.4f}, {metric_fb:.4f}, {metric_lfc:.4f}, {metric_rfc:.4f})"
                f"\nbest mean dice: {best_metric:.4f}    "
            )
            # reset the status for next validation round
            dice_metric.reset()
            dice_metric_batch.reset()

if name == "main":
main()

I have tried several methods in the tutorial, but they have failed with same KeyError

[wyli]

thanks for reporting, could you try removing ToTensord as well as , meta_data[0] in save image calls, if it's based on monai 0.9.1. Which tutorial are you referring to? it might need some update to use the latest API.

[jjhhan]

Thanks a lot , It works. I ignored the transform operation when I looked at the SaveImage usage, it's all my fault😊