Merge pull request #398 from Sofia-Mouchtaris/unetr

A new transformer-based segmentation arch (`UNETR`) has been added
mlcommons · Apr 12, 2022 · 3e35ff0 · 3e35ff0
2 parents def9014 + ae2df32
commit 3e35ff0
Show file tree

Hide file tree

Showing 3 changed files with 580 additions and 2 deletions.
diff --git a/GANDLF/models/__init__.py b/GANDLF/models/__init__.py
@@ -31,6 +31,7 @@
 from .sdnet import SDNet
 from .MSDNet import MSDNet
 from .brain_age import brainage
+from .unetr import unetr
 
 # defining dict for models - key is the string and the value is the transform object
 global_models_dict = {
@@ -52,6 +53,7 @@
     "deep_resunet": deep_resunet,
     "light_resunet": light_resunet,
     "light_resunet_multilayer": light_resunet_multilayer,
+    "unetr": unetr,
     "fcn": fcn,
     "uinc": uinc,
     "vgg": vgg19,