Patch dataset based on labels.

[YerePhy]

Hi,

I would like to know if there is a way of replicating the behaviour of torchio.data.LabelSampler in MONAI. If not, it would be nice to implement this feature.

[wyli]

I'm not familiar with torchio, looks like it's a weighted cropping function. There are quite a few transforms in monai to achieve this, such as https://docs.monai.io/en/stable/transforms.html#randweightedcropd

[YerePhy]

Thanks for your answer, @wyli. I've explored the variety of samplers that MONAI has, however I can't get a patch dataset where every sample is a patch extracted from subject, could you help?

First of all, I am working with dict data:

class BratsModalities(Enum):
    T1 = 'T1'
    T1CE = 'T1CE'
    T2 = 'T2'
    FLAIR = 'FLAIR'
    SEG = 'SEG'

# read samples
base_path = Path('/home/yere/Repositories/brain_tumor/datasets/dataset_concat')
ext = '.nii.gz'
subjects = [
    {modality.value.lower(): subject.joinpath(modality.value + ext) for modality in BratsModalities} 
    for subject in base_path.iterdir()
]
subjects

>>> [{'t1': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/550/T1.nii.gz'),
  't1ce': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/550/T1CE.nii.gz'),
  't2': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/550/T2.nii.gz'),
  'flair': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/550/FLAIR.nii.gz'),
  'seg': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/550/SEG.nii.gz')},
 {'t1': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/497/T1.nii.gz'),
  't1ce': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/497/T1CE.nii.gz'),
  't2': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/497/T2.nii.gz'),
  'flair': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/497/FLAIR.nii.gz'),
  'seg': PosixPath('/home/yere/Repositories/brain_tumor/datasets/dataset_concat/497/SEG.nii.gz')},
...]

I've tried two approaches in order to get a patch dataset:

1. To add the patch sampler at the final of the transform:

tranformations = transforms.Compose(
    [
        transforms.LoadImaged(keys={'t1', 'flair', 't2', 't1ce', 'seg'}, reader='NibabelReader'),
        transforms.AddChanneld(keys={'t1', 'flair', 't2', 't1ce', 'seg'}),
        transforms.ConcatItemsd(keys={'t1', 'flair', 't2', 't1ce'}, name='input'),
        transforms.ConvertToMultiChannelBasedOnBratsClassesd(keys={'seg'}),
        transforms.NormalizeIntensityd(keys={'input'}, channel_wise=True, nonzero=True),
        transforms.RandCropByLabelClassesd(
            keys=['seg', 'input'], 
            label_key='seg',
            spatial_size=(64, 64, 1), 
            num_samples=2,
            num_classes=3,
            ratios=[0.2, 0.2, 0.4]
)
    ]
)
ds = data.Dataset(subjects[:3], transform=tranformations)

for item in data.DataLoader(ds):
    print(item['input'].shape)

>> torch.Size([2, 4, 64, 64, 1])
torch.Size([2, 4, 64, 64, 1])
torch.Size([2, 4, 64, 64, 1])

Problem: this way the monai.data.DataLoader stacks the patches.

2. Using monai.data.Dataset with monai.data.PatchDataset and a patch_func

patch_transformations = transforms.Compose(transforms=[
    transforms.SelectItemsd(keys={'seg', 'input', 'flair'})
])
tranformations = transforms.Compose(
    [
        transforms.LoadImaged(keys={'t1', 'flair', 't2', 't1ce', 'seg'}, reader='NibabelReader'),
        transforms.AddChanneld(keys={'t1', 'flair', 't2', 't1ce', 'seg'}),
        transforms.ConcatItemsd(keys={'t1', 'flair', 't2', 't1ce'}, name='input'),
        transforms.ConvertToMultiChannelBasedOnBratsClassesd(keys={'seg'}),
        transforms.NormalizeIntensityd(keys={'input'}, channel_wise=True, nonzero=True)#,
    ]
)
patch_func = transforms.RandCropByLabelClassesd(
            keys=['seg', 'input'], 
            label_key='seg',
            spatial_size=(64, 64, 1), 
            num_samples=2,
            num_classes=3,
            ratios=[0.2, 0.2, 0.4]
)
ds = data.Dataset(subjects[:3], transform=tranformations)
patch_ds = data.PatchDataset(ds, transform=None,  patch_func=patch_func, samples_per_image=2)
patch_dl = data.DataLoader(patch_ds, batch_size=1)

This way, something weird happens: the 'input' and 'seg' patches in patch_ds have different center...:

for i in range(len(patch_ds)):
    print(
        f"INPUT PATCH CENTER: {patch_ds[i]['input_transforms'][0]['extra_info']['center']} |",
        f"SEG PATCH CENTER: {patch_ds[i]['seg_transforms'][0]['extra_info']['center']}"
    )
    
>>> INPUT PATCH CENTER: [99, 122, 96] | SEG PATCH CENTER: [104, 111, 77]
INPUT PATCH CENTER: [119, 114, 69] | SEG PATCH CENTER: [96, 125, 88]
INPUT PATCH CENTER: [156, 173, 109] | SEG PATCH CENTER: [144, 157, 115]
INPUT PATCH CENTER: [143, 169, 109] | SEG PATCH CENTER: [156, 164, 119]
INPUT PATCH CENTER: [69, 148, 96] | SEG PATCH CENTER: [94, 160, 111]
INPUT PATCH CENTER: [99, 166, 76] | SEG PATCH CENTER: [75, 153, 92]

However, if I pass patch_ds to a monai.data.DataLoader the patches have the same center... this behaviour is very weird:

for item in patch_dl:
    print(
        f"INPUT PATCH CENTER: {item['input_transforms'][0]['extra_info']['center']} |",
        f"SEG PATCH CENTER: {item['seg_transforms'][0]['extra_info']['center']}"
    )
 
>>> INPUT PATCH CENTER: [tensor([101]), tensor([122]), tensor([70])] | SEG PATCH CENTER: [tensor([101]), tensor([122]), tensor([70])]
INPUT PATCH CENTER: [tensor([119]), tensor([108]), tensor([91])] | SEG PATCH CENTER: [tensor([119]), tensor([108]), tensor([91])]
INPUT PATCH CENTER: [tensor([147]), tensor([158]), tensor([117])] | SEG PATCH CENTER: [tensor([147]), tensor([158]), tensor([117])]
INPUT PATCH CENTER: [tensor([156]), tensor([159]), tensor([111])] | SEG PATCH CENTER: [tensor([156]), tensor([159]), tensor([111])]
INPUT PATCH CENTER: [tensor([56]), tensor([127]), tensor([99])] | SEG PATCH CENTER: [tensor([56]), tensor([127]), tensor([99])]
INPUT PATCH CENTER: [tensor([72]), tensor([152]), tensor([89])] | SEG PATCH CENTER: [tensor([72]), tensor([152]), tensor([89])] 

Problem: monai.data.PatchDataset analog patches ('input', 'seg') have different center.

[YerePhy]

As far as I know if you use a cropper that yiels a list of patches or a batch of them, you must use monai.data.Dataloader in order to decollate them and pop one patch at a time (not a batch of them). That mimics torchIO Queue behaviour.

[wyli]

cool examples, you are almost there, patch_ds[i] will trigger the randomisation every time, see also #2827

for i in range(len(patch_ds)):
    item_i = patch_ds[i]
    print(
        f"INPUT PATCH CENTER: {item_i['input_transforms'][0]['extra_info']['center']} |",
        f"SEG PATCH CENTER: {item_i['seg_transforms'][0]['extra_info']['center']}"
    )

[YerePhy]

@wyli, thanks!

One related question, is it possible to use monai.transforms.RandCropByLabelClassesd after monai.transforms.ConvertToMultiChannelBasedOnBratsClassesd or should be used before because monai.transforms.ConvertToMultiChannelBasedOnBratsClassesd introduces class overlapping? If this is the case it would be nice to pointing it out in the documentation. In addition, parameters like indices have non-sense if the classes overlap because, if I am not missing something the indices are not well defined since one pixel/voxel could belong to one or more classes in general.

[wyli]

These two approaches for patch-based analysis happen to answer your question @jxchen01 #4406 (reply in thread)