SaveImage error for multilabel Swin UNet: 'Tensor' object has no attribute 'affine'

[esther1262]

Hi, I ran a swin UNet for 3 labels segmentation. Everything was fine including inference on test set. But when I tried to save the predicted results, I ran into issues. Can anyone help? I included the codes seems to be most relevant.

feature_size = 24
batch_dim = (96,96,32)

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        EnsureTyped(keys=["image", "label"]),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        #NormalizeIntensityd(keys='image'),
        RandScaleIntensityd(keys=["image"], prob=1, factors=(2,3)),
        RandShiftIntensityd(keys=["image"], prob=1, offsets=(500,1000)),
        RandBiasFieldd(keys="image", coeff_range=(0.1,0.3), prob=0.5),
        RandRotated(keys=["image", "label"], prob=1, range_z=0.1, mode=['bilinear','nearest']), 
        RandZoomd(keys=["image", "label"], prob=1, min_zoom=0.9, max_zoom=1.1, mode=['area','nearest']),
        RandCropByPosNegLabeld(keys=["image", "label"], image_key="image", label_key="label", spatial_size=batch_dim,
                               pos=1, neg=1,num_samples=2, image_threshold=0, allow_smaller = True),
        ResizeWithPadOrCropd(keys=["image", "label"], spatial_size=batch_dim, mode='constant')
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        EnsureTyped(keys=["image", "label"]),
        CropForegroundd(keys=["image", "label"], source_key="image")
    ]
)

# main model
    model = SwinUNETR(
        img_size=batch_dim, # the third dim has to be at least 32
        in_channels=1,
        out_channels=4,
        feature_size=feature_size # need to be divisible by 12
    ).to(device)


# transformation on test set
test_org_transforms = Compose(
    [    LoadImaged(keys=["image", "label"]),
         EnsureChannelFirstd(keys=["image", "label"]),
         EnsureTyped(keys=["image", "label"]),
    ]
)
test_files = load_decathlon_datalist(datasets, True, "test") # only use the test set
test_ds = Dataset(data=test_files, transform=test_org_transforms)
test_loader = DataLoader(test_ds, batch_size=1)

post_transforms_pred = Compose([
    AsDiscrete(argmax=True, to_onehot=4), FillHoles(),
    SaveImage(output_dir="./results/allimages_multilabel_swin_ver6", output_postfix="seg", resample=False, separate_folder=False)
    ])
post_transforms_label = Compose([
    AsDiscrete(to_onehot=4), FillHoles()
    ])

dice_metric_test = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)

# function for inference on test set
with torch.no_grad():
    for test_data in test_loader:
        test_inputs, test_labels = (test_data["image"].to(device), test_data["label"].to(device))
        test_outputs = sliding_window_inference(test_inputs, batch_dim, batch_size, model)
        
        test_labels_list = decollate_batch(test_labels)
        test_labels_convert = [post_transforms_label(test_label_tensor) for test_label_tensor in test_labels_list]
        test_outputs_list = decollate_batch(test_outputs)
        test_output_convert = [post_transforms_pred(test_pred_tensor) for test_pred_tensor in test_outputs_list]
        dice_metric_test(y_pred=test_output_convert, y=test_labels_convert)

Error:

RuntimeError: SaveImage cannot find a suitable writer for results/allimages_multilabel_swin_ver6/0_seg.nii.gz.
    Please install the writer libraries, see also the installation instructions:
    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies.
   The current registered writers for .nii.gz: (<class 'monai.data.image_writer.NibabelWriter'>, <class 'monai.data.image_writer.ITKWriter'>).
Traceback (most recent call last):
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/transforms/io/array.py", line 445, in __call__
    writer_obj.set_metadata(meta_dict=meta_data, **self.meta_kwargs)
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/data/image_writer.py", line 591, in set_metadata
    self.data_obj, self.affine = self.resample_if_needed(
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/data/image_writer.py", line 275, in resample_if_needed
    affine, *_ = convert_data_type(output_array.affine, output_type=orig_type)  # type: ignore
AttributeError: 'Tensor' object has no attribute 'affine'

Traceback (most recent call last):
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/transforms/io/array.py", line 445, in __call__
    writer_obj.set_metadata(meta_dict=meta_data, **self.meta_kwargs)
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/data/image_writer.py", line 431, in set_metadata
    self.data_obj, self.affine = self.resample_if_needed(
  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/data/image_writer.py", line 275, in resample_if_needed
    affine, *_ = convert_data_type(output_array.affine, output_type=orig_type)  # type: ignore
AttributeError: 'Tensor' object has no attribute 'affine'

[KumoLiu]

Hi @esther1262, could you please check your MONAI version and maybe use 1.2 or the weekly version to try it again?
Hope it helps, thanks!

[esther1262]

Hi @KumoLiu SaveImage worked fine if I didnt include set_track_meta(False). So I guess I have to change it to True before inference. However, there's another issue. My ground truth labelmap has labels 0=background, 1,2,3,4 for different brain regions. Using the codes above, the predicted labelmaps have 1=background, 0=otherwise. How to fix it?

test_outputs_list[0].shape = torch.Size([5, 256, 256, 16])

This discussion Project-MONAI/MONAI#4280 suggested SplitChannel or SplitDim.

When I used SplitDim(dim=1), 16 images were saved and overrode the previous one. Similar issues with SplitChannel(channel_dim=3).

[KumoLiu]

Hi @esther1262, that makes sense, without track_meta, a tensor doesn't have an affine attribute.

My ground truth labelmap has labels 0=background, 1,2,3,4 for different brain regions. Using the codes above, the predicted labelmaps have 1=background, 0=otherwise. How to fix it?
test_outputs_list[0].shape = torch.Size([5, 256, 256, 16])

How were you able to get “1=background, 0=otherwise” this information from shape?
If you want save the image for each image, I think you should specify the dim to 0 since dim0 is the channel dim.

Thanks!

[esther1262]

Hi @KumoLiu Using the codes above but without SplitChannel or SplitDim, I were able to save the predicted labels. I then looked at the data array. It's all 0s and 1s, where the 1s were the areas that should be labelled as background. Previously when I ran a binary segmentation, the predicted labels also had the same issues, but at least I could manually switch the 0s and 1s. This time the areas that should be labelled as 1,2,3,4s were all placed 0s. I cannot differentiate which is which. My goal is to save the label (possibly with just one channel) but different values (0 to 4) representing different regions.
If I used SplitDim(dim=0) and SaveImage in the same Compose(), one binary label for one channel was saved, and the new file would replace the previous file. Or if I removed to_onehot in Compose(), I could save a multi-label mask but the calculation of dice coefficients would be just one ground mean rather than for each label. So I guess the easiest solution is to create two post_transform for two tasks. Please advise if you have better options. Thanks.

[KumoLiu]

Hi @esther1262, if you want to save the label with only one channel, you can remove to_onehot in AsDiscrete.
And if you want to get the dice for each channel, you can change the reduction to mean_batch and remember to set num_classes(always including the background).

DiceMetric(include_background=False, reduction="mean_batch", get_not_nans=False, num_classes=5)

Hope it helps, thanks!