Training segmentation algorithm for different sized images

[vikashg]

Hi all,
I am trying to train a segmentation algorithm for a dataset which has 3D images with varying number of slices ranging from 40s to 200. So it is not a good idea to resample images and I was trying to use UNetR referring to the following example.

The logic behind using the above example is that I could use patch-based training. I have a couple of questions,

1. I wanted to save the patches and visually confirm if they are alright. I did so using the following transform chain

train_transforms = Compose([LoadImaged(keys=["image", "label"]), 
                                               EnsureChannelFirstd(keys=["image", "label"]), 
                                               Orientation(keys=["image", "label"], axcodes="RAS"),
                                               CropForegroundd(keys=["image", "label"], source_key="image),
                                               RandCropByPosNegLabeld(keys=["image", "label"], label_key="label", spatial_size=(96, 96, 96),
                                                                                               pos=1, neg=1, num_samples=6)])

# Saving the patches
for _data in data_list:
     a = train_transforms(_data)
     for i in range(6):
        SaveImage(output_dir='./output_dir', output_postfix='img', resample=False)(a[i]['image'])
        SaveImage(output_dir='./output_dir', output_postfix='label', resample=False)(a[i]['image'])

Also, Does the transform Spacingd has any effect on training? Because if images are from different scanners, they have different image spacings. Can all of them be normalized to some arbitrary spacing if it is just for visual inspection.

2. In order to avoid the "cropped area larger than image size error", I decided to train on images that have a number of slices larger than 96. During training, I get the following warning
   
   which is not good for training I would think and the results were not good indeed. I was wondering how can I avoid this kind of warning.

3. Lastly, I would assume that it is possible to reduce the patch size on the model

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)

I think I had an earlier discussion where @tangy5 mentioned that the image patch size needs to be 96X96X96. else I can pad it. But can I not have a model with smaller image patch size just a model which is natively trained on images with some other patch size ?

[KumoLiu]

Hi @vikashg, thanks for your interest here.

• For the first question, it's better to add Spacingd to ensure all data with the same spacing otherwise the physical scale of the data would be different. Of course, if you only want to visualize, either is ok. However, if you are checking the data sent into the network, it would be better to add it.
• For the second one, I think it's due to the label you sent to the network don't have any foreground, it's better to check the data instead of avoiding the warning.
• For the final one, the data sent into the network can be of any size, the only thing you have to pay attention to is the downsampling times you have to prevent the data from being too small and not having any information left at the last level of downsampling. In addition, we require the data to be of the same size then they could be in a batch, which can also speed up the training.
  Hope it can help you, thanks!

[vikashg]

Thanks @KumoLiu I will try these changes you mentioned.
Vikash