From 5b25641dd6b3a47b108ca767827cad9cce0f060e Mon Sep 17 00:00:00 2001 From: Jeff Stout Date: Wed, 22 Jul 2026 16:29:49 -0400 Subject: [PATCH 1/3] Updating weights / minor version Setting check on the hugging face config.json file --- MEGnet/megnet_init.py | 65 +++++++++++++++++++++++--------- MEGnet/tests/test_megnet_init.py | 35 +++++++++++++++++ pyproject.toml | 2 +- 3 files changed, 84 insertions(+), 18 deletions(-) create mode 100644 MEGnet/tests/test_megnet_init.py diff --git a/MEGnet/megnet_init.py b/MEGnet/megnet_init.py index 6c21d34..35e8710 100644 --- a/MEGnet/megnet_init.py +++ b/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.1' + + +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') diff --git a/MEGnet/tests/test_megnet_init.py b/MEGnet/tests/test_megnet_init.py new file mode 100644 index 0000000..76bf0f3 --- /dev/null +++ b/MEGnet/tests/test_megnet_init.py @@ -0,0 +1,35 @@ +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_v2_1(monkeypatch, tmp_path): + weights_path = _set_config_path(monkeypatch, tmp_path) + weights_path.mkdir() + 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_v2_1(monkeypatch, tmp_path): + weights_path = _set_config_path(monkeypatch, tmp_path) + weights_path.mkdir() + 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 diff --git a/pyproject.toml b/pyproject.toml index 411f712..1f2ba63 100644 --- a/pyproject.toml +++ b/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'] From d0580dfb62d4a0bc4cb9b40bf91f6494536c67ef Mon Sep 17 00:00:00 2001 From: Jeff Stout Date: Wed, 22 Jul 2026 17:52:10 -0400 Subject: [PATCH 2/3] Updating model to new weight inputs Check that the model is of v2.2 or redownload from HF --- MEGnet/megnet_init.py | 2 +- MEGnet/prep_inputs/ICA.py | 2 +- MEGnet/prep_inputs/convert_keras_model.sh | 172 +++++++++++++++++++++ MEGnet/prep_inputs/tests/test_ica2input.py | 2 +- 4 files changed, 175 insertions(+), 3 deletions(-) create mode 100755 MEGnet/prep_inputs/convert_keras_model.sh diff --git a/MEGnet/megnet_init.py b/MEGnet/megnet_init.py index 35e8710..7ec0d54 100644 --- a/MEGnet/megnet_init.py +++ b/MEGnet/megnet_init.py @@ -15,7 +15,7 @@ _megnet_path = MEGnet.__path__[0] weights_path = op.join(_megnet_path, 'model_v2k3') config_path = op.join(weights_path, 'config.json') -min_model_version = 'v2.1' +min_model_version = 'v2.2' def _version_tuple(version): diff --git a/MEGnet/prep_inputs/ICA.py b/MEGnet/prep_inputs/ICA.py index 43f4d5e..a37ca3d 100755 --- a/MEGnet/prep_inputs/ICA.py +++ b/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: diff --git a/MEGnet/prep_inputs/convert_keras_model.sh b/MEGnet/prep_inputs/convert_keras_model.sh new file mode 100755 index 0000000..9a9462b --- /dev/null +++ b/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 [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')" diff --git a/MEGnet/prep_inputs/tests/test_ica2input.py b/MEGnet/prep_inputs/tests/test_ica2input.py index d757a00..4ea6768 100644 --- a/MEGnet/prep_inputs/tests/test_ica2input.py +++ b/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 # ============================================================================= From ca69a37cdd090aa3e45c65a4a709705256a845ad Mon Sep 17 00:00:00 2001 From: Jeff Stout Date: Wed, 22 Jul 2026 18:15:55 -0400 Subject: [PATCH 3/3] Updated for current weights --- MEGnet/prep_inputs/tests/test_ica2input.py | 4 ++-- MEGnet/tests/test_megnet_init.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/MEGnet/prep_inputs/tests/test_ica2input.py b/MEGnet/prep_inputs/tests/test_ica2input.py index 4ea6768..66fae04 100644 --- a/MEGnet/prep_inputs/tests/test_ica2input.py +++ b/MEGnet/prep_inputs/tests/test_ica2input.py @@ -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')) diff --git a/MEGnet/tests/test_megnet_init.py b/MEGnet/tests/test_megnet_init.py index 76bf0f3..88f2019 100644 --- a/MEGnet/tests/test_megnet_init.py +++ b/MEGnet/tests/test_megnet_init.py @@ -17,18 +17,20 @@ def test_check_weights_requires_config_json(monkeypatch, tmp_path): assert megnet_init._check_weights() is False -def test_check_weights_rejects_model_version_not_greater_than_v2_1(monkeypatch, tmp_path): +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_v2_1(monkeypatch, tmp_path): +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)