Update weights w version check

MEGnet/megnet_init.py

@@ -7,18 +7,57 @@
 """
 
 
+import json
+import re
 import MEGnet
 import os, os.path as op
 
 _megnet_path = MEGnet.__path__[0]
-weights_path = op.join(_megnet_path, 'model_v2_k3')
+weights_path = op.join(_megnet_path, 'model_v2k3')
+config_path = op.join(weights_path, 'config.json')
+min_model_version = 'v2.2'
+
+
+def _version_tuple(version):
+    version_match = re.match(r'^v?(\d+(?:\.\d+)*)$', version)
+    if not version_match:
+        return None
+    return tuple(int(i) for i in version_match.group(1).split('.'))
+
 
 def _check_weights():
-    if op.exists(weights_path):
-        return True
-    else:
+    if not op.exists(config_path):
+        return False
+
+    try:
+        with open(config_path, encoding='utf-8') as fid:
+            config = json.load(fid)
+    except (OSError, json.JSONDecodeError):
+        return False
+
+    model_version = config.get('model_version')
+    if not isinstance(model_version, str):
         return False
-
+
+    model_version_tuple = _version_tuple(model_version)
+    min_model_version_tuple = _version_tuple(min_model_version)
+    if model_version_tuple is None or min_model_version_tuple is None:
+        return False
+
+    return model_version_tuple > min_model_version_tuple
+
+
+def _download_weights():
+    from huggingface_hub import snapshot_download
+
+    snapshot_download(
+        repo_id='jstout211/MEGnetV2',
+        local_dir=_megnet_path,
+        local_dir_use_symlinks=False,
+        allow_patterns=["model_v2k3/*"],
+        force_download=True
+    )
+
 
 def main():
     """
@@ -28,20 +67,12 @@ def main():
     if _check_weights():
         print('Model weights present - check successful')
     else:
-        print(f'''Model weights were not found in:  
+        print(f'''Model weights are missing or out of date in:
               {weights_path}
-              Performing download from huggingface repository''')
-
-        # Download the data
-        from huggingface_hub import snapshot_download
-
+              Pulling newest weights from huggingface repository''')
+
         try:
-            snapshot_download(
-                repo_id='jstout211/MEGnetV2',
-                local_dir= _megnet_path,
-                local_dir_use_symlinks=False,  
-                allow_patterns=["model_v2k3/*"]
-            )
+            _download_weights()
         except BaseException as e:
             print('Could not download the weights for classification')
             print('This is likely an issue with network access to the huggingface repository')

MEGnet/prep_inputs/ICA.py

@@ -814,7 +814,7 @@ def classify_ica(results_dir=None, outbasename=None, filename=None):
     import keras
     model_path = op.join(MEGnet.__path__[0] ,  'model_v2k3/model_v2.keras') 
     # This is set to use CPU in initial import
-    kModel=keras.models.load_model(model_path)
+    kModel=keras.models.load_model(model_path, compile=False)
 
     #Set output names
     if outbasename != None:

MEGnet/prep_inputs/convert_keras_model.sh

@@ -0,0 +1,172 @@
+#!/usr/bin/env bash
+#
+# convert_keras_model.sh
+#
+# Converts a legacy Keras/TensorFlow SavedModel directory
+# (containing assets/, keras_metadata.pb, saved_model.pb, variables/)
+# into a Keras 3 native ".keras" model file.
+#
+# Usage:
+#   ./convert_keras_model.sh /path/to/model_folder [output_name.keras]
+#
+# Requirements:
+#   - A Python environment where the model can be loaded, typically
+#     tf-keras (legacy Keras 2 shim) since raw SavedModel loading isn't
+#     supported directly by keras>=3. tf-keras is used specifically to
+#     load the old format before re-saving in the new one.
+#
+#     pip install tensorflow tf-keras keras
+#
+set -euo pipefail
+
+# ---- Argument parsing -------------------------------------------------
+if [ $# -lt 1 ]; then
+    echo "Usage: $0 <path_to_saved_model_dir> [output_name.keras]" >&2
+    exit 1
+fi
+
+MODEL_DIR="$1"
+OUTPUT_NAME="${2:-converted_model.keras}"
+# Derive the .h5 sibling name from OUTPUT_NAME (same basename, .h5 extension)
+OUTPUT_BASENAME="${OUTPUT_NAME%.keras}"
+OUTPUT_H5_NAME="${OUTPUT_BASENAME}.h5"
+
+# ---- Sanity checks ------------------------------------------------------
+if [ ! -d "$MODEL_DIR" ]; then
+    echo "Error: '$MODEL_DIR' is not a directory." >&2
+    exit 1
+fi
+
+if [ ! -f "$MODEL_DIR/saved_model.pb" ]; then
+    echo "Error: '$MODEL_DIR' does not look like a TensorFlow SavedModel" \
+         "(missing saved_model.pb)." >&2
+    exit 1
+fi
+
+if [ ! -d "$MODEL_DIR/variables" ]; then
+    echo "Error: '$MODEL_DIR' does not look like a TensorFlow SavedModel" \
+         "(missing variables/ directory)." >&2
+    exit 1
+fi
+
+# Resolve to an absolute path so the Python snippet is unambiguous.
+MODEL_DIR_ABS="$(cd "$MODEL_DIR" && pwd)"
+OUTPUT_PATH="$(pwd)/$OUTPUT_NAME"
+OUTPUT_H5_PATH="$(pwd)/$OUTPUT_H5_NAME"
+
+echo "Source SavedModel dir : $MODEL_DIR_ABS"
+echo "Output Keras 3 file   : $OUTPUT_PATH"
+echo "Output H5 file        : $OUTPUT_H5_PATH"
+echo
+
+# ---- Check for required Python packages --------------------------------
+python3 - <<'PYCHECK'
+import importlib
+import importlib.util
+import sys
+
+missing = []
+for pkg in ("tensorflow", "tf_keras", "keras"):
+    if importlib.util.find_spec(pkg) is None:
+        missing.append(pkg)
+
+if missing:
+    print(f"Missing required Python packages: {', '.join(missing)}", file=sys.stderr)
+    print("Either activate an environment that already has these installed,", file=sys.stderr)
+    print("or install them with:", file=sys.stderr)
+    print("    pip install tf-keras keras", file=sys.stderr)
+    sys.exit(1)
+PYCHECK
+
+# ---- Run the actual conversion ------------------------------------------
+python3 - "$MODEL_DIR_ABS" "$OUTPUT_PATH" "$OUTPUT_H5_PATH" <<'PYCONVERT'
+import sys
+import os
+
+model_dir = sys.argv[1]
+output_path = sys.argv[2]
+output_h5_path = sys.argv[3]
+
+# tf_keras provides the legacy Keras 2 loader capable of reading
+# TF SavedModel-format models (with keras_metadata.pb).
+import tf_keras as legacy_keras
+
+# Skip compiling the model on load. Compilation requires reconstructing
+# the optimizer/loss/metrics (e.g. custom TF Addons objects), which we
+# don't need just to migrate the architecture + weights to Keras 3.
+print(f"Loading legacy SavedModel from: {model_dir}")
+legacy_model = legacy_keras.models.load_model(model_dir, compile=False)
+print("Model loaded successfully (compile=False).")
+
+# The object above is a tf_keras.Functional/Sequential instance. Saving it
+# directly with .save() embeds 'module': 'tf_keras.src.engine...' in the
+# config, which Keras 3 cannot deserialize (it only knows 'keras.*' /
+# 'keras.src.*' module paths). To produce a genuinely native Keras 3
+# model, we rebuild the architecture using the real `keras` package from
+# the legacy model's config, patching module paths, then transfer weights.
+import keras
+import json
+
+print(f"\nRebuilding architecture as native Keras 3 model "
+      f"(keras version: {keras.__version__})...")
+
+legacy_config = legacy_model.get_config()
+
+def _fix_modules(obj):
+    """Recursively rewrite tf_keras module/class refs to native keras ones
+    so keras 3's deserializer can resolve every layer/object in the config.
+    Also fixes known config-shape differences between tf_keras and keras 3
+    (e.g. BatchNormalization's `axis` stored as a list instead of an int)."""
+    if isinstance(obj, dict):
+        if obj.get("module", "").startswith("tf_keras"):
+            obj["module"] = "keras.layers" if "layers" in obj.get("module", "") else "keras"
+        if obj.get("class_name") == "Functional":
+            obj["module"] = "keras"
+            obj["registered_name"] = None
+        if obj.get("class_name") == "BatchNormalization":
+            axis = obj.get("config", {}).get("axis")
+            if isinstance(axis, list) and len(axis) == 1:
+                obj["config"]["axis"] = axis[0]
+        for v in obj.values():
+            _fix_modules(v)
+    elif isinstance(obj, list):
+        for item in obj:
+            _fix_modules(item)
+    return obj
+
+legacy_config = _fix_modules(legacy_config)
+
+# Reconstruct using native Keras 3's Functional/Sequential deserializer.
+if legacy_model.__class__.__name__ == "Sequential":
+    model = keras.Sequential.from_config(legacy_config)
+else:
+    model = keras.Model.from_config(legacy_config)
+
+# Transfer weights by name to be robust to any minor ordering differences.
+model.set_weights(legacy_model.get_weights())
+print("Weights transferred to native Keras 3 model.")
+model.summary()
+
+if not output_path.endswith(".keras"):
+    output_path += ".keras"
+
+model.save(output_path)
+print(f"\nConversion complete. Saved Keras 3 model to: {output_path}")
+print("Note: model was loaded with compile=False, so it has no optimizer/")
+print("loss/metrics attached. Re-compile with model.compile(...) before training.")
+
+# Also save a legacy HDF5 (.h5) copy. Keras 3 still supports writing the
+# H5 format via the same save() call when given an .h5 extension.
+if not output_h5_path.endswith(".h5"):
+    output_h5_path += ".h5"
+
+model.save(output_h5_path)
+print(f"Also saved legacy H5 model to: {output_h5_path}")
+PYCONVERT
+
+echo
+echo "Done. You can now load the model in Keras 3 with:"
+echo "    import keras"
+echo "    model = keras.models.load_model('$OUTPUT_NAME')"
+echo "or the legacy H5 file with:"
+echo "    model = keras.models.load_model('$OUTPUT_H5_NAME')"

MEGnet/prep_inputs/tests/test_ica2input.py

@@ -31,7 +31,7 @@
 import keras
 from MEGnet.megnet_utilities import fPredictChunkAndVoting_parrallel
 model_path = op.join(MEGnet.__path__[0] ,  'model_v2k3/model_v2.keras')    # << May want to change this to function
-kModel=keras.models.load_model(model_path)
+kModel=keras.models.load_model(model_path, compile=False)
 
 from numpy.testing import assert_almost_equal
 # =============================================================================
@@ -134,8 +134,8 @@ def test_classify_ica():
     savemat(ts_fname, {'arrICATimeSeries':ica_ts})
     # Classify the data vectors
     ica_dict = classify_ica(results_dir=results_dir, filename=ctf_filename)
-    assert np.all(ica_dict['classes']==[1, 0, 0, 0, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
-    assert np.all(ica_dict['bads_idx']==[0,4,5])    
+    assert np.all(ica_dict['classes']==[0, 0, 0, 0, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    assert np.all(ica_dict['bads_idx']==[4,5])    
 
 def get_inputs(dirname):
     classID = np.load(op.join(dirname, 'cl.npy'))

MEGnet/tests/test_megnet_init.py

@@ -0,0 +1,37 @@
+import json
+import os.path as op
+
+from MEGnet import megnet_init
+
+
+def _set_config_path(monkeypatch, tmp_path):
+    weights_path = tmp_path / 'model_v2k3'
+    monkeypatch.setattr(megnet_init, 'weights_path', str(weights_path))
+    monkeypatch.setattr(megnet_init, 'config_path', str(weights_path / 'config.json'))
+    return weights_path
+
+
+def test_check_weights_requires_config_json(monkeypatch, tmp_path):
+    _set_config_path(monkeypatch, tmp_path).mkdir()
+
+    assert megnet_init._check_weights() is False
+
+
+def test_check_weights_rejects_model_version_not_greater_than_min_version(monkeypatch, tmp_path):
+    weights_path = _set_config_path(monkeypatch, tmp_path)
+    weights_path.mkdir()
+    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1')
+    with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid:
+        json.dump({'model_version': 'v2.1'}, fid)
+
+    assert megnet_init._check_weights() is False
+
+
+def test_check_weights_accepts_model_version_greater_than_min_version(monkeypatch, tmp_path):
+    weights_path = _set_config_path(monkeypatch, tmp_path)
+    weights_path.mkdir()
+    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1')
+    with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid:
+        json.dump({'model_version': 'v2.2'}, fid)
+
+    assert megnet_init._check_weights() is True

pyproject.toml

@@ -22,7 +22,7 @@ classifiers=[
 dependencies = [
     'mne>1.10', 'numpy', 'scipy', 'pandas', 'munch', 'nibabel', 'joblib',  'torch', 'keras>3.0', 'scikit-learn', 'huggingface_hub>=0.24.0'
 ]
-version = "0.3.3"
+version = "0.3.4"
 
 [project.optional-dependencies]
 dev = ['stabilized-ica']