DataLoader for 2D images of different sizes

[hermancollin]

Hi,

I am trying to do segmentation on 2D microscopy images. Usually, these images are pretty large, and the images in my dataset have various sizes. I am following the 2D-unet tutorial but the dataloader cannot properly load the images: the monai.utils.misc.first() function raises an error. I don't want to resize the images because it will degrade the quality of the data (e.g. manual annotations ground-truths).

So for training, an ideal scenario would be to crop the images into 256x256 non-overlapping contiguous patches. I don't mind zero padding a few patches to avoid resizing. However, I can't really find a way to do it. For example, the following is almost what I want, using GridPatchd to transform the images into multiple patches:

train_transforms = Compose(
        [
            LoadImaged(keys=["im", "seg-ax", "seg-my"]),
            EnsureChannelFirstd(keys=["im", "seg-ax", "seg-my"]),
            NormalizeIntensityd(keys="im"),
            # for affine and elastic transforms, adapted from default-SEM-model/model_seg_rat_axon-myelin_sem/model_seg_rat_axon-myelin_sem.json
            # see https://github.com/axondeepseg/default-SEM-model
            RandAffined(
                keys=["im", "seg-ax", "seg-my"], 
                prob=1.0, 
                rotate_range=np.pi/64, 
                scale_range=0.05,
                translate_range=(10, 10),
                padding_mode="zeros",
                device=device
            ),
            Rand2DElasticd(
                keys=["im", "seg-ax", "seg-my"],
                prob=0.5,
                spacing=(30, 30),
                magnitude_range=(1, 2),
                padding_mode="zeros",
                device=device,
            ),
            GridPatchd(keys=["im", "seg-ax", "seg-my"], patch_size=[256, 256])
        ]
    )
    check_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
    check_loader = DataLoader(check_ds, batch_size=2, num_workers=4, collate_fn=list_data_collate)
    check_data = monai.utils.misc.first(check_loader)
    print(check_data["im"].shape, check_data["seg-ax"].shape, check_data["seg-my"])

But monai.utils.misc.first fails with something like this:

E: stack expects each tensor to be equal size, but got [20, 1, 256, 256] at entry 0 and [9, 1, 256, 256] at entry 1

I understand list_data_collate fails here because the GridPatch transform returns a different number of patches if the images have different sizes, but I'm not sure how to fix this.

Similar to #678 but I am really trying to avoid resizing the images to a common size because my annotations are large and detailed and I'm concerned it would degrade them. Would you have any recommendation on best practices for such a case? Would the best solution be to write my own collate function to pass to DataLoader?

[KumoLiu]

Hi @hermancollin, maybe you could try to use GridPatchDataset without GridPatchd.
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/grid_dataset.py#L133
Hope it can help you, thanks!

[hermancollin]

@KumoLiu thank you for the quick response. Indeed, this is exactly what I was looking for. I still have a problem though. I kept the same transforms as described in my original post (excluding the GridPatchd transform) and I now have this with GridPatchDataset:

patch_iterator = monai.data.PatchIterd(keys=["im", "seg-ax", "seg-my"], patch_size=(256, 256), mode='constant')
check_ds = monai.data.GridPatchDataset(
    data=train_files, 
    patch_iter=patch_iterator,
    transform=train_transforms
)
check_loader = DataLoader(check_ds, batch_size=2, num_workers=1, collate_fn=list_data_collate)
check_data = monai.utils.misc.first(check_loader)

However, I think the transforms fed to GridPatchDataset are applied patch-wise, but I am using the LoadImaged transform to load the images. When GridPatchDataset tries to apply the transforms to the patches, I get this error:

Original Traceback (most recent call last):
  File "/home/herman/.local/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/herman/.local/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 34, in fetch
    data.append(next(self.dataset_iter))
  File "/home/herman/.local/lib/python3.10/site-packages/monai/data/grid_dataset.py", line 191, in __iter__
    for patch, *others in self.patch_iter(image):
  File "/home/herman/.local/lib/python3.10/site-packages/monai/data/grid_dataset.py", line 116, in __call__
    original_spatial_shape = d[first(self.keys)].shape[1:]
AttributeError: 'str' object has no attribute 'shape'

So my understanding is that I need to load the images first in a preliminary dataloader, then feed them to GridPatchDataset, similarly to the Volume-level preprocessing in this tutorial: https://github.com/Project-MONAI/tutorials/blob/main/modules/2d_slices_from_3d_training.ipynb. This brings me back to the original problem, because my images all have different sizes. Do you have any insights on where I should go from here?

[KumoLiu]

Hi @hermancollin, you could add a Dataset before. Something like:

trans = LoadImaged(keys="img", ensure_channel_first=True)
ds = Dataset(data=data, transform=trans)

patch_iter = PatchIterd(keys='img', patch_size=(250, 250), start_pos=(0, 0))
patch_intensity = RandShiftIntensityd(keys='img', offsets=1.0, prob=1.0)

# construct the dataset
ds = GridPatchDataset(data=ds,
                      patch_iter=patch_iter,
                      transform=patch_intensity)
# use the grid patch dataset
for item in DataLoader(ds, batch_size=4, num_workers=2):
    print("patch size:", item[0]['img'].shape)
    print("coordinates:", item[1])

Hope it can help you, thanks!

[hermancollin]

@KumoLiu thank you. I ended up doing something like this, putting all my transformations in the preliminary Dataset and then using the GridPatchDataset. Now this part works.