Merge pull request qubvel#43 from qubvel/feature-new-backbones

New backbones

requirements.txt

@@ -1,3 +1,3 @@
 keras>=2.1.4
 scikit-image
-image-classifiers==0.1.0rc0
+image-classifiers==0.2.0

segmentation_models/backbones/__init__.py

@@ -1,5 +1,75 @@
-from .inception_resnet_v2 import InceptionResNetV2
-from .inception_v3 import InceptionV3
+from classification_models import Classifiers
+from classification_models import resnext
 
-from .backbones import get_backbone
-from .preprocessing import get_preprocessing
+from . import inception_resnet_v2 as irv2
+from . import inception_v3 as iv3
+
+# replace backbones with others, which have corrected padding mode in first pooling
+Classifiers._models.update({
+    'inceptionresnetv2': [irv2.InceptionResNetV2, irv2.preprocess_input],
+    'inceptionv3': [iv3.InceptionV3, iv3.preprocess_input],
+    'resnext50': [resnext.ResNeXt50, resnext.models.preprocess_input],
+    'resnext101': [resnext.ResNeXt101, resnext.models.preprocess_input],
+})
+
+DEFAULT_FEATURE_LAYERS = {
+
+    # List of layers to take features from backbone in the following order:
+    # (x16, x8, x4, x2, x1) - `x4` mean that features has 4 times less spatial
+    # resolution (Height x Width) than input image.
+
+    # VGG
+    'vgg16': ('block5_conv3', 'block4_conv3', 'block3_conv3', 'block2_conv2', 'block1_conv2'),
+    'vgg19': ('block5_conv4', 'block4_conv4', 'block3_conv4', 'block2_conv2', 'block1_conv2'),
+
+    # ResNets
+    'resnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'resnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'resnet50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'resnet101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'resnet152': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+
+    # ResNeXt
+    'resnext50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'resnext101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+
+    # Inception
+    'inceptionv3': (228, 86, 16, 9),
+    'inceptionresnetv2': (594, 260, 16, 9),
+
+    # DenseNet
+    'densenet121': (311, 139, 51, 4),
+    'densenet169': (367, 139, 51, 4),
+    'densenet201': (479, 139, 51, 4),
+
+    # SE models
+    'seresnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'seresnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
+    'seresnet50': (233, 129, 59, 4),
+    'seresnet101': (522, 129, 59, 4),
+    'seresnet152': (811, 197, 59, 4),
+    'seresnext50': (1065, 577, 251, 4),
+    'seresnext101': (2442, 577, 251, 4),
+    'senet154': (6837, 1614, 451, 12),
+
+    # Mobile Nets
+    'mobilenet': ('conv_pw_11_relu', 'conv_pw_5_relu', 'conv_pw_3_relu', 'conv_pw_1_relu'),
+    'mobilenetv2': ('block_13_expand_relu', 'block_6_expand_relu', 'block_3_expand_relu', 'block_1_expand_relu'),
+
+}
+
+
+def get_names():
+    return list(DEFAULT_FEATURE_LAYERS.keys())
+
+
+def get_feature_layers(name, n=5):
+    return DEFAULT_FEATURE_LAYERS[name][:n]
+
+
+def get_backbone(name, *args, **kwargs):
+    return Classifiers.get_classifier(name)(*args, **kwargs)
+
+
+def get_preprocessing(name):
+    return Classifiers.get_preprocessing(name)

segmentation_models/fpn/model.py

@@ -1,25 +1,8 @@
 from .builder import build_fpn
-from ..backbones import get_backbone
+from ..backbones import get_backbone, get_feature_layers
 from ..utils import freeze_model
 from ..utils import legacy_support
 
-DEFAULT_FEATURE_PYRAMID_LAYERS = {
-    'vgg16': ('block5_conv3', 'block4_conv3', 'block3_conv3'),
-    'vgg19': ('block5_conv4', 'block4_conv4', 'block3_conv4'),
-    'resnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet152': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnext50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnext101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'inceptionv3': (228, 86, 16),
-    'inceptionresnetv2': (594, 260, 16),
-    'densenet121': (311, 139, 51),
-    'densenet169': (367, 139, 51),
-    'densenet201': (479, 139, 51),
-}
-
 old_args_map = {
     'freeze_encoder': 'encoder_freeze',
     'fpn_layers': 'encoder_features',
@@ -83,7 +66,7 @@ def FPN(backbone_name='vgg16',
                             include_top=False)
 
     if encoder_features == 'default':
-        encoder_features = DEFAULT_FEATURE_PYRAMID_LAYERS[backbone_name]
+        encoder_features = get_feature_layers(backbone_name, n=3)
 
     upsample_rates = [2] * len(encoder_features)
     last_upsample = 2 ** (5 - len(encoder_features))

segmentation_models/linknet/model.py

@@ -1,24 +1,7 @@
 from .builder import build_linknet
 from ..utils import freeze_model
 from ..utils import legacy_support
-from ..backbones import get_backbone
-
-DEFAULT_SKIP_CONNECTIONS = {
-    'vgg16': ('block5_conv3', 'block4_conv3', 'block3_conv3', 'block2_conv2'),
-    'vgg19': ('block5_conv4', 'block4_conv4', 'block3_conv4', 'block2_conv2'),
-    'resnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet152': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnext50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnext101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'inceptionv3': (228, 86, 16, 9),
-    'inceptionresnetv2': (594, 260, 16, 9),
-    'densenet121': (311, 139, 51, 4),
-    'densenet169': (367, 139, 51, 4),
-    'densenet201': (479, 139, 51, 4),
-}
+from ..backbones import get_backbone, get_feature_layers
 
 old_args_map = {
     'freeze_encoder': 'encoder_freeze',
@@ -84,7 +67,7 @@ def Linknet(backbone_name='vgg16',
                             include_top=False)
 
     if encoder_features == 'default':
-        encoder_features = DEFAULT_SKIP_CONNECTIONS[backbone_name]
+        encoder_features = get_feature_layers(backbone_name, n=4)
 
     model = build_linknet(backbone,
                           classes,

segmentation_models/pspnet/model.py

@@ -1,33 +1,17 @@
 from .builder import build_psp
 from ..utils import freeze_model
 from ..utils import legacy_support
-from ..backbones import get_backbone
-
-PSP_BASE_LAYERS = {
-    'vgg16': ('block5_conv3', 'block4_conv3', 'block3_conv3'),
-    'vgg19': ('block5_conv4', 'block4_conv4', 'block3_conv4'),
-    'resnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnet152': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnext50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'resnext101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1'),
-    'inceptionv3': (228, 86, 16),
-    'inceptionresnetv2': (594, 260, 16),
-    'densenet121': (311, 139, 51),
-    'densenet169': (367, 139, 51),
-    'densenet201': (479, 139, 51),
-}
+from ..backbones import get_backbone, get_feature_layers
 
 
 def _get_layer_by_factor(backbone_name, factor):
+    feature_layers = get_feature_layers(backbone_name, n=3)
     if factor == 4:
-        return PSP_BASE_LAYERS[backbone_name][-1]
+        return feature_layers[-1]
     elif factor == 8:
-        return PSP_BASE_LAYERS[backbone_name][-2]
+        return feature_layers[-2]
     elif factor == 16:
-        return PSP_BASE_LAYERS[backbone_name][-3]
+        return feature_layers[-3]
     else:
         raise ValueError('Unsupported factor - `{}`, Use 4, 8 or 16.'.format(factor))
 

segmentation_models/unet/model.py

@@ -1,29 +1,12 @@
 from .builder import build_unet
 from ..utils import freeze_model
 from ..utils import legacy_support
-from ..backbones import get_backbone
-
-DEFAULT_SKIP_CONNECTIONS = {
-    'vgg16': ('block5_conv3', 'block4_conv3', 'block3_conv3', 'block2_conv2', 'block1_conv2'),
-    'vgg19': ('block5_conv4', 'block4_conv4', 'block3_conv4', 'block2_conv2', 'block1_conv2'),
-    'resnet18': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),  # check 'bn_data'
-    'resnet34': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnet152': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnext50': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'resnext101': ('stage4_unit1_relu1', 'stage3_unit1_relu1', 'stage2_unit1_relu1', 'relu0'),
-    'inceptionv3': (228, 86, 16, 9),
-    'inceptionresnetv2': (594, 260, 16, 9),
-    'densenet121': (311, 139, 51, 4),
-    'densenet169': (367, 139, 51, 4),
-    'densenet201': (479, 139, 51, 4),
-}
+from ..backbones import get_backbone, get_feature_layers
 
 old_args_map = {
     'freeze_encoder': 'encoder_freeze',
     'skip_connections': 'encoder_features',
-    'upsample_rates': None, # removed
+    'upsample_rates': None,  # removed
     'input_tensor': None,  # removed
 }
 
@@ -80,7 +63,7 @@ def Unet(backbone_name='vgg16',
                             include_top=False)
 
     if encoder_features == 'default':
-        encoder_features = DEFAULT_SKIP_CONNECTIONS[backbone_name]
+        encoder_features = get_feature_layers(backbone_name, n=4)
 
     model = build_unet(backbone,
                        classes,

tests/test_models.py

@@ -9,10 +9,20 @@
 from segmentation_models import Linknet
 from segmentation_models import PSPNet
 from segmentation_models import FPN
-from segmentation_models.backbones import backbones as bkb
+from segmentation_models import backbones as sm_backbones
 
 
-BACKBONES = list(bkb.backbones.keys())
+def get_backbones():
+    is_travis = os.environ.get('TRAVIS', False)
+    exclude = ['senet154']
+    backbones = sm_backbones.get_names()
+
+    if is_travis:
+        backbones = [b for b in backbones if b not in exclude]
+    return backbones
+
+
+BACKBONES = get_backbones()
 
 
 def _select_names(names):