Skip to content

Commit

Permalink
Amend documentation for generate_model_image
Browse files Browse the repository at this point in the history
  • Loading branch information
FrankWhoee committed Mar 30, 2019
1 parent 3b77248 commit ff693e6
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions CNN-visualiser.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,17 @@
from matplotlib import pyplot


def generate_model_image(model_path):
def generate_model_image(model_path, filename):
"""
Generates a flow chart of the model to model.png
Generates a flow chart of the model.
:param model_path: Path of the Keras model to be loaded. Expects string input.
:param filename: Path of file to be generated and downloaded. Expects string input.
:return: None
"""
print("Loading model...")
model = load_model(model_path)
print("Model loaded. Plotting model...")
plot_model(model, to_file='model.png', show_shapes=True)
plot_model(model, to_file=filename, show_shapes=True)
print("Plotting complete. File is ready at model.png")


Expand Down Expand Up @@ -58,3 +59,4 @@ def visualize_feature_maps(model_path, layer):
pyplot.axes().set_aspect('equal', 'datalim')
pyplot.pcolormesh(layers[:, :, 0])
pyplot.show()

0 comments on commit ff693e6

Please sign in to comment.