segmentation has different orientation from original volume

[tommydino93]

Hi :)
I am working on an adapted version of the BRATS 3D segmentation script. Training works smoothly and the validation dice looks fine. Now, I wanted run inference on the test set and save the segmentation masks to disk.

I’ve been following tutorial1 and tutorial2

However, it looks like the segmentations are correct but have a different orientation from the original volumes (see image)

What am I doing wrong? Thanks a lot in advance! here is my pseudo code:

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # if available, set GPU to use during training loop; otherwise, use cpu
    random.seed(10)  # set random seed for reproducibility

    # -------------------------------------- create test dataset
    test_dir = os.path.join(ds_path, "imagesTs")

    # create list of dictionaries
    test_files = [os.path.join(test_dir, x) for x in os.listdir(test_dir)]
    test_files = [{"image": volume_name} for volume_name in test_files]

    test_transforms = Compose([LoadImaged(keys=["image"]),
                              Spacingd(keys=["image"], pixdim=(1.0, 1.0, 1.0), mode="bilinear"),
                              Orientationd(keys=["image"], axcodes="RAS"),
                              EnsureChannelFirstd(keys=["image"]),
                              NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                              EnsureTyped(keys=["image"])])
    # create test data loader
    test_ds = Dataset(data=test_files, transform=test_transforms)
    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=nb_workers, pin_memory=torch.cuda.is_available())

    # -------------------------------------- create model
    model = SegResNet(blocks_down=(1, 2, 2, 4),
                      blocks_up=(1, 1, 1),
                      init_filters=16,
                      in_channels=1,
                      out_channels=1,
                      dropout_prob=0.2).to(device)

    postprocessing_transforms = Compose([EnsureTyped(keys="pred"),
                                         Invertd(keys="pred",
                                                 transform=test_transforms,
                                                 orig_keys="image",
                                                 meta_keys="pred_meta_dict",
                                                 orig_meta_keys="image_meta_dict",
                                                 meta_key_postfix="meta_dict",
                                                 nearest_interp=False,
                                                 to_tensor=True),
                                         Activationsd(keys="pred", sigmoid=True),
                                         AsDiscreted(keys="pred", threshold=0.5),
                                         SaveImaged(keys="pred", meta_keys="pred_meta_dict", output_dir=out_dir, output_postfix="seg", separate_folder=False)])

    # enable cuDNN benchmark
    torch.backends.cudnn.benchmark = True

    # --------------------------------- begin inference
    # saver = SaveImaged(out_dir, separate_folder=False, output_postfix="seg")  # will be used to save the segmentation masks to disk
    model.load_state_dict(torch.load(training_weights_path, map_location=device))  # load training weights
    model.eval()  # set model in inference mode
    with torch.no_grad():  # set all requires_grad flags to False
        for test_data in test_loader:  # loop over batches
            test_inputs = test_data["image"].to(device)  # type: torch.Tensor # extract volumes
            test_data["pred"] = inference(test_inputs, val_amp, model)  # compute inference
            test_data = [postprocessing_transforms(i) for i in decollate_batch(test_data)]  # apply post-processing to output tensors

[rijobro]

And what does your visualisation code look like? If you superpose image and pred, then pred will have been inverted, but image won't. So I'm wondering if that could be causing the issue you're seeing.

[tommydino93]

Hi @rijobro!
Thanks for your answer. The image is just a snapshot that I took from fsleyes where I loaded the two .nii.gz volumes.

To double check I also plotted them in matplotlib, but it still looks like the orientations are different

.

[rijobro]

Ok, let's get the matplotlib sorted first. What does the code for the matplotlib visualisation look like? Again, we need to be careful that we invert either both or neither of pred and image.

[rijobro]

So just:

test_data = next(iter(test_loader))
test_data["image"] = test_data["image"].to(device)
test_data["pred"] = model(test_data["image"])[0]  # remove batch (requires batch_size==1)
# visualise both "image" and "pred" without any inverse
blended = blend_images(test_data["image"], test_data["pred"])
...

[tommydino93]

So, the way I created the plot with matplotlib was simply loading the saved mask from disk with:

import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt

orig_volume = nib.load("/media/newuser/7EF9B56B27259371/BRATS_dataset/Task01_Brain_Tumour_Reduced/imagesTs/BRATS_487.nii.gz")
orig_volume_np = np.asanyarray(orig_volume.dataobj)

segm_volume = nib.load("/media/newuser/7EF9B56B27259371/BRATS_dataset/out_inference/out_brats_inference_Jan_26_2022/BRATS_487_seg.nii.gz")
segm_volume_np = np.asanyarray(segm_volume.dataobj)

plt.figure(1)
slice = 79
plt.imshow(orig_volume_np[:, :, slice], cmap="gray")
plt.imshow(segm_volume_np[:, :, slice], alpha=0.5)
plt.show()

but you are right, if I plot them before thresholding and before saving to disk, they are aligned! I tried with:

for test_data in test_loader:  # loop over batches
    test_inputs = test_data["image"].to(device)  # type: torch.Tensor # extract volumes
    test_data["pred"] = inference(test_inputs, val_amp, model)  # compute inference
    
    slice = 79
    test_inputs_np = test_inputs.cpu().detach().numpy()[0, 0, :, :, slice]
    test_pred_np = test_data["pred"].cpu().detach().numpy()[0, 0, :, :, slice]

    plt.figure(1)
    plt.imshow(test_inputs_np, cmap="gray")
    plt.figure(2)
    plt.imshow(test_pred_np, cmap="hot")

So I guess the problem is in the saving process (i.e. when I call postprocessing_transforms)

[rijobro]

Ok great! Next step is to invert the predicted image and compare it to the same image with minimal transforms, e.g.:

fname = "/media/newuser/7EF9B56B27259371/BRATS_dataset/Task01_Brain_Tumour_Reduced/imagesTs/BRATS_487.nii.gz"
data = test_transforms({"image": fname})
data["image"] = data["image"].to(device)
data["pred"] = model(image[None])

# note that we're not using SaveImaged yet!
postprocessing_transforms = Compose([
    EnsureTyped(keys="pred"),
    Invertd(keys="pred",
                  transform=test_transforms,
                  orig_keys="image",
                  meta_keys="pred_meta_dict",
                  orig_meta_keys="image_meta_dict",
                  meta_key_postfix="meta_dict",
                  nearest_interp=True,   # this should be true by the way for segmentations
                  to_tensor=True),
    Activationsd(keys="pred", sigmoid=True),
    AsDiscreted(keys="pred", threshold=0.5),
])
data = [postprocessing_transforms(i) for i in decollate_batch(test_data)]

# Load original image -- no Orientationd
fresh_img_transforms = Compose([
    LoadImaged(keys=["image"]),
    Spacingd(keys=["image"], pixdim=(1.0, 1.0, 1.0), mode="bilinear"),
    EnsureChannelFirstd(keys=["image"]),
    NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
    EnsureTyped(keys=["image"])
])
fresh_img = fresh_img_transforms({"image": fname})["image]

# visualise fresh_img and data["pred"]
...

[tommydino93]

Thanks again for your answer! I slightly adapted your code and visualized the inverted prediction which is indeed different from the original.

But by looking at your code I realized what was probably wrong in my initial script.

I had:

for test_data in test_loader:  # loop over batches
    test_inputs = test_data["image"].to(device)  # type: torch.Tensor # extract volumes
    test_data["pred"] = inference(test_inputs, val_amp, model)  # compute inference
    test_data = [postprocessing_transforms(i) for i in decollate_batch(test_data)]  # apply post-processing to output tensors

while now I adjusted to

for test_data in test_loader:  # loop over batches
    test_data["image"] = test_data["image"].to(device)  # type: torch.Tensor # extract volumes
    test_data["pred"] = inference(test_data["image"], val_amp, model)  # compute inference
    test_data = [postprocessing_transforms(i) for i in decollate_batch(test_data)]  # apply post-processing to output tensors

I guess I had to explicitly set test_data["image"] = test_data["image"].to(device) in order to have the proper keys in postprocessing_transforms

Could it make sense? I now checked the masks that are saved to disk and they are ok

[rijobro]

Hm, I find it strange that this would make a difference. The only difference between the top and the bottom examples is that the dictionary entry for image is on two different devices.

I wonder if you changed some other small lines of code along the way that could have also fixed your problem? Feel free to stop here if you've got things working as you need.

If you're interested in debugging, you could do some things like:

1. Switch between the code snippets in your previous comment, but don't change anything else in the code. Does the problem persist?
2. After inference, do test_data["image"] = test_data["image"].cpu() -- is it the difference in devices between image and pred that's causing the error?

[tommydino93]

You are right. There are actually several other things I changed. I won't have time to dig which one was the correct one, but I will just list them here in case someone has the same problem

1. I modified the transformation (which I believe is the main modification):

# ------------ BEFORE ------------
test_transforms = Compose([LoadImaged(keys=["image"]),
                              Spacingd(keys=["image"], pixdim=(1.0, 1.0, 1.0), mode="bilinear"),
                              Orientationd(keys=["image"], axcodes="RAS"),
                              EnsureChannelFirstd(keys=["image"]),
                              NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                              EnsureTyped(keys=["image"])])

# ------------ AFTER ------------
median_shape = [144, 176, 144]  # I removed empty rows, cols and slice around the brain, so I reshape to a median shape
median_padded_shape = [max(median_shape[0:2]), max(median_shape[0:2]), median_shape[2]]

test_transforms = Compose([LoadImaged(keys="image"),
                               EnsureChannelFirstd(keys="image"),
                               Resized(keys="image", spatial_size=tuple(median_shape), mode=["area"]),  # resize all volumes to median shape
                               SpatialPadd(keys="image", spatial_size=tuple(median_padded_shape), method="symmetric"),  # zero padding
                               Orientationd(keys="image", axcodes="RAS"),
                               NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                               EnsureTyped(keys="image")])

2. I noticed there was a bug in the way I loaded volumes and labels which probably led the model to learn weird segmentations. If not adding sorted, for many subjects I accidentally mixed volumes and labels cause os.listdir does not always guarantee the same order

# ------------ BEFORE ------------
train_dir = os.path.join(ds_path, "imagesTr")
train_labels_dir = os.path.join(ds_path, "labelsTr")
# create list of dictionaries
all_image_files = [os.path.join(train_dir, x) for x in os.listdir(train_dir)]
all_label_files = [os.path.join(train_labels_dir, x) for x in os.listdir(train_labels_dir)]

# ------------ AFTER ------------
train_dir = os.path.join(ds_path, "imagesTr")
train_labels_dir = os.path.join(ds_path, "labelsTr")
# create list of dictionaries
all_image_files = [os.path.join(train_dir, x) for x in sorted(os.listdir(train_dir))]
all_label_files = [os.path.join(train_labels_dir, x) for x in sorted(os.listdir(train_labels_dir))]

Thank you very much again for your time