Error implementing ResizeWithPadOrCrop

[jmenville11]

Apologies if this is an obvious question, I'm new to both Python and Monai. Using the BRATS tutorial as a base, I'm building a multi-label brain segmentation model. However, the tutorial uses CenterSpatialCropd to limit the ROI to (48,48,48). I'd like to instead use Monai's ResizeWithPadOrCrop transform to pad my images to the size of the largest image's volume (for now I just chose a standard size of (200,200,200). However, when I try to execute the following the kernel keeps dying. Is there an issue with how I'm implementing the transform?

# Setup imports

!python -c "import monai" || pip install -q "monai-weekly[nibabel, tqdm]"
!python -c "import matplotlib" || pip install -q matplotlib
%matplotlib inline

import os
import shutil
import tempfile
import glob
import torch
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from monai.apps import DecathlonDataset
from monai.data import CacheDataset, DataLoader, decollate_batch, Dataset, NiftiSaver
from monai.config import print_config
from monai.data import DataLoader
from monai.losses import DiceLoss
from monai.metrics import DiceMetric
from monai.networks.nets import UNet
from monai.transforms import (
    Activations,
    AsChannelFirstd,
    AsDiscrete,
    CenterSpatialCropd,
    Compose,
    LoadImaged,
    MapTransform,
    NormalizeIntensityd,
    Orientationd,
    RandFlipd,
    RandScaleIntensityd,
    RandShiftIntensityd,
    RandSpatialCropd,
    ResizeWithPadOrCropd,
    Spacingd,
    ToTensord,
    ToTensor,    
    RandCropByPosNegLabeld,
    AddChanneld,
)
from monai.utils import set_determinism

print_config()


# Establish directory
data_dir = Path("directory")
os.chdir(data_dir)


# Establish file paths
train_images = sorted(glob.glob(os.path.join(data_dir, "train", "Images", "*_FLAIR.nii")))
train_labels = sorted(glob.glob(os.path.join(data_dir, "train", "Labels", "*_label.nii")))
val_images = sorted(glob.glob(os.path.join(data_dir, "val", "Images", "*_FLAIR.nii")))
val_labels = sorted(glob.glob(os.path.join(data_dir, "val", "Labels", "*_label.nii")))
test_images = sorted(glob.glob(os.path.join(data_dir, "test", "Images", "*_FLAIR.nii")))

print('Training images found: ',len(train_images))
print('Training segmentations found: ',len(train_labels))
print('Validation images found: ',len(val_images))
print('Validation segmentations found: ',len(val_labels))


# Zip labels and images into dictionaries for both training and validation:
train_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]

val_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(val_images, val_labels)
]


# Set determinism
set_determinism(seed=0)


# Convert brain tumor labels:
class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Convert labels to multi channels:
    label 1 is the hematoma (HE)
    label 2 is the edema (ED)
    The possible classes are ED or HE

    """
    def __call__(self, data):
        d = dict(data)
        for key in self.keys:
            result = []
            # label 1 is HE (hematoma)
            result.append(d[key] == 1)
            # label 2 is ED (edema)
            result.append(d[key] == 2)
            d[key] = np.stack(result, axis=0).astype(np.float32)
        return d


# Setup transforms for training and validation
train_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(2, 2, 2),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(
            keys=["image", "label"],
            axcodes="PLS"
        ), 
        ResizeWithPadOrCropd(keys=["image", "label"],
            spatial_size=(200, 200, 200),
            mode=("constant"),
        ),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        ToTensord(keys=["image", "label"]),
    ]
)

val_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(2, 2, 2),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(
            keys=["image", "label"],
            axcodes="PLS"
        ), 
        ResizeWithPadOrCropd(keys=["image", "label"],
            spatial_size=(200, 200, 200),
            mode=("constant"),
        ),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        ToTensord(keys=["image", "label"]),
    ]
)

[Nic-Ma]

Hi @jmenville11 ,

Thanks for your interest here.
Could you please help paste your error message here?
Maybe something related to out of memory? You can decrease the spatial_size or batch size, etc.

Thanks.

[ericspod]

The docstring has been updated for UNet but maybe a little more explanation would help. The stride parameter states the stride for kernels that are downsampling or upsampling data with padding added to ensure the output's spatial dimensions are even fractions/multiples of the stride. The idea was to typically use a stride of 2 when down/upsampling so that the output dimensions are either half that of the input when downsampling, and twice when upsampling.

This causes problems if the input isn't some multiple of a power of 2. If you have an input of spatial size (100,100,100) and stride values of (2,2,2), the output spatial dimensions when downsampling will be (50,50,50), then (25,25,25), then (12,12,12). However that last size will be upsampled to (24,24,24) which can't be concatenated with the input coming from the skip connection which is (25,25,25). So if you have N numbers of layers in the network your inputs must have spatial dimensions that are all multiples of 2**N.

Your inputs of shape (464, 512, 256) I would pad to (512,512,256) and this will permit a network with up to 7 layers, though I would recommend no more than 5. Here is a very simple demonstration of what I'd suggest:

import monai
import monai.transforms as mt
import torch

inp=torch.rand(1, 464, 512, 256)  # test input tensor

# simple network definition of 4 layers
net=monai.networks.nets.UNet(3,1,2,(8,26,32,64),(2,2,2))

# this demonstrates using the SpatialPad transform
trans=mt.Compose([
    mt.SpatialPad((512,512,256)),
    mt.ToTensor()
])
    
inp_trans=trans(inp)  # pretend this was done in a dataset/dataloader

pred=net(inp_trans[None])  # [None] adds a batch dimension
print(pred.shape)  # torch.Size([1, 2, 512, 512, 256])

[Nic-Ma]

Hi @ericspod ,

Thanks for your detailed explanation, let me try to enhance the doc-string ASAP.

[Nic-Ma]

Hi @ericspod , I enhanced the doc-string with your explanation in PR: #2659, could you please help review it?

Thanks in advance.

[jmenville11]

Thank you both, this is enormously helpful!! @Nic-Ma you were correct, it was a memory issue.

[Nic-Ma]

Hi @jmenville11 , glad to see your good progress!

Thanks.