To turn pretrained Auto3dseg segresnet model into a MONAI Label custom app

[Hirriririir]

Hello,

Great toolkits first. I have trained a segmentation model for thigh muscles using the Auto3dseg pipeline. The segresnet model turned out to perform the best compared to the other three algorithms. However, I don't want to use the ensemble model because it appears larger than I need. Therefore, I want to add my pretrained segresnet model into Monai Label.

I have referenced the [MONAI Label Radiology App - OHIF Spleen Segmentation Example] tutorial to start the Monai Label service. Additionally, I used the README.md file called [DeepLearning models for Radiology use-case(s)] in the Radiology app to build a thigh muscle app based on the code provided by the spleen seg app and the segmenter.py provided by Auto3dseg/segresnet_0/scripts.

The model is perfect and the Monai Label server indeed ran successfully in 3D Slicer (monailabel start_server --app apps/radiology --studies datasets/Thigh_muscles --conf models segmentation_muscle --conf skip_trainers true),but the results are not as expected. I believe this might be caused by the wrong transforms I configured in the apps/radiology/lib/configs/segmentation_muscle.py and apps/radiology/lib/infers/segmentation_muscle.py.

Could you please help me? Thank you!

I attached my Auto3dseg segresnet config files as below:

[Auto3dseg segresnet] training.log
[Auto3dseg segresnet] hyper_parameters.yaml.txt

Monai label configs: segmentation_muscle.py

import logging
import os
from typing import Any, Dict, Optional, Union

import lib.infers
import lib.trainers
from monai.networks.nets import SegResNetDS

from monailabel.interfaces.config import TaskConfig
from monailabel.interfaces.tasks.infer_v2 import InferTask
from monailabel.interfaces.tasks.train import TrainTask
from monailabel.utils.others.generic import download_file, strtobool

logger = logging.getLogger(__name__)


class SegmentationMuscle(TaskConfig):
    def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **kwargs):
        super().init(name, model_dir, conf, planner, **kwargs)

        # Labels
        self.labels = {
            "SA": 1,
            "RF": 2,
            "VL": 3,
            "VI": 4,
            "VM": 5,
            "AM": 6,
            "GR": 7,
            "BL": 8,
            "ST": 9,
            "SM": 10,
            "BB": 11,
        }

        # Model Files
        self.path = [
            os.path.join(self.model_dir, f"pretrained_{name}.pt"),  # pretrained
            os.path.join(self.model_dir, f"{name}.pt"),  # published
        ]

        # Download PreTrained Model
        if strtobool(self.conf.get("use_pretrained_model", "true")):
            url = f"{self.conf.get('pretrained_path', self.PRE_TRAINED_PATH)}"
            url = f"{url}/radiology_segmentation_muscle.pt"
            download_file(url, self.path[0])

        self.target_spacing = (0.7813000082969666, 0.7813000082969666, 4.0)  # target space for image
        self.roi_size = (336, 336, 88)

        # Network
        self.network = SegResNetDS(
            spatial_dims=3,
            norm = 'INSTANCE', 
            in_channels=1,
            out_channels=12,
            init_filters=32,
            blocks_down=(1, 2, 2, 4, 4),
            #blocks_up=(1, 1, 1),
            dsdepth=4,
            resolution =[0.7813000082969666, 0.7813000082969666, 4.0],
        )

    def infer(self) -> Union[InferTask, Dict[str, InferTask]]:
        task: InferTask = lib.infers.SegmentationMuscle(
            path=self.path,
            network=self.network,
            roi_size=self.roi_size,
            #target_spacing=self.target_spacing,
            labels=self.labels,
            preload=strtobool(self.conf.get("preload", "false")),
        )
        return task

    def trainer(self) -> Optional[TrainTask]:
        output_dir = os.path.join(self.model_dir, self.name)
        load_path = self.path[0] if os.path.exists(self.path[0]) else self.path[1]

        task: TrainTask = lib.trainers.SegmentationMuscle(
            model_dir=output_dir,
            network=self.network,
            roi_size=self.roi_size,
            target_spacing=self.target_spacing,
            load_path=load_path,
            publish_path=self.path[1],
            description="Train vertebra segmentation Model",
            labels=self.labels,
        )
        return task

Monai label infer: segmentation_muscle.py

from typing import Callable, Sequence

from monai.inferers import Inferer, SlidingWindowInferer
from monai.transforms import (
    Activationsd,
    AsDiscreted,
    EnsureChannelFirstd,
    EnsureTyped,
    LoadImaged,
    NormalizeIntensityd,
    ScaleIntensityRanged,
    Spacingd,
    SpatialPadd
)

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


class SegmentationMuscle(BasicInferTask):
    """
    This provides Inference Engine for pre-trained spleen segmentation (UNet) model over MSD Dataset.
    """

    def __init__(
        self,
        path,
        network=None,
        type=InferType.SEGMENTATION,
        labels=None,
        dimension=3,
        description="A pre-trained model for volumetric (3D) segmentation of the spleen from CT image",
        **kwargs,
    ):
        super().__init__(
            path=path,
            network=network,
            type=type,
            labels=labels,
            dimension=dimension,
            description=description,
            **kwargs,
        )


    def pre_transforms(self, data=None) -> Sequence[Callable]:
        return [
            LoadImaged(keys="image"),
            EnsureTyped(keys="image", device=data.get("device") if data else None),
            EnsureChannelFirstd(keys="image"),
            Spacingd(keys="image", mode="bilinear", pixdim=[0.7813000082969666, 0.7813000082969666, 4.0]),
                     #min_pixdim=[0.7262306388071416, 0.7262306249726234, 1.0],
                     #max_pixdim=[1.6035165011882808, 1.6035165011882808, 26.00048770904541],
                     #allow_missing_keys=True,),
            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            #ScaleIntensityRanged(keys="image", a_min=173.66656895101505, a_max=3272.910755018248, b_min=-1, b_max=1, clip=False),
        ]

    def inferer(self, data=None) -> Inferer:
        return SlidingWindowInferer(roi_size=(336, 336, 88),
                                    sw_batch_size=1,
                                    overlap=0.625,
                                    mode="gaussian",
                                    cache_roi_weight_map=False,
                                    progress=False)

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

[KumoLiu]

Hi @Hirriririir, thanks for your interest here.
First I want to make sure I understand your process. You use Auto3dseg to train a model for segmenting thigh muscles and you want to use this model to do inference in the model label, right?
Your question is that you can't infer well in the MONAI label, right? If I understand correctly, I think you only need to ensure that you have the same pre-processing and post-processing between auto3dseg and MONAI label app. Could you please double confirm that part?
Hope it can help you, thanks!

[Hirriririir]

Hi KumoLiu, thanks for your reply!

Yes, all your assumptions are correct. The Auto3dseg segresnet model performs well with my dataset (average dice above 0.8). It's the problem of the transformations.

I am a noob in deep learning. I have tried to reverse the pre-processing and post-processing configuration of the Auto3dseg segresnet model in the MONAI label app to the best of my ability, but it did not work well. Could you help me, at your convenience? Thanks

[Auto3dseg segresnet script] infer.py.txt
[Auto3dseg segresnet script] segmenter.py.txt

[KumoLiu]

Hi @myron, could you please help take a look at this question, I think you are more familiar with this part. Is it possible to provide some guides for @Hirriririir to apply segresnet in MONAILabel apps, e.g. a guide to use the same transforms that the segmenter in Auto3DSeg segresnet has built? What do you think?
Thanks in advance!

[myron]

Hi @Hirriririir, I'm glad you found Segresnet model useful! Try these debugging steps.

1. Ensure it outputs a correct segmentation results from commandline first. Run
   python /pathto/segmenter.py --config_file=/pathto/configs/hyper_parameters.yaml --infer#enabled=True

set correct path to "segmenter.py" (that you used for training) and the path to "hyper_parameters.yaml" (default is in configs folder). And in your "datalist.json" file with a list of training images, add a new key "testing" with a list of images to run inference on (same structure as training set, but only the "image" key is necessary)

this will run inference and save the result .nii.gz files in the "predictions" folder.
Open those files in 3D slicer (as labels) to ensure they look correct

2. Your screenshot looks "too colorful", this artifact can happen if you try to open an "image" as a "segmentation label", since intensity image has many intensity levels, they create a "colorful" incorrect segmentation. Just make sure you don't send the CT as your segmentation somehow.

3. is your image 3D MRI or CT, your input transform indicate it's MRI, but if it's CT, then replace
   NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),

with 2 lines

ScaleIntensityRanged(keys="image", a_min=intensity_bounds[0], a_max=intensity_bounds[1], b_min=-1, b_max=1, clip=False)
Lambdad(keys="image", func=lambda x: torch.sigmoid(x))

you can find intensity_bounds values in hyper_parameters.yaml (they were auto detected)
cheers!

[Hirriririir]

Hi myron, thanks for your reply!

The 1 step you mentioned, I had already run. The SegResNet performed very well on the testing MRI dataset in the "datalist.json" file.

However, the problem seems to be in the MONAI label inference configuration I wrote. I wrote this configuration based on the "hyper_parameters.yaml" provided by the SegResNet Auto3dseg bundle. Do you think these pre and post transforms are correct? Thanks in advance!"

 def pre_transforms(self, data=None) -> Sequence[Callable]:
        return [
            LoadImaged(keys="image"),
            EnsureTyped(keys="image", device=data.get("device") if data else None),
            EnsureChannelFirstd(keys="image"),
            Spacingd(keys="image", mode="bilinear", pixdim=[0.7813000082969666, 0.7813000082969666, 4.0]),
                     #min_pixdim=[0.7262306388071416, 0.7262306249726234, 1.0],
                     #max_pixdim=[1.6035165011882808, 1.6035165011882808, 26.00048770904541],
                     #allow_missing_keys=True,),
            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            #ScaleIntensityRanged(keys="image", a_min=173.66656895101505, a_max=3272.910755018248, b_min=-1, b_max=1, clip=False),
        ]

    def inferer(self, data=None) -> Inferer:
        return SlidingWindowInferer(roi_size=(336, 336, 88),
                                    sw_batch_size=1,
                                    overlap=0.625,
                                    mode="gaussian",
                                    cache_roi_weight_map=False,
                                    progress=False)

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