Merge pull request #65 from SonanceAI/feat/DAT-899

One-line trainer UNETR++ (DAT-899)

Author: luandalmazo

datamint/lightning/__init__.py

@@ -8,9 +8,11 @@
     SemanticSegmentation2DTrainer,
     SemanticSegmentation3DTrainer,
     SegmentationTrainer,
+    VolumeSegmentationTrainer,
     UNetPPTrainer,
     DeepLabV3PlusTrainer,
     TransUNetTrainer,
+    UNETRPPTrainer,
 )
 
 __all__ = [
@@ -21,7 +23,9 @@
     "SemanticSegmentation2DTrainer",
     "SemanticSegmentation3DTrainer",
     "SegmentationTrainer",
+    "VolumeSegmentationTrainer",
     "UNetPPTrainer",
     "DeepLabV3PlusTrainer",
     "TransUNetTrainer",
+    "UNETRPPTrainer",
 ]

datamint/lightning/trainers/__init__.py

@@ -8,15 +8,19 @@
 from .specialized.unetpp import UNetPPTrainer
 from .specialized.deeplabv3plus import DeepLabV3PlusTrainer
 from .specialized.transunet import TransUNetTrainer
+from .vol_seg_trainer import VolumeSegmentationTrainer
+from .specialized.unetrpp import UNETRPPTrainer
 
 __all__ = [
     "BaseTrainer",
     "SegmentationTrainer",
     "SemanticSegmentation2DTrainer",
     "SemanticSegmentation3DTrainer",
+    "VolumeSegmentationTrainer",
     "UNetPPTrainer",
     "DeepLabV3PlusTrainer",
     "TransUNetTrainer",
+    "UNETRPPTrainer",
     "ClassificationTrainer",
     "ImageClassificationTrainer",
 ]

datamint/lightning/trainers/lightning_modules/__init__.py

@@ -1,6 +1,6 @@
 from .base import DatamintLightningModule
 from .segmentation_module import SegmentationModule
-from .segmentation_modules import SMPSegmentationModule, UNetPPModule, DeepLabV3PlusModule, TransUNetModule
+from .segmentation_modules import SMPSegmentationModule, UNetPPModule, DeepLabV3PlusModule, TransUNetModule, UNETRPPModule
 from .classification_module import ClassificationModule
 
-__all__ = ["DatamintLightningModule", "SegmentationModule", "SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule", "TransUNetModule", "ClassificationModule"]
+__all__ = ["DatamintLightningModule", "SegmentationModule", "SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule", "TransUNetModule", "UNETRPPModule", "ClassificationModule"]

datamint/lightning/trainers/lightning_modules/segmentation_modules/__init__.py

@@ -2,5 +2,6 @@
 from .unetpp import UNetPPModule
 from .deeplabv3plus import DeepLabV3PlusModule
 from .transunet import TransUNetModule
+from .unetrpp import UNETRPPModule
 
-__all__ = ["SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule", "TransUNetModule"]
+__all__ = ["SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule", "TransUNetModule", "UNETRPPModule"]