The problem with the accuracy of training from scratch

[Minxiangliu]

My version: from docker projectmonai/monailabel:0.4.2

3D Slicer version: 5.0.3

I referred to the suggestions here to add my own segmentation model. Since my data belongs to PET CT images, it is not implemented in the current pre-training model, so I plan to train a new model by myself. At present, the modified code can run normally in 3D Slicer. I currently want to try the performance of the model, so I prepare 5 sets of labeled data and 1 set of unlabeled data. When training, 4 sets of training and 1 set of verification. Finally, after repeated execution 3 times, an accuracy rate of about 49% can be obtained, but the results displayed when performing inference seem to be much less similar to 50%.

1. Where might my problem be?
2. Also, is it possible to get better results with deepedit than with segmentation?

Thanks in advance.

Server Executio: monailabel start_server --app apps/radiology --studies HNC-datasets --conf models segmentation_hnc

Here are the datasets, code, and training results used.

Results:

configs/segmentation_hnc.py

    def trainer(self) -> Optional[TrainTask]:
        output_dir = os.path.join(self.model_dir, self.name)
        task: TrainTask = lib.trainers.SegmentationHNC(
            model_dir=output_dir,
            network=self.network,
            description="Train HNC Segmentation Model",
            target_spacing=(4.,4.,2.57),
            padd_size=(64,64,80),
            crop_size=(48,48,48),
            load_path=self.path[0],
            publish_path=self.path[1],
            labels=self.labels,
            config={"pretrained":strtobool(self.conf.get("use_pretrained_model", "false")),
                    "max_epochs":50,
                    'train_batch_size':8},
        )
        return task
    
    def infer(self) -> Union[InferTask, Dict[str, InferTask]]:
        task: InferTask = lib.infers.SegmentationHNC(
            path=self.path,
            network=self.network,
            type=InferType.SEGMENTATION,
            labels=self.labels,
            dimension=3,
            description='A pre-trained model for volumetric (3D) segmentation of the HNC from PET CT image',
            target_spacing=(4.,4.,2.57),
            padd_size=(64,64,80),
            crop_size=(48,48,48),
            preload=strtobool(self.conf.get("preload", "false")),
        )
        return task

trainers/segmentation_hnc.py

class SegmentationHNC(BasicTrainTask):
    def __init__(
        self,
        model_dir,
        network,
        description,
        target_spacing,
        padd_size,
        crop_size,
        labels,
        **kwargs,
    ):
        self._network = network
        self.padd_size = padd_size
        self.crop_size = crop_size
        self.target_spacing = target_spacing
        self.labels = labels
        self._transforms = [
            LoadImaged(keys=['image','label']),
            EnsureChannelFirstd(keys=['image','label']),
            Orientationd(keys=['image','label'], axcodes='RAS'),
            Spacingd(keys=['image','label'], pixdim=self.target_spacing, mode=('bilinear', 'nearest')),
            Rotate90d(keys=['image','label'], spatial_axes=(0,1)),
            Flipd(keys=['image','label'], spatial_axis=1),
            ScaleIntensityRanged(keys=['image'], a_min=0, a_max=30000,
                                 b_min=0.0, b_max=1.0, clip=True),
            CropForegroundd(keys=['image','label'],
                            source_key='image', 
                            select_fn=lambda x:x>0.05, 
                            margin=0, 
                            return_coords=False),
            SpatialPadd(keys=['image','label'], spatial_size=self.padd_size),
            EnsureTyped(keys=['image','label'], device='cuda')
        ]
        super().__init__(model_dir, description, **kwargs)

    def network(self, context: Context):
        return self._network

    def optimizer(self, context: Context):
        return torch.optim.Adam(context.network.parameters(), lr=0.0001)

    def loss_function(self, context: Context):
        return DiceCELoss(to_onehot_y=True, softmax=True, squared_pred=True, batch=True)

    def train_pre_transforms(self, context: Context):
        train_transforms = self._transforms.copy()
        train_transforms.extend([
            RandScaleIntensityd(keys=['image'], factors=0.2, prob=1.0),
            RandShiftIntensityd(keys=['image'], offsets=(-0.8, 0.8), prob=1.0),
            RandCropByPosNegLabeld(keys=['image', 'label'],
                                   image_key='image',
                                   label_key='label',
                                   spatial_size=self.crop_size,
                                   pos=1,
                                   neg=1,
                                   num_samples=4,
                                   image_threshold=0),
            RandFlipd(keys=['image','label'], spatial_axis=[0], prob=0.5),
            RandFlipd(keys=['image','label'], spatial_axis=[1], prob=0.5),
            RandFlipd(keys=['image','label'], spatial_axis=[2], prob=0.5),
            RandRotate90d(keys=['image','label'], prob=0.5),
            SelectItemsd(keys=['image','label'])
        ])
        return train_transforms

    def train_post_transforms(self, context: Context):
        return [
            ToTensord(keys=("pred", "label")),
            Activationsd(keys="pred", softmax=True),
            AsDiscreted(
                keys=("pred", "label"),
                argmax=(True, False),
                to_onehot=True,
                n_classes=len(self.labels),
            ),
        ]

    def val_pre_transforms(self, context: Context):
        val_transforms = self._transforms.copy()
        return val_transforms

    def val_inferer(self, context: Context):
        return SlidingWindowInferer(roi_size=(64, 64, 64), sw_batch_size=1, overlap=0.25)

infers/segmentation_hnc.py

class SegmentationHNC(InferTask):
    def __init__(
        self,
        path,
        network,
        type,
        labels,
        dimension,
        description,
        target_spacing,
        padd_size,
        crop_size,
        **kwargs,
    ):
        self.padd_size = padd_size
        self.crop_size = crop_size
        self.target_spacing = target_spacing
        self.labels = labels
        super().__init__(
            path=path,
            network=network,
            type=type,
            labels=self.labels,
            dimension=dimension,
            description=description,
            **kwargs,
        )

    def pre_transforms(self, data=None) -> Sequence[Callable]:
        return [
            LoadImaged(keys='image'),
            EnsureChannelFirstd(keys='image'),
            Orientationd(keys='image', axcodes='RAS'),
            Spacingd(keys='image', pixdim=self.target_spacing, mode='bilinear'),
            Rotate90d(keys='image', spatial_axes=(0,1)),
            Flipd(keys='image', spatial_axis=1),
            ScaleIntensityRanged(keys='image', a_min=0, a_max=30000,
                                 b_min=0.0, b_max=1.0, clip=True),
            CropForegroundd(keys='image',
                            source_key='image', 
                            select_fn=lambda x:x>0.05, 
                            margin=0, 
                            return_coords=False),
            SpatialPadd(keys='image', spatial_size=self.padd_size),
            EnsureTyped(keys='image', device=data.get("device") if data else None)
        ]

    def inferer(self, data=None) -> Inferer:
        return SlidingWindowInferer(roi_size=(64, 64, 64), sw_batch_size=1, overlap=0.25)

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

[diazandr3s]

Hi @Minxiangliu,

Thanks for the detailed discussion.

I've checked the files you attached and one of the main issues is that not all the HNC volumes have the same number of labels - some of them have only one label. In the most recent MONAI Label version, we've proposed a transform to deal with these cases: https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/trainers/segmentation.py#L77

The second issue I saw is that in the segmentation model there is no need to declare a background label. It is only needed for the DeepEdit model - This is because the DeepEdit model simulates clicks also for the background.

Another issue is that you weren't using the most recent MONAI Label version. Please fetch the latest docker image.

Here you can find a model I've trained on the dataset you attached: https://drive.google.com/file/d/19JQNDx57RVqg9qqI6xOPw7ZzqWs-Agep/view?usp=sharing

And here is a video showing how that works:

hnc_project.mp4

Hope that helps,

[Minxiangliu]

Hi, @diazandr3s

Thank you for your help, first of all I have to clarify why I don't use the new docker first, the latest version has an error a few days ago, and the error occurs after executing the monailabel command, but the current first version is available.

I have modified the configs, trainers, infers files according to your suggestion(use custom transformation), and also successfully trained, but the verification result shows that it can reach 80%, but after the inference, it seems that the location of the cancer cell is not correctly found, may I ask this Is it normal?

The possible reason is that the inference setting is wrong, or the number of training samples is too small to cause overfitting?

Training result:

Results file: https://drive.google.com/file/d/1__5twHLRt99bNxWBU-PGfzQywAx_hF4P/view?usp=sharing

[Minxiangliu]

After removing Rotate90d and Flipd in pre_transforms, and retraining the results, it seems that the position of the mask is opposite to the direction of the cancer cell city?

[diazandr3s]

Thanks for the update, @Minxiangliu.

Thank you for your help, first of all I have to clarify why I don't use the new docker first, the latest version has an error a few days ago, and the error occurs after executing the monailabel command, but the current first version is available.

If there is an error on the latest docker container, please do let us know. There shouldn't be any issue with that.

I have modified the configs, trainers, infers files according to your suggestion(use custom transformation), and also successfully trained, but the verification result shows that it can reach 80%, but after the inference, it seems that the location of the cancer cell is not correctly found, may I ask this Is it normal?

No, it is not normal. The restored transform (https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/infers/localization_spine.py#L84) resamples the prediction into the original space/orientation as I showed in the video I uploaded in my previous post.

After checking the files you sent, I see you applied these pre-transforms during inference:

Rotate90d(keys='image', spatial_axes=(0,1)),
Flipd(keys='image', spatial_axis=1),

I think that's the issue. In inference you shouldn't apply these rotation and flipping transforms. Please leave these transforms only: https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/infers/localization_spine.py#L58-L66

Additionally, you also applied these:

CropForegroundd(keys='image',
                source_key='image', 
                select_fn=lambda x:x>0.05, 
                margin=0, 
                return_coords=False),
SpatialPadd(keys='image', spatial_size=self.padd_size),

In inference, you shouldn't crop the foreground as there are no labels. I'm not sure that thresholding is helpful in inference. Let the model predict from the whole image.

Padding in inference isn't needed as you're using the sliding window inferrer: https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/infers/localization_spine.py#L69

[Minxiangliu]

Hi @diazandr3s ,
Thanks for the detailed explanation.
After your prompt, I found that it is indeed in the inferred post transforms that the predicted mask is not transformed back to the original orientation. Since my data is not uniform in direction, amount, and size, I have to do some crop and pad in pre transforms to achieve consistent size. Finally I added the inverse transformation to the post transforms of the inference, and finally output the result in SqueezeDimd.

Final code in lib/infers:

    def post_transforms(self, data=None) -> Sequence[Callable]:
        _pre_trans = self.pre_trans.copy()
        _pre_trans.append(EnsureTyped(keys='image', device=data.get("device") if data else None))

        largest_cc = False if not data else data.get("largest_cc", False)
        applied_labels = list(self.labels.values()) if isinstance(self.labels, dict) else self.labels
        t = [
            EnsureTyped(keys='pred', device=data.get('device') if data else None),
            Activationsd(keys='pred', softmax=True),
            Invertd(keys='pred',
                    transform=Compose(_pre_trans),
                    orig_keys='image',
                    nearest_interp=False,
                    to_tensor=True,
                    device=data.get('device') if data else None),
            AsDiscreted(keys='pred', argmax=True),
        ]
        if largest_cc:
            t.append(KeepLargestConnectedComponentd(keys='pred', applied_labels=applied_labels))
        t.append(BinaryMaskd(keys='pred'))
        t.append(SqueezeDimd(keys='pred', dim=0))
        return t

Resluts:

Compared to the previous results, it is indeed a bit more normal.
I don't know if I'm right. In order to maximize the use of monai's transforms, will it be faster and more accurate to use Invertd directly in the post transforms of the inference?

[diazandr3s]

Thanks for letting us know, @Minxiangliu
Happy to see this working.

I don't know if I'm right. In order to maximize the use of monai's transforms, will it be faster and more accurate to use Invertd directly in the post transforms of the inference?

Isn't that what you're showing? We also do this in the infer post_transforms: https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/infers/localization_spine.py#L84 - we use the Restored transform