GaNDLF/GANDLF/models/__init__.py at master · mlcommons/GaNDLF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# Import the different model architectures and functions
from .unet import unet, resunet
from .light_unet import light_unet, light_resunet
from .unet_multilayer import unet_multilayer, resunet_multilayer
from .light_unet_multilayer import light_unet_multilayer, light_resunet_multilayer
from .deep_unet import deep_unet, deep_resunet
from .uinc import uinc
from .fcn import fcn
from .vgg import vgg11, vgg13, vgg16, vgg19
from .densenet import densenet121, densenet169, densenet201, densenet264
from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152, resnet200
from .dynunet_wrapper import dynunet_wrapper
from .efficientnet import (
    efficientnetB0,
    efficientnetB1,
    efficientnetB2,
    efficientnetB3,
    efficientnetB4,
    efficientnetB5,
    efficientnetB6,
    efficientnetB7,
)
from .imagenet_vgg import (
    imagenet_vgg11,
    imagenet_vgg11_bn,
    imagenet_vgg13,
    imagenet_vgg13_bn,
    imagenet_vgg16,
    imagenet_vgg16_bn,
    imagenet_vgg19,
    imagenet_vgg19_bn,
)
from .imagenet_unet import imagenet_unet_wrapper
from .sdnet import SDNet
from .MSDNet import MSDNet
from .brain_age import brainage
from .unetr import unetr
from .transunet import transunet
from .modelBase import ModelBase

# Define a dictionary of model architectures and corresponding functions
global_models_dict = {
    # Types of unet
    "unet": unet,
    "unet_multilayer": unet_multilayer,
    "resunet": resunet,
    "resunet_multilayer": resunet_multilayer,
    "residualunet": resunet,  # Alias for "resunet"
    "residualunet_multilayer": resunet_multilayer,  # Alias for "resunet_multilayer"
    "deepunet": deep_unet,
    "lightunet": light_unet,
    "lightunet_multilayer": light_unet_multilayer,
    "deep_unet": deep_unet,  # Alias for "deepunet"
    "light_unet": light_unet,  # Alias for "lightunet"
    "light_unet_multilayer": light_unet_multilayer,  # Alias for "lightunet_multilayer"
    "deepresunet": deep_resunet,
    "lightresunet": light_resunet,
    "lightresunet_multilayer": light_resunet_multilayer,
    "deep_resunet": deep_resunet,  # Alias for "deepresunet"
    "light_resunet": light_resunet,  # Alias for "lightresunet"
    "light_resunet_multilayer": light_resunet_multilayer,  # Alias for "lightresunet_multilayer"
    "unetr": unetr,
    "transunet": transunet,
    "uinc": uinc,
    # UNet models with imagenet support from segmentation_models.pytorch
    "imagenet_unet": imagenet_unet_wrapper,
    # Additional segmentation model
    "fcn": fcn,
    # VGG models
    "vgg": vgg19,
    "vgg11": vgg11,
    "vgg13": vgg13,
    "vgg16": vgg16,
    "vgg19": vgg19,
    # VGG models with imagenet support
    "imagenet_vgg11": imagenet_vgg11,
    "imagenet_vgg11_bn": imagenet_vgg11_bn,
    "imagenet_vgg13": imagenet_vgg13,
    "imagenet_vgg13_bn": imagenet_vgg13_bn,
    "imagenet_vgg16": imagenet_vgg16,
    "imagenet_vgg16_bn": imagenet_vgg16_bn,
    "imagenet_vgg19": imagenet_vgg19,
    "imagenet_vgg19_bn": imagenet_vgg19_bn,
    # DenseNet models
    "densenet": densenet264,
    "densenet121": densenet121,
    "densenet169": densenet169,
    "densenet201": densenet201,
    "densenet264": densenet264,
    # ResNet models
    "resnet18": resnet18,
    "resnet34": resnet34,
    "resnet50": resnet50,
    "resnet101": resnet101,
    "resnet152": resnet152,
    "resnet200": resnet200,
    # EfficientNet models
    "efficientnetb0": efficientnetB0,
    "efficientnetb1": efficientnetB1,
    "efficientnetb2": efficientnetB2,
    "efficientnetb3": efficientnetB3,
    "efficientnetb4": efficientnetB4,
    "efficientnetb5": efficientnetB5,
    "efficientnetb6": efficientnetB6,
    "efficientnetb7": efficientnetB7,
    "dynunet": dynunet_wrapper,
    # Custom models
    "msdnet": MSDNet,
    "brain_age": brainage,
    "sdnet": SDNet,
}


def get_model(params: dict) -> ModelBase:
    """
    Function to get the model definition.

    Args:
        params (dict): The parameters' dictionary.

    Returns:
        model (ModelBase): The model definition.
    """
    chosen_model = params["model"]["architecture"].lower()
    assert (
        chosen_model in global_models_dict
    ), f"Could not find the requested model '{params['model']['architecture']}'"
    return global_models_dict[chosen_model](parameters=params)