cellfinder 1.1.3__py3-none-any.whl → 1.3.0__py3-none-any.whl

cellfinder/__init__.py

@@ -1,24 +1,33 @@
+import os
 from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
 
+# Check cellfinder is installed
 try:
     __version__ = version("cellfinder")
 except PackageNotFoundError as e:
     raise PackageNotFoundError("cellfinder package not installed") from e
 
-# If tensorflow is not present, tools cannot be used.
+# If Keras is not present, tools cannot be used.
 # Throw an error in this case to prevent invocation of functions.
 try:
-    TF_VERSION = version("tensorflow")
+    KERAS_VERSION = version("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
+    raise PackageNotFoundError(
+        f"cellfinder tools cannot be invoked without Keras. "
+        f"Please install Keras with a backend into your environment "
+        f"to use cellfinder tools. "
+        f"For more information on Keras backends, please see "
+        f"https://keras.io/getting_started/#installing-keras-3."
+        f"For more information on brainglobe, please see "
+        f"https://github.com/brainglobe/brainglobe-meta#readme."
+    ) from e
+
+
+# Set the Keras backend to torch
+os.environ["KERAS_BACKEND"] = "torch"
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
-__author__ = "Adam Tyson, Christian Niedworok, Charly Rousseau"
 __license__ = "BSD-3-Clause"
+
+DEFAULT_CELLFINDER_DIRECTORY = Path.home() / ".brainglobe" / "cellfinder"

cellfinder/core/classify/classify.py

@@ -1,10 +1,11 @@
 import os
+from datetime import datetime
 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
-from tensorflow import keras
 
 from cellfinder.core import logger, types
 from cellfinder.core.classify.cube_generator import CubeGeneratorFromFile
@@ -48,9 +49,9 @@ def main(
         callbacks = None
 
     # Too many workers doesn't increase speed, and uses huge amounts of RAM
-    workers = get_num_processes(
-        min_free_cpu_cores=n_free_cpus, n_max_processes=max_workers
-    )
+    workers = get_num_processes(min_free_cpu_cores=n_free_cpus)
+
+    start_time = datetime.now()
 
     logger.debug("Initialising cube generator")
     inference_generator = CubeGeneratorFromFile(
@@ -63,6 +64,8 @@ def main(
         cube_width=cube_width,
         cube_height=cube_height,
         cube_depth=cube_depth,
+        use_multiprocessing=False,
+        workers=workers,
     )
 
     model = get_model(
@@ -73,10 +76,9 @@ def main(
     )
 
     logger.info("Running inference")
+    # 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,
     )
@@ -91,6 +93,11 @@ def main(
         cell.type = predictions[idx] + 1
         points_list.append(cell)
 
+    time_elapsed = datetime.now() - start_time
+    print(
+        "Classfication complete - all points done in : {}".format(time_elapsed)
+    )
+
     return points_list
 
 

cellfinder/core/classify/cube_generator.py

@@ -2,13 +2,13 @@ from pathlib import Path
 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 tensorflow.keras.utils import Sequence
 
 from cellfinder.core import types
 from cellfinder.core.classify.augment import AugmentationParameters, augment
@@ -40,7 +40,7 @@ class CubeGeneratorFromFile(Sequence):
         background_array: types.array,
         voxel_sizes: Tuple[int, int, int],
         network_voxel_sizes: Tuple[int, int, int],
-        batch_size: int = 16,
+        batch_size: int = 64,
         cube_width: int = 50,
         cube_height: int = 50,
         cube_depth: int = 20,
@@ -56,7 +56,14 @@ class CubeGeneratorFromFile(Sequence):
         translate: Tuple[float, float, float] = (0.05, 0.05, 0.05),
         shuffle: bool = False,
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
     ):
+        # 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
@@ -218,10 +225,10 @@ class CubeGeneratorFromFile(Sequence):
 
         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
+            return images, batch_labels.astype(np.float32)
         elif self.extract:
             batch_info = self.__get_batch_dict(cell_batch)
             return images, batch_info
@@ -252,7 +259,8 @@ class CubeGeneratorFromFile(Sequence):
                 (number_images,)
                 + (self.cube_height, self.cube_width, self.cube_depth)
                 + (self.channels,)
-            )
+            ),
+            dtype=np.float32,
         )
 
         for idx, cell in enumerate(cell_batch):
@@ -337,7 +345,7 @@ class CubeGeneratorFromDisk(Sequence):
         signal_list: List[Union[str, Path]],
         background_list: List[Union[str, Path]],
         labels: Optional[List[int]] = None,  # only if training or validating
-        batch_size: int = 16,
+        batch_size: int = 64,
         shape: Tuple[int, int, int] = (50, 50, 20),
         channels: int = 2,
         classes: int = 2,
@@ -350,7 +358,14 @@ class CubeGeneratorFromDisk(Sequence):
         translate: Tuple[float, float, float] = (0.2, 0.2, 0.2),
         train: bool = False,  # also return labels
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
     ):
+        # 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
@@ -410,10 +425,10 @@ class CubeGeneratorFromDisk(Sequence):
 
         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
+            return images, batch_labels.astype(np.float32)
         else:
             return images
 
@@ -424,7 +439,8 @@ class CubeGeneratorFromDisk(Sequence):
     ) -> np.ndarray:
         number_images = len(list_signal_tmp)
         images = np.empty(
-            ((number_images,) + self.im_shape + (self.channels,))
+            ((number_images,) + self.im_shape + (self.channels,)),
+            dtype=np.float32,
         )
 
         for idx, signal_im in enumerate(list_signal_tmp):
@@ -433,7 +449,7 @@ class CubeGeneratorFromDisk(Sequence):
                 images, idx, signal_im, background_im
             )
 
-        return images.astype(np.float16)
+        return images
 
     def __populate_array_with_cubes(
         self,

cellfinder/core/classify/resnet.py

@@ -1,9 +1,11 @@
 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 (
+    KerasTensor as Tensor,
+)
+from keras import Model
+from keras.initializers import Initializer
+from keras.layers import (
     Activation,
     Add,
     BatchNormalization,
@@ -14,7 +16,7 @@ from tensorflow.keras.layers import (
     MaxPooling3D,
     ZeroPadding3D,
 )
-from tensorflow.keras.optimizers import Adam, Optimizer
+from keras.optimizers import Adam, Optimizer
 
 #####################################################################
 # Define the types of ResNet
@@ -113,7 +115,7 @@ def non_residual_block(
     activation: str = "relu",
     use_bias: bool = False,
     bn_epsilon: float = 1e-5,
-    pooling_padding: str = "same",
+    pooling_padding: str = "valid",
     axis: int = 3,
 ) -> Tensor:
     """
@@ -131,6 +133,7 @@ def non_residual_block(
     )(x)
     x = BatchNormalization(axis=axis, epsilon=bn_epsilon, name="conv1_bn")(x)
     x = Activation(activation, name="conv1_activation")(x)
+
     x = MaxPooling3D(
         max_pool_size,
         strides=strides,

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,31 @@ 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 = []

cellfinder/core/detect/detect.py

@@ -22,6 +22,7 @@ from typing import Callable, List, Optional, Sequence, Tuple, TypeVar
 import numpy as np
 from brainglobe_utils.cells.cells import Cell
 from brainglobe_utils.general.system import get_num_processes
+from numba import set_num_threads
 
 from cellfinder.core import logger, types
 from cellfinder.core.detect.filters.plane import TileProcessor
@@ -157,6 +158,13 @@ def main(
         )
     n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
     n_ball_procs = max(n_processes - 1, 1)
+
+    # we parallelize 2d filtering, which typically lags behind the 3d
+    # processing so for n_ball_procs 2d filtering threads, ball_z_size will
+    # typically be in use while the others stall waiting for 3d processing
+    # so we can use those for other things, such as numba threading
+    set_num_threads(max(n_ball_procs - int(ball_z_size), 1))
+
     start_time = datetime.now()
 
     (
@@ -236,7 +244,10 @@ def main(
         # then 3D filtering has finished. As batches of planes are filtered
         # by the 3D filter, it releases the locks of subsequent 2D filter
         # processes.
-        cells = mp_3d_filter.process(async_results, locks, callback=callback)
+        mp_3d_filter.process(async_results, locks, callback=callback)
+
+        # it's now done filtering, get results with pool
+        cells = mp_3d_filter.get_results(worker_pool)
 
     time_elapsed = datetime.now() - start_time
     logger.debug(