Using MONAI Auto3DSeg model as pretrained model for MONAI Label

[klchtsk]

Dear community,

I'm trying to use the aorta model (MONAI Auto3dSeg pretrained models, ./.MONAIAuto3DSeg/models/aorta-v1.1.0/model.pt' as an input pretrained model for MONAI Label: /apps/radiology/model/pretrained_segmentation.pt.

The server starts: monailabel start_server --app ./apps/radiology --studies ./Documents/SlicerDICOMDatabase/ --conf models segmentation

When being used in MONAI Label segmentation app (renamed as pretrained_segmentation.pt), the result is chaotic and different from Auto3DSeg.

Auto3DSeg result:

Same model used in MONAI Label - it's just a block of everything:

I have also kept training the Auto3Dseg model on 3 CT angiographies. Programatically it all works, but the results are unfavourable.

Is it possible to use Auto3DSeg models as pretrained models for segmentation app of MONAI Label or these models have different network architectures?

Thank you!

[lassoan]

You can find the MONAIAuto3DSeg inference script here: https://github.com/lassoan/SlicerMONAIAuto3DSeg/blob/main/MONAIAuto3DSeg/Scripts/auto3dseg_segresnet_inference.py Probably MONAILabel misses some steps.

[klchtsk]

Dear Andras,
Thank you for your answer! I have found a discussion: Project-MONAI/tutorials#1714 and implemented the changes in code, suggested by @diazandr3s, with slight modification to use CPU and label only aorta.

It works perfectly and segments the aorta in MONAI Label, as it would have been done in Auto3DSeg, with the aorta.pt file, which is the ./.MONAIAuto3DSeg/models/aorta-v1.1.0/model.pt file from Auto3DSeg.

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

import lib.infers
import lib.trainers
from monai.networks.nets import SegResNet, SegResNetDS
from monai.utils import optional_import
from monai.bundle import ConfigParser

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

_, has_cp = optional_import("cupy")
_, has_cucim = optional_import("cucim")

logger = logging.getLogger(__name__)


class Segmentation(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 = {
            "aorta": 1,
        }

        # 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", "false")):
            url = f"{self.conf.get('pretrained_path', self.PRE_TRAINED_PATH)}"
            url = f"{url}/radiology_icrh_segresnet.pt"
            download_file(url, self.path[0])

        self.target_spacing = (1.5, 1.5, 1.5)  # target space for image
        # Setting ROI size - This is for the image padding
        self.roi_size = (128, 128, 128)

        # Network
        checkpoint = torch.load(os.path.join(os.path.join(self.model_dir, 'aorta.pt')), map_location="cpu")
        config = checkpoint["config"]
        state_dict = checkpoint["state_dict"]
        model = ConfigParser(config["network"]).get_parsed_content()
        model.load_state_dict(state_dict, strict=True)
        self.network = model.to(device=torch.device("cpu"), memory_format=torch.channels_last_3d)

    def infer(self) -> Union[InferTask, Dict[str, InferTask]]:
        task: InferTask = lib.infers.Segmentation(
            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")),
            config={"largest_cc": True if has_cp and has_cucim else 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.Segmentation(
            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 Segmentation Model",
            labels=self.labels,
        )
        return task

infer file was updated too:

from typing import Callable, Sequence

from lib.transforms.transforms import GetCentroidsd
from monai.inferers import Inferer, SlidingWindowInferer, SlidingWindowInfererAdapt
from monai.transforms import (
    Activationsd,
    AsDiscreted,
    EnsureChannelFirstd,
    EnsureTyped,
    GaussianSmoothd,
    KeepLargestConnectedComponentd,
    LoadImaged,
    Lambdad,
    ScaleIntensityRanged,
    Orientationd,
    ScaleIntensityd,
    Spacingd,
)

from monailabel.interfaces.tasks.infer_v2 import InferType
from monailabel.tasks.infer.basic_infer import BasicInferTask
from monailabel.transform.post import Restored
import torch
import numpy as np


class Segmentation(BasicInferTask):
    """
    This provides Inference Engine for pre-trained Segmentation (SegResNet) model.
    """

    def __init__(
        self,
        path,
        network=None,
        target_spacing=(1.0, 1.0, 1.0),
        type=InferType.SEGMENTATION,
        labels=None,
        dimension=3,
        description="A pre-trained model for volumetric (3D) Segmentation from CT image",
        **kwargs,
    ):
        super().__init__(
            path=path,
            network=network,
            type=type,
            labels=labels,
            dimension=dimension,
            description=description,
            load_strict=False,
            **kwargs,
        )
        self.target_spacing = target_spacing

    def pre_transforms(self, data=None) -> Sequence[Callable]:
        t = [
            LoadImaged(keys="image", ensure_channel_first=True, dtype=None, allow_missing_keys=True, image_only=False),
            EnsureTyped(keys="image", device=data.get("device") if data else None),
            Orientationd(keys="image", axcodes="RAS"),
            Spacingd(
                keys="image",
                pixdim=list(self.target_spacing),
                mode="bilinear",
                dtype=torch.float,
                min_pixdim=np.array(self.target_spacing) * 0.75,
                max_pixdim=np.array(self.target_spacing) * 1.25,
                allow_missing_keys=True,
            ),
            ScaleIntensityRanged(keys="image",
                                 a_min=-1.186875225438012,
                                 a_max=141.96979081895617,
                                 b_min=-1,
                                 b_max=1,
                                 clip=False),
            Lambdad(keys="image", func=lambda x: torch.sigmoid(x))
        ]
        return t

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

    def inverse_transforms(self, data=None):
        return []

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

        if data and data.get("largest_cc", False):
            t.append(KeepLargestConnectedComponentd(keys="pred"))
        t.extend(
            [
                Restored(
                    keys="pred",
                    ref_image="image",
                    config_labels=self.labels if data.get("restore_label_idx", False) else None,
                ),
                GetCentroidsd(keys="pred", centroids_key="centroids"),
            ]
        )
        return t

But after this implementation I'm facing a PyTorch error when I try to keep learning the model, [AttributeError: 'list' object has no attribute 'dim']. I'm looking for a workaround. There should be a way to convert the data.

[klchtsk]

I have solved my issue and now I have a new one!
[ThreadPoolExecutor-0_0] [INFO] (monailabel.utils.async_tasks.utils:83) - Return code: -9

[diazandr3s]

Hi @klchtsk,

Current status is: MONAI Label only runs inference on Auto3DSeg models. To train the Auto3DSeg, you have to follow the Auto3DSeg instructions.
I've also comment here: Project-MONAI/tutorials#1714 (reply in thread)

Hope this helps,