Overlapping Classes in BRATS Segmentation Task

[amida47]

in the brats_segmentation_3d.ipynb tutorial, you made a transformation from the provided target classes in the brats dataset in the function ConvertToMultiChannelBasedOnBratsClassesd, so we went from (peritumoral edema, GD-enhancing tumor, necrotic and non-enhancing tumor core) to (Tumor core, Whole tumor, Enhancing tumor), the problem I have with this is that it would create an overlapping classes issue, for example a pixel that belongs to the whole tumor class is definitely going to belong also to either the tumor core or the enhancing tumor class, so wouldn't this confuse the model and make it harder to learn?

[KumoLiu]

Hi @amida47, thanks for your interest here, actually, this shouldn't be an issue. The transformation changes the label channels to match the BRATS evaluation script. The fact that classes overlap is taken into account when training the model and during inference.
Consider this scenario. If a voxel belongs to the "whole tumor" class, it cannot belong to the "tumor core" or "enhancing tumor" class at the same time. This is because each pixel's class is mutually exclusive in this scenario. So, during training, the model will learn the appropriate representations for distinguishing voxels as either "whole tumor", "tumor core", or "enhancing tumor".
Many medical segmentation tasks have similar setups where one could argue that there is a hierarchical relationship between classes. For example, "tumour" would be a superset of "necrotic tumour tissue" and "enhancing tumour". Yet, the models seem to perform fairly well in practice on these tasks.
The key point is that when people manually annotate these images, they don't label a single voxel with multiple classes. Rather, they decide the most suitable class for each voxel. The model will learn to mimic this behavior.
Remember, the loss function, such as multi-class cross entropy or Dice with softmax, during training ensures that only one class is predicted per voxel. These classes are mutually exclusive even though conceptually one might consider a certain hierarchy or overlap among them. They are not simultaneous. In the output of the model, each voxel will have probabilities assigned to each class, and the class with the highest probability will be assigned to that voxel.
The metric calculations during model evaluation also take this into account, and they are done on a voxel-by-voxel basis. This design is not a problem but rather a standard, practical approach for such kind of segmentation tasks in medical image analysis.

Thanks,
Yun

[amida47]

hey @KumoLiu .
thank for taking the time to answer my question.
I am sorry but I couldn't quite understand your answer, first we both agree that creating a target labels with tumorCore, wholeTumor and enhancingTumor will definitely create a hierarchy between the classes. so just to verify this idea I wrote this code, I hope it's correct:

from util import get_MyIterableDataset
import torch
config = {
    'num_classes': 4,
    'batch_size': 50,
    'num_of_workers': 2,
    'learning_rate': 1e-4,
    'epochs': 2,
}
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


train_dataset, test_dataset = get_MyIterableDataset(
    'ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData'
)

train_loader = torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=config['batch_size'], 
    num_workers=config['num_of_workers']
)
val_loader = torch.utils.data.DataLoader(
    test_dataset, 
    batch_size=config['batch_size'], 
    num_workers=config['num_of_workers']
)
nb_slices = 0
nb_batches = 0
for inputs, target in train_loader:
    nb_batches += 1
    # collapse the channel dimension so if two or more pixels are labeled with multiple classes,
    # they are summed and sum should be greater than 1
    target = torch.sum(target, dim=1)
    
    # if the sum is greater than 1, then the pixel is labeled with multiple classes
    res = target > 1
    # collapse the spatial dimensions
    res = torch.sum(res, dim=(1,2))
    
    # if the collapsed slice has a pixel labeled with multiple classes, then the slice 
    # then the sum of the slice should be greater than 0
    res = res > 0
    
    # the number of slices with pixels labeled with multiple classes
    nb_slices += torch.sum(res)
print(f"Number of batches: {nb_batches}")
print(nb_slices)

so the code only iterates over the 10 first examples of the brats dataset, it first takes the batch of segmentation target, then collapse the channel dim by sum [batch_size, 4, 240, 240] --> [batch_size, 240, 240], this way if a pixel is labeled with more than one class the resulting slice would have pixels with a value greater than one, given this a check for these pixels that have a value greater than 1, assign them 1 if yes otherwise 0, finally I collapse the spatial dims by sum [batch_size, 240, 240] --> [batch_size], the logic here is if a slice didn't have a pixel labeled with multiple classes then the sum should be equal to zero, using this I look for slices that have a sum greater than one and then by using another sum I get the number of slices that did have a pixel labeled with more than one class.
in this example I ended up with 268 slices out of the 1000 slices processed. this just to prove that indeed some pixels are labeled with more than one class.