MAGDI Segmentation Models 3D

This python package named magdi_segmentation_models_3d responsible for providing custom Hugging Face compatible models for 3D image segmentation for the project MAGDI.

Hugging Face Custom Models

Documentation on Hugging Face: https://huggingface.co/docs/transformers/en/custom_models

Examples: https://github.com/huggingface/transformers/tree/main/src/transformers/models

mednext

MedNeXt implementation from monai wrapped as Hugging Face model.

References:

• 10.48550/arXiv.2303.09975

Usage example:

from magdi_segmentation_models_3d import MedNeXtModel, MedNeXtConfig,

MedNeXtForImageSegmentation, MedNeXtImageProcessor

MedNeXtConfig.register_for_auto_class()
MedNeXtModel.register_for_auto_class("AutoModel")
MedNeXtForImageSegmentation.register_for_auto_class(
    "AutoModelForImageSegmentation"
)
MedNeXtImageProcessor.register_for_auto_class("AutoImageProcessor")

mednext_config = MedNeXtConfig(
    variant='B',
    spatial_dims=3,
    in_channels=1,
    out_channels=5,
    kernel_size=3,
    deep_supervision=False,
)
mednext_model = MedNeXtForImageSegmentation(mednext_config)
processor = MedNeXtImageProcessor()

nnunetresenc

ResidualEncoderUNet from dynamic_network_architectures.architectures.unet wrapped as Hugging Face model. This architecture is also being used by nnUNet https://github.com/MIC-DKFZ/nnUNet.

References:

• 10.48550/arXiv.1809.10486
• 10.48550/arXiv.2404.09556

Usage example:

from magdi_segmentation_models_3d import nnUNetResEncConfig, nnUNetResEncModel,

nnUNetResEncForImageSegmentation, nnUNetResEncImageProcessor

nnUNetResEncConfig.register_for_auto_class()
nnUNetResEncModel.register_for_auto_class("AutoModel")
nnUNetResEncForImageSegmentation.register_for_auto_class(
    "AutoModelForImageSegmentation"
)
nnUNetResEncImageProcessor.register_for_auto_class("AutoImageProcessor")

nnunet_config = nnUNetResEncConfig(
    variant="B",  # only B supported yet
    in_channels=1,
    out_channels=5,
    enable_deep_supervision=False,
)
nnunet_model = nnUNetResEncForImageSegmentation(nnunet_config)
processor = nnUNetResEncImageProcessor()

stunet

STU-Net from https://github.com/Ziyan-Huang/STU-Net wrapped as Hugging Face model.

References:

• 10.48550/arXiv.2304.06716

Usage example:

from magdi_segmentation_models_3d import STUNetConfig, STUNetModel,

STUNetForImageSegmentation, STUNetImageProcessor

STUNetConfig.register_for_auto_class()
STUNetModel.register_for_auto_class("AutoModel")
STUNetForImageSegmentation.register_for_auto_class(
    "AutoModelForImageSegmentation"
)
STUNetImageProcessor.register_for_auto_class("AutoImageProcessor")

stu_net_config = STUNetConfig(
    variant='B',
    in_channels=1,
    out_channels=5,
    kernel_size=[[3, 3, 3]] * 6,
    deep_supervision=True,
)
stu_net_model = STUNetForImageSegmentation(stu_net_config)
processor = STUNetImageProcessor()

swinunetrv2

swinUNETRV2 implementation from monai wrapped as Hugging Face model.

References:

• 10.48550/arXiv.2201.01266

Usage example:

from magdi_segmentation_models_3d import SwinUNETRv2Config, SwinUNETRv2Model,

SwinUNETRv2ForImageSegmentation, SwinUNETRv2ImageProcessor

SwinUNETRv2Config.register_for_auto_class()
SwinUNETRv2Model.register_for_auto_class("AutoModel")
SwinUNETRv2ForImageSegmentation.register_for_auto_class(
    "AutoModelForImageSegmentation"
)
SwinUNETRv2ImageProcessor.register_for_auto_class("AutoImageProcessor")
swin_unetr_v2_config = SwinUNETRv2Config(
    in_channels=1,
    out_channels=5,
    depths=(2, 2, 2, 2),
    num_heads=(3, 6, 12, 24),
    feature_size=48,
    patch_size=2,
    window_size=7,
    drop_rate=0.2,
    attn_drop_rate=0.2,
    dropout_path_rate=0.2,
    spatial_dims=3,
)
swinunetrv2_model = SwinUNETRv2ForImageSegmentation(swin_unetr_v2_config)

processor = SwinUNETRv2ImageProcessor()

unet

Enhanced version of U-Net - Residual U-Net - implementation from monai wrapped as Hugging Face model.

References:

• https://link.springer.com/chapter/10.1007/978-3-030-12029-0_40

Usage example:

from magdi_segmentation_models_3d import UnetConfig, UnetModel,

UnetForImageSegmentation, UnetImageProcessor

UnetConfig.register_for_auto_class()
UnetModel.register_for_auto_class("AutoModel")
UnetForImageSegmentation.register_for_auto_class(
    "AutoModelForImageSegmentation"
)
UnetImageProcessor.register_for_auto_class("AutoImageProcessor")

unet_config = UnetConfig(
    in_channels=1,
    out_channels=5,
    channels=(64, 128, 256, 512, 1024),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    spatial_dims=3,
    dropout=0.2,
)
unet_model = UnetForImageSegmentation(unet_config)
processor = UnetImageProcessor()