Segmentation mask not showing after inference when using monai label in OHIF

[JackTheSturgeon]

Hello,

I have trained a model to segment blood vessels using a monai UNetR model on brain MRI. I have done so in jupyter notebook separately from monai.
My task is to integrate that pretrained model so that it works with monai in OHIF to be able to adjust the segmentations and retrain the model using monai.

How did I do so far:
I am working in the radiology app and created my own config, infers and trainers files from a copy of the already existing segmentation spleen that I modified.
I added my pretrained model, saved using torch.jit.save, to the model folder of the application.
I also created my own Inferer to adapt the inference to my model. I am doing inference on different patches of the MRI that I merge at the end.

My problem:
When I run my model in OHIF, it finishes, I see the label added to OHIF but the segmentation mask is empty.
I already verified that my inferer correctly outputs what I am expecting and it is. I have a 3D tensor with ones where there are blood vessels and zeros where its the background. I even saved the masks as images to be sure that it works.
I don't know what to do to be able to see my segmentation mask properly in monai and don't understand where the problem is coming from.

This is why i am calling for help here, thank you for reading.

Here is my code, I can provide more information if needed.

`

Infers file for my unetr

from typing import Callable, Sequence, Union

import torch

import lib.monai_inferers

from monai.apps.deepedit.transforms import (
AddGuidanceFromPointsDeepEditd,
AddGuidanceSignalDeepEditd,
DiscardAddGuidanced,
ResizeGuidanceMultipleLabelDeepEditd,
)
from monai.inferers import Inferer, SimpleInferer
from monai.transforms import (
Activationsd,
AsDiscreted,
EnsureChannelFirstd,
EnsureTyped,
MapTransform,
LoadImaged,
Orientationd,
Resized,
ScaleIntensityRanged,
SqueezeDimd,
ToNumpyd,
Activations,
AsDiscrete,
Compose,
Invertd,
NormalizeIntensityd,
RandFlipd,
RandScaleIntensityd,
RandShiftIntensityd,
RandSpatialCropd,
Spacingd,
CropForegroundd,
RandRotate90d,
RandCropByPosNegLabeld
)

from monailabel.interfaces.tasks.infer_v2 import InferType
from monailabel.tasks.infer.basic_infer import BasicInferTask
from monailabel.transform.post import Restored

class SegmentationUnetrCorticalVessels(BasicInferTask):

def __init__(
    self,
    path,
    network=None,
    type=InferType.SEGMENTATION,
    labels=None,
    dimension=3,
    spatial_size=(64, 64, 64),
    target_spacing=(1.0, 1.0, 1.0),
    number_intensity_ch=1,
    description="Cortical Vessel Segmentation using the UnetR model",
    **kwargs,
):
    super().__init__(
        path=path,
        network=network,
        type=type,
        labels=labels,
        dimension=dimension,
        description=description,
        # input_key="image",
        # output_label_key="pred",
        # output_json_key="result",
        **kwargs,
    )

    self.spatial_size = spatial_size
    self.target_spacing = target_spacing
    self.number_intensity_ch = number_intensity_ch
    self.num_samples=8

def pre_transforms(self, data=None):
    return [
LoadImaged(keys=("image")),
EnsureChannelFirstd(keys="image"),
EnsureTyped(keys=("image")),
Spacingd(keys="image", pixdim=self.target_spacing, allow_missing_keys=True),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
]

def inferer(self, data=None) -> Inferer:
    return lib.monai_inferers.unetrInferer(spatial_size=self.spatial_size)

def post_transforms(self, data=None) -> Sequence[Callable]:
    return [
        EnsureTyped(keys="pred"),
        Activationsd(keys="pred", softmax=True),
        AsDiscreted(keys="pred", argmax=True),
        Restored(keys="pred", ref_image="image"),
    ]`

`

Custom Inferer

import os
import numpy as np
import torch
from tqdm import tqdm

import matplotlib.pyplot as plt

from monai.inferers import Inferer, Splitter, SlidingWindowSplitter, AvgMerger
from monai.transforms import Activations, AsDiscrete, Compose

class unetrInferer(Inferer):

def __init__(self, threshold=0.5, spatial_size=(64,64,64)) -> None:
    super().__init__()
    self.threshold=threshold
    self.spatial_size=spatial_size
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def __call__(self, inputs, network, *args, **kwargs) -> torch.Tensor:
    
    device = self.device
    
    splitter = SlidingWindowSplitter(self.spatial_size, overlap=0, device=device)
    splits = splitter(inputs)

    padded_shape = tuple(inputs.shape[0:2]) + splitter.get_padded_shape(inputs)
    cropped_shape = inputs.shape

    post_trans = Compose([Activations(sigmoid=True), 
                           AsDiscrete(threshold=self.threshold)])

    counter=0
    values=[]
    locations=[]
    for split in tqdm(splits):
        split_pred=post_trans(network(split[0].to(device)))
        values.append(split_pred)
        locations.append(split[1])
        counter+=1

    merger=AvgMerger(merged_shape=padded_shape, 
                     cropped_shape=cropped_shape, 
                     device=device)

    for i,j in zip(values, locations):  
        merger.aggregate(i, j)
    
    output=merger.finalize()
    return output`

Here is one slice of the segmentation mask:

`

Model and Settings

    self.target_spacing = (1.0, 1.0, 1.0)
    self.roi_size = (224,224,128)
    self.spatial_size=(64,64,64)

    # Network
    self.network = UNETR(
        in_channels=1,
        out_channels=1,
        img_size=self.spatial_size,
        feature_size=16,
        hidden_size=768,
        mlp_dim=3072,
        num_heads=12,
        pos_embed="perceptron",
        norm_name="instance",
        res_block=True,
        dropout_rate=0.2,
    )`

[diazandr3s]

Hi @JackTheSturgeon,

Thanks for the detailed discussion.
Have you tried running inference using 3DSlicer? Once you start the MONAI Label server, run inference using 3DSlicer and see if you get a mask.
I just want to discard issues from the model.

Let us know

[JackTheSturgeon]

Hello @diazandr3s,

Thank you for the quick response, I took my time because I realized that maybe my problem was about the format of my output.
Which was effectively the problem.
I answered to my problem in the comments as precisely as I could.

[JackTheSturgeon]

I resolved the problem. It was about the shape and values of my inference output tensor.

The output of the inference from the model has to be of shape:

• [1, 2, x, y, z] with x, y, z being the dimensions of the input images.
  • 2 for two channels with one of them being the negative of the other (the negative is the 1st channel)
  • the first dimension is added because for some reason the data acquired in the post transform function is squeezed. So I had to add this additionnal dimension.
• Non thresholded Tensor
  • a softmax or sigmoid is already used in the post transform to obtain the mask

My precise thought process to get to the solution:

When i did inference with the custom model on OHIF, the segmentation mask obtained was empty.
However, I saved the segmentation masks infered by the model outside of OHIF and I could see that they were not empty, they actually looked like this:

Which is what I was expecting. However, when analysing the inference output of the segmentation spleen model I could see two differences:

• Tensor shape was [2, x, y, z], mine was [1, x, y, z]
• Tensor values were in [-15, 18] while mine were in [0, 1]

What it means:

→ There is supposed to be two channels on the inference output (which are the positive and negative predictions embedded on the real image)
What one slice of each channel looks like for the spleen segmentation:

→ The thresholding of these images to obtain the segmentation mask is not done at the inference level but at the post transform level.

For any person that had the same problem, if you need more details or something is not clear please ask i will update my post.