Question about Configuring UNetr

[vikashg]

Hi all,
I am trying to do a 3D segmentation based on the UNetr example in the tutorials https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/unetr_btcv_segmentation_3d.ipynb.
Here is some info about my data.

1. Number of classes = 10 + background
2. Size varies as 512 X 512 X (89-110) depending on the number of slices.

I am going through the paper to understand the architecture but looking at the code, I cannot figure out where do we specify the number of labels in the Network architecture. Is it num_heads ? or we do not need to specify this ?

The second part I am confused about is the parameter img_size in the UNETR function. As I see in the tutorial that img_size in UNETR function and the spatial_size in RandCropByPosNegLabeld are same (96, 96, 96). Do they need to be same?

I think I need to reduce the spatial size to something other than 96 as some of my images do not have 96 slices?

Can someone throw some light ?
Thanks
Vikash

[tangy5]

HI @vikashg, thanks for the interests of the work.

1. To define a different number of output classes: you could change the out_channels, e.g., in the tutorial, there are 13 organs and 1 background, so the out_channels=14 here.

model = UNETR(
    in_channels=1,
    out_channels=14,
    img_size=(96, 96, 96),
    feature_size=16,
    hidden_size=768,
    mlp_dim=3072,
    num_heads=12,
    pos_embed="perceptron",
    norm_name="instance"
    res_block=True,
    dropout_rate=0.0,
).to(device)

2. Yes, this is the patch size we use that cropped from CT volume, the (96x96x96) is the crop size in transform om 'RandCropByPosNegLabeld'. the 'img_size' is the input size of the model, they should be the same in this scenario.
   You can try smaller size of (64x64x64). You can also use 'SpatialPadd' to pad your images to reach the 96 size.

[vikashg]

Thanks @tangy5 That would have been my next question about how to modify the neural network but you answered that.
Vikash

[vikashg]

Hi @tangy5,
So I resampled all my images to a size (512, 512, 108) and used the same Unetr as in the tutorial. But I get this error on RandCropByPosNegLabeld for the crop size being larger than the image size but that is not the case.
Any thoughts on what is going on?

[tangy5]

The input image is all (512, 512, 108), which is good. Can you check whether there are Spacingd transform? This step will resample the input volume to a specific resolution, then the dimensions will change, I guess the z axis, after resample, there are less than 96 slices.

In detail, this transform: you can remove this if your images resolution are consistent.

        Spacingd(
            keys=["image", "label"],
            pixdim=(1.5, 1.5, 2.0),
            mode=("bilinear", "nearest"),
        ),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-175,
            a_max=250,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),

[vikashg]

Yes that was the issue, I removed the spacing transform and it is working. I think we should update the notebook with some information about image size, spacing issues as well as number of labels as out_channels.

Thanks a lot for your help @tangy5