SW-FastEdit Integration into MONAI (GPU based DeepEdit rework)

[matt3o]

Hey guys!

We reworked the DeepEdit code within the scope of my Masters thesis and I want to ask if we shall reintegrate that code into MONAI. Sadly the DeepEdit code in MONAI tutorials is completely broken and does not work, at least the last time I was using it, so even more incentive to provide a new start into interactive segmentation.
This is the corresponding paper, which I will also present at the ISBI conference in May 2024: https://arxiv.org/pdf/2311.14482.pdf
The code can be found here: https://github.com/matt3o/AutoPET2-Submission

A few advantages: Fully GPU based distance transform (now also in MONAI > 1.3.0 itself), Sliding-Window approach for resource-reduction and higher Dice score, fixed a lot of GPU leaks on the way (see my issues about that topic) and generally quite a lot of optimizations to make the code real-time fast. This means an iteration of SW-FastEdit on a full AutoPET volume takes 5 seconds on an A100 with 80Gb of GPU.
Also, as the papers shows, we added quite some new click generation strategies and stopping criterions.
(Bonus: disabling the interactive part, this code did fine enough to let us win the third place in this years AutoPET II challenge on robustness across domains)

Long story short, I would love to integrate the code into MONAI for others to use - if that is desired.
There would be more or less three steps, that I would see as of now. However different approaches may of course be chosen as well.

1. Integration of the core functionality into MONAI, most importantly the new transforms and reworked interaction class. Those could then be used in MONAILabel.
2. Adding the rest of the code and scripts to MONAI tutorials and provide an interactive tutorial on how to use the code
3. Add integrations to MONAILabel such that the code can be started there. A working implementation can already be found in the monailabel folder, which has been used for our user study.

Currently the code has one major limitation, it is fully Nvidia based since I started doing a lot of transforms in cupy. I could look into transforming it for CPU based architectures as well, if that would be necessary for the integration into MONAI.

What are your thoughts on that?

Cheers,
Matthias (and @Zrrr1997)

Ps: Especially @wyli and @diazandr3s , since I was writing with you previously on this topic

[diazandr3s]

Hi @matt3o,

This is great! Congratulations on this work! I think the community will benefit a lot from this.

Sadly the DeepEdit code in MONAI tutorials is completely broken and does not work, at least the last time I was using it, so even more incentive to provide a new start into interactive segmentation.

This is strange. I've recently tested it and it is working as intended. Can you please be more specific about what's not working?

With regards to your work, I'd suggest we include this model in the radiology app of MONAI Label - would you like to create a PR for this? You could include all custom transforms here temporarily until we make sure they all work.

BTW, we don't need to replace DeepEdit. They both can be there.

I quickly checked the monailabel code and I didn't see a trainer file. Is this only for inference? Have you tested on CT images?

Please let us know.

Great work!!

[matt3o]

@diazandr3s
Coming back to your comments about what exactly is broken:

Trying to reproduce the settings from the paper I ran into the following issues:
First of all the tutorial code, which does not respect the --max_train_interactions and --max_train_interactions flags [Bug 1], crashes as soon the interactions are increased to more than 1 click [Bug 2]. You can do so by manually adding:

        iteration_update=Interaction(
            ....
            max_interactions=10,
        ),

The error message is

[2023-09-24 10:28:21.078][ERROR](<https://www.notion.so/ignite.engine.engine.SupervisedTrainer>) - Current run is terminating due to exception: applying transform <monai.apps.deepedit.transforms.AddRandomGuidanceDeepEditd object at 0x7f9624d3ee20>
[2023-09-24 10:28:21.078][ERROR](<https://www.notion.so/ignite.engine.engine.SupervisedTrainer>) - Exception: applying transform <monai.apps.deepedit.transforms.AddRandomGuidanceDeepEditd object at 0x7f9624d3ee20>
Traceback (most recent call last):
File "/home/matteo/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/transform.py", line 141, in apply_transform
return _apply_transform(transform, data, unpack_items, lazy, overrides, log_stats)
File "/home/matteo/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/transform.py", line 98, in _apply_transform
return transform(data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(data)
File "/home/matteo/anaconda3/envs/monai/lib/python3.9/site-packages/monai/apps/deepedit/transforms.py", line 585, in call
self.guidance[key_label] = [j for j in tmp_gui if -1 not in j]
TypeError: 'NoneType' object is not iterable

• -> This tutorial is thus only working in the default setting. The transforms could be fixed easily by adding a None check like this: https://github.com/Zrrr1997/Guiding-The-Guidance/blob/main/user_guidance_study/utils/transforms.py#L395C25-L395C25

[matt3o]

(This is from my notes but I would guess that this is still up to date. I did log the clicks at some point and realized they were too many)
The second issue is even more problematic: The DeepEdit code creates the wrong amount of clicks. I cannot prove that on the tutorial since for the reasons above it does not work.
For 0 clicks, all is good, no guidance is provided, or more precisely it gets removed in the Interaction class (# zero out input guidance channels section).
However, for anything greater than 0 clicks, there is one click too many. This happens only for non-background labels. So e.g. 1 click on tumor actually results in 2 clicks on tumor. For 10 clicks we get 11 clicks instead.
This is due to the AddInitialSeedPointMissingLabelsd adding a click in the beginning and then for each interaction one more click is added.

[diazandr3s]

Hi @matt3o,

I'm sorry I'm just coming back to this. It's been a hectic time with another project.

Regarding this:

First of all the tutorial code, which does not respect the --max_train_interactions and --max_train_interactions flags [Bug 1]

You're right. The train.py script doesn't obey these flags because they aren't passed from the train.py args script to the Interaction classes (eval and train) (https://github.com/Project-MONAI/tutorials/blob/main/deepedit/ignite/train.py#L304-L310). So the Interaction classes use the arg set by default (1): https://github.com/Project-MONAI/MONAI/blob/dev/monai/apps/deepedit/interaction.py#L52

As you mentioned, it should be solved by just passing the arg to the Interaction class.

With regards to this:

crashes as soon the interactions are increased to more than 1 click [Bug 2]

I'd like to reproduce the error you're facing. Can you please confirm the number of labels you were using for this test? Did you use the default ones on the same dataset? https://github.com/Project-MONAI/tutorials/blob/main/deepedit/ignite/train.py#L411

It seem this is an error related to the way the labels are passed.

I was wondering if there is any blocker in MONAI Label that you'll need to be fixed for your PR?

Thanks again for your comments.

[diazandr3s]

Thnaks again for your comments, @matt3o.
This PR should solve this issue.

[matt3o]

@diazandr3s Very nice, thanks a lot! Also sorry for the late respone, was rather busy as well.