Add efficientnets (qubvel#116)(qubvel#118)

* Add efficientnet backbones (B0-B3) * Remove `input_tensor` kwarg * Update README
AloneGu · Jun 7, 2019 · f3b174b · f3b174b
1 parent d513142
commit f3b174b
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 21 deletions.
diff --git a/README.rst b/README.rst
@@ -155,19 +155,20 @@ PSPNet        FPN
 
 .. table:: 
 
-    ===========  ===== 
-    Type         Names
-    ===========  =====
-    VGG          ``'vgg16' 'vgg19'``
-    ResNet       ``'resnet18' 'resnet34' 'resnet50' 'resnet101' 'resnet152'``
-    SE-ResNet    ``'seresnet18' 'seresnet34' 'seresnet50' 'seresnet101' 'seresnet152'``
-    ResNeXt      ``'resnext50' 'resnext101'``
-    SE-ResNeXt   ``'seresnext50' 'seresnext101'``
-    SENet154     ``'senet154'``
-    DenseNet     ``'densenet121' 'densenet169' 'densenet201'`` 
-    Inception    ``'inceptionv3' 'inceptionresnetv2'``
-    MobileNet    ``'mobilenet' 'mobilenetv2'``
-    ===========  =====
+    =============  ===== 
+    Type           Names
+    =============  =====
+    VGG            ``'vgg16' 'vgg19'``
+    ResNet         ``'resnet18' 'resnet34' 'resnet50' 'resnet101' 'resnet152'``
+    SE-ResNet      ``'seresnet18' 'seresnet34' 'seresnet50' 'seresnet101' 'seresnet152'``
+    ResNeXt        ``'resnext50' 'resnext101'``
+    SE-ResNeXt     ``'seresnext50' 'seresnext101'``
+    SENet154       ``'senet154'``
+    DenseNet       ``'densenet121' 'densenet169' 'densenet201'`` 
+    Inception      ``'inceptionv3' 'inceptionresnetv2'``
+    MobileNet      ``'mobilenet' 'mobilenetv2'``
+    EfficientNet   ``'efficientnetb0' 'efficientnetb1' 'efficientnetb2' 'efficientnetb3'``
+    =============  =====
 
 .. epigraph::
     All backbones have weights trained on 2012 ILSVRC ImageNet dataset (``encoder_weights='imagenet'``). 

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,5 @@
 keras>=2.2.0
-keras_applications>=1.0.7
+keras_applications==1.0.7
 scikit-image
-image-classifiers==0.2.0
+image-classifiers==0.2.0
+efficientnet>=0.0.3
diff --git a/segmentation_models/backbones/__init__.py b/segmentation_models/backbones/__init__.py
@@ -1,3 +1,4 @@
+import efficientnet as eff
 from classification_models import Classifiers
 from classification_models import resnext
 
@@ -14,6 +15,18 @@
     'resnext101': [resnext.ResNeXt101, resnext.models.preprocess_input],
     'mobilenet': [mbn.MobileNet, mbn.preprocess_input],
     'mobilenetv2': [mbn2.MobileNetV2, mbn2.preprocess_input],
+
+    'efficientnetb0': [eff.EfficientNetB0, eff.preprocess_input],
+    'efficientnetb1': [eff.EfficientNetB1, eff.preprocess_input],
+    'efficientnetb2': [eff.EfficientNetB2, eff.preprocess_input],
+    'efficientnetb3': [eff.EfficientNetB3, eff.preprocess_input],
+
+    # weights are not released
+#     'efficientnetb4': [eff.EfficientNetB4, eff.preprocess_input],
+#     'efficientnetb5': [eff.EfficientNetB5, eff.preprocess_input],
+#     'efficientnetb6': [eff.EfficientNetB6, eff.preprocess_input],
+#     'efficientnetb7': [eff.EfficientNetB7, eff.preprocess_input],
+
 })
 
 DEFAULT_FEATURE_LAYERS = {
@@ -59,7 +72,19 @@
     # 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'),
-
+
+    # EfficientNets
+    'efficientnetb0': (169, 77, 47, 17),
+    'efficientnetb1': (246, 122, 76, 30),
+    'efficientnetb2': (246, 122, 76, 30),
+    'efficientnetb3': (278, 122, 76, 30),
+
+    # weights are not released
+#     'efficientnetb4': (342, 154, 92, 30),
+#     'efficientnetb5': (419, 199, 121, 43),
+#     'efficientnetb6': (483, 231, 137, 43),
+#     'efficientnetb7': (592, 276, 166, 56),
+
 }
 
 

diff --git a/segmentation_models/fpn/model.py b/segmentation_models/fpn/model.py
@@ -11,13 +11,13 @@
     'interpolation': 'final_interpolation',
     'upsample_rates': None,  # removed
     'last_upsample': None,  # removed
+    'input_tensor': None, # removed
 }
 
 
 @legacy_support(old_args_map)
 def FPN(backbone_name='vgg16',
         input_shape=(None, None, 3),
-        input_tensor=None,
         classes=21,
         activation='softmax',
         encoder_weights='imagenet',
@@ -61,7 +61,6 @@ def FPN(backbone_name='vgg16',
 
     backbone = get_backbone(backbone_name,
                             input_shape=input_shape,
-                            input_tensor=input_tensor,
                             weights=encoder_weights,
                             include_top=False)
 

diff --git a/segmentation_models/linknet/model.py b/segmentation_models/linknet/model.py
@@ -62,7 +62,6 @@ def Linknet(backbone_name='vgg16',
 
     backbone = get_backbone(backbone_name,
                             input_shape=input_shape,
-                            input_tensor=None,
                             weights=encoder_weights,
                             include_top=False)
 

diff --git a/segmentation_models/pspnet/model.py b/segmentation_models/pspnet/model.py
@@ -84,7 +84,6 @@ def PSPNet(backbone_name='vgg16',
 
     backbone = get_backbone(backbone_name,
                             input_shape=input_shape,
-                            input_tensor=None,
                             weights=encoder_weights,
                             include_top=False)
 

diff --git a/segmentation_models/unet/model.py b/segmentation_models/unet/model.py
@@ -58,7 +58,6 @@ def Unet(backbone_name='vgg16',
 
     backbone = get_backbone(backbone_name,
                             input_shape=input_shape,
-                            input_tensor=None,
                             weights=encoder_weights,
                             include_top=False)