From 8953ceca72b9f8bfac3f252c8abe533e5083c032 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:36:54 +0000
Subject: [PATCH 01/17] remove pytest-lazy-fixture as dev dependency and skip
 test (with WG temp fix)

---
 pyproject.toml                                |  1 -
 tests/core/test_integration/test_detection.py | 10 ++++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 92cad582..96006dfe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,7 +46,6 @@ dev = [
     "pre-commit",
     "pyinstrument",
     "pytest-cov",
-    "pytest-lazy-fixture",
     "pytest-mock",
     "pytest-qt",
     "pytest-timeout",
diff --git a/tests/core/test_integration/test_detection.py b/tests/core/test_integration/test_detection.py
index ed289d10..037b63ef 100644
--- a/tests/core/test_integration/test_detection.py
+++ b/tests/core/test_integration/test_detection.py
@@ -43,15 +43,17 @@ def background_array():
 
 
 # FIXME: This isn't a very good example
+@pytest.mark.skip()
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    "n_free_cpus",
+    "free_cpus",
     [
-        pytest.lazy_fixture("no_free_cpus"),
-        pytest.lazy_fixture("run_on_one_cpu_only"),
+        pytest.param("no_free_cpus", id="No free CPUs"),
+        pytest.param("run_on_one_cpu_only", id="One CPU"),
     ],
 )
-def test_detection_full(signal_array, background_array, n_free_cpus):
+def test_detection_full(signal_array, background_array, free_cpus, request):
+    n_free_cpus = request.getfixturevalue(free_cpus)
     cells_test = main(
         signal_array,
         background_array,

From a1acce09a984dedb89933422786000511fc9e129 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Mon, 5 Feb 2024 16:18:58 +0000
Subject: [PATCH 02/17] change tensorflow dependency for cellfinder

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 96006dfe..d8a5315c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,8 +30,8 @@ dependencies = [
     "scikit-image",
     "scikit-learn",
     # See https://github.com/brainglobe/cellfinder-core/issues/103 for < 2.12.0 pin
-    "tensorflow-macos>=2.5.0,<2.12.0; platform_system=='Darwin' and platform_machine=='arm64'",
-    "tensorflow>=2.5.0,<2.12.0; platform_system!='Darwin' or platform_machine!='arm64'",
+    "tensorflow-macos>=2.5.0; platform_system=='Darwin' and platform_machine=='arm64'",
+    "tensorflow>=2.5.0; platform_system!='Darwin' or platform_machine!='arm64'",
     "tifffile",
     "tqdm",
 ]

From 6846ed25116774499b6c3edee24ca3a1552d732b Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:42:35 +0000
Subject: [PATCH 03/17] replace keras imports from tensorflow to just keras
 imports

---
 cellfinder/core/classify/classify.py       |  2 +-
 cellfinder/core/classify/cube_generator.py |  6 +++---
 cellfinder/core/classify/resnet.py         |  8 ++++----
 cellfinder/core/classify/tools.py          | 22 +++++++++++-----------
 cellfinder/core/train/train_yml.py         |  2 +-
 5 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
index 86b709f5..e91d32b8 100644
--- a/cellfinder/core/classify/classify.py
+++ b/cellfinder/core/classify/classify.py
@@ -4,7 +4,7 @@
 import numpy as np
 from brainglobe_utils.cells.cells import Cell
 from brainglobe_utils.general.system import get_num_processes
-from tensorflow import keras
+import keras
 
 from cellfinder.core import logger, types
 from cellfinder.core.classify.cube_generator import CubeGeneratorFromFile
diff --git a/cellfinder/core/classify/cube_generator.py b/cellfinder/core/classify/cube_generator.py
index f4c67515..6520bbe3 100644
--- a/cellfinder/core/classify/cube_generator.py
+++ b/cellfinder/core/classify/cube_generator.py
@@ -8,7 +8,7 @@
 from brainglobe_utils.general.numerical import is_even
 from scipy.ndimage import zoom
 from skimage.io import imread
-from tensorflow.keras.utils import Sequence
+from keras.utils import Sequence
 
 from cellfinder.core import types
 from cellfinder.core.classify.augment import AugmentationParameters, augment
@@ -220,7 +220,7 @@ def __getitem__(
 
         if self.train:
             batch_labels = [cell.type - 1 for cell in cell_batch]
-            batch_labels = tf.keras.utils.to_categorical(
+            batch_labels = keras.utils.to_categorical(
                 batch_labels, num_classes=self.classes
             )
             return images, batch_labels
@@ -414,7 +414,7 @@ def __getitem__(
 
         if self.train and self.labels is not None:
             batch_labels = [self.labels[k] for k in indexes]
-            batch_labels = tf.keras.utils.to_categorical(
+            batch_labels = keras.utils.to_categorical(
                 batch_labels, num_classes=self.classes
             )
             return images, batch_labels
diff --git a/cellfinder/core/classify/resnet.py b/cellfinder/core/classify/resnet.py
index ed313642..85158bd0 100644
--- a/cellfinder/core/classify/resnet.py
+++ b/cellfinder/core/classify/resnet.py
@@ -1,9 +1,9 @@
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
 from tensorflow import Tensor
-from tensorflow.keras import Model
-from tensorflow.keras.initializers import Initializer
-from tensorflow.keras.layers import (
+from keras import Model
+from keras.initializers import Initializer
+from keras.layers import (
     Activation,
     Add,
     BatchNormalization,
@@ -14,7 +14,7 @@
     MaxPooling3D,
     ZeroPadding3D,
 )
-from tensorflow.keras.optimizers import Adam, Optimizer
+from keras.optimizers import Adam, Optimizer
 
 #####################################################################
 # Define the types of ResNet
diff --git a/cellfinder/core/classify/tools.py b/cellfinder/core/classify/tools.py
index 2d5c44b2..77f19072 100644
--- a/cellfinder/core/classify/tools.py
+++ b/cellfinder/core/classify/tools.py
@@ -1,9 +1,10 @@
 import os
-from typing import List, Optional, Sequence, Tuple, Union
+from collections.abc import Sequence
+from typing import List, Optional, Tuple, Union
 
+import keras
 import numpy as np
-import tensorflow as tf
-from tensorflow.keras import Model
+from keras import Model
 
 from cellfinder.core import logger
 from cellfinder.core.classify.resnet import build_model, layer_type
@@ -17,8 +18,7 @@ def get_model(
     inference: bool = False,
     continue_training: bool = False,
 ) -> Model:
-    """
-    Returns the correct model based on the arguments passed
+    """Returns the correct model based on the arguments passed
     :param existing_model: An existing, trained model. This is returned if it
     exists
     :param model_weights: This file is used to set the model weights if it
@@ -30,29 +30,29 @@ def get_model(
     by using the default one
     :param continue_training: If True, will ensure that a trained model
     exists. E.g. by using the default one
-    :return: A tf.keras model
+    :return: A keras model
 
     """
     if existing_model is not None or network_depth is None:
         logger.debug(f"Loading model: {existing_model}")
-        return tf.keras.models.load_model(existing_model)
+        return keras.models.load_model(existing_model)
     else:
         logger.debug(f"Creating a new instance of model: {network_depth}")
         model = build_model(
-            network_depth=network_depth, learning_rate=learning_rate
+            network_depth=network_depth, learning_rate=learning_rate,
         )
         if inference or continue_training:
             logger.debug(
-                f"Setting model weights according to: {model_weights}"
+                f"Setting model weights according to: {model_weights}",
             )
             if model_weights is None:
-                raise IOError("`model_weights` must be provided")
+                raise OSError("`model_weights` must be provided")
             model.load_weights(model_weights)
         return model
 
 
 def make_lists(
-    tiff_files: Sequence, train: bool = True
+    tiff_files: Sequence, train: bool = True,
 ) -> Union[Tuple[List, List], Tuple[List, List, np.ndarray]]:
     signal_list = []
     background_list = []
diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 63e97e00..68c843a6 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -324,7 +324,7 @@ def run(
 
     suppress_tf_logging(tf_suppress_log_messages)
 
-    from tensorflow.keras.callbacks import (
+    from keras.callbacks import (
         CSVLogger,
         ModelCheckpoint,
         TensorBoard,

From 2a34fda8fb852862b70a0b797f65a76918044db4 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:48:18 +0000
Subject: [PATCH 04/17] add keras import and reorder

---
 cellfinder/core/classify/classify.py       | 2 +-
 cellfinder/core/classify/cube_generator.py | 4 ++--
 cellfinder/core/classify/resnet.py         | 2 +-
 cellfinder/core/classify/tools.py          | 6 ++++--
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
index e91d32b8..93605176 100644
--- a/cellfinder/core/classify/classify.py
+++ b/cellfinder/core/classify/classify.py
@@ -1,10 +1,10 @@
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
+import keras
 import numpy as np
 from brainglobe_utils.cells.cells import Cell
 from brainglobe_utils.general.system import get_num_processes
-import keras
 
 from cellfinder.core import logger, types
 from cellfinder.core.classify.cube_generator import CubeGeneratorFromFile
diff --git a/cellfinder/core/classify/cube_generator.py b/cellfinder/core/classify/cube_generator.py
index 6520bbe3..fb936dbf 100644
--- a/cellfinder/core/classify/cube_generator.py
+++ b/cellfinder/core/classify/cube_generator.py
@@ -2,13 +2,13 @@
 from random import shuffle
 from typing import Dict, List, Optional, Tuple, Union
 
+import keras
 import numpy as np
-import tensorflow as tf
 from brainglobe_utils.cells.cells import Cell, group_cells_by_z
 from brainglobe_utils.general.numerical import is_even
+from keras.utils import Sequence
 from scipy.ndimage import zoom
 from skimage.io import imread
-from keras.utils import Sequence
 
 from cellfinder.core import types
 from cellfinder.core.classify.augment import AugmentationParameters, augment
diff --git a/cellfinder/core/classify/resnet.py b/cellfinder/core/classify/resnet.py
index 85158bd0..de6e720f 100644
--- a/cellfinder/core/classify/resnet.py
+++ b/cellfinder/core/classify/resnet.py
@@ -1,6 +1,5 @@
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
-from tensorflow import Tensor
 from keras import Model
 from keras.initializers import Initializer
 from keras.layers import (
@@ -15,6 +14,7 @@
     ZeroPadding3D,
 )
 from keras.optimizers import Adam, Optimizer
+from tensorflow import Tensor
 
 #####################################################################
 # Define the types of ResNet
diff --git a/cellfinder/core/classify/tools.py b/cellfinder/core/classify/tools.py
index 77f19072..3bf48876 100644
--- a/cellfinder/core/classify/tools.py
+++ b/cellfinder/core/classify/tools.py
@@ -39,7 +39,8 @@ def get_model(
     else:
         logger.debug(f"Creating a new instance of model: {network_depth}")
         model = build_model(
-            network_depth=network_depth, learning_rate=learning_rate,
+            network_depth=network_depth,
+            learning_rate=learning_rate,
         )
         if inference or continue_training:
             logger.debug(
@@ -52,7 +53,8 @@ def get_model(
 
 
 def make_lists(
-    tiff_files: Sequence, train: bool = True,
+    tiff_files: Sequence,
+    train: bool = True,
 ) -> Union[Tuple[List, List], Tuple[List, List, np.ndarray]]:
     signal_list = []
     background_list = []

From 4a0e886b0ad24253c9388ee4a6e7871fd55ecc8b Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 11:32:22 +0000
Subject: [PATCH 05/17] add keras and TF 2.16 to pyproject.toml

---
 pyproject.toml | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index d8a5315c..39941d6d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,9 +29,11 @@ dependencies = [
     "numpy",
     "scikit-image",
     "scikit-learn",
+    "keras", #
+    "tf-nightly==2.16.0.dev20240101", # pinning to same TF as Keras 3.0
     # See https://github.com/brainglobe/cellfinder-core/issues/103 for < 2.12.0 pin
-    "tensorflow-macos>=2.5.0; platform_system=='Darwin' and platform_machine=='arm64'",
-    "tensorflow>=2.5.0; platform_system!='Darwin' or platform_machine!='arm64'",
+    # "tensorflow-macos>=2.5.0; platform_system=='Darwin' and platform_machine=='arm64'",
+    # "tensorflow>=2.5.0; platform_system!='Darwin' or platform_machine!='arm64'",
     "tifffile",
     "tqdm",
 ]
@@ -124,7 +126,6 @@ commands = python -m pytest -v --color=yes
 deps =
     pytest
     pytest-cov
-    pytest-lazy-fixture
     pytest-mock
     pytest-timeout
     # Even though napari is a requirement for cellfinder.napari, we have to
@@ -141,4 +142,5 @@ passenv =
     XAUTHORITY
     NUMPY_EXPERIMENTAL_ARRAY_FUNCTION
     PYVISTA_OFF_SCREEN
+    # export KERAS_BACKEND="tensorflow" ?
 """

From 5e9e41e96687c0e8749ec6e1f54292597a8897bd Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 11:32:58 +0000
Subject: [PATCH 06/17] comment out TF version check for now

---
 cellfinder/__init__.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/cellfinder/__init__.py b/cellfinder/__init__.py
index 9971f648..30fb135c 100644
--- a/cellfinder/__init__.py
+++ b/cellfinder/__init__.py
@@ -7,18 +7,18 @@
 
 # If tensorflow is not present, tools cannot be used.
 # Throw an error in this case to prevent invocation of functions.
-try:
-    TF_VERSION = version("tensorflow")
-except PackageNotFoundError as e:
-    try:
-        TF_VERSION = version("tensorflow-macos")
-    except PackageNotFoundError as e:
-        raise PackageNotFoundError(
-            f"cellfinder tools cannot be invoked without tensorflow. "
-            f"Please install tensorflow into your environment to use cellfinder tools. "
-            f"For more information, please see "
-            f"https://github.com/brainglobe/brainglobe-meta#readme."
-        ) from e
+# try:
+#     TF_VERSION = version("tensorflow")  # replace by keras?
+# except PackageNotFoundError as e:
+#     try:
+#         TF_VERSION = version("tensorflow-macos")
+#     except PackageNotFoundError as e:
+#         raise PackageNotFoundError(
+#             f"cellfinder tools cannot be invoked without tensorflow. "
+#             f"Please install tensorflow into your environment to use cellfinder tools. "
+#             f"For more information, please see "
+#             f"https://github.com/brainglobe/brainglobe-meta#readme."
+#         ) from e
 
 __author__ = "Adam Tyson, Christian Niedworok, Charly Rousseau"
 __license__ = "BSD-3-Clause"

From 750171c90fbf9296158975549ff81c5476d357b6 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 13:04:39 +0000
Subject: [PATCH 07/17] change checkpoint filename for compliance with keras 3.
 remove use_multiprocessing=False from fit() as it is no longer an input.
 test_train() passing

---
 cellfinder/core/train/train_yml.py        | 34 +++++++++++++++--------
 tests/core/test_integration/test_train.py |  2 +-
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 68c843a6..9a42253a 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -386,15 +386,15 @@ def run(
             labels=labels_test,
             batch_size=batch_size,
             train=True,
-        )
+        )  # PyDataset
 
         # for saving checkpoints
-        base_checkpoint_file_name = "-epoch.{epoch:02d}-loss-{val_loss:.3f}.h5"
+        base_checkpoint_file_name = "-epoch.{epoch:02d}-loss-{val_loss:.3f}"
 
     else:
         logger.info("No validation data selected.")
         validation_generator = None
-        base_checkpoint_file_name = "-epoch.{epoch:02d}.h5"
+        base_checkpoint_file_name = "-epoch.{epoch:02d}"
 
     training_generator = CubeGeneratorFromDisk(
         signal_train,
@@ -404,7 +404,9 @@ def run(
         shuffle=True,
         train=True,
         augment=not no_augment,
-    )
+    )  # PyDataset - in Keras 3.0, multiprocessing should be specified here
+    # (rather than in input to fit() for ex)
+    # by default: use_multiprocessing=False
     callbacks = []
 
     if tensorboard:
@@ -419,10 +421,18 @@ def run(
         callbacks.append(tensorboard)
 
     if not no_save_checkpoints:
+        # Keras 3.0: filepath name needs to end with
+        # - ".weights.h5" when save_weights_only=True or
+        # - ".keras" when checkpoint saving the whole model (default)
         if save_weights:
-            filepath = str(output_dir / ("weight" + base_checkpoint_file_name))
+            filepath = str(
+                output_dir
+                / ("weight" + base_checkpoint_file_name + ".weights.h5")
+            )
         else:
-            filepath = str(output_dir / ("model" + base_checkpoint_file_name))
+            filepath = str(
+                output_dir / ("model" + base_checkpoint_file_name + ".keras")
+            )
 
         checkpoints = ModelCheckpoint(
             filepath,
@@ -431,25 +441,27 @@ def run(
         callbacks.append(checkpoints)
 
     if save_progress:
-        filepath = str(output_dir / "training.csv")
-        csv_logger = CSVLogger(filepath)
+        csv_filepath = str(output_dir / "training.csv")
+        csv_logger = CSVLogger(csv_filepath)
         callbacks.append(csv_logger)
 
     logger.info("Beginning training.")
+    # Keras 3.0: for model the `use_multiprocessing` input does not exist
+    # anymore, it is set in `training_generator` instead and it is False
+    # by default
     model.fit(
         training_generator,
         validation_data=validation_generator,
-        use_multiprocessing=False,
         epochs=epochs,
         callbacks=callbacks,
     )
 
     if save_weights:
         logger.info("Saving model weights")
-        model.save_weights(str(output_dir / "model_weights.h5"))
+        model.save_weights(str(output_dir / "model.weights.h5"))
     else:
         logger.info("Saving model")
-        model.save(output_dir / "model.h5")
+        model.save(str(output_dir / "model.keras"))
 
     logger.info(
         "Finished training, " "Total time taken: %s",
diff --git a/tests/core/test_integration/test_train.py b/tests/core/test_integration/test_train.py
index 0d4f5c38..8cfc8c31 100644
--- a/tests/core/test_integration/test_train.py
+++ b/tests/core/test_integration/test_train.py
@@ -35,5 +35,5 @@ def test_train(tmpdir):
     sys.argv = train_args
     train_run()
 
-    model_file = os.path.join(tmpdir, "model.h5")
+    model_file = os.path.join(tmpdir, "model.keras")
     assert os.path.exists(model_file)

From 04f182b5ac1b5c89fdd7ae08291041d398c52c39 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 13:57:58 +0000
Subject: [PATCH 08/17] add multiprocessing parameters to cube generator
 constructor and remove from fit() signature (keras3 change)

---
 cellfinder/core/classify/classify.py       |  7 +++--
 cellfinder/core/classify/cube_generator.py | 30 ++++++++++++++++++++++
 cellfinder/core/train/train_yml.py         |  6 ++---
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
index 93605176..c6554acc 100644
--- a/cellfinder/core/classify/classify.py
+++ b/cellfinder/core/classify/classify.py
@@ -63,6 +63,8 @@ def main(
         cube_width=cube_width,
         cube_height=cube_height,
         cube_depth=cube_depth,
+        use_multiprocessing=True,
+        workers=workers,
     )
 
     model = get_model(
@@ -73,10 +75,11 @@ def main(
     )
 
     logger.info("Running inference")
+    # in Keras 3.0, multiprocessing params are now specified in the generator
     predictions = model.predict(
         inference_generator,
-        use_multiprocessing=True,
-        workers=workers,
+        # use_multiprocessing=True,
+        # workers=workers,
         verbose=True,
         callbacks=callbacks,
     )
diff --git a/cellfinder/core/classify/cube_generator.py b/cellfinder/core/classify/cube_generator.py
index fb936dbf..a10551e8 100644
--- a/cellfinder/core/classify/cube_generator.py
+++ b/cellfinder/core/classify/cube_generator.py
@@ -56,7 +56,28 @@ def __init__(
         translate: Tuple[float, float, float] = (0.05, 0.05, 0.05),
         shuffle: bool = False,
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
+        # use_multiprocessing: bool = False,
+        # workers: int = 3, # 3 from max num workers
     ):
+        # Multiprocessing parameters
+        # Option 1
+        # use_multiprocessing will set _use_multiprocessing too,
+        # see https://github.com/keras-team/keras/blob/5bc8488c0ea3f43c70c70ebca919093cd56066eb/keras/trainers/data_adapters/py_dataset_adapter.py#L134
+        # self.use_multiprocessing = use_multiprocessing
+        # self.workers = workers
+        # ----------
+        # Option 2:
+        # call constructor of superclass to avoid redundance instead?
+        # Sequence.__init__(self, use_multiprocessing, workers)
+        # ----------
+        # Option 3:
+        # pass any additional arguments not specified in signature to the
+        # constructor of the superclass (e.g.: `use_multiprocessing` or
+        # `workers`)
+        super().__init__(*args, **kwargs)
+
         self.points = points
         self.signal_array = signal_array
         self.background_array = background_array
@@ -352,7 +373,16 @@ def __init__(
         translate: Tuple[float, float, float] = (0.2, 0.2, 0.2),
         train: bool = False,  # also return labels
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
+        # use_multiprocessing: bool = False,
+        # workers: int = 3, # 3 from max num workers
     ):
+        # pass any additional arguments not specified in signature to the
+        # constructor of the superclass (e.g.: `use_multiprocessing` or
+        # `workers`)
+        super().__init__(*args, **kwargs)
+
         self.im_shape = shape
         self.batch_size = batch_size
         self.labels = labels
diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 9a42253a..3a22cd38 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -386,7 +386,7 @@ def run(
             labels=labels_test,
             batch_size=batch_size,
             train=True,
-        )  # PyDataset
+        )  # PyDataset, validation_generator.use_multiprocessing=False; assert?
 
         # for saving checkpoints
         base_checkpoint_file_name = "-epoch.{epoch:02d}-loss-{val_loss:.3f}"
@@ -446,7 +446,7 @@ def run(
         callbacks.append(csv_logger)
 
     logger.info("Beginning training.")
-    # Keras 3.0: for model the `use_multiprocessing` input does not exist
+    # Keras 3.0: for model.fit() the `use_multiprocessing` input does not exist
     # anymore, it is set in `training_generator` instead and it is False
     # by default
     model.fit(
@@ -454,7 +454,7 @@ def run(
         validation_data=validation_generator,
         epochs=epochs,
         callbacks=callbacks,
-    )
+    )  # training_generator.use_multiprocessing = False
 
     if save_weights:
         logger.info("Saving model weights")

From b7e8c09e64a43baf5f9b510463b9ade9deb530a3 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 15:27:38 +0000
Subject: [PATCH 09/17] apply temp garbage collector fix

---
 tests/core/conftest.py                        | 21 -------
 tests/core/test_integration/test_detection.py | 59 +++++++++++++++----
 2 files changed, 46 insertions(+), 34 deletions(-)

diff --git a/tests/core/conftest.py b/tests/core/conftest.py
index f05ec88a..83778a9f 100644
--- a/tests/core/conftest.py
+++ b/tests/core/conftest.py
@@ -1,4 +1,3 @@
-import os
 from typing import Tuple
 
 import numpy as np
@@ -9,26 +8,6 @@
 from cellfinder.core.tools.prep import DEFAULT_INSTALL_PATH
 
 
-@pytest.fixture(scope="session")
-def no_free_cpus() -> int:
-    """
-    Set number of free CPUs so all available CPUs are used by the tests.
-    """
-    return 0
-
-
-@pytest.fixture(scope="session")
-def run_on_one_cpu_only() -> int:
-    """
-    Set number of free CPUs so tests can use exactly one CPU.
-    """
-    cpus = os.cpu_count()
-    if cpus is not None:
-        return cpus - 1
-    else:
-        raise ValueError("No CPUs available.")
-
-
 @pytest.fixture(scope="session")
 def download_default_model():
     """
diff --git a/tests/core/test_integration/test_detection.py b/tests/core/test_integration/test_detection.py
index 037b63ef..1a43dca5 100644
--- a/tests/core/test_integration/test_detection.py
+++ b/tests/core/test_integration/test_detection.py
@@ -43,17 +43,45 @@ def background_array():
 
 
 # FIXME: This isn't a very good example
-@pytest.mark.skip()
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    "free_cpus",
+    "cpus_to_leave_available",
     [
-        pytest.param("no_free_cpus", id="No free CPUs"),
-        pytest.param("run_on_one_cpu_only", id="One CPU"),
+        pytest.param(0, id="Leave no CPUS free"),
+        pytest.param(-1, id="Only use one CPU"),
     ],
 )
-def test_detection_full(signal_array, background_array, free_cpus, request):
-    n_free_cpus = request.getfixturevalue(free_cpus)
+def test_detection_full(
+    signal_array, background_array, cpus_to_leave_available: int
+):
+    """
+    cpus_to_leave_available is interpreted as follows:
+
+    - For values >=0, this is the number of CPUs to leave available
+    to the system when running this test.
+    - For values <0, this is HOW MANY CPUS to request be used to
+    run the test.
+
+    In each case, we check that we will be running on at least one CPU,
+    and not requesting more CPUs than the system can provide.
+    """
+    # Determine the number of CPUs to leave available
+    system_cpus = os.cpu_count()
+    # How many CPUs do we want to leave free?
+    if cpus_to_leave_available >= 0:
+        n_free_cpus = cpus_to_leave_available
+    else:
+        # Number of CPUs to keep free is <0, interpret as
+        # number of CPUs _to use_. Thus;
+        # n_free_cpus = system_cpus - |cpus_to_leave_available|
+        n_free_cpus = system_cpus - abs(cpus_to_leave_available)
+    # Check that there are enough CPUs
+    if not 0 <= n_free_cpus < system_cpus:
+        RuntimeError(
+            f"Not enough CPUS available (you want to leave {n_free_cpus} "
+            f"available, but there are only {system_cpus} on the system)."
+        )
+
     cells_test = main(
         signal_array,
         background_array,
@@ -81,10 +109,13 @@ def test_detection_full(signal_array, background_array, free_cpus, request):
 
 
 def test_detection_small_planes(
-    signal_array, background_array, no_free_cpus, mocker
+    signal_array,
+    background_array,
+    mocker,
+    cpus_to_leave_free: int = 0,
 ):
     # Check that processing works when number of planes < number of processes
-    nproc = get_num_processes(no_free_cpus)
+    nproc = get_num_processes(cpus_to_leave_free)
     n_planes = 2
 
     # Don't want to bother classifying in this test, so mock classifcation
@@ -101,11 +132,13 @@ def test_detection_small_planes(
         background_array[0:n_planes],
         voxel_sizes,
         ball_z_size=5,
-        n_free_cpus=no_free_cpus,
+        n_free_cpus=cpus_to_leave_free,
     )
 
 
-def test_callbacks(signal_array, background_array, no_free_cpus):
+def test_callbacks(
+    signal_array, background_array, cpus_to_leave_free: int = 0
+):
     # 20 is minimum number of planes needed to find > 0 cells
     signal_array = signal_array[0:20]
     background_array = background_array[0:20]
@@ -130,7 +163,7 @@ def detect_finished_callback(points):
         detect_callback=detect_callback,
         classify_callback=classify_callback,
         detect_finished_callback=detect_finished_callback,
-        n_free_cpus=no_free_cpus,
+        n_free_cpus=cpus_to_leave_free,
     )
 
     np.testing.assert_equal(planes_done, np.arange(len(signal_array)))
@@ -148,13 +181,13 @@ def test_floating_point_error(signal_array, background_array):
         main(signal_array, background_array, voxel_sizes)
 
 
-def test_synthetic_data(synthetic_bright_spots, no_free_cpus):
+def test_synthetic_data(synthetic_bright_spots, cpus_to_leave_free: int = 0):
     signal_array, background_array = synthetic_bright_spots
     detected = main(
         signal_array,
         background_array,
         voxel_sizes,
-        n_free_cpus=no_free_cpus,
+        n_free_cpus=cpus_to_leave_free,
     )
     assert len(detected) == 8
 

From 2b0b01baab4ff84c0049f3114c3b1d70f97efa17 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 15:29:13 +0000
Subject: [PATCH 10/17] skip troublesome test

---
 tests/core/test_integration/test_detection.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/core/test_integration/test_detection.py b/tests/core/test_integration/test_detection.py
index 1a43dca5..39cc738c 100644
--- a/tests/core/test_integration/test_detection.py
+++ b/tests/core/test_integration/test_detection.py
@@ -43,6 +43,7 @@ def background_array():
 
 
 # FIXME: This isn't a very good example
+@pytest.mark.skip(reason="resolved in PR369")
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "cpus_to_leave_available",

From 71d112720d2804e4b26cb7266f8429c5e0c05d21 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 16:37:32 +0000
Subject: [PATCH 11/17] skip running tests on CI on windows

---
 .github/workflows/test_and_deploy.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index 79e765b4..002dc0b6 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -45,8 +45,8 @@ jobs:
         include:
         - os: macos-latest
           python-version: "3.10"
-        - os: windows-latest
-          python-version: "3.10"
+        # - os: windows-latest
+        #   python-version: "3.10"
 
     steps:
       # Cache the tensorflow model so we don't have to remake it every time

From 2e90028013731472aef0f4cc536ef1f95e425bcd Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:23:49 +0000
Subject: [PATCH 12/17] remove commented out TF check

---
 cellfinder/__init__.py | 15 ---------------
 1 file changed, 15 deletions(-)

diff --git a/cellfinder/__init__.py b/cellfinder/__init__.py
index 30fb135c..9d60c02f 100644
--- a/cellfinder/__init__.py
+++ b/cellfinder/__init__.py
@@ -5,20 +5,5 @@
 except PackageNotFoundError as e:
     raise PackageNotFoundError("cellfinder package not installed") from e
 
-# If tensorflow is not present, tools cannot be used.
-# Throw an error in this case to prevent invocation of functions.
-# try:
-#     TF_VERSION = version("tensorflow")  # replace by keras?
-# except PackageNotFoundError as e:
-#     try:
-#         TF_VERSION = version("tensorflow-macos")
-#     except PackageNotFoundError as e:
-#         raise PackageNotFoundError(
-#             f"cellfinder tools cannot be invoked without tensorflow. "
-#             f"Please install tensorflow into your environment to use cellfinder tools. "
-#             f"For more information, please see "
-#             f"https://github.com/brainglobe/brainglobe-meta#readme."
-#         ) from e
-
 __author__ = "Adam Tyson, Christian Niedworok, Charly Rousseau"
 __license__ = "BSD-3-Clause"

From db441c9043a6c1d01596a049d78a3944e0862778 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:44:34 +0000
Subject: [PATCH 13/17] clean commented out code. Explicitly pass
 use_multiprocessing=False (as before)

---
 cellfinder/core/classify/classify.py       |  4 +---
 cellfinder/core/classify/cube_generator.py | 16 ----------------
 cellfinder/core/train/train_yml.py         | 20 ++++++++++----------
 pyproject.toml                             |  6 +-----
 4 files changed, 12 insertions(+), 34 deletions(-)

diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
index c6554acc..2193ac08 100644
--- a/cellfinder/core/classify/classify.py
+++ b/cellfinder/core/classify/classify.py
@@ -75,11 +75,9 @@ def main(
     )
 
     logger.info("Running inference")
-    # in Keras 3.0, multiprocessing params are now specified in the generator
+    # in Keras 3.0 multiprocessing params are specified in the generator
     predictions = model.predict(
         inference_generator,
-        # use_multiprocessing=True,
-        # workers=workers,
         verbose=True,
         callbacks=callbacks,
     )
diff --git a/cellfinder/core/classify/cube_generator.py b/cellfinder/core/classify/cube_generator.py
index a10551e8..17509061 100644
--- a/cellfinder/core/classify/cube_generator.py
+++ b/cellfinder/core/classify/cube_generator.py
@@ -58,21 +58,7 @@ def __init__(
         interpolation_order: int = 2,
         *args,
         **kwargs,
-        # use_multiprocessing: bool = False,
-        # workers: int = 3, # 3 from max num workers
     ):
-        # Multiprocessing parameters
-        # Option 1
-        # use_multiprocessing will set _use_multiprocessing too,
-        # see https://github.com/keras-team/keras/blob/5bc8488c0ea3f43c70c70ebca919093cd56066eb/keras/trainers/data_adapters/py_dataset_adapter.py#L134
-        # self.use_multiprocessing = use_multiprocessing
-        # self.workers = workers
-        # ----------
-        # Option 2:
-        # call constructor of superclass to avoid redundance instead?
-        # Sequence.__init__(self, use_multiprocessing, workers)
-        # ----------
-        # Option 3:
         # pass any additional arguments not specified in signature to the
         # constructor of the superclass (e.g.: `use_multiprocessing` or
         # `workers`)
@@ -375,8 +361,6 @@ def __init__(
         interpolation_order: int = 2,
         *args,
         **kwargs,
-        # use_multiprocessing: bool = False,
-        # workers: int = 3, # 3 from max num workers
     ):
         # pass any additional arguments not specified in signature to the
         # constructor of the superclass (e.g.: `use_multiprocessing` or
diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 3a22cd38..49a76808 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -386,7 +386,8 @@ def run(
             labels=labels_test,
             batch_size=batch_size,
             train=True,
-        )  # PyDataset, validation_generator.use_multiprocessing=False; assert?
+            use_multiprocessing=False,
+        )
 
         # for saving checkpoints
         base_checkpoint_file_name = "-epoch.{epoch:02d}-loss-{val_loss:.3f}"
@@ -404,9 +405,8 @@ def run(
         shuffle=True,
         train=True,
         augment=not no_augment,
-    )  # PyDataset - in Keras 3.0, multiprocessing should be specified here
-    # (rather than in input to fit() for ex)
-    # by default: use_multiprocessing=False
+        use_multiprocessing=False,
+    )
     callbacks = []
 
     if tensorboard:
@@ -421,9 +421,10 @@ def run(
         callbacks.append(tensorboard)
 
     if not no_save_checkpoints:
-        # Keras 3.0: filepath name needs to end with
+        # In Keras 3.0, the name of the checkpoint needs to end with
         # - ".weights.h5" when save_weights_only=True or
-        # - ".keras" when checkpoint saving the whole model (default)
+        # - ".keras" when saving the whole model
+        # see https://keras.io/api/callbacks/model_checkpoint/#modelcheckpoint-class
         if save_weights:
             filepath = str(
                 output_dir
@@ -446,15 +447,14 @@ def run(
         callbacks.append(csv_logger)
 
     logger.info("Beginning training.")
-    # Keras 3.0: for model.fit() the `use_multiprocessing` input does not exist
-    # anymore, it is set in `training_generator` instead and it is False
-    # by default
+    # Keras 3.0: `use_multiprocessing` input is set in the
+    # `training_generator` instead (False by default)
     model.fit(
         training_generator,
         validation_data=validation_generator,
         epochs=epochs,
         callbacks=callbacks,
-    )  # training_generator.use_multiprocessing = False
+    )
 
     if save_weights:
         logger.info("Saving model weights")
diff --git a/pyproject.toml b/pyproject.toml
index 39941d6d..5c73d68c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,11 +29,8 @@ dependencies = [
     "numpy",
     "scikit-image",
     "scikit-learn",
-    "keras", #
+    "keras",
     "tf-nightly==2.16.0.dev20240101", # pinning to same TF as Keras 3.0
-    # See https://github.com/brainglobe/cellfinder-core/issues/103 for < 2.12.0 pin
-    # "tensorflow-macos>=2.5.0; platform_system=='Darwin' and platform_machine=='arm64'",
-    # "tensorflow>=2.5.0; platform_system!='Darwin' or platform_machine!='arm64'",
     "tifffile",
     "tqdm",
 ]
@@ -142,5 +139,4 @@ passenv =
     XAUTHORITY
     NUMPY_EXPERIMENTAL_ARRAY_FUNCTION
     PYVISTA_OFF_SCREEN
-    # export KERAS_BACKEND="tensorflow" ?
 """

From 0a719b81a5bbb5d59c06cc6e49e59875e66a5cb9 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:44:57 +0000
Subject: [PATCH 14/17] remove str conversion before model.save

---
 cellfinder/core/train/train_yml.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 49a76808..f9c60f6d 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -458,10 +458,10 @@ def run(
 
     if save_weights:
         logger.info("Saving model weights")
-        model.save_weights(str(output_dir / "model.weights.h5"))
+        model.save_weights(output_dir / "model.weights.h5")
     else:
         logger.info("Saving model")
-        model.save(str(output_dir / "model.keras"))
+        model.save(output_dir / "model.keras")
 
     logger.info(
         "Finished training, " "Total time taken: %s",

From fc5ad6941b0cc6329c35bcdad28e6c9d8f0d2196 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:48:44 +0000
Subject: [PATCH 15/17] raise test_detection error for sonarcloud happy

---
 tests/core/test_integration/test_detection.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/core/test_integration/test_detection.py b/tests/core/test_integration/test_detection.py
index 39cc738c..d4d89745 100644
--- a/tests/core/test_integration/test_detection.py
+++ b/tests/core/test_integration/test_detection.py
@@ -78,7 +78,7 @@ def test_detection_full(
         n_free_cpus = system_cpus - abs(cpus_to_leave_available)
     # Check that there are enough CPUs
     if not 0 <= n_free_cpus < system_cpus:
-        RuntimeError(
+        raise RuntimeError(
             f"Not enough CPUS available (you want to leave {n_free_cpus} "
             f"available, but there are only {system_cpus} on the system)."
         )

From e7b2d51f56a8d422075dc6064896c97343bb8293 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:20:55 +0000
Subject: [PATCH 16/17] skip running tests on windows on CI

---
 .github/workflows/test_and_deploy.yml | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index 002dc0b6..1f0ec49a 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -41,12 +41,10 @@ jobs:
         # Run all supported Python versions on linux
         os: [ubuntu-latest]
         python-version: ["3.9", "3.10"]
-        # Include one windows, one macos run
+        # Include one macos run
         include:
         - os: macos-latest
           python-version: "3.10"
-        # - os: windows-latest
-        #   python-version: "3.10"
 
     steps:
       # Cache the tensorflow model so we don't have to remake it every time

From d88080bb8ef5644d6e4c814836749e3025593685 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:23:06 +0000
Subject: [PATCH 17/17] remove filename comment and small edits

---
 cellfinder/core/train/train_yml.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index f9c60f6d..dd83ebc5 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -421,10 +421,6 @@ def run(
         callbacks.append(tensorboard)
 
     if not no_save_checkpoints:
-        # In Keras 3.0, the name of the checkpoint needs to end with
-        # - ".weights.h5" when save_weights_only=True or
-        # - ".keras" when saving the whole model
-        # see https://keras.io/api/callbacks/model_checkpoint/#modelcheckpoint-class
         if save_weights:
             filepath = str(
                 output_dir
@@ -448,7 +444,7 @@ def run(
 
     logger.info("Beginning training.")
     # Keras 3.0: `use_multiprocessing` input is set in the
-    # `training_generator` instead (False by default)
+    # `training_generator` (False by default)
     model.fit(
         training_generator,
         validation_data=validation_generator,