Add random crop transform

[mathpluscode]

🚀 Feature

Add random crop as a transformer apart from the LabelSampler.

Motivation

I'm doing patch based segmentation. My images are relatively large 512x512x512 and I only take patches of size 128x128x64.
I'm now currently applying random affine on the entire image, but this process is too slow. Therefore, it would be nice to apply a random crop to get images of size 256x256x128, then apply a random affine, then extract patches.

Or is there any other better way to reduce the compute here?

[fepegar]

Hi, @mathpluscode. Nice to see you here. Would the example in #205 (comment) work?

If the reading time is another bottleneck, we could implement a reader that uses NiBabel's memmap to perform the random crop at reading time. Please let me know.

[mathpluscode]

Hi, @mathpluscode. Nice to see you here. Would the example in #205 (comment) work?

If the reading time is another bottleneck, we could implement a reader that uses NiBabel's memmap to perform the random crop at reading time. Please let me know.

Hi @fepegar, thanks for the quick reply! Yes this example sounds enough for my application. Will try it and let you know if it works still!

[mathpluscode]

Hi, @mathpluscode. Nice to see you here. Would the example in #205 (comment) work?

If the reading time is another bottleneck, we could implement a reader that uses NiBabel's memmap to perform the random crop at reading time. Please let me know.

Based on your example code, the following code works for me.
Feel free to close this issue. @fepegar

"""Module for custom transforms."""
from typing import Tuple

import torchio as tio


class RandomCrop:
    """Random cropping on subject."""

    def __init__(self, roi_size: Tuple):
        """Init.

        Args:
            roi_size: cropping size.
        """
        self.sampler = tio.data.LabelSampler(patch_size=roi_size, label_name="label")

    def __call__(self, subject: tio.Subject) -> tio.Subject:
        """Use patch sampler to crop.

        Args:
            subject: subject having image and label.

        Returns:
            cropped subject
        """
        for patch in self.sampler(subject=subject, num_patches=1):
            return patch

[fepegar]

Nice usage of samplers! Thanks, @mathpluscode.

[jxchen01]

any plan of adding this to torchio as a native transformation?

[fepegar]

@jxchen01 it would be nice to have this, yes. That's why I never closed the issue. Using a WeightedSampler (the parent of LabelSampler) might make training slow, but using a UniformSampler would probably be fast. Something similar to Yunguan's code #570 (comment).

Would you like to submit a PR?

Another way to do this would be using the translation kwarg of RandomAffine, but I think the implementation with the sampler would be faster.

[jxchen01]

hmm... interesting. I want to give it shot and try to make a PR.

[jxchen01]

@fepegar While I look more in the details, I have a slightly different idea for implementation:

https://github.com/fepegar/torchio/blob/main/torchio/transforms/preprocessing/spatial/crop_or_pad.py#L165

The current method for CropOrPad, when cropping, always do a center cropping. What if we add another args, like random: boolean = False? Basically, if the user wants to make the cropping (or even padding) random, we can do that. Otherwise, it is center cropping or padding by default. The only adjustment to support random cropping besides center cropping is to adjusting the cropping parameters cropping_params = self._get_six_bounds_parameters(cropping) like this:

cropping_params = self._get_six_bounds_parameters(cropping)
random_x = randint(-cropping[0] // 2, cropping[0] // 2)
random_y = randint(-cropping[1] // 2, cropping[1] // 2)
random_z = randint(-cropping[2] // 2, cropping[2] // 2)
cropping_params = [
    cropping_params[0] + random_x,
    cropping_params[1] - random_x,
    cropping_params[2] + random_y,
    cropping_params[3] - random_y,
    cropping_params[4] + random_z,
    cropping_params[5] - random_z,
]

what do you think?

[fepegar]

Sure, I guess that could work too, and might be faster. But to be consistent with the rest of the library, it should be a new transform called RandomCrop that uses CropOrPad under the hood, I think.

[jxchen01]

Sure, I guess that could work too, and might be faster. But to be consistent with the rest of the library, it should be a new transform called RandomCrop that uses CropOrPad under the hood, I think.

Just to make sure I understand your suggestion correctly. The TODOs are:

• implement this minor change in CropOrPad to support cropping at a random location
• implement a wrapper RandomCrop which calls CropOrPad to use the random cropping

Even though the CropOrPad is updated, we don't expect users to us CropOrPad to do the random cropping, but using RandomCrop instead, for the sake of consistent naming. Sounds correct?

[fepegar]

I was rather thinking of leaving CropOrPad as is and creating a RandomCrop that computes the needed cropping argument for CropOrPad.

[fepegar]

Hey, @jxchen01. I've seen your PR. When would you need random padding? Why not simply implementing random cropping as in #570 (comment)? Moreover, why not using patch-based training instead?

[jxchen01]

Why random padding is needed?

It is necessary for classification problems. You can find one example here: https://github.com/AllenCell/image_classifier_3d/blob/master/image_classifier_3d/data_loader/universal_loader.py#L224

In the example above, we need to build a 3D mitotic classifier from microscopy images, where we need to CropOrPad all 3D images into the same size (padding is more common, more than 95% of the cases probably. We set the target size relatively large. Only those very large images will be cropped, since we don't want to loss any information if possible). Different from natural scene images, re-sizing the images is not possible, because re-sizing will change the morphology of the cells. Training with batches of images with different sizes is possible in theory, but our empirical experiments showed that random padding into the same size achieved much better performance. In the example above, we implemented a random CropOrPad in a "naive" way, where is a little bit cumbersome. I think this PR gives a better shot and will be super helpful for future works.

Why not patch-based training instead?

As you mentioned earlier, #570 (comment), as a standalone transformation, a more native implementation would be faster, instead of depending on Sampling-based pipeline. I believe the solution proposed by @mathpluscode is a more or less workaround when the native solution does not exist.