How to apply same transform to image and label using 'Dataset'?

[ivychancheers]

Hi, I am a newbie for using monai and PyTorch. I am trying to reconstruct an image by a paired pixel-wise supervised learning and a simple U-net. I have transformed and loaded the image and label image (or we can say it as the groundtruth image) in this way:

train_transform = Compose(
    [   LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CenterSpatialCropd(keys=["image", "label"], roi_size=[32, 128, 128]),
        RandFlipd(keys=["image", "label"], prob=1, spatial_axis=0),
        NormalizeIntensityd(keys=["image", "label"], nonzero=True, channel_wise=True),
        Rand3DElasticd(keys=["image", "label"], prob=1, spatial_size=(32, 128, 128), 
                       magnitude_range=(5, 6),
                       rotate_range=(np.pi / 4,),scale_range=(0.2, 0.2),
                       sigma_range=(5,6), translate_range=(10, 10),
                       padding_mode="border").set_random_state(seed=123),
        ToTensord(keys=["image", "label"]),
    ])

train_ds = Dataset(data=train_files, transform=train_transform)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)

When I plot the "image" and "label" in the train_ds, they are transformed in a different way... Also, the orders of images and labels are changed.

1. Is it possible to apply same transformation on the training images and labels with Dataset? Or we must use ArrayDataset?
2. Can we use set_random_state for this purpose?

Thank you very much for your kind help.

[rijobro]

You're applying Rand3DElasticd to your image, but not your label (and with scaling of 0.2 your image and label will be vastly different).

For any transforms that will require interpolation, you'll want to use nearest neighbour interpolation with your label. So you would use a list of different interpolations for the different keys, e.g., mode=["trilinear", "nearest"].

[rijobro]

ValueError: 'trilinear' is not a valid GridSampleMode

Sorry, try "bilinear", which is the default.

But still the "image" and "label" in the train_ds are transformed in a different way...

Different in what way? Do they have same shapes? I think you'll have to remove transformations, check they match, and then gradually add the transformations back in one by one.

[rijobro]

I would also delete the set_random_state attached to your Rand3DElasticd. If you want, you can use set_determinism at the start of your script.

[ivychancheers]

Yes I have removed set_random_state , and used set_determinism at the start of the script.

train_transform = Compose(
    [  
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CenterSpatialCropd(keys=["image", "label"], roi_size=[72, 368, 368]),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
        ToTensord(keys=["image", "label"]),
    ])

Even only with RandFlipd, the image and label still have different flips.

[wyli]

related to this #1063? if not, would be great to have more information to help replicate your issue

[rijobro]

I would also remove Orientationd whilst debugging -- there's no guarantee that your input images and label have the same orientation according to their metadata. If you just have:

data = {"image": "/path/to/image.nii", "label": "/path/to/label.nii"}
transform = Compose([
    LoadImaged(keys=["image", "label"]),
])
loaded_data = transform(data)
im = loaded_data["image"]
lbl = loaded_data["label"]

then can im and lbl be superposed? If yes, then add in transforms one by one until they no longer superpose.

[ivychancheers]

Sorry we have found my problem in plotting the images. I have used

train_ds = Dataset(data=train_files, transform=train_transform)
print(f"image shape: {train_ds['image'].shape}")
plt.figure("image", (15, 6))
plt.subplot(1, 2, 1)
plt.imshow(train_ds[0]["image"][0,25,:, :])
plt.subplot(1, 2, 2)
plt.imshow(train_ds[0]["label"][0,25, :, :])
plt.show() 

In the above I apply transform in the train_files each time when I plot train_ds separately. Instead, I should plot in this way:

train_ds = Dataset(data=train_files, transform=train_transform)
plot_train=train_ds[0]
print(f"image shape: {plot_train['image'].shape}")
plt.figure("image", (15, 6))
plt.subplot(1, 2, 1)
plt.imshow(plot_train["image"][0,25,:, :])
plt.subplot(1, 2, 2)
plt.imshow(plot_train["label"][0,25, :, :])
plt.show() 

Then the train_transform would be only applied once, and we could observe the same transformation in the image and label.

Thanks a lot! :)