How to perform effective 3D segmentation of small targets in medical imaging.

[Minxiangliu]

Thank you for the powerful tools provided by MONAI for training medical imaging. I greatly enjoy using them in my research. Currently, I am working on segmenting two targets (T: primary cancer cells, N: metastatic cancer cells) in PET CT brain images.

I am facing two challenges:
a. Due to the significant size variation of cancer cells in the brain, ranging from large to small, the model struggles to effectively segment the smaller targets.
b. Both T and N belong to cancer cells, but their definitions are somewhat different. The current model has difficulty distinguishing between them effectively.

I am using the SwinUNETR model for training with 5-fold cross-validation on 98 sets of data (78 for training, 20 for testing). The Dice scores are approximately around 0.45.

I would appreciate some suggestions for improvement. Thanks in advance!

The following describes the content of the data:

All images have the same spacing in their metadata. I used the original spacing during the transform process.
The sizes of the images are as follows:

• Max and Min Width = 82, 58
• Max and Min Height = 91, 71
• Max and Min Slice = 86, 46

Training Transform:

tr_transform = Compose([
            LoadImaged(keys=['image','label'], reader='ITKReader', image_only=False),
            EnsureChannelFirstd(keys=['image','label']),
            Orientationd(keys=['image','label'], axcodes='RAS'),
            Spacingd(keys=['image','label'], pixdim=(2.734, 2.734, 2.79), mode=('bilinear', 'nearest')),
            SpatialPadd(keys=['image','label'], spatial_size=(96,96,96)),
            EnsureTyped(keys=['image','label'], device='cuda'),
            RandFlipd(keys=['image','label'], prob=0.5, spatial_axis=0),
            RandFlipd(keys=['image','label'], prob=0.5, spatial_axis=1),
            RandFlipd(keys=['image','label'], prob=0.5, spatial_axis=2),
            RandSpatialCropd(
                keys=['image','label'], 
                roi_size=(32,32,32), 
                max_roi_size=(64,64,64), 
                random_size=True,
                random_center=True),
            SpatialPadd(keys=['image','label'], spatial_size=(64,64,64)),
            NormalizeIntensityd(keys='image', nonzero=True, channel_wise=True),
            RandScaleIntensityd(keys='image', factors=0.1, prob=1.0),
            RandShiftIntensityd(keys='image', offsets=0.1, prob=1.0), 
        ])

Test Transform:

te_transform = Compose([
            LoadImaged(keys=['image','label'], reader='ITKReader', image_only=False),
            EnsureChannelFirstd(keys=['image','label']),
            Orientationd(keys=['image','label'], axcodes='RAS'),
            Spacingd(keys=['image','label'], pixdim=(2.734, 2.734, 2.79), mode=('bilinear', 'nearest')),
            NormalizeIntensityd(keys='image', nonzero=True, channel_wise=True),
            EnsureTyped(keys=['image','label'], device='cuda')
        ])

Post Transform:

post_pred = Compose([AsDiscrete(argmax=True, to_onehot=3)])
post_label = Compose([AsDiscrete(to_onehot=3)])

Define Loss and Metric

dice_metric = DiceMetric(include_background=False, reduction='mean_batch')
loss = DiceCELoss(to_onehot_y=True, softmax=True, squared_pred=True, batch=True, include_background=True)

Define Model:

model = SwinUNETR(
            img_size=(64,64,64),
            in_channels=1,
            out_channels=3,
            feature_size=48,
            drop_rate=0.2,
            attn_drop_rate=0.2,
            dropout_path_rate=0.2,
            use_checkpoint=True)

Define training parameters:

max_epochs=1000
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)
EarlyStopping(monitor=dice_metric, mode='max', patience=100)

Define validation methods:

sliding_window_inference(image, roi_size=(64,64,64), sw_batch_size=4, model=model, overlap=0.5)

Results:

Visualization of validation data results:

[KumoLiu]

Hi @Minxiangliu, thanks for your interest here.
For the imbalanced issue, I suggest you use some sampling tips, such as cropping the patches based on the pos and neg ratio.

MONAI/monai/transforms/croppad/array.py

Line 1049 in 782c1e6

class RandCropByPosNegLabel(Randomizable, TraceableTransform, LazyTransform, MultiSampleTrait):

You can also try to add some weight in the loss function.

MONAI/monai/losses/focal_loss.py

Line 86 in 782c1e6

weight: weights to apply to the voxels of each class. If None no weights are applied.

Hope it helps, thanks!

[Minxiangliu]

Hi @KumoLiu ,
Thank you very much for your response. I will give it a try.
Additionally, I would like to inquire about how to adjust the parameters if I want to deepen and enlarge the SwinUNETR model.

[KumoLiu]

Hi @Minxiangliu, you could refer to these discussions.
#6664
#5459
Thanks!

[Minxiangliu]

Hi @KumoLiu ,
I attempted to use RandCropByPosNegLabeld on the data, but encountered the following message.

message: Num foregrounds XXX, Num backgrounds 0, unable to generate class balanced samples, setting `pos_ratio` to 1.

Can I ignore this? To reiterate, the values in my labels range from 0 to 2.

Additionally, after using RandCropByPosNegLabeld(4 batch), the visualized results all look the same, Is this the expected outcome?
Thank you very much for your assistance!

Code:

from monai.data import Dataset, DataLoader
from monai.transforms import (
    Compose,
    LoadImaged,
    EnsureChannelFirstd,
    Orientationd,
    Spacingd,
    RandCropByPosNegLabeld
    )

def getDataLoader(filelist:list, trans:list):
    dataset = Dataset(data=filelist, transform=Compose(trans))
    return DataLoader(dataset, batch_size=1, shuffle=False, drop_last=False, num_workers=0)

trans = [
    LoadImaged(keys=['image','label'], reader='ITKReader', image_only=False),
    EnsureChannelFirstd(keys=['image','label']),
    Orientationd(keys=['image','label'], axcodes='RAS'),
    MaskTrans(keys='label'),
    Spacingd(keys=['image','label'], pixdim=(2.734, 2.734, 2.79), mode=('bilinear', 'nearest')),
    RandCropByPosNegLabeld(
        keys=['image', 'label'],
        label_key='label',
        spatial_size=(64,64,64),
        pos=1,
        neg=1,
        num_samples=4,
        image_key='label',
        image_threshold=0,
        allow_smaller=True)
]
filelist = [{'image': 'image_01.nii.gz', 'label':'label_01.nii.gz'}]
dataloader = getDataLoader(filelist, trans)
for batch in dataloader:
    image, label = batch['image'].squeeze(), batch['label'].squeeze()

[KumoLiu]

Hi @Minxiangliu,

• for the warning message, it means for the current case, the pixel num of your background is 0. Possible reasons such as your roi size is too small and only the foreground is included.
• This also doesn't make sense for the issue of all four RandCropByPosNegLabeld outputs being the same. Perhaps you should check your input.
• I looked at the data again, and you have a small percentage of foreground, so you could add weight to the loss, since the penalty for misclassified foreground should be higher, giving more contribution in the loss.

Thanks!

[Minxiangliu]

Hi @KumoLiu ,
Sorry to bother you again, regarding the weight issue, Suppose I have two classes, 0 for background, and 1 for foreground. I want to assign a higher penalty to the foreground. Is the following code correct?

from monai.losses import FocalLoss
import torch

weight = [0.2, 0.8]
loss = FocalLoss(include_background=True, to_onehot_y=True, use_softmax=True, weight=weight)

logits = torch.tensor([[1, 0], [0, 1], [1, 0]], dtype=torch.float32)
label = torch.tensor([0, 1, 0], dtype=torch.int64)

print("focal loss from monai: ", loss(logits, label))

[KumoLiu]

Yes. Looks good.