Multi-class 3D segmentation visualization in Patch-Based training

[Nastii22]

Hello,

I'm working on a multi-class 3D segmentation project which contains 3 classes, including background.
I've utilized Medical image segmentation with TorchIO, MONAI & PyTorch Lightning the example provided in GitHub repository (https://github.com/fepegar/torchio-notebooks/blob/main/notebooks/TorchIO_MONAI_PyTorch_Lightning.ipynb). The key difference in my approach is that I'm training the model in a patch-based manner as provided in TorchIO: a tutorial, GitHub repository (https://github.com/fepegar/torchio-notebooks/blob/main/notebooks/TorchIO_tutorial.ipynb).

Could you help me understand how to modify the testing part of the patch-based method to fit multi-class segmentation? Specifically, I need to visualize and check the model output for each class alongside the input image and the corresponding labels for each class.

Any assistance would be greatly appreciated.

[fepegar]

Hi, @Nastii22. Do you need help with visualization only, or something else?

[fepegar]

If you share an image, label and prediction, we'll be able to help better.

[Nastii22]

Hi @fepegar thanks for coming back to me so quickly.
The dataset I'm using consists of CT images of the brain, which contain two labels: infarction and hemorrhagic areas. Considering the background, it's a multi-class 3D segmentation project. For easier follow-up, I put my code here.

Data

def change_img_to_label_path(path):
    """
    Replaces imagesTr with labelsTr
    """
    parts = list(path.parts) 
    parts[parts.index("imagesTr")] = "labelsTr"  
    return Path(*parts) 

path = Path("/Data/imagesTr")
subjects_paths = list(path.glob("PAT_*"))
subjects = []

for subject_path in subjects_paths:
    label_path = change_img_to_label_path(subject_path)
    subject = tio.Subject({"CT":tio.ScalarImage(subject_path), "Label":tio.LabelMap(label_path)})
    subjects.append(subject)    
dataset = tio.SubjectsDataset(subjects)

# Training and validation split
num_subjects = len(dataset)
num_train_subjects = int(train_val_ratio * num_subjects)
num_val_subjects = num_subjects - num_train_subjects
num_split_subjects = num_train_subjects, num_val_subjects
train_subjects, val_subjects = torch.utils.data.random_split(subjects, num_split_subjects)

# Preprocessing
process = tio.Compose([
            tio.ToCanonical(),
            tio.Resample('CT'),
            tio.CopyAffine('CT'),
            tio.CropOrPad((512, 512, 35)),   
            tio.EnsureShapeMultiple(4),  
            tio.RescaleIntensity((-1, 1)), 
            tio.OneHot(3),
            ])

augmentation = tio.Compose([
                tio.RandomAnisotropy(p=0.25),  
                tio.OneOf({
                    tio.RandomAffine():0.8,
                    tio.RandomElasticDeformation():0.2,
                }),
                tio.RandomFlip(),
                tio.RandomNoise(p=0.25),
                tio.RandomGamma(p=0.5),
                ])

train_transform = tio.Compose([process, augmentation])
val_transform = process

train_dataset = tio.SubjectsDataset(train_subjects, transform=train_transform)
val_dataset = tio.SubjectsDataset(val_subjects, transform=val_transform)

# Patching
sampler = tio.data.LabelSampler(patch_size, label_name="Label", label_probabilities={0:0.2, 1:0.3, 2:0.5})

train_patches_queue = tio.Queue(
     train_dataset,
     max_length=40,   
     samples_per_volume=5, 
     sampler=sampler,
     num_workers=4
    )

val_patches_queue = tio.Queue(
     val_dataset,
     max_length=40,
     samples_per_volume=5,
     sampler=sampler,
     num_workers=4
    )

train_loader = torch.utils.data.DataLoader(train_patches_queue, batch_size=train_batch_size, num_workers=0)
val_loader = torch.utils.data.DataLoader(val_patches_queue, batch_size=val_batch_size, num_workers=0)

The below images are sample subjects of training dataset.

Training

# Training step
class Segmenter(pl.LightningModule):
    def __init__(self, model, learning_rate, loss_function, optimizer_name, scheduler_name):
        super().__init__()
        self.save_hyperparameters()
        self.model = self.hparams.model
        self.lr = self.hparams.learning_rate
        self.loss_fn = self.hparams.loss_function
        self.optimizer_name = self.hparams.optimizer_name
        self.scheduler_name = self.hparams.scheduler_name

       self.training_step_outputs = [] 
        self.validation_step_outputs = []  
         
        
    def forward(self, data):
        pred = self.model(data)
        return pred
          
    def training_step(self, batch, batch_idx):       
        img = batch["Image"]["data"]
        mask = batch["Label"]["data"]   
        mask = mask.long()    
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        self.training_step_outputs.append(loss)
        self.log("Train Loss", loss, on_step=True, on_epoch=True, logger=True)
        return loss

   def on_train_epoch_end(self):
        train_epoch_average = torch.stack(self.training_step_outputs).mean()
        self.logger.experiment.add_scalars('loss/epoch', {'Train': train_epoch_average}, global_step=self.current_epoch)
        self.training_step_outputs.clear()
        
    def validation_step(self, batch, batch_idx):        
        img = batch["Image"]["data"]
        mask = batch["Label"]["data"]   
        mask = mask.long()  
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        preds = [self.post_pred(i) for i in decollate_batch(pred)]
        masks = [self.post_label(i) for i in decollate_batch(mask)]
        self.validation_step_outputs.append(loss)     
        self.log("Val Loss", loss, on_step=True, on_epoch=True, logger=True)
        return loss

   def on_validation_epoch_end(self):
        val_epoch_average = torch.stack(self.validation_step_outputs).mean()
        self.logger.experiment.add_scalars('loss/epoch', {'Val':val_epoch_average}, global_step=self.current_epoch)
        self.validation_step_outputs.clear()
    
    def configure_optimizers(self):
        optimizer = self.optimizer_name(self.parameters(), lr = self.lr)
        lr_scheduler = {
        "scheduler": self.scheduler_name(optimizer, gamma=0.99),
        "interval": "epoch",
        "monitor": "val_loss",
        }
        return [optimizer], [lr_scheduler]

# ###
unet = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=3, 
    channels=(16, 32, 64),
    strides=(2, 2),
)

if __name__ == "__main__":
       
    model = Segmenter(
                model=unet,
                learning_rate=1e-4,
                loss_function=DiceLoss(include_background=False, to_onehot_y=False, softmax=True,  
                reduction='mean', smooth_nr=1e-5, smooth_dr=1e-5,  weight=[2, 10]),
                optimizer_name=torch.optim.AdamW,
                scheduler_name=torch.optim.lr_scheduler.ExponentialLR
    )

checkpoint_callback = ModelCheckpoint(
    monitor='Val Loss',
    save_top_k=10,  
    mode='min')
lr_monitor = LearningRateMonitor()
callbacks = [checkpoint_callback, lr_monitor]
logger = TensorBoardLogger(save_dir="./logs")
trainer = pl.Trainer(devices=1, accelerator="gpu",
                         logger=logger,
                         log_every_n_steps=1,
                         check_val_every_n_epoch=1,
                         callbacks=callbacks,
                         max_epochs=num_epochs

                        )
trainer.fit(model, train_loader, val_loader)

Evaluation
My question is related to the evaluation part, which I am not sure about, as the visualization shows something messed up.

# Evaluation on one validation subject

model = Segmenter.load_from_checkpoint("/checkpoints/epoch=88-step=18690.ckpt")
model = model.eval()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device);

# Patch Aggregation
# Select a validation subject and extract the images and segmentation for evaluation
IDX = 4
subject = val_dataset[IDX]
imgs = subject["CT"]["data"]
mask = subject["Label"]["data"]
# GridSampler
grid_sampler = tio.inference.GridSampler(val_dataset[IDX], patch_size=(96, 96, 8), patch_overlap=(2, 2, 2))
# GridAggregator
aggregator = tio.inference.GridAggregator(grid_sampler)
patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=val_batch_size)
# Prediction
CHANNELS_DIMENSION = 1
with torch.no_grad():
    for patches_batch in patch_loader:
        input_tensor = patches_batch['CT']["data"].to(device)  
        locations = patches_batch[tio.LOCATION]  
        probabilities = model(input_tensor).softmax(dim=CHANNELS_DIMENSION) 
        aggregator.add_batch(probabilities, locations)  
# Extract the volume prediction
output_tensor = aggregator.get_output_tensor()

# Post-Processing
post_label = Compose([EnsureType("tensor", device="cpu"), AsDiscrete(argmax=True, to_onehot=3)])(mask)
post_pred = Compose([EnsureType("tensor", device="cpu"), AsDiscrete(argmax=True, to_onehot=3)])(output_tensor)

The image below is the output (prediction mask), which is incorrect and should only depict the segmented portion. I would be thankful if you could please help me with the issue.

[romainVala]

Hi

the code looks fine to me, the aggregation should work. (patch over_lap of 2 voxel is small, but this is not the problem)
may be the problem comes from somewhere else

[Nastii22]

Hi @romainVala , thank you for your time and concern.
It is really strange for me, and I've been trying to solve it, but unfortunately, I have not succeeded. I would be grateful if you could share any thoughts or suggestions you might have.

[romainVala]

Does the model learn something, I guess the training loss is going down ? (what the end DICE score on train and validation ?)

I had strange results in the past, when I forget the softmax. but here it looks fine, there are no sofmax in the model, you add it before computing the training loss and also after the prediction, .... all good

Not sure if necessary, but in my code, I use

prediction = aggregator.get_output_tensor().unsqueeze(0)

what is the shape of your output_tensor ?
I am not used with AsDiscrete so not sure if it matter or not ...

[Nastii22]

Dear @romainVala , thanks for coming back to me so quickly.

The image below is loss/epoch curves for train and validation over 100 epochs.

I edited the post-processing part because I had already done the OneHot transformation in the preprocessing, so I think there was no need to do it again. But the image of the output (prediction mask) remained the same.

The mask.shape is torch.Size([3, 512, 512, 36]) and torch.unique(mask) is tensor([0, 1], dtype=torch.int32).
The output_tensor.shape is torch.Size([3, 512, 512, 36]) and torch.unique(output_tensor) is tensor([4.0821e-24, 5.4153e-24, 2.1569e-23, ..., 1.0000e+00, 1.0000e+00, 1.0000e+00]).

As Mean Dice of Monai, DiceMetric, expects y_pred to have binarized predictions, I changed post processing as below:

# Post-Processing
post_label = mask
transform =  AsDiscrete(threshold=0.5)
post_pred = transform(output_tensor)

Now, torch.unique(post_pred) is metatensor([0., 1.]).
But there is not any change on the visualization and it is the same as before.

For Dice score, I put my code and results here for easier follow-up.

Dice score calculation for the specific validation subject

IDX = 4
subject = val_dataset[IDX]
mask = subject["Label"]["data"]
GT_mask = mask.float().unsqueeze(0)  # add batch dimension

transform = AsDiscrete(threshold=0.5)
PRED_mask = (transform(output_tensor)).unsqueeze(0)  # add batch dimension

dice_metric = DiceMetric(include_background=False, reduction='mean', get_not_nans=True, ignore_empty=True, num_classes=3)
dice = dice_metric(PRED_mask.to(device), GT_mask.to(device))

metric, not_nans_count = dice_metric.aggregate()
print(metric)
print(not_nans_count)

classes = ['class1', 'class2']  
for i, row in enumerate(dice):
    for j, element in enumerate(row):
        print(f"{classes[j]}: {element.item()}")

GT_mask shape: torch.Size([1, 3, 512, 512, 36])
unique values of GT_mask: tensor([0., 1.])

PRED_mask shape: torch.Size([1, 3, 512, 512, 36])
unique values of PRED_mask: metatensor([0., 1.])

tensor([0.0103])
tensor(1.)

dice shape: torch.Size([1, 2])
class1: 0.010345708578824997
class2: nan

Evaluation for all validation data

dice_scores = []  
non_nan_counts = []  

for PAT_idx, PAT in enumerate(val_dataset):
  
    with torch.no_grad(): 
        img = PAT["CT"]["data"]
        mask = PAT["Label"]["data"]
       
        Gt_mask = mask.unsqueeze(0)  # add batch dimension 
        
        transform = AsDiscrete(threshold=0.5)
        pred = transform(output_tensors_list[PAT_idx])
        pred_mask = pred.unsqueeze(0)  # add batch dimension 
        
         # Calculate Dice score
        dice_metric = DiceMetric(include_background=False, reduction='mean', get_not_nans=True, ignore_empty=True, num_classes=3)
        dice = dice_metric(pred_mask, GT_mask)

        dice_scores.append(dice)  
        
        # calculate mean dice score and count of non-nan values for each patient
        metric, not_nans_count = dice_metric.aggregate()

        print(f"Mean Dice Score for patient {PAT_idx}: {metric}")
        print(f"Non-nan Count for patient {PAT_idx}: {not_nans_count}")
        
# convert the list of dice scores to a tensor        
dice_scores_tensor = torch.stack(dice_scores)       


# calculate mean dice score across all patients and classes
mean_dice_score = torch.nanmean(dice_scores_tensor, dim=(0, 1)) 

for class_idx, dice_score in enumerate(mean_dice_score):
    print(f"Mean Dice score for class {class_idx}: {dice_score.item()}")

Mean Dice score for class 0: 0.006322517991065979
Mean Dice score for class 1: nan

Thank you in advance for your assistance.

[Nastii22]

Hi @romainVala , @fepegar

I wanted to follow up on my previous question. Could you kindly provide an update on my query? Your assistance is crucial for my project's progress, and I would greatly appreciate any guidance or solutions you can offer.

Thank you in advance for your time and support.