Fixed loading from pretrained model

JanPalasek · Apr 26, 2020 · 026cbd6 · 026cbd6
1 parent c9316fe
commit 026cbd6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/captcha_detection/captcha_network.py b/src/captcha_detection/captcha_network.py
@@ -69,12 +69,12 @@ def __init__(self, image_shape, classes: int, image_preprocess_pipeline, label_p
 
             self._model = tf.keras.Model(inputs=input, outputs=output)
         else:
-            self._model = tf.saved_model.load(args.pretrained_model)
+            self._model = tf.keras.models.load_model(args.pretrained_model)
 
         if args.weights_file is not None:
             self._model.load_weights(args.weights_file)
 
-        print(f"Total layers: {len(self._model.layers)}")
+        # print(f"Total layers: {len(self._model.layers)}")
         if args.remove_layers:
             # remove classification header and add new one
             input = self._model.layers[0].input

diff --git a/src/test.py b/src/test.py
@@ -25,7 +25,7 @@
 if __name__ == "__main__":
     # Parse arguments
     parser = argparse.ArgumentParser()
-    parser.add_argument("--weights_file", default="src/captcha_detection/model.h5", type=str,
+    parser.add_argument("--weights_file", default=None, type=str,
                         help="Path to file that contains pre-trained weights.")
     parser.add_argument("--pretrained_model", default=None, type=str)
     parser.add_argument("--freeze_layers", default=0, type=int,