Make Keras 2 available for tf 16.1

caiman/__init__.py

@@ -1,5 +1,17 @@
 #!/usr/bin/env python
 
+# Find keras, depending on tensorflow version
+import os
+try:
+    import tensorflow.keras as keras
+except ModuleNotFoundError:
+    try:
+        # workaround to continue using Keras 2 with tensorflow >= 2.16
+        os.environ["TF_USE_LEGACY_KERAS"] = "1"
+        import tf_keras as keras
+    except ModuleNotFoundError:
+        keras = None
+
 import pkg_resources
 from caiman.base.movies import movie, load, load_movie_chain, _load_behavior, play_movie
 from caiman.base.timeseries import concatenate

caiman/components_evaluation.py

@@ -269,12 +269,12 @@ def evaluate_components_CNN(A,
     if not isGPU and 'CAIMAN_ALLOW_GPU' not in os.environ:
         print("GPU run not requested, disabling use of GPUs")
         os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-    try:
+    if caiman.keras is not None:
         os.environ["KERAS_BACKEND"] = "tensorflow"
-        from tensorflow.keras.models import model_from_json
+        model_from_json = caiman.keras.models.model_from_json
         use_keras = True
         logging.info('Using Keras')
-    except (ModuleNotFoundError):
+    else:
         use_keras = False
         logging.info('Using Tensorflow')
 

caiman/source_extraction/cnmf/online_cnmf.py

@@ -323,14 +323,9 @@ def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
             self.tf_in = None
             self.tf_out = None
         else:
-            try:
-                from tensorflow.keras.models import model_from_json
+            if caiman.keras is not None:
                 logging.info('Using Keras')
-                use_keras = True
-            except(ModuleNotFoundError):
-                use_keras = False
-                logging.info('Using Tensorflow')
-            if use_keras:
+                model_from_json = caiman.keras.models.model_from_json
                 path = self.params.get('online', 'path_to_model').split(".")[:-1]
                 json_path = ".".join(path + ["json"])
                 model_path = ".".join(path + ["h5"])
@@ -342,6 +337,7 @@ def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
                 self.tf_in = None
                 self.tf_out = None
             else:
+                logging.info('Using Tensorflow')
                 path = self.params.get('online', 'path_to_model').split(".")[:-1]
                 model_path = '.'.join(path + ['h5', 'pb'])
                 loaded_model = load_graph(model_path)

caiman/source_extraction/volpy/mrcnn/model.py

@@ -15,13 +15,14 @@
 import multiprocessing
 import numpy as np
 import tensorflow as tf
-import tensorflow.keras as keras
-import tensorflow.keras.backend as K
-import tensorflow.keras.layers as KL
-import tensorflow.keras.layers as KE
-import tensorflow.keras.utils as KU
+
+from caiman import keras
+K = keras.backend
+KE = KL = keras.layers
+KU = keras.utils
+KM = keras.models
+
 from tensorflow.python.eager import context
-import tensorflow.keras.models as KM
 
 from ..mrcnn import utils
 

caiman/tests/test_tensorflow.py

@@ -5,14 +5,15 @@
 
 from caiman.paths import caiman_datadir
 from caiman.utils.utils import load_graph
+from caiman import keras
 
 #FIXME A lot of what this did is no longer relevant
 
-try:
+if keras is not None:
     os.environ["KERAS_BACKEND"] = "tensorflow"
-    from tensorflow.keras.models import model_from_json
+    model_from_json = keras.models.model_from_json
     use_keras = True
-except(ModuleNotFoundError):
+else:
     import tensorflow as tf
     use_keras = False
 

caiman/utils/nn_models.py

@@ -8,13 +8,20 @@
 import numpy as np
 import os
 import tensorflow as tf
-from tensorflow.keras.layers import Input, Dense, Reshape, Layer, Activation
-from tensorflow.keras.models import Model
-from tensorflow.keras.optimizers import Adam
-from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, LearningRateScheduler
-from tensorflow.keras import backend as K
-from tensorflow.keras.initializers import Constant, RandomUniform
-from tensorflow.keras.utils import Sequence
+from caiman import keras
+K = keras.backend
+
+for submodule, names in {
+    'layers': ['Input', 'Dense', 'Reshape', 'Layer', 'Activation'],
+    'models': ['Model'],
+    'optimizers': ['Adam'],
+    'callbacks': ['ModelCheckpoint', 'EarlyStopping', 'LearningRateScheduler'],
+    'initializers': ['Constant', 'RandomUniform'],
+    'utils': ['Sequence']
+}.items():
+    for name in names:
+        globals()[name] = getattr(getattr(keras, submodule), name)
+
 import time
 
 import caiman.base.movies