Does how does RandSpatialCrop in monai transforms behave? Its cropping different 3D patches from image and segmentation...

[ummadiviany]

I working on a 3D binary segmentation problem the training images and labels are of 320x320x40. I want to randomly crop the 100x100x20 3D patch and same respective patch from the label as well. I am using the following code. I have an issue with RandSpatialCrop. I want to know how to use array transforms and especially Rand array transforms. Can you please help me.

import torch
import torch.nn as nn
import nibabel as nib
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from glob import glob
from monai.transforms import (ScaleIntensityRange, Compose,  AddChannel, ToTensor,
                              RandCropByPosNegLabel, RandAxisFlip, EnsureChannelFirst, 
                              CropForeground, Spacing, EnsureType, Resize, RandSpatialCrop)
from monai.utils import set_determinism

torch.manual_seed(2000)
set_determinism(seed=2000)

class Load3D(Dataset):
    def __init__(self, img_paths : list = None, label_paths : list = None,
                    img_transform = None, label_transform = None) -> None:

        super(Load3D).__init__()
        self.img_paths = img_paths
        self.label_paths = label_paths
        self.img_transform = img_transform
        self.label_transform = label_transform
    
    def __len__(self) -> int:
        return len(self.img_paths)

    def preprocess(self, img, label):
        """
        This function is used to preprocess the data.
        1. Rotates the image by 270 degrees in the x-y plane.
        2. Keep the original intensity values from 0 to 99.9 % 
        2. Scales the intensity range of the image to [0, 1].
        """
        img = np.flip(np.rot90(img, k=3, axes = (0,1)), axis = 1)
        label = np.flip(np.rot90(label, k=3, axes = (0,1)), axis = 1)

        a_max = np.percentile(img, 99.9)
        SIR = ScaleIntensityRange(a_min = 0, a_max = a_max, b_min = 0, b_max = 1, clip = True)
        img = SIR(img)

        return img, label

    def __getitem__(self, idx):
        print(f"Loading {self.img_paths[idx]} and {self.label_paths[idx]}")
        img = nib.load(self.img_paths[idx]).get_fdata()
        label = nib.load(self.label_paths[idx]).get_fdata()
        img, label = self.preprocess(img, label)

        if self.img_transform is not None:
            img = self.img_transform(img)
        if self.label_transform is not None:
            label = self.label_transform(label)
        return img, label



if __name__ == "__main__":
    data_dir = "../../MIS/datasets/promise12/"
    img_paths = glob(data_dir + "imagesTr/*.nii.gz")
    label_paths = glob(data_dir + "labelsTr/*.nii.gz")
    img_paths.sort()
    label_paths.sort()
    # Removing the file with  15 slices ; Atleast 16 slices should be present
    img_paths.pop(13)
    label_paths.pop(13)

    roi_size= (256, 256, 16)
    train_transform = Compose(
        [
            AddChannel(),
            RandSpatialCrop(roi_size= roi_size, random_center = True, random_size=False),
            RandAxisFlip(prob = 0.5),
            ToTensor()
        ]
    )

    test_transform = Compose(
        [
            AddChannel(),
            ToTensor()
        ]
    )

    train_set = Load3D(img_paths = img_paths[:45], label_paths = label_paths[:45], 
                        img_transform = train_transform, label_transform = train_transform)
    test_set = Load3D(img_paths = img_paths[45:], label_paths = label_paths[45:],
                        img_transform = test_transform, label_transform = test_transform)

    print(f"Length of train set: {len(train_set)}")
    print(f"Length of test set: {len(test_set)}")

    train_loader = DataLoader(train_set, batch_size = 1, shuffle = True)
    test_loader = DataLoader(test_set, batch_size = 1, shuffle = False)


    imgs,labels = next(iter(train_loader))
    print(f"Image shape : {imgs.shape}")
    print(f"Label shape : {labels.shape}")

    slice = 7
    plt.figure(figsize=(5*2,5*1))
    plt.subplot(1,3,1)
    plt.imshow(imgs[0,0,:,:,slice], cmap='gray')
    plt.axis('off')
    plt.title('Image')
    plt.subplot(1,3,2)
    plt.imshow(labels[0,0,:,:,slice], cmap='gray')
    plt.axis('off')
    plt.title('Label')
    plt.subplot(1,3,3)
    plt.imshow(imgs[0,0,:,:,slice], cmap='gray')
    plt.imshow(labels[0,0,:,:,slice], 'jet', alpha=0.1)
    plt.axis('off')
    plt.title('Overlay')
    plt.show()

---------------------------------Output------------------------------------------------------
Length of train set: 45
Length of test set: 4

Image shape : torch.Size([1, 1, 256, 256, 16])
Label shape : torch.Size([1, 1, 256, 256, 16])

---

The image crop and label crops are from different, they are not associated. How to use RandSpatialCrop properly??

[Nic-Ma]

Hi @ummadiviany ,

Thanks for your interest here.
You can refer to this example:
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/torch/unet_training_array.py#L72
The ImageDataset handles the image & label logic with array transforms.
You can also use the ArrayDataset and additional LoadImage transform:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/dataset.py#L1151

Thanks.