Troubleshooting Inversion of Prediction Mask to Match Original Image Dimensions After Transformations

[hyeonLIB]

Hi, I'm encountering an issue where i cannot correctly invert a prediction mask back to the original image shape after applying a series of transformations including 'CropForegoundd', 'Orientationd' and 'Spacingd'. Despite attempting to use 'Invertd' transform on the prediction mask, the output does not match the original image dimensions. Below is the relevant portion of my code:

import os
import glob
import nibabel as nib
import numpy as np
import torch
from monai.transforms import (
    EnsureChannelFirstd,
    Compose,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    ScaleIntensityRanged,
    Spacingd,
    Activationsd,
    AsDiscreted,
    Invertd,
)
from monai.data import (
    DataLoader,
    Dataset
)
val_transforms = Compose([
    LoadImaged(keys=["image"]),
    EnsureChannelFirstd(keys=["image"]),
    ScaleIntensityRanged(
        keys=["image"],
        a_min=-57,
        a_max=164,
        b_min=0.0,
        b_max=1.0,
        clip=True,
    ),
    CropForegroundd(keys=["image"], source_key="image"),
    Orientationd(keys=["image"], axcodes="RAS"),
    Spacingd(keys=["image"], pixdim=(1.5, 1.5, 2.0), mode="bilinear"),
])
pred_transform = Compose([
    Activationsd(keys="pred", softmax=True),
    AsDiscreted(keys="pred", argmax=True),
    Invertd(keys="pred", transform=val_transforms, orig_keys="image", nearest_interp=True)
])

input_dir = '../database/1/input'
images = sorted(glob.glob(os.path.join(input_dir, "*.nii.gz")))[:1]
data_dicts = [{"image":image_name} for image_name in images]
image = nib.load(data_dicts[0]['image']).get_fdata()
print('original image shape : ', image.shape) # OUTPUT -> original image shape :  (512, 512, 193)

dataset = Dataset(data=data_dicts, transform=val_transforms)
loader = DataLoader(dataset, batch_size=1, num_workers=4)
sample = next(iter(loader))
print(sample.keys()) # OUTPUT -> dict_keys(['image', 'image_meta_dict', 'foreground_start_coord', 'foreground_end_coord'])
print('transformed image shape : ',sample['image'].shape) # OUTPUT -> transformed image shape :  torch.Size([1, 1, 267, 201, 289])
print(type(sample['image'])) # OUTPUT -> <class 'monai.data.meta_tensor.MetaTensor'>

pred_dict = { 
    'pred': torch.zeros(sample['image'].shape),
    'image_meta_dict': sample['image_meta_dict'],
    'foreground_start_coord': sample['foreground_start_coord'],
    'foreground_end_coord': sample['foreground_end_coord']
}
print('predicted mask shape : ', pred_dict['pred'].shape) # OUTPUT -> predicted mask shape :  torch.Size([1, 1, 267, 201, 289])
inverted_dict = pred_transform(pred_dict)
print('inverted image shape : ', inverted_dict['pred'].shape) # OUTPUT -> inverted image shape :  torch.Size([1, 1, 267, 201, 289])

I have confirmed the dimensions of the original image, the transformed image and the prediction mask. However after applying 'Invertd', the dimensions of the inverted prediction mask do not revert to the original image's dimensions.

Original image shape: (512,512,193)
Transformed image shape: torch.Size([1,1,267,201,289)]
Predicted mask shape before inversion: torch.Size([1,1,267,201,289)]
Inverted mask shape (expected original shape): torch.Size([1,1,267,201,289)] - which is incorrect.

Could you please advice on what might be going wrong here, or if there are additional steps or considerations needed to ensure the Invertd transform accurately reverts the prediction mask to the original image dimensions?

The purpose of the process, measuring each volume accurately and visualizing the mask on the 3D slicer. I would like to convert the prediction mask to the label map which can display the segmentation result on 3D slicer.

Thank you for your help.

[KumoLiu]

Hi @hyeonLIB, I don't think the pred in your "pred_dict" has been done any transforms.
You can refer to the example here.

for val_data in val_org_loader:
        val_inputs = val_data["image"].to(device)
        val_data["pred"] = inference(val_inputs)
        val_data = [post_transforms(i) for i in decollate_batch(val_data)]

Or you can specify the meta information and applied_operations in this line.
'pred': torch.zeros(sample['image'].shape)
Like 'pred': MetaTensor(torch.zeros(sample['image'].shape), meta=sample['image'].meta, applied_operations=sample['image'].applied_operations)
Thanks.

[hyeonLIB]

Thank you for your suggestion. I solved the issue with the codes down below.

post_transforms = Compose(
    [
        Activationsd(keys="pred", softmax=True),
        Invertd(
            keys="pred",  
            transform=val_transforms,
            orig_keys="image",  
            nearest_interp=False,
            to_tensor=True,
        ),
        AsDiscreted(keys="pred", argmax=True),
        ]
    )

def predict(data, model, device):
    inputImage = data["image"].to(device)
    outputImage = sliding_window_inference(inputImage, (96,96,96), 4, model)
    return outputImage

with torch.no_grad():
    for data in loader:
        outputImage = predict(data, model, device)
        data['pred'] = outputImage
        data = [post_transforms(i) for i in decollate_batch(data)]
        print('orginal image shape: ', nib.load(data[0]['image_meta_dict']['filename_or_obj']).get_fdata().shape)
        print('pre_transformed image shape: ',data[0]['image'].shape)
        pred_img = data[0]['pred']
        print('post_transformed pred shape: ',pred_img.shape)

# The result I got is 
orginal image shape:  (512, 512, 193)
pre_transformed image shape:  torch.Size([1, 267, 201, 289])
post_transformed pred shape:  torch.Size([1, 512, 512, 193])