Integration of SlidingWindow DeepEdit Approach

[matt3o]

Hey guys!

As I have been writing a few times, I am working on a Sliding Window rework of DeepEdit. What it became is a major rework mostly of the transforms to use torch and thus run on the GPU (without running into OOMs, the hardest part of this the work). Plus a cupy based edt distance transform implementation. This code also fixes a bug in the initial implementation which always produced one click too much, i.e. 11 clicks instead of 10.
Note that this code is currently only tested on AutoPET, some functionality supported by the initial code like 2D volume support has been dropped or at least not been checked.

We (@Zrrr1997 and me) decided to release the code even before my master thesis is written such that I have enough time to integrate into MONAI if that is desired. The code and a ton of additional information can be found here:
https://git.scc.kit.edu/updcv/sliding-window-based-interactive-segmentation-of-volumetric-medical-images

If you are interested and if it's possible I would love to integrate it into MONAI until the end of September. What I would like to know is what you think would be the best approach? From converting this to a tutorial to integrating the transforms into the official MONAI transforms or keeping it as a separate package, whereever you see fit.
The code also showcases some features which I have rarely or never seen in tutorials like the MONAI Checkpoinloader or the WorkflowProfiler.

Based on this code I created a MONAILabel model in the radiology sample app, which I will try to push to the official repository as well.

More extensive results will follow as well as a user study and probably a paper. Also expect a few final changes, including a move to ScaleIntensityRangePercentilesd, where I still have to verify the metric impact. What is planned too is to verify the pretrained model (AutoPET) on a different dataset, probably HECKTOR.

If there are any questions, I'll be happy to answer them here or personally in a call.

Cheers from Germany ,
Matthias Hadlich

@diazandr3s

[wyli]

thanks! the distance transform based on cucim looks good, If you have time to make the wrappers part of utility functions in monai.transforms, that would be a possible step towards #1332. the deepedit workflow could be considered at monailabel cc @diazandr3s @SachidanandAlle

[matt3o]

Sure I can add that. Can you describe how those should be added?

What what in theory also be possible is to add the distance transforms as an acutal transforms, e.g. DistanceTransformEDTd which creates distance transformed tensors and puts them in key_edt. And analouge to the code then a SelectPointd. Not sure though if that is what you had in your mind.

[wyli]

sure, perhaps we can start with a utility function that is compatible with both numpy and cupy inputs, like this:

MONAI/monai/transforms/utils.py

Lines 1014 to 1026 in a86c0e0

	def get_largest_connected_component_mask(
	img: NdarrayTensor, connectivity: int | None = None, num_components: int = 1
	) -> NdarrayTensor:
	"""
	Gets the largest connected component mask of an image.
	Args:
	img: Image to get largest connected component from. Shape is (spatial_dim1 [, spatial_dim2, ...])
	connectivity: Maximum number of orthogonal hops to consider a pixel/voxel as a neighbor.
	Accepted values are ranging from 1 to input.ndim. If ``None``, a full
	connectivity of ``input.ndim`` is used. for more details:
	https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label.
	num_components: The number of largest components to preserve.

then creating the transforms could be the next step such as

MONAI/monai/transforms/post/array.py

Lines 236 to 238 in a86c0e0

	class KeepLargestConnectedComponent(Transform):
	"""
	Keeps only the largest connected component in the image.

[matt3o]

Awesome, I'll check that out and come back to you in a few days!

[SachidanandAlle]

Based on this code I created a MONAILabel model in the radiology sample app, which I will try to push to the official repository as well.
How about if you can share it so that we can test/verify/compare.. @diazandr3s ?

[matt3o]

Yes I have just added the code under the monailabel folder. Note that this code is not yet cleaned up / does not fully match the way MONAILabel uses configs. I can of course fix that in case you are interested in adding this to the MONAILabel code.
Not sure it adds anything fancy, but the code and the transforms are rather fast. We will use it for PET only annotation DeepEdit style, so 10 clicks on tumor/ background plus a model update after every round of clicks. (Similar to the way the model is trained)
I added a pretrained weight file too. Note that this is trained on PET Scans from AutoPET, so you will have to select those too.

[diazandr3s]

It is great to see this contribution, @matt3o.

We could integrate this version into the main MONAI Label repo as a new model within the Radiology folder.

Just out of curiosity, is this version for binary segmentation, right? Have you tried multilabel segmentation?

In any case, I'd suggest we put the code in MONAI Label format so others can also benefit from this. I know a couple of use cases already.
We can work together on this. I'm happy to join a video call if needed.

Great contribution!

[matt3o]

@diazandr3s Thanks!
So far the code is only tested for binary segmentation. The implementation should work for multilabel as well but I have not tested it.
Since you're interested I'll go clean up the code and make it work similar to what MONAILabel does.
I'll come back to you as soon as I'm done and then we can discuss the next steps.

[diazandr3s]

Superb! Please let me know. Thanks again!