Small Label regions & RandCropByLabelClassesd

[dyollb]

I am using the RandCropByLabelClassesd to crop sample patches of size [96]*3.

When I am trying to train a network with a single class (e.g. hippocampus), and I am using some spatial augmentation e.g. RandRotated and RandZoomd I frequently end up getting errors like this:

My transforms are ordered as follows:

        xforms = [
            LoadImaged(keys=keys, reader="itkreader"),
            EnsureChannelFirstd(keys="image"),
            Orientationd(keys=keys, axcodes="RAS"),
            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            CropForegroundd(keys=keys, source_key="image"),
            EnsureTyped(keys=keys, dtype=np.float32, device=torch.device(self.device)),
            AddChanneld(keys="label"),
            RandFlipd(keys=keys, prob=0.2, spatial_axis=0),
            RandFlipd(keys=keys, prob=0.2, spatial_axis=1),
            RandFlipd(keys=keys, prob=0.2, spatial_axis=2),
            RandRotated(keys=keys, prob=0.2, range_z=0.4, mode=mode),
            RandRotated(keys=keys, prob=0.2, range_y=0.4, mode=mode),
            RandRotated(keys=keys, prob=0.2, range_x=0.4, mode=mode),
            RandZoomd(keys=keys, prob=0.2, min_zoom=0.8, max_zoom=1.3, mode=mode),
            RandCropByLabelClassesd(
                keys=keys,
                label_key="label",
                image_key="image",
                spatial_size=self.spatial_size,
                num_classes=self.num_classes,
                num_samples=4,
            )
        ]

I think the rotation moves the hippocampus out of the image region and then RandCropByLabelClassesd fails to find any samples and divides by zero.

I decided to crop at the end, since cropping before doing rotation/zooming would add "padding".

Can you recommend a different ordering, or workaround?

• do RandCropByLabelClassesd with larger roi before rotation/zoom and finally CenterSpatialCropd to [96]*3
• is there any way to sample the rotation center by label class?
• another workaround I can imagine would "fix" the issue is to add a head mask (or brain mask) class, which almost certainly always gets samples. This would prevent the divide by zero issue.

[dyollb]

Debugging the issue I found that the thresholded image in RandCropByLabelClassesd was responsible for my problem.

In summary

• it had nothing to do with the rotation/zoom transforms
• for some reason (tbd) my image>image_threshold (where image_threshold=0) did not overlap with the mask for my labels:
  label_flat = img_flat & label_flat if img_flat is not None else label_flat is ZERO
• therefore monai.transforms.utils.map_classes_to_indices returned an empty list of indices
• monai.transforms.utils.generate_label_classes_crop_centers then tries to sample with probabilities:
  classes = rand_state.choice(len(ratios_), size=num_samples, p=np.asarray(ratios_) / np.sum(ratios_))
• but ratios_ is ZERO, so probabilities are NaN

Workaround/fix

set image_key=None, i.e. RandCropByLabelClassesd(keys=keys, label_key="label", image_key=None, ...)

[saruarlive]

In the RandCropByLabelClassesd class (Monai 0.9.0) default value of image_key is None, right!

[dyollb]

yes, but I initially was using image_key="image" since I thought it generally made sense.