Weird predicted segmentation maps for BraTS 2021 dataset

[soumbane]

Hi,

I am using MONAI's Swin UNETR architecture for prediction of segmentation masks using BraTS 2021 dataset. However, I get these weird segmentation maps predicted even though the dice score is very high (above 90%).

[soumbane]

I just want to do binary segmentation: background (channel 0) and whole tumor (channel 1). So, the transformation is:

For BraTS 2021 Dataset

import os
import json
import torch
from torch.utils.data import Dataset
from monai import transforms
from monai.config import KeysCollection
import numpy as np

class ConvertToTwoChannelsd(transforms.MapTransform):
def init(self, keys: KeysCollection):
super().init(keys)

def __call__(self, data):
    d = dict(data)
    for key in self.keys:
        img = d[key]
        if img.shape[0] > 1:
            # If the input is already multi-channel, we assume it's in the correct format
            continue
        
        # Create background channel (class 0)
        # Background is now only the original background (0)
        background = (img == 0)
        
        # Create whole tumor channel (class 1)
        # Whole tumor now includes original WT (2), TC (1), and ET (4)
        whole_tumor = (img == 1) | (img == 2) | (img == 4)
        
        # Stack the channels
        d[key] = torch.stack([background.float(), whole_tumor.float()], dim=0)
    return d

class BraTS2021Dataset(Dataset):
def init(self, config_path, mode='train', val_split=135):
config_path = os.path.abspath(config_path)
if not os.path.exists(config_path):
raise FileNotFoundError(f"Config file not found: {config_path}")

    with open(config_path, 'r') as f:
        config = json.load(f)
    
    self.dataset_path = config['dataset_path']
    self.mri_seqs = config['mri_seqs']
    self.label = config['label']
    self.file_type = config['file_type']
    self.filename = config['filename']
    self.roi_size = tuple(config['roi_size'])  # Ensure it's a tuple of 3 values
    self.num_samples = config['num_crops_per_image']
    self.is_random_crops = config['is_random_crops']
    self.mode = mode

    self.folders = [item for item in os.listdir(self.dataset_path) if item.startswith("BraTS2021_0")]
    self.valid_folders = []
    self.test_folders = []
    required_subfolders = set(self.mri_seqs + [self.label])
    
    for folder in self.folders:
        folder_path = os.path.join(self.dataset_path, folder)
        subfolders = set(os.listdir(folder_path))
        if required_subfolders.issubset(subfolders):
            self.valid_folders.append(folder)
        elif set(self.mri_seqs).issubset(subfolders):
            self.test_folders.append(folder)

    self.valid_folders = sorted(self.valid_folders)
    self.test_folders = sorted(self.test_folders)
    
    # Split the dataset
    np.random.seed(42)  # for reproducibility
    np.random.shuffle(self.valid_folders)
    self.val_folders = self.valid_folders[:val_split]
    self.train_folders = self.valid_folders[val_split:]

    if mode == 'train':
        self.folders = self.train_folders
        self.transforms = self.get_train_transforms()
    elif mode == 'val':
        self.folders = self.val_folders
        self.transforms = self.get_val_transforms()
    elif mode == 'test':
        self.folders = self.test_folders
        self.transforms = self.get_test_transforms()
    else:
        raise ValueError("Mode must be either 'train', 'val', or 'test'")

def get_train_transforms(self):
    common_transforms = [
        transforms.LoadImaged(keys=self.mri_seqs + [self.label]),
        transforms.EnsureChannelFirstd(keys=self.mri_seqs + [self.label]),
        transforms.Orientationd(keys=self.mri_seqs + [self.label], axcodes="RAS"),
        transforms.NormalizeIntensityd(keys=self.mri_seqs, nonzero=True, channel_wise=True),
        transforms.Spacingd(keys=self.mri_seqs + [self.label], pixdim=(1.0, 1.0, 1.0), mode=("bilinear",) * len(self.mri_seqs) + ("nearest",)),
        transforms.RandFlipd(keys=self.mri_seqs + [self.label], spatial_axis=[0], prob=0.5),
        transforms.RandFlipd(keys=self.mri_seqs + [self.label], spatial_axis=[1], prob=0.5),
        transforms.RandFlipd(keys=self.mri_seqs + [self.label], spatial_axis=[2], prob=0.5),
        transforms.RandRotate90d(keys=self.mri_seqs + [self.label], prob=0.5, max_k=3),
        transforms.RandScaleIntensityd(keys=self.mri_seqs, factors=0.1, prob=0.5),
        transforms.RandShiftIntensityd(keys=self.mri_seqs, offsets=0.1, prob=0.5),
        ConvertToTwoChannelsd(keys=[self.label]),
        transforms.EnsureTyped(keys=self.mri_seqs + [self.label]),
    ]

    if self.is_random_crops:
        random_crop_transform = [
            transforms.RandCropByPosNegLabeld(
                keys=self.mri_seqs + [self.label],
                label_key=self.label,
                spatial_size=self.roi_size,
                pos=1,
                neg=0,
                num_samples=self.num_samples,
                image_key=self.mri_seqs[0],
                image_threshold=0
            )
        ]
        return transforms.Compose(common_transforms[:5] + random_crop_transform + common_transforms[5:])
    else:
        return transforms.Compose(common_transforms)

def get_val_transforms(self):
    return transforms.Compose([
        transforms.LoadImaged(keys=self.mri_seqs + [self.label]),
        transforms.EnsureChannelFirstd(keys=self.mri_seqs + [self.label]),
        transforms.Orientationd(keys=self.mri_seqs + [self.label], axcodes="RAS"),
        transforms.NormalizeIntensityd(keys=self.mri_seqs, nonzero=True, channel_wise=True),
        ConvertToTwoChannelsd(keys=[self.label]),
        transforms.EnsureTyped(keys=self.mri_seqs + [self.label]),
    ])

def get_test_transforms(self):
    return transforms.Compose([
        transforms.LoadImaged(keys=self.mri_seqs),
        transforms.EnsureChannelFirstd(keys=self.mri_seqs),
        transforms.Orientationd(keys=self.mri_seqs, axcodes="RAS"),
        transforms.NormalizeIntensityd(keys=self.mri_seqs, nonzero=True, channel_wise=True),
        transforms.EnsureTyped(keys=self.mri_seqs),
    ])
    
def __len__(self):
    if self.mode == 'train' and self.is_random_crops:
        return len(self.folders) * self.num_samples
    else:
        return len(self.folders)

def __getitem__(self, idx):
    if self.mode == 'train' and self.is_random_crops:
        folder_idx = idx // self.num_samples
        crop_idx = idx % self.num_samples
    else:
        folder_idx = idx
        crop_idx = 0

    patient_folder = self.folders[folder_idx]
    patient_path = os.path.join(self.dataset_path, patient_folder)
    
    if self.mode == 'test':
        data_dict = {key: os.path.join(patient_path, key, self.file_type, self.filename) 
                     for key in self.mri_seqs}
    else:
        data_dict = {key: os.path.join(patient_path, key, self.file_type, self.filename) 
                     for key in self.mri_seqs + [self.label]}
    
    data = self.transforms(data_dict)
    
    if self.mode == 'train' and self.is_random_crops:
        data = data[crop_idx]
    
    images = torch.cat([data[key] for key in self.mri_seqs], dim=0)
    
    if self.mode == 'test':
        return {'image': images, 'patient_id': patient_folder}
    else:
        label = data[self.label].squeeze(1)
        return {'image': images, 'label': label}

And the architecture is:
import torch
from monai.networks.nets import SwinUNETR

model = SwinUNETR(
img_size=config.roi_size,
in_channels=4,
out_channels=2,
feature_size=12,
drop_rate=0.0,
attn_drop_rate=0.0,
dropout_path_rate=0.0,
use_checkpoint=False,
)

[soumbane]

This is my validate script:
def validate(epoch):
model.eval()
dice_meter = AverageMeter()
hd95_meter = AverageMeter()
val_loss_meter = AverageMeter()
start_time = time.time()

with torch.no_grad():
    for idx, batch_data in enumerate(val_loader):
        data, target = batch_data["image"].to(device), batch_data["label"].to(device)
        print(f'Val. Batch_data image data shape: {batch_data["image"].shape}')
        print(f'Val. Batch_data Target label shape: {batch_data["label"].shape}')
        show_image_sample(batch_data)
        show_label_sample(batch_data)

logits = model_inferer(data)

        logits = sliding_window_inference(data, (128,128,128), 1, model, overlap=0.5)
        print(f'Logits shape: {logits.shape}')
        
        loss = criterion(logits, target)
        val_loss_meter.update(loss.item(), n=config.val_batch_size)

        val_labels_list = decollate_batch(target)
        print(f'Length of GT-labels: {len(val_labels_list)}')
        print(f'Val. output GT-Labels decollated shape: {val_labels_list[0].shape}')
        
        print(f'Showing output GT-Labels')
        show_sample(val_labels_list[0])
        
        val_outputs_list = decollate_batch(logits)
        print(f'Length of Predictions: {len(val_outputs_list)}')
        print(f'Val. output predictions decollated shape: {val_outputs_list[0].shape}')

        val_output_convert = [post_pred(post_sigmoid(val_pred_tensor)) for val_pred_tensor in val_outputs_list]
        
        print(f'Val. output converted predictions shape: {val_output_convert[0].shape}')
        
        print(f'Showing output Predictions')
        show_sample(val_output_convert[0])
        
        dice_metric.reset()
        hd95_metric.reset()
        dice_metric(y_pred=val_output_convert, y=val_labels_list)
        hd95_metric(y_pred=val_output_convert, y=val_labels_list)
        dice, dice_not_nans = dice_metric.aggregate()
        print(f'Dice score: {dice}')
        print(f'Dice not nans: {dice_not_nans}')
        
        raise ValueError("Stop Validation")
        
        hd95, hd95_not_nans = hd95_metric.aggregate()
        dice_meter.update(dice.cpu().numpy(), n=dice_not_nans.cpu().numpy())
        hd95_meter.update(hd95.cpu().numpy(), n=hd95_not_nans.cpu().numpy())

        print(
            f"Val {epoch + 1}/{config.max_epochs} {idx}/{len(val_loader)}",
            f"loss: {val_loss_meter.avg:.4f}",
            f"dice: {dice_meter.avg}",
            f"hd95: {hd95_meter.avg}",
            f"time {time.time() - start_time:.2f}s"
        )
        start_time = time.time()

return dice_meter.avg, hd95_meter.avg, val_loss_meter.avg

with the post transformations as: post_sigmoid = Activations(sigmoid=True)
post_pred = AsDiscrete(argmax=False, threshold=0.5)
dice_metric = DiceMetric(include_background=True, reduction=MetricReduction.MEAN_BATCH, get_not_nans=True)

[soumbane]

I am using most of the code from here: "https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb" but with custom transforms for binary segmentations

[KumoLiu]

Hi @soumbane, I would suggest you check your data after transforms first, please ensure all the data have ensure channel first.
And also perhaps try the notebook without modification first and then adapt with your own data.

Thanks.

[soumbane]

yes channel first is already there. does the trainer code here have some problems? "import time
import torch
import numpy as np
from monai.losses import DiceCELoss
from monai.inferers import sliding_window_inference
from monai.transforms import AsDiscrete, Activations
from monai.metrics import DiceMetric
from monai.utils.enums import MetricReduction
from monai.data import decollate_batch
from functools import partial
from .utils import AverageMeter, save_checkpoint, load_checkpoint
from .configuration import config

class Trainer:
def init(self, model, train_loader, val_loader, config, logger=None, resume_from=None):
self.model = model
self.train_loader = train_loader
self.val_loader = val_loader
self.config = config
self.logger = logger
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)

    self.criterion = DiceCELoss(include_background=True, to_onehot_y=False, softmax=True, sigmoid=False, lambda_dice=1.0, lambda_ce=1.0)
    self.optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
    self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=config.max_epochs)

    self.post_sigmoid = Activations(softmax=True)
    self.post_pred = AsDiscrete(argmax=True, to_onehot=2)
    self.dice_metric = DiceMetric(include_background=True, reduction=MetricReduction.MEAN_BATCH, get_not_nans=True)

    self.model_inferer = partial(
        sliding_window_inference,
        roi_size=config.roi_size,
        sw_batch_size=config.sw_batch_size,
        predictor=model,
        overlap=config.infer_overlap,
    )

    self.start_epoch = 0
    self.best_metric = -1
    self.best_metrics_all_classes = None
    
    if resume_from:
        self.model, self.optimizer, self.scheduler, self.start_epoch, self.best_metric, self.best_metrics_all_classes = load_checkpoint(
            self.model, self.optimizer, self.scheduler, self.config, resume_from
        )

    if self.logger is not None:
        self.logger.log_model(self.model)

def train_epoch(self, epoch):
    self.model.train()
    loss_meter = AverageMeter()
    start_time = time.time()

    for idx, batch_data in enumerate(self.train_loader):
        data, target = batch_data["image"].to(self.device), batch_data["label"].to(self.device)
        self.optimizer.zero_grad()
        logits = self.model(data)
        loss = self.criterion(logits, target)
        loss.backward()
        self.optimizer.step()
        loss_meter.update(loss.item(), n=self.config.train_batch_size)
        
        print(
            f"Epoch {epoch + 1}/{self.config.max_epochs} {idx}/{len(self.train_loader)}",
            f"loss: {loss_meter.avg:.4f}",
            f"time {time.time() - start_time:.2f}s"
        )
        start_time = time.time()

    return loss_meter.avg

def validate(self, epoch):
    self.model.eval()
    dice_meter = AverageMeter()
    val_loss_meter = AverageMeter()
    start_time = time.time()

    with torch.no_grad():
        for idx, batch_data in enumerate(self.val_loader):
            data, target = batch_data["image"].to(self.device), batch_data["label"].to(self.device)
            logits = self.model_inferer(data)
            loss = self.criterion(logits, target)
            val_loss_meter.update(loss.item(), n=self.config.val_batch_size)
            
            val_labels_list = decollate_batch(target)
            val_outputs_list = decollate_batch(logits)

            val_output_convert = [self.post_pred(self.post_sigmoid(val_pred_tensor)) for val_pred_tensor in val_outputs_list]
            self.dice_metric.reset()
            self.dice_metric(y_pred=val_output_convert, y=val_labels_list)
            dice, dice_not_nans = self.dice_metric.aggregate()
            dice_meter.update(dice.cpu().numpy(), n=dice_not_nans.cpu().numpy())

            print(
                f"Val {epoch + 1}/{self.config.max_epochs} {idx}/{len(self.val_loader)}",
                f"loss: {val_loss_meter.avg:.4f}",
                f"dice: {dice_meter.avg}",
                f"time {time.time() - start_time:.2f}s"
            )
            start_time = time.time()

    return dice_meter.avg, val_loss_meter.avg

def train(self):
    print("Training starting...")
    best_metric = self.best_metric
    best_metric_epoch = self.start_epoch
    best_metrics_all_classes = self.best_metrics_all_classes

    for epoch in range(self.start_epoch, self.config.max_epochs):
        print(f"Epoch {epoch + 1}/{self.config.max_epochs}")
        train_loss = self.train_epoch(epoch)
        
        if self.logger is not None:
            self.logger.log_train_epoch(train_loss, epoch + 1)
        
        if (epoch + 1) % self.config.val_every == 0:
            val_metric, val_loss = self.validate(epoch)
            val_metric_mean = np.mean(val_metric)
            if self.logger is not None:
                self.logger.log_validation(val_loss, val_metric, epoch + 1)

            if val_metric_mean > best_metric:
                best_metric = val_metric_mean
                best_metric_epoch = epoch + 1
                best_metrics_all_classes = val_metric
                save_checkpoint(self.model, self.optimizer, self.scheduler, epoch, 
                                best_acc=best_metric, 
                                best_metrics_all_classes=best_metrics_all_classes,
                                config=self.config)
                
                if self.logger is not None:
                    self.logger.log_best_metrics(best_metric, best_metrics_all_classes, best_metric_epoch)
        
            print(
                f"Current epoch: {epoch + 1} current mean dice: {val_metric_mean:.4f}"
                f"\nCurrent dice - bg: {val_metric[0]:.4f}, wt: {val_metric[1]:.4f}"
                f"\nBest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
            )
            
            if best_metrics_all_classes is not None:
                print(f"Best dice - bg: {best_metrics_all_classes[0]:.4f}, wt: {best_metrics_all_classes[1]:.4f}")
        
        self.scheduler.step()
        
        if self.logger is not None:
            self.logger.log_lr(self.optimizer.param_groups[0]['lr'], epoch + 1)

    print(f"Training completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
    if best_metrics_all_classes is not None:
        print(f"Best dice scores - bg: {best_metrics_all_classes[0]:.4f}, wt: {best_metrics_all_classes[1]:.4f}")
    
    return best_metric, best_metrics_all_classes"

[soumbane]

This issue is solved