Update model loading to be compliant with new versions of tensorflow

sarwaveifrproc/config.yaml

@@ -1,5 +1,6 @@
-model_intraburst: 'model_intraburst.keras'
-model_interburst: 'model_interburst.keras'
+model_version: '1.0.0'
+model_intraburst: 'model_intraburst.h5'
+model_interburst: 'model_interburst.h5'
 scaler_intraburst: 'scaler_intraburst.npy'
 scaler_interburst: 'scaler_interburst.npy'
 bins_intraburst: 'bins_intraburst'

sarwaveifrproc/main.py

@@ -4,19 +4,21 @@
 import sys
 import glob
 import numpy as np
-from sarwaveifrproc.utils import get_output_safe, load_config, load_models, process_files
 
 def parse_args():
 
     parser = argparse.ArgumentParser(description="Generate a L2 WAVE product from a L1B or L1C SAFE.")
-
+    
     # Define arguments
     parser.add_argument("--input_path", required=True, help="l1b or l1c safe path or listing path (.txt file).")
     parser.add_argument("--save_directory", required=True, help="where to save output data.")
     parser.add_argument("--product_id", required=True, help="3 digits ID representing the processing options. Ex: E00.")
+
+    parser.add_argument("--config_path", default=None, help="custom config path for model prediction.")
 
     # Group related arguments under 'model' and 'bins'
     model_group = parser.add_argument_group('model', 'Arguments related to the neural models')
+    model_group.add_argument("--model_version", required=False, help="models version.")
     model_group.add_argument("--model_intraburst", required=False, help="neural model path to predict sea states on intraburst data.")
     model_group.add_argument("--model_interburst", required=False, help="neural model path to predict sea states on interburst data.")
 
@@ -41,6 +43,7 @@ def setup_logging(verbose=False):
     level = logging.DEBUG if verbose else logging.INFO
     logging.basicConfig(level=level, format=fmt, datefmt='%d/%m/%Y %H:%M:%S', force=True)
 
+
 def get_files(dir_path, listing):
 
     fn = []
@@ -59,13 +62,23 @@ def main():
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # Suppress TensorFlow INFO and WARNING messages
     os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # Hide CUDA devices
 
+    #  Functions are loaded here otherwise it takes too long to display help
+    from sarwaveifrproc.utils import get_output_safe, load_config, load_models_and_scalers, process_files
+
     logging.info("Loading configuration file...")
-    conf = load_config()
 
     input_path = args.input_path
     save_directory = args.save_directory
     product_id = args.product_id
+
+    conf = load_config(args.config_path)
+
+    model_version = args.model_version or conf['model_version']
     predicted_variables = args.predicted_variables or conf['predicted_variables']
+
+    # ========================= # 
+    # === PROCESS A LISTING === #
+    # ========================= # 
 
     if input_path.endswith('.txt'):
         files = np.loadtxt(input_path, dtype=str)
@@ -90,14 +103,17 @@ def main():
         'bins_intraburst': args.bins_intraburst or conf['bins_intraburst'],
         'bins_interburst': args.bins_interburst or conf['bins_interburst'],
         }
-        model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst = load_models(paths, predicted_variables)
+        model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst = load_models_and_scalers(model_version, paths, predicted_variables)
         logging.info('Models loaded.')
 
         logging.info('Processing files...')
         for f, output_safe in zip(files, output_safes):
             process_files(f, output_safe, model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst, predicted_variables, product_id)
 
-
+    # ============================= # 
+    # === PROCESS A SINGLE FILE === #
+    # ============================= # 
+
     else:
         logging.info("Checking if output safe already exists...")
         output_safe = get_output_safe(input_path, save_directory, product_id)
@@ -115,7 +131,7 @@ def main():
             'bins_intraburst': args.bins_intraburst or conf['bins_intraburst'],
             'bins_interburst': args.bins_interburst or conf['bins_interburst'],
         }
-        model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst = load_models(paths, predicted_variables)
+        model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst = load_models_and_scalers(model_version, paths, predicted_variables)
         logging.info('Models loaded.')
 
         logging.info('Processing files...')

sarwaveifrproc/models.py

@@ -0,0 +1,48 @@
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Dense, Concatenate, Input, Dropout
+
+def get_model(version):
+    """
+    Retrieves a neural network model builder function based on the specified version.
+
+    Parameters:
+    - version (str): A string indicating the version of the model to retrieve.
+
+    Returns:
+    - model_builder (function): A function that builds the specified version of the neural network model.
+    """
+    model_builders = {'1.0.0': build_model_v1_0_0,
+                      '1.0.1': build_model_v1_0_0}
+
+    return model_builders[version]
+
+
+def build_model_v1_0_0(nb_classes):
+    """
+    Builds a neural network model corresponding to the version v1.0.0.
+
+    Parameters:
+    - nb_classes (list or tuple): A list or tuple containing the number of output classes for each output layer.
+
+    Returns:
+    - model (keras.Model): A compiled Keras model with the specified architecture.
+
+    Architecture:
+    - Input layer: 24 features.
+    - 10 Dense layers, each with 1024 neurons and ReLU activation.
+    - Output layer: Concatenation of output layers with a number of neurons specified by nb_classes.
+
+    """
+    input_shape = (24, )
+    X_input = Input(input_shape)
+
+    X = Dense(1024, activation='relu')(X_input)
+    for i in range(9):
+        X = Dense(1024, activation='relu')(X)
+
+    X_output = Concatenate()([Dense(n)(X) for n in nb_classes])
+
+    model = Model(inputs=X_input, outputs=X_output)
+    model.compile()
+
+    return model

sarwaveifrproc/utils.py

@@ -10,6 +10,7 @@
 import os
 from datetime import datetime
 from sarwaveifrproc.l2_wave import generate_l2_wave_product
+from sarwaveifrproc.models import get_model
 
 def get_safe_date(safe):
     """
@@ -96,13 +97,18 @@ def get_output_filename(l1x_path, output_safe, tail='e00'):
     savepath = os.path.join(output_safe, final_filename)
     return savepath
 
-def load_config():
+def load_config(path=None):
     """
+    Parameters:
+        path (str): Optional. Path to a given config.
 
     Returns:
         conf: dict
     """
-    local_config_path = os.path.join(os.path.dirname(sarwaveifrproc.__file__), 'localconfig.yaml')
+    if path:
+        local_config_path = path
+    else:
+        local_config_path = os.path.join(os.path.dirname(sarwaveifrproc.__file__), 'localconfig.yaml')
 
     if os.path.exists(local_config_path):
         config_path = local_config_path
@@ -115,15 +121,16 @@ def load_config():
     return conf
 
 
-def load_models(paths, predicted_variables):
+def load_models_and_scalers(version, paths, predicted_variables):
     """
     Loads models, scalers, and bins necessary for prediction.
 
     Parameters:
+        version (str): Model version to use.
         paths (dict): Dictionary containing paths to model files, scaler files, and bin files.
             Keys:
-                - 'model_intraburst': Path to the intraburst model file.
-                - 'model_interburst': Path to the interburst model file.
+                - 'model_intraburst': Path to the intraburst model directory containing both model architecture and weights.
+                - 'model_interburst': Path to the interburst model directory containing both model architecture and weights.
                 - 'scaler_intraburst': Path to the intraburst scaler file.
                 - 'scaler_interburst': Path to the interburst scaler file.
                 - 'bins_intraburst': Path to the intraburst bins directory.
@@ -137,25 +144,49 @@ def load_models(paths, predicted_variables):
             - scaler_interburst (RobustScaler): Interburst scaler loaded from the provided path.
             - bins_intraburst (dict): Dictionary containing intraburst bins for each predicted variable.
             - bins_interburst (dict): Dictionary containing interburst bins for each predicted variable.
-
     """    
     # Unpack paths
     path_model_intraburst, path_model_interburst, path_scaler_intraburst, path_scaler_interburst, path_bins_intraburst, path_bins_interburst = paths.values()
 
-    # Load models and scalers using paths provided
-    model_intraburst = tf.keras.models.load_model(path_model_intraburst, compile=False)
+    # Load scalers using paths provided
     scaler_intraburst_array = np.load(path_scaler_intraburst)
     scaler_intraburst = RobustScaler(medians=scaler_intraburst_array[0], iqrs=scaler_intraburst_array[1])    
 
-    model_interburst = tf.keras.models.load_model(path_model_interburst, compile=False)
     scaler_interburst_array = np.load(path_scaler_interburst)
     scaler_interburst = RobustScaler(medians=scaler_interburst_array[0], iqrs=scaler_interburst_array[1]) 
 
+    # Load bins
     bins_intraburst = {f'{var}': np.load(os.path.join(path_bins_intraburst, f'bins_{var}.npy')) for var in predicted_variables}
     bins_interburst = {f'{var}': np.load(os.path.join(path_bins_interburst, f'bins_{var}.npy')) for var in predicted_variables}
+
+    # Get number of classes to load models
+    nb_classes_intraburst = [len(bins_intraburst[var])-1 for var in predicted_variables]
+    nb_classes_interburst = [len(bins_interburst[var])-1 for var in predicted_variables]
+
+    # Load models
+    model_intraburst = load_model(version, path_model_intraburst, nb_classes_intraburst)
+    model_interburst = load_model(version, path_model_interburst, nb_classes_interburst)
 
     return model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst
 
+
+def load_model(version, path, nb_classes):
+    """
+    This function builds the model architecture and reads the model weights from an HDF5 file, and reconstructs the Keras model.
+
+    Parameters
+        version (str): Neural model version to use.
+        path (str): The path where the model weights are stored.
+        nb_classes (list of int): The number of output classes of the neural model. 
+
+    Returns
+    (tf.keras.Model) The reconstructed Keras model with the architecture and weights loaded.
+    """
+    model = get_model(version)(nb_classes)
+    model.load_weights(path)
+
+    return model
+
 
 def process_files(input_safe, output_safe, model_intraburst, model_interburst, scaler_intraburst, scaler_interburst, bins_intraburst, bins_interburst, predicted_variables, product_id):
     """