How to deal with multi class (one-file-)label-masks?

[desserdmi]

Hi guys,

I have a very similar problem with 3D segmentation task as already discussed in issue #415, but the solution does not work for me..

My input is a 3D image (CT scan) and multi-class label segmentation mask containing 3 labels: background, bleeding and edema, labeled as [0,1,2].

I am trying to apply 3D spleen segmentation tutorial and I have already changed the parameter to softmax=True, but I still get the error: IndexError: index 5 is out of bounds for dimension 4.

Does anybody maybe have an idea?

Any help is appreciated and thank you very much in advance!

[Nic-Ma]

Hi @desserdmi ,

Thanks for your interest here, I didn't see some problem in your description.
Could you please help share your test program for further debug?

Thanks in advance.

[desserdmi]

Define model, loss function and optimizer

device = torch.device("cpu")
model = UNet(
dimensions=3,
in_channels=1,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean")

Run the training

max_epochs = 10
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=False, n_classes=2)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=False, n_classes=2)])

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}, "
f"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():
        metric_sum = 0.0
        metric_count = 0
        for val_data in val_loader:
            val_inputs, val_labels = (
                val_data["image"].to(device),
                val_data["label"].to(device),
            )
            roi_size = (160, 160, 160)
            sw_batch_size = 4
            val_outputs = sliding_window_inference(
                val_inputs, roi_size, sw_batch_size, 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)

        # aggregate the final mean dice result
        metric = dice_metric.aggregate().item()
        # reset the status for next validation round
        dice_metric.reset()

        metric_values.append(metric)
        if metric > best_metric:
            best_metric = metric
            best_metric_epoch = epoch + 1
            torch.save(model.state_dict(), os.path.join(
                root_dir, "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"at epoch: {best_metric_epoch}"
        )

results in error:

epoch 1/10

---

RuntimeError Traceback (most recent call last)
in
22 optimizer.zero_grad()
23 outputs = model(inputs)
---> 24 loss = loss_function(outputs, labels)
25 loss.backward()
26 optimizer.step()

~/miniconda3/envs/edema/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []

~/miniconda3/envs/edema/lib/python3.8/site-packages/monai/losses/dice.py in forward(self, input, target)
130 warnings.warn("single channel prediction, to_onehot_y=True ignored.")
131 else:
--> 132 target = one_hot(target, num_classes=n_pred_ch)
133
134 if not self.include_background:

~/miniconda3/envs/edema/lib/python3.8/site-packages/monai/networks/utils.py in one_hot(labels, num_classes, dtype, dim)
84
85 o = torch.zeros(size=sh, dtype=dtype, device=labels.device)
---> 86 labels = o.scatter_(dim=dim, index=labels.long(), value=1)
87
88 return labels

RuntimeError: index 3 is out of bounds for dimension 1 with size 3

Thanks you very much!

[Nic-Ma]

Hi @desserdmi ,

Could you please help share your transforms for label?
I think the issue is that: your transforms didn't ensure the label to be channel first shape, or your label is already One-Hot format so no need to_onehot_y =True.

Thanks.

[desserdmi]

Sure! here are my transforms:

train_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Spacingd(keys=["image", "label"], pixdim=(
1.5, 1.5, 2), mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(
keys=["image"], a_min=0, a_max=80,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),

    RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=1.0, spatial_size=(32, 32, 32),
             rotate_range=(0, np.pi/30, 0), translate_range=(10, 10, 10), padding_mode='zeros'),
    EnsureTyped(keys=["image", "label"]),
]

)
val_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Spacingd(keys=["image", "label"], pixdim=(
1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(
keys=["image"], a_min=0, a_max=80,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),
EnsureTyped(keys=["image", "label"]),
]
)

[Nic-Ma]

Hi @desserdmi ,

Could you please help print out the shape of label before loss function?
I think using EnsureChannelFirstd instead of AddChanneld may fix your issue.

Thanks.

[desserdmi]

Thank you very much for your response! Indeed, it seems that my images were of different shapes. Now, I have insured that both images and labels are of the same shape and orientation
and here are the tensor shapes:
image: torch.Size([2, 1, 256, 256, 256])
label: torch.Size([2, 1, 256, 256, 256]).

But there seems still to be a shape problem directly at the start of the training:
AssertionError: ground truth has different shape (torch.Size([2, 1, 256, 256, 256])) from input (torch.Size([2, 3, 256, 256, 256]))

I do not really understand why and at which step the shape of label tensor changes from [2, 1, 256, 256, 256] to [2, 3, 256, 256, 256]

Are there any transforms missing?

Thank you very much!

[Nic-Ma]

Hi @desserdmi ,

Seems you need to set to_onehot_y =True in the loss function?
Could you please help share the full program and error log to check?

Thanks in advance.

[desserdmi]

Hope the link works:

https://colab.research.google.com/drive/1CE50SrSY-z6hRlnihJBkZH1HRLpeo-wr?usp=sharing

Error log:

AssertionError Traceback (most recent call last)
in
25 optimizer.zero_grad()
26 outputs = model(inputs)
---> 27 loss = loss_function(outputs, labels)
28 loss.backward()
29 optimizer.step()

~/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []

~/miniconda3/lib/python3.9/site-packages/monai/losses/dice.py in forward(self, input, target)
141
142 if target.shape != input.shape:
--> 143 raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
144
145 # reducing only spatial dimensions (not batch nor channels)

AssertionError: ground truth has different shape (torch.Size([2, 1, 256, 256, 256])) from input (torch.Size([2, 3, 256, 256, 256]))

---

setting to_onehot_y =True results in error:
RuntimeError: index 3 is out of bounds for dimension 1 with size 3

[Nic-Ma]

Hi @desserdmi ,

I think you must set to_onehot_y=True as this is a multi classes task.
I think maybe your label values are [1, 2, 3] instead of [0, 1, 2]? Could you please help verify it?
If yes, you can use the MapLabelValued transform to convert it to [0, 1, 2].

Thanks.

[desserdmi]

Hi,
thank you very much! Indeed, this seems to be the reason for the error..
I just realized that my labels are inconsistent over the dataset..

e.g.:
label-mask 1 = [0, 1, 2];
label-mask 2 = [0, 1, 2, 3]
label-mask 3 = [0, 1, 3]

is there any option to handle this label data inconsistency?

Thanks again for your advice!

[Nic-Ma]

No problem.
In your case, I think maybe you should output 4 classes, some labels may miss 1 or 2 classes, that's common.

Thanks.