diff --git a/MEGnet/megnet_init.py b/MEGnet/megnet_init.py index 7ec0d54..844a321 100644 --- a/MEGnet/megnet_init.py +++ b/MEGnet/megnet_init.py @@ -14,6 +14,7 @@ _megnet_path = MEGnet.__path__[0] weights_path = op.join(_megnet_path, 'model_v2k3') +model_path = op.join(weights_path, 'model_v2.keras') config_path = op.join(weights_path, 'config.json') min_model_version = 'v2.2' @@ -26,7 +27,7 @@ def _version_tuple(version): def _check_weights(): - if not op.exists(config_path): + if not op.isfile(model_path) or not op.isfile(config_path): return False try: @@ -44,7 +45,7 @@ def _check_weights(): if model_version_tuple is None or min_model_version_tuple is None: return False - return model_version_tuple > min_model_version_tuple + return model_version_tuple >= min_model_version_tuple def _download_weights(): diff --git a/MEGnet/prep_inputs/ICA.py b/MEGnet/prep_inputs/ICA.py index a37ca3d..c591205 100755 --- a/MEGnet/prep_inputs/ICA.py +++ b/MEGnet/prep_inputs/ICA.py @@ -34,7 +34,7 @@ from mne.viz.utils import _setup_vmin_vmax, _get_cmap, plt_show from scipy.io import savemat import PIL.Image -import MEGnet +from MEGnet import megnet_init from MEGnet.megnet_utilities import fPredictChunkAndVoting_parrallel import functools @@ -49,6 +49,19 @@ # Helper Functions # ============================================================================= +def _require_model_weights(): + """Return the model path after checking its presence and version.""" + if megnet_init._check_weights(): + return megnet_init.model_path + + raise RuntimeError( + 'MEGnet model weights are missing or incompatible. Expected ' + f'{megnet_init.model_path} and a config.json model_version at least ' + f'{megnet_init.min_model_version}. Run `megnet_init` to download ' + 'compatible model weights.' + ) + + # function to transform Cartesian coordinates to spherical coordinates # theta = azimuth # phi = elevation @@ -656,9 +669,9 @@ def circle_plot(circle_pos=None, data=None, out_fname=None): #return matrix_out -def main(filename, outbasename=None, mains_freq=60.0, +def main(filename, results_dir, outbasename=None, mains_freq=60.0, save_preproc=False, save_ica=False, seedval=0, - results_dir=None, filename_raw=None, do_assess_bads=False, + filename_raw=None, do_assess_bads=False, bad_channels=[]): ''' Perform all of the steps to preprocess the ica maps: @@ -674,6 +687,8 @@ def main(filename, outbasename=None, mains_freq=60.0, filename : str or Raw MNE data object Path to file + results_dir : str / path + Path to output directory filename_raw : str Required for MEGIN datasets Path to file @@ -688,12 +703,12 @@ def main(filename, outbasename=None, mains_freq=60.0, Save the ica output seedval : Int Set the numpy random seed - results_dir : str / path - Path to output directory do_assess_bads : Bool Assess bad channels if not already done ''' + _require_model_weights() + if (type(filename) == str) | (type(filename) == PosixPath): raw = read_raw(filename) elif type(filename) in raw_typelist: @@ -806,13 +821,14 @@ def classify_ica(results_dir=None, outbasename=None, filename=None): ''' from scipy.io import loadmat + model_path = _require_model_weights() + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 os.environ["CUDA_VISIBLE_DEVICES"] = "" os.environ["KERAS_BACKEND"] = "torch" 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, compile=False) @@ -830,6 +846,10 @@ def classify_ica(results_dir=None, outbasename=None, filename=None): arrSP = np.stack([loadmat(i)['array'] for i in arrSP_fnames]) preds, probs = fPredictChunkAndVoting_parrallel(kModel, arrTS, arrSP) meg_rest_ica_classes = preds.argmax(axis=1) + np.save( + op.join(results_dir, 'megnet_classification.npy'), + meg_rest_ica_classes, + ) ica_comps_toremove = [index for index, value in enumerate(meg_rest_ica_classes) if value in [1, 2, 3]] return {'classes':meg_rest_ica_classes, 'bads_idx': ica_comps_toremove} @@ -899,7 +919,8 @@ def cmdline(): standard_args = parser.add_argument_group('standard') standard_args.add_argument('-filename', help='Path to MEG dataset') - standard_args.add_argument('-results_dir', help='Path to save the results') + standard_args.add_argument( + '-results_dir', required=True, help='Path to save the results') standard_args.add_argument('-line_freq', help='{60,50} Hz - AC electric frequency') @@ -942,4 +963,3 @@ def cmdline(): - diff --git a/MEGnet/tests/test_ica_model_weights.py b/MEGnet/tests/test_ica_model_weights.py new file mode 100644 index 0000000..e2f2c2e --- /dev/null +++ b/MEGnet/tests/test_ica_model_weights.py @@ -0,0 +1,89 @@ +import pytest +import numpy as np + +from MEGnet.prep_inputs import ICA + + +def test_main_requires_results_dir(): + with pytest.raises(TypeError, match='results_dir'): + ICA.main('input.fif') + + +def test_cmdline_requires_results_dir(monkeypatch): + monkeypatch.setattr( + 'sys.argv', + ['ICA.py', '-filename', 'input.fif', '-line_freq', '60'], + ) + + with pytest.raises(SystemExit) as error: + ICA.cmdline() + + assert error.value.code == 2 + + +def test_require_model_weights_returns_compatible_model(monkeypatch, tmp_path): + model_path = tmp_path / 'model_v2k3' / 'model_v2.keras' + + monkeypatch.setattr(ICA.megnet_init, 'model_path', str(model_path)) + monkeypatch.setattr(ICA.megnet_init, '_check_weights', lambda: True) + + assert ICA._require_model_weights() == str(model_path) + + +def test_require_model_weights_rejects_invalid_install(monkeypatch, tmp_path): + model_path = tmp_path / 'model_v2k3' / 'model_v2.keras' + monkeypatch.setattr(ICA.megnet_init, 'model_path', str(model_path)) + monkeypatch.setattr(ICA.megnet_init, '_check_weights', lambda: False) + + with pytest.raises(RuntimeError, match='Run `megnet_init`'): + ICA._require_model_weights() + + +def test_classify_ica_saves_classification_vector(monkeypatch, tmp_path): + results_root = tmp_path / 'results' + output_dir = results_root / 'sample' + output_dir.mkdir(parents=True) + (output_dir / 'ICATimeSeries.mat').touch() + for index in range(1, 21): + (output_dir / f'component{index}.mat').touch() + + class FakeKerasModels: + @staticmethod + def load_model(model_path, compile=False): + return object() + + class FakeKeras: + models = FakeKerasModels + + monkeypatch.setattr(ICA, '_require_model_weights', lambda: 'model_v2.keras') + monkeypatch.setitem(__import__('sys').modules, 'keras', FakeKeras) + monkeypatch.setattr( + ICA, + 'fPredictChunkAndVoting_parrallel', + lambda model, arrTS, arrSP: ( + np.array([ + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + ]), + None, + ), + ) + + def fake_loadmat(path): + if str(path).endswith('ICATimeSeries.mat'): + return {'arrICATimeSeries': np.ones((10, 3))} + return {'array': np.ones((2, 2, 3))} + + monkeypatch.setattr('scipy.io.loadmat', fake_loadmat) + + result = ICA.classify_ica( + results_dir=str(results_root), + outbasename='sample', + filename='ignored.fif', + ) + + classification_path = output_dir / 'megnet_classification.npy' + assert classification_path.is_file() + np.testing.assert_array_equal(np.load(classification_path), np.array([0, 1, 2])) + np.testing.assert_array_equal(result['classes'], np.array([0, 1, 2])) diff --git a/MEGnet/tests/test_megnet_init.py b/MEGnet/tests/test_megnet_init.py index 88f2019..250c566 100644 --- a/MEGnet/tests/test_megnet_init.py +++ b/MEGnet/tests/test_megnet_init.py @@ -1,37 +1,55 @@ import json import os.path as op +import pytest + from MEGnet import megnet_init -def _set_config_path(monkeypatch, tmp_path): +def _set_model_paths(monkeypatch, tmp_path): weights_path = tmp_path / 'model_v2k3' monkeypatch.setattr(megnet_init, 'weights_path', str(weights_path)) + monkeypatch.setattr(megnet_init, 'model_path', str(weights_path / 'model_v2.keras')) 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() + weights_path = _set_model_paths(monkeypatch, tmp_path) + weights_path.mkdir() + (weights_path / 'model_v2.keras').touch() + + assert megnet_init._check_weights() is False + + +def test_check_weights_requires_model_file(monkeypatch, tmp_path): + weights_path = _set_model_paths(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 False -def test_check_weights_rejects_model_version_not_greater_than_min_version(monkeypatch, tmp_path): - weights_path = _set_config_path(monkeypatch, tmp_path) +def test_check_weights_rejects_model_version_below_min_version(monkeypatch, tmp_path): + weights_path = _set_model_paths(monkeypatch, tmp_path) weights_path.mkdir() - monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1') + (weights_path / 'model_v2.keras').touch() + monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.2') 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) +@pytest.mark.parametrize('model_version', ['v2.2', 'v2.3']) +def test_check_weights_accepts_compatible_model_version( + monkeypatch, tmp_path, model_version): + weights_path = _set_model_paths(monkeypatch, tmp_path) weights_path.mkdir() - monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1') + (weights_path / 'model_v2.keras').touch() + monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.2') with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid: - json.dump({'model_version': 'v2.2'}, fid) + json.dump({'model_version': model_version}, fid) assert megnet_init._check_weights() is True