Add Weights for B4 and B5 models, bump version and setup.py

titu1994 · Jun 19, 2019 · 156df01 · 156df01
1 parent b3d3905
commit 156df01
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 29 deletions.
diff --git a/keras_efficientnets/__init__.py b/keras_efficientnets/__init__.py
@@ -1,4 +1,4 @@
 from keras_efficientnets.efficientnet import *
 from keras_efficientnets.config import BlockArgs
 
-__version__ = '0.1.3'
+__version__ = '0.1.4'
diff --git a/keras_efficientnets/efficientnet.py b/keras_efficientnets/efficientnet.py
@@ -479,32 +479,32 @@ def EfficientNet(input_shape,
                     cache_subdir='models')
             model.load_weights(weights_path)
 
-        # elif default_size == 380:
-        #     if include_top:
-        #         weights_path = get_file(
-        #             'efficientnet-b4.h5',
-        #             "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b4.h5",
-        #             cache_subdir='models')
-        #     else:
-        #         weights_path = get_file(
-        #             'efficientnet-b4_notoph5',
-        #             "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b4_notop.h5",
-        #             cache_subdir='models')
-        #     model.load_weights(weights_path)
-        #
-        # elif default_size == 456:
-        #     if include_top:
-        #         weights_path = get_file(
-        #             'efficientnet-b5.h5',
-        #             "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b5.h5",
-        #             cache_subdir='models')
-        #     else:
-        #         weights_path = get_file(
-        #             'efficientnet-b5_notop.h5',
-        #             "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b5_notop.h5",
-        #             cache_subdir='models')
-        #     model.load_weights(weights_path)
-        #
+        elif default_size == 380:
+            if include_top:
+                weights_path = get_file(
+                    'efficientnet-b4.h5',
+                    "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b4.h5",
+                    cache_subdir='models')
+            else:
+                weights_path = get_file(
+                    'efficientnet-b4_notoph5',
+                    "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b4_notop.h5",
+                    cache_subdir='models')
+            model.load_weights(weights_path)
+
+        elif default_size == 456:
+            if include_top:
+                weights_path = get_file(
+                    'efficientnet-b5.h5',
+                    "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b5.h5",
+                    cache_subdir='models')
+            else:
+                weights_path = get_file(
+                    'efficientnet-b5_notop.h5',
+                    "https://github.com/titu1994/keras-efficientnets/releases/download/v0.1/efficientnet-b5_notop.h5",
+                    cache_subdir='models')
+            model.load_weights(weights_path)
+
         # elif default_size == 528:
         #     if include_top:
         #         weights_path = get_file(
@@ -822,7 +822,7 @@ def EfficientNetB3(input_shape=None,
 
 def EfficientNetB4(input_shape=None,
                    include_top=True,
-                   weights=None,
+                   weights='imagenet',
                    input_tensor=None,
                    pooling=None,
                    classes=1000,
@@ -892,7 +892,7 @@ def EfficientNetB4(input_shape=None,
 
 def EfficientNetB5(input_shape=None,
                    include_top=True,
-                   weights=None,
+                   weights='imagenet',
                    input_tensor=None,
                    pooling=None,
                    classes=1000,

diff --git a/setup.py b/setup.py
@@ -73,6 +73,13 @@ def run(self):
         os.system('git tag v{0}'.format(get_version()))
         os.system('git push --tags')
 
+        try:
+            self.status('Removing build artifacts...')
+            rmtree(os.path.join(base_path, 'build'))
+            rmtree(os.path.join(base_path, 'keras_efficientnets.egg-info'))
+        except OSError:
+            pass
+
         sys.exit()
 
     def initialize_options(self):