Add Swin-UNETR to Keras-Hub

[innat]

Is your feature request related to a problem? Please describe.

Swin-UNETR was originally designed for 3D medical image segmentation, using Swin Transformers for effective feature extraction. However, there is currently no implementation of this model available on Keras Hub.

MONAI provides Swin-UNETR for both 2D and 3D segmentation tasks, offering a well-optimized framework for medical image analysis. Having Swin-UNETR to Keras Hub would make it easier for users to apply it with torch backend in MONAI workflows.

Describe the solution you'd like

Paper of Swin-UNETR: https://arxiv.org/abs/2201.01266
Cited by (until now): 1341

Additional context

• The backbone of Swin-UNETR is Video Swin Transformer model with few modifications.
• Having Keras 3 implementation, it would be possible to run it with MONAI with torch backend.
• In order to run it with Keras 3 workflows with all backend (instead of MONAI), several utilities methos will be required, i.e.
  • sliding_window_inference
  • Augmentation methods, i.e. ScaleIntensityRanged, CropForegroundd, RandCropByPosNegLabeld, Spacingd, etc. These are from MONAI.
  • Segmentation metrics: DICE for 3D cases.
• I am open to contributing to this effort if the Keras team considers it an important and high-priority feature.

[mattdangerw]

Do we have some indication of how many downloads/users this model has today? I would think that it is lower priority than, say, qwen, which we are working on adding now. But I don't have a clear read on it's popularit.

Is the main task a segmentation task essentially? How well does fit with the high level abstractions we have in Keras today?

Assigning to @divyashreepathihalli for thoughts, she's been maintaining our model wish list.

[innat]

@mattdangerw
I don’t have specific data on the number of users, downloads, or the overall popularity of this model. One potential way to gauge its impact could be through citations. To my knowledge, this model was originally developed by NVIDIA for 3D medical imaging tasks. Through MONAI, it can also be applied to both 2D and 3D segmentation tasks, particularly in the medical imaging domain.

Currently, I’ve successfully converted the model from MONAI to Keras 3. A key requirement for this model is the use of a Video Swin Transformer as the backbone (for 3D applications) combined with UNETR’s decoder. While the model can be used with MONAI pipelines, pure PyTorch, or PyTorch Lightning when using the Torch backend. Supporting it across all backends in Keras will require implementing several additional methods, i.e. sliding_window_inference, resample depth for volumetric data etc.