# ================================================================================ #
# Authors: Fabio Frazao and Oliver Kirsebom #
# Contact: fsfrazao@dal.ca, oliver.kirsebom@dal.ca #
# Organization: MERIDIAN (https://meridian.cs.dal.ca/) #
# Team: Data Analytics #
# Project: ketos #
# Project goal: The ketos library provides functionalities for handling #
# and processing acoustic data and applying deep neural networks to sound #
# detection and classification tasks. #
# #
# License: GNU GPLv3 #
# #
# This program is free software: you can redistribute it and/or modify #
# it under the terms of the GNU General Public License as published by #
# the Free Software Foundation, either version 3 of the License, or #
# (at your option) any later version. #
# #
# This program is distributed in the hope that it will be useful, #
# but WITHOUT ANY WARRANTY; without even the implied warranty of #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the #
# GNU General Public License for more details. #
# #
# You should have received a copy of the GNU General Public License #
# along with this program. If not, see <https://www.gnu.org/licenses/>. #
# ================================================================================ #
""" 'neural_networks.dev_utils.export' module within the ketos library
This module contains utilities for saving ketos models in various formats.
"""
import ketos.data_handling.parsing as pa
from tensorflow.saved_model import save as save_pb
# from ketos.audio.audio_loader import audio_repres_dict
from ketos.utils import ensure_dir
from zipfile import ZipFile
import warnings
import json
import shutil
import os
import inspect
import numpy as np
from pathlib import Path, PurePath
def get_export_function(output_name):
r""" Get the appropriate export function for the chosen format.
The function automatically detects the output format from the filename
extension,
* '.pb' or '.PB': protobuf
* '.kt' or '.KT': ketos
* '.ktpb' or '.KTPB': ketos-protobuf (compatible with PAMGuard)
For any other extension, the ketos export function is used.
Args:
output_name: str
The name of the exported model.
Returns:
: func
The export function
"""
if output_name[-3:].lower() == '.pb':
return export_to_protobuf
elif output_name[-3:].lower() == '.kt':
return export_to_ketos
elif output_name[-5:].lower() == '.ktpb':
return export_to_ketos_protobuf
else:
return export_to_ketos
[docs]def export_to_ketos_protobuf(model, output_name, audio_repr, input_shape=None,
tmp_folder="tmp_export_folder", overwrite=True,
duration=None, backward_compat=True, **kwargs):
r""" Export a ketos model to Ketos-Protobuf format.
Saving your ketos model in Ketos-Protobuf format makes it easier to share it with
collaborators and use it with other software applications.
In particular, the output file generated by this function can be loaded directly into
PAMGuard, an open-source and widely adopted application for passive acoustic monitoring (PAM).
If the output directory does not already exist, it is automatically created.
The function generates a zipped archive containing,
* the tensorflow model in protobuf format (model/model.pb)
* the audio representation (audio_repr.json)
* the ketos model recipe (recipe.json)
The user is free to specify the extension of the output file, but we recommend using \*.ktpb
as this will allow the file to be recognized and loaded into PAMGuard.
A warning will be printed if the method is unable to infer values for the parameters
`duration` and `step`. The model will be saved, but it will not be possible to load
it into PAMGuard.
Args:
model:
The ketos model to be exported. Usually created by one of the Interface classes found
in ketos.neural_networks (e.g.: ResNetInterface)
output_name: str
The name of the exported model. Must have the extension \*.ktpb to ensure that it can
be loaded into PAMGuard.
input_shape: list or tuple.
The input shape expected by the model. It can be represented by a tuple or list of four
elements: (number of intances, width, height, number of channels). The number of instances
and number of channels are commonly 1, and the width and height are usually the number of
time and frequency bins in a spectrogram, respectively. This, however, can vary with
the model in question.
audio_repr: dict or str
Audio representation. For example,
>>> audio_repr = {"type": "MagSpectrogram",
... "rate": "1000 Hz",
... "window": "0.256 s",
... "step": "0.032 s",
... "freq_min": "0 Hz",
... "freq_max": "500 Hz",
... "window_func": "hamming",
... "transforms": [{"name":"normalize"}]
... }
It is also possible to specify the path to a json file containing the audio representation.
tmp_folder: str
The name of a temporary folder created during the model conversion. It will be deleted
upon sucessful execution. If the folder already exists, a 'FileExistsError will be thrown,
unless 'overwrite' is set to True.
overwrite: bool
If True and the folder specified in 'tmp_folder' exists, the folder will be overwritten.
duration: float
Duration in seconds of the input sample. If not specified, the duration is extracted from
the audio representation, or, if not available there, it is computed as `step * input_shape[0]`,
provided that `input_shape` has been specified.
backward_compat: bool
Ensure backward compatibility with ketos versions 4.2.1 and older
Raises:
AssertionError if the input shape cannot be inferred.
"""
if input_shape != None and len(input_shape) == 4 and input_shape[0] == 1 and backward_compat:
print("Warning: ketos is using the following interpretation of the input_shape:" \
+ f"{input_shape} = (number of intances, width, height, number of channels). " \
+ "If this is not the correct interpretation, you should call the " \
+ "export_to_ketos_protobuf function with backward_compat=False.")
input_shape = input_shape[1:]
# if the path to a json file containing the audio representation(s)
# has been specified, load its contents
if isinstance(audio_repr, (str, PurePath)):
audio_repr = pa.load_audio_representation(audio_repr, return_unparsed=True)
if os.path.exists(tmp_folder):
if not overwrite:
raise FileExistsError("{} already exists. If you want to overwrite it set the \
'overwrite' argument to True.".format(tmp_folder))
else:
shutil.rmtree(tmp_folder)
# ensure that the audio representation is specified as a nested directory
if 'type' in audio_repr.keys():
audio_repr = {'spectrogram': audio_repr}
# check that one of the representations is named 'spectrogram'
# if this is not the case, issue a warning
if 'spectrogram' not in audio_repr.keys():
warnings.warn("audio representation named `spectrogram` not found. You will not be " \
"able to load the saved model into PAMGuard.", category=UserWarning)
# if the output directory does not already exist, create it
ensure_dir(output_name)
# loop over audio representations
for name, params in audio_repr.items():
# parse audio representation, if it is not already parsed
if pa.is_encoded(params):
params = pa.parse_audio_representation(params)
# if the user did not specify the sample duration, attempt to infer it
if duration == None:
duration = _infer_duration(params, input_shape)
if duration != None:
print(f'Inferred sample duration of {duration}s')
if duration == None:
warnings.warn("ketos was unable to infer the sample duration. You will not be " \
"able to load the saved model into PAMGuard. If you intend to load your model " \
"into PAMGuard, you must specify the sample duration.", category=UserWarning)
else:
params['duration'] = duration
# if the user did not specify an input shape, attempt to infer it
if input_shape == None:
input_shape = _infer_shape(params)
if input_shape != None:
print(f'Inferred input shape of {input_shape}')
assert input_shape != None, 'ketos was unable to infer the input shape. Use the input_shape ' \
+ 'argument to specify the input shape expected by the model.'
if 'step' not in params.keys():
warnings.warn("ketos was unable to infer the step size. You will not be " \
"able to load the saved model into PAMGuard. If you intend to load your model " \
"into PAMGuard, you must specify the step size.", category=UserWarning)
# if model has not been built, build it
X = model._transform_input(np.ones(shape=input_shape))
if not model.model.built:
model.run_on_instance(X[0])
params['dtype'] = model.model.dtype
params['input_ndims'] = model.model.layers[0].input_spec.min_ndim
params['input_shape'] = X.shape
# encode audio representation parameters before saving to file
params = pa.encode_audio_representation(params)
# create tmp folder
os.makedirs(tmp_folder)
# save recipe to tmp folder
recipe_path = os.path.join(tmp_folder, 'recipe.json')
model.save_recipe_file(recipe_path)
# save model to tmp folder
model_path = os.path.join(tmp_folder, 'model')
save_pb(obj=model.model, export_dir=model_path)
# save audio representation to tmp folder
audio_repr_path = os.path.join(tmp_folder, "audio_repr.json")
with open(audio_repr_path, 'w') as json_repr:
json.dump(audio_repr, json_repr)
# now, archive the contents of the tmp folder
with ZipFile(output_name, 'w') as zip:
zip.write(model_path, "model")
for root, dirs, files in os.walk(model_path):
renamed_root = root.replace(model_path, "model")
for d in dirs:
zip.write(os.path.join(root,d), os.path.join(renamed_root,d))
for f in files:
zip.write(os.path.join(root,f),os.path.join(renamed_root,f))
zip.write(recipe_path, "recipe.json")
zip.write(audio_repr_path, "audio_repr.json")
# tidy up
shutil.rmtree(tmp_folder)
def _infer_duration(params, input_shape):
""" Helper function for :func:`export_to_ketos_protobuf`.
The function first searches the `params` dictionary for an item
with key `duration`. If found, the value of this item is returned.
If not found, the function searches for the key `step` and, if found,
returns the product `params['step'] * input_shape[0]`.
If the function is unable to infer the duration through one of
these two approaches, it returns a None value.
Args:
params: dict
Parsed audio representation.
input_shape: tuple or list
Shape of an input sample.
Returns:
: float
Duration in seconds.
"""
if 'duration' in params.keys():
return params['duration']
elif input_shape != None and 'step' in params.keys():
return params['step'] * input_shape[0]
else:
return None
def _infer_shape(params):
""" Helper function for :func:`export_to_ketos_protobuf`.
Args:
params: dict
Parsed audio representation.
Returns:
: tuple
Inferred shape. If the shape could not be inferred,
a None value is returned.
"""
return params['type'].infer_shape(**params)
[docs]def export_to_protobuf(model, output_name, **kwargs):
r""" Export a ketos model to Protobuf format (\*.pb).
See also the related fuction :func:`export_to_ketos_protobuf`.
Note that the model must be built before it can be saved.
Args:
model:
The ketos model to be exported. Usually created by one of the Interface
classes found in ketos.neural_networks (e.g.: ResNetInterface)
output_name: str
Folder where the exported model will be saved
"""
assert model.model.built, "The model must be built. Call model.run_on_instance() on a sample input"
# if the output directory does not already exist, create it
ensure_dir(output_name)
save_pb(obj=model.model, export_dir=output_name)
def export_to_ketos(model, output_name, checkpoint_name="cp-0000", audio_repr=None, custom_module=None,
tmp_folder="tmp_export_folder", overwrite=True, metadata=None, extra=None, **kwargs):
r""" Export a ketos model to ketos format (\*.kt).
If the output directory does not already exist, it is automatically created.
The function generates a zipped archive containing,
* the model weights (\*.ckpt)
* the audio representation (audio_repr.json)
* the ketos model recipe (recipe.json)
Args:
model:
The ketos model to be exported. Usually created by one of the Interface classes found
in ketos.neural_networks (e.g.: ResNetInterface)
output_name: str
The name of the exported model. Should have the extension \*.kt
checkpoint_name: str
The name of the checkpoint file that the model weights will be saved to. Optional.
audio_repr: dict or str
Optional audio representation dictionary. If passed, it will be added to the \*.kt file.
For example,
>>> audio_repr = {"spectrogram": {
... "type": "MagSpectrogram",
... "rate": "1000 Hz",
... "window": "0.256 s",
... "step": "0.032 s",
... "freq_min": "0 Hz",
... "freq_max": "500 Hz",
... "window_func": "hamming",
... "transforms": [{"name":"normalize"}]}
... }
It is also possible to specify the path to a json file containing the audio representation.
custom_module: str
Path to a custom module. The custom module can contain anynumber of files. Common files are user defined nn architecture,
audio representation, input transform function and output transform function.
All files inside the module will be saved to the ketos output file and can be later accessed to build the model for
re-training or inference.
tmp_folder: str
The name of a temporary folder created during the model conversion. It will be deleted
upon sucessful execution. If the folder already exists, a 'FileExistsError will be thrown,
unless 'overwrite' is set to True.
overwrite: bool
If True and the folder specified in 'tmp_folder' exists, the folder will be overwritten.
metadata: dict
Optional metadata dictionary. If passed, it will be added to the \*.kt output archive file.
extra: str, list(str)
Full path to one or several additional files to be saved to the output \*.kt archive file.
"""
if not isinstance(extra, list) and extra is not None:
extra = [extra]
tmp_folder = Path(tmp_folder)
# create tmp folder
if tmp_folder.exists():
if not overwrite:
raise FileExistsError("{} already exists. If you want to overwrite it set the \
'overwrite' argument to True.".format(tmp_folder))
else:
shutil.rmtree(tmp_folder)
recipe = model._extract_recipe_dict()
# save audio representation to tmp folder
if audio_repr is not None:
if isinstance(audio_repr, str) and Path(audio_repr).suffix == ".json":
audio_repr = pa.load_audio_representation(audio_repr, return_unparsed=False)
if isinstance(audio_repr, dict) and pa.is_encoded(audio_repr):
audio_repr = pa.parse_audio_representation(audio_repr)
audio_repr = pa.encode_audio_representation(audio_repr)
if custom_module is not None:
custom_module = Path(custom_module).resolve()
# save checkpoints to tmp folder
# Unfortunetly, the current implementation of save_eights and the ZipFile library does not allow to save the model weights directly to the zipfile
weights_path = (tmp_folder / "checkpoints" / checkpoint_name)
weights_path.parent.mkdir(parents=True, exist_ok=True)
model.model.save_weights(weights_path)
# save to zip archive
with ZipFile(output_name, 'w') as archive:
# save recipe
recipe = json.dumps(recipe) # converts the dictionary as a json_string
archive.writestr('recipe.json', recipe)
# save audio representation
if audio_repr is not None:
audio_repr = json.dumps(audio_repr) # properly converts to json string
archive.writestr('audio_repr.json', audio_repr)
# save metadata
if metadata is not None:
metadata = json.dumps(metadata)
archive.writestr("metadata.json", metadata)
# save checkpoints
checkpoints = weights_path.parent.glob('*')
for checkpoint in checkpoints:
archive.write(checkpoint, Path("checkpoints", checkpoint.name))
# Save custom module to the archive
if custom_module is not None:
files = custom_module.rglob('*')
for file in files:
archive.write(file, "custom/" + file.name)
# save any extra contents
if extra is not None:
for extra_file in extra:
archive.write(extra_file, Path(extra_file).name)
# tidy up
shutil.rmtree(tmp_folder)