Issue with visualization of output while excluding the background in the loss function

[NastaranVB]

Hello,
I am doing multi class 3D segmentation which contains 3 classes including background. I have used Network, Loss and Metric from Monai. When excluding the background in the loss function and defining it as "DiceLoss(include_background=False, to_onehot_y=True, softmax=True)," upon visualizing the prediction, the output (prediction mask) displays an image as attached, which is incorrect and should only depict the segmented portion. Could you please help me with the issue?
Thank you in advance.

[KumoLiu]

Hi @NastaranVB,
The parameter include_background in the DiceLoss function does not affect how the data is visualized. Instead, it indicates whether the background should be considered when calculating the loss. If it's set to False, this means the background is not included in the computation of the loss.
When it comes to visualization, if you want to view a specific channel, you should directly specify it in your visualization function or code. For instance, if you're using matplotlib for visualization, you would select the channel when calling imshow(). If the output of your model is a multiple channel tensor output, you could simply index it like output[channel_number] to get the specific channel you want to visualize.

Hope it helps, thanks.

[NastaranVB]

Hi @KumoLiu, thanks for coming back to me so quickly.

For visualization, I had done output[channel_number], as you also have mentioned but the image was as I attached before.
I am doing my segmentation task based on Spleen 3D tutorial, as below.

class Net(pytorch_lightning.LightningModule):
    def __init__(self):
        super().__init__()
        self._model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=3,
            channels=(16, 32, 64, 128, 256),
            strides=(2, 2, 2, 2),
       )
        self.loss_function = DiceLoss(include_background=False, to_onehot_y=True, softmax=True)
        
     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"]       
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        self.log("Train Loss", loss, on_step=True, on_epoch=True, logger=True)
        return loss
        
    def validation_step(self, batch, batch_idx):        
        img = batch["Image"]["data"]
        mask = batch["Label"]["data"]     
        pred = self(img)
        loss = self.loss_fn(pred, mask)       
        self.log("Val Loss", loss, on_step=True, on_epoch=True, logger=True) 
        return loss
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self._model.parameters(), 1e-4)
        return optimizer

And the model was trained in a patch wise manner as following:

imgs = val_dataset[IDX]["Image"]["data"]
mask = val_dataset[IDX]["Label"]["data"]
grid_sampler = tio.inference.GridSampler(val_dataset[IDX], patch_size=patch_size, patch_overlap=patch_overlap)
aggregator = tio.inference.GridAggregator(grid_sampler)
patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=val_batch_size)
with torch.no_grad():
    for patches_batch in patch_loader:
        input_tensor = patches_batch['Image']["data"].to(device)  
        locations = patches_batch[tio.LOCATION]  
        pred = model(input_tensor)  
        aggregator.add_batch(pred, locations)
output_tensor = aggregator.get_output_tensor()

output_class_1= output_tensor[1]

When I use pytorch CrossEntropyLoss (without Monai), the code works and gives the proper image of prediction mask slices. So it would be nice if you could tell me if there is anything that I am doing wrong with the Monai part?
Thank you in advance.