Segmentation model only learning background and organ label not the cancer label

[faisalahm3d]

I am new to the medical imaging domain and MONAI. As part of my learning, I attempted to create a segmentation pipeline for CT images with three labels: background (0), liver (1), and cancer (2). Unfortunately, the model is only learning for the background and liver, while the validation loss for cancer always returns 'NaN'. Could you please help me identify where I may have made a mistake in the segmentation pipeline from the code provided?

Thank you very much.
`

data_dir = '/home/Task03_Liver/imagesTr'
labels_dir = '/home/Task03_Liver/labelsTr'

images = sorted(
glob.glob(os.path.join(data_dir, "liver*.nii.gz")))
labels = sorted(
glob.glob(os.path.join(data_dir, "liver*.nii.gz")))

print('No. of images:', len(images), ' labels:', len(labels))

img = nib.load(images[0]).get_fdata()
lbl = nib.load(labels[0]).get_fdata()
print('\nimg shape:', img.shape, ' lbl shape:', lbl.shape)
print('img intensity min.:', np.min(img), ' max.:', np.max(img), ' unique labels:', np.unique(lbl))

data_dicts = [
{"image": image_name, "label": label_name}
for image_name, label_name in zip(images, labels)
]

train_files, val_files, test_files = data_dicts[:93], data_dicts[93:119], data_dicts[119:131]

print('train files:', len(train_files), ' val files:', len(val_files), ' test files:', len(test_files))

train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
ScaleIntensityRanged(
keys=["image"], a_min=-1000, a_max=1000,
b_min=0.0, b_max=1.0, clip=True,
),

    # Change labels from 255 to 1
    ScaleIntensityRanged(
        keys=["label"], a_min=0, a_max=255,
        b_min=0.0, b_max=1.0, clip=True,
    ),
    CropForegroundd(keys=["image", "label"], source_key="image"),
    Orientationd(keys=["image", "label"], axcodes="PLS"),
    Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),

    # Randomly crop a patch from both image and label files
    RandSpatialCropd(
        keys=["image", "label"],
        roi_size=[128, 128, 128],
        random_size=False,
    ),

    # Data augmentation transforms
    RandAffined(
        keys=['image', 'label'],
        mode=('bilinear', 'nearest'),
        prob=1.0, spatial_size=(128, 128, 128),
        rotate_range=(0, 0, np.pi / 15),
        scale_range=(0.1, 0.1, 0.1)),
    RandFlipd(keys=["image", "label"], prob=0.25, spatial_axis=0),
    RandFlipd(keys=["image", "label"], prob=0.25, spatial_axis=1),
    RandFlipd(keys=["image", "label"], prob=0.25, spatial_axis=2),
    RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
    RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
])

val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
ScaleIntensityRanged(
keys=["image"], a_min=-1000, a_max=1000,
b_min=0.0, b_max=1.0, clip=True,
),
ScaleIntensityRanged(
keys=["label"], a_min=0, a_max=255,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),
Orientationd(keys=["image", "label"], axcodes="PLS"),
Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
])

train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True)

val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0)
val_loader = DataLoader(val_ds, batch_size=1)

device = torch.device("cuda:0")

model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=3, # 2 if using Softmax activation, 1 if using Sigmoid
channels=(32, 64, 128, 256, 320, 320),
strides=(1, 2, 2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(device)

loss_function = DiceLoss(to_onehot_y=True, softmax=True)

loss_function = DiceCELoss(to_onehot_y=True, softmax=True)

optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean",)
dice_metric2 = DiceMetric(include_background=True, reduction="none")

from monai.inferers import sliding_window_inference

Define hyperparameters

max_epochs = 500
val_interval = 2
best_metric = -1
classwise_best_metrics = None
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []

Define post-processing transforms

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

Training loop

for epoch in range(max_epochs):
print("-" * 10)
print(f"epoch {epoch + 1}/{max_epochs}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)

    loss.backward()
    optimizer.step()

    epoch_loss += loss.item()
    print(f"{step}/{len(train_ds) // train_loader.batch_size}, train_loss: {loss.item():.4f}")

epoch_loss /= step
epoch_loss_values.append(epoch_loss)

print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

if (epoch + 1) % val_interval == 0:
    model.eval()
    with torch.no_grad():
        for val_data in val_loader:
            val_inputs, val_labels = val_data["image"].to(device), val_data["label"].to(device),
            val_outputs = sliding_window_inference(val_inputs, (128, 128, 128), 1, model)

            val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
            val_labels = [post_label(i) for i in decollate_batch(val_labels)]

            # compute metric for current iteration
            dice_metric(y_pred=val_outputs, y=val_labels)
            dice_metric2(y_pred=val_outputs, y=val_labels)

        # aggregate the final mean dice result
        metric = dice_metric.aggregate().item()
        classwise_metric = dice_metric2.aggregate().mean(dim=0)

        # reset the status for next validation round
        dice_metric.reset()
        dice_metric2.reset()

        metric_values.append(metric)
        if metric > best_metric:
            best_metric = metric
            classwise_best_metrics = classwise_metric
            best_metric_epoch = epoch + 1
            torch.save(model.state_dict(), "../best_metric_model.pth")
            print("saved new best metric model")
        print(
            f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
            f"\nbest mean dice: {best_metric:.4f} "
            f"\nclasswise best mean dice: {classwise_best_metrics}"
            f"at epoch: {best_metric_epoch}"
        )

print(
f"train completed \nbest_metric: {best_metric:.4f} \n classwise best mean dice: {classwise_best_metrics}"
f"\nat epoch: {best_metric_epoch}")

`