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

cellfinder/__init__.py

@@ -1,27 +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/main.py

@@ -1,23 +1,13 @@
-"""
-N.B imports are within functions to prevent tensorflow being imported before
-it's warnings are silenced
-"""
-
 import os
 from typing import Callable, List, Optional, Tuple
 
 import numpy as np
 from brainglobe_utils.cells.cells import Cell
-from brainglobe_utils.general.logging import suppress_specific_logs
 
 from cellfinder.core import logger
 from cellfinder.core.download.download import model_type
 from cellfinder.core.train.train_yml import depth_type
 
-tf_suppress_log_messages = [
-    "multiprocessing can interact badly with TensorFlow"
-]
-
 
 def main(
     signal_array: np.ndarray,
@@ -28,7 +18,7 @@ def main(
     trained_model: Optional[os.PathLike] = None,
     model_weights: Optional[os.PathLike] = None,
     model: model_type = "resnet50_tv",
-    batch_size: int = 32,
+    batch_size: int = 64,
     n_free_cpus: int = 2,
     network_voxel_sizes: Tuple[int, int, int] = (5, 1, 1),
     soma_diameter: int = 16,
@@ -58,13 +48,11 @@ def main(
         Called every time a plane has finished being processed during the
         detection stage. Called with the plane number that has finished.
     classify_callback : Callable[int], optional
-        Called every time tensorflow has finished classifying a point.
+        Called every time a point has finished being classified.
         Called with the batch number that has just finished.
     detect_finished_callback : Callable[list], optional
         Called after detection is finished with the list of detected points.
     """
-    suppress_tf_logging(tf_suppress_log_messages)
-
     from cellfinder.core.classify import classify
     from cellfinder.core.detect import detect
     from cellfinder.core.tools import prep
@@ -98,7 +86,7 @@ def main(
     if not skip_classification:
         install_path = None
         model_weights = prep.prep_model_weights(
-            model_weights, install_path, model, n_free_cpus
+            model_weights, install_path, model
         )
         if len(points) > 0:
             logger.info("Running classification")
@@ -120,17 +108,4 @@ def main(
             )
         else:
             logger.info("No candidates, skipping classification")
-
     return points
-
-
-def suppress_tf_logging(tf_suppress_log_messages: List[str]) -> None:
-    """
-    Prevents many lines of logs such as:
-    "2019-10-24 16:54:41.363978: I tensorflow/stream_executor/platform/default
-    /dso_loader.cc:44] Successfully opened dynamic library libcuda.so.1"
-    """
-    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-
-    for message in tf_suppress_log_messages:
-        suppress_specific_logs("tensorflow", message)

cellfinder/core/tools/prep.py

@@ -9,9 +9,7 @@ from pathlib import Path
 from typing import Optional
 
 from brainglobe_utils.general.config import get_config_obj
-from brainglobe_utils.general.system import get_num_processes
 
-import cellfinder.core.tools.tf as tf_tools
 from cellfinder.core import logger
 from cellfinder.core.download.download import (
     DEFAULT_DOWNLOAD_DIRECTORY,
@@ -26,20 +24,13 @@ def prep_model_weights(
     model_weights: Optional[os.PathLike],
     install_path: Optional[os.PathLike],
     model_name: model_type,
-    n_free_cpus: int,
 ) -> Path:
-    n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
-    prep_tensorflow(n_processes)
+    # prepare models (get default weights or provided ones)
     model_weights = prep_models(model_weights, install_path, model_name)
 
     return model_weights
 
 
-def prep_tensorflow(max_threads: int) -> None:
-    tf_tools.set_tf_threads(max_threads)
-    tf_tools.allow_gpu_memory_growth()
-
-
 def prep_models(
     model_weights_path: Optional[os.PathLike],
     install_path: Optional[os.PathLike],

cellfinder/core/tools/system.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import keras
 from brainglobe_utils.general.exceptions import CommandLineInputError
 
 
@@ -80,3 +81,12 @@ def memory_in_bytes(memory_amount, unit):
         )
     else:
         return memory_amount * 10 ** supported_units[unit]
+
+
+def force_cpu():
+    """
+    Forces the CPU to be used, even if a GPU is available
+    """
+    keras.src.backend.common.global_state.set_global_attribute(
+        "torch_device", "cpu"
+    )

cellfinder/core/train/train_yml.py

@@ -22,7 +22,10 @@ from brainglobe_utils.general.numerical import (
     check_positive_float,
     check_positive_int,
 )
-from brainglobe_utils.general.system import ensure_directory_exists
+from brainglobe_utils.general.system import (
+    ensure_directory_exists,
+    get_num_processes,
+)
 from brainglobe_utils.IO.cells import find_relevant_tiffs
 from brainglobe_utils.IO.yaml import read_yaml_section
 from fancylog import fancylog
@@ -33,11 +36,6 @@ from cellfinder.core import logger
 from cellfinder.core.classify.resnet import layer_type
 from cellfinder.core.download.download import DEFAULT_DOWNLOAD_DIRECTORY
 
-tf_suppress_log_messages = [
-    "sample_weight modes were coerced from",
-    "multiprocessing can interact badly with TensorFlow",
-]
-
 depth_type = Literal["18", "34", "50", "101", "152"]
 
 models: Dict[depth_type, layer_type] = {
@@ -318,11 +316,7 @@ def run(
     save_progress=False,
     epochs=100,
 ):
-    from cellfinder.core.main import suppress_tf_logging
-
-    suppress_tf_logging(tf_suppress_log_messages)
-
-    from tensorflow.keras.callbacks import (
+    from keras.callbacks import (
         CSVLogger,
         ModelCheckpoint,
         TensorBoard,
@@ -339,7 +333,6 @@ def run(
         model_weights=model_weights,
         install_path=install_path,
         model_name=model,
-        n_free_cpus=n_free_cpus,
     )
 
     yaml_contents = parse_yaml(yaml_file)
@@ -361,6 +354,7 @@ def run(
 
     signal_train, background_train, labels_train = make_lists(tiff_files)
 
+    n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
     if test_fraction > 0:
         logger.info("Splitting data into training and validation datasets")
         (
@@ -387,15 +381,17 @@ def run(
             labels=labels_test,
             batch_size=batch_size,
             train=True,
+            use_multiprocessing=False,
+            workers=n_processes,
         )
 
         # 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,
@@ -405,6 +401,8 @@ def run(
         shuffle=True,
         train=True,
         augment=not no_augment,
+        use_multiprocessing=False,
+        workers=n_processes,
     )
     callbacks = []
 
@@ -421,9 +419,14 @@ def run(
 
     if not no_save_checkpoints:
         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,
@@ -432,25 +435,26 @@ 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: `use_multiprocessing` input is set in the
+    # `training_generator` (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(output_dir / "model.weights.h5")
     else:
         logger.info("Saving model")
-        model.save(output_dir / "model.h5")
+        model.save(output_dir / "model.keras")
 
     logger.info(
         "Finished training, " "Total time taken: %s",

cellfinder/napari/curation.py

@@ -54,7 +54,7 @@ class CurationWidget(QWidget):
         self.save_empty_cubes = save_empty_cubes
         self.max_ram = max_ram
         self.voxel_sizes = [5, 2, 2]
-        self.batch_size = 32
+        self.batch_size = 64
         self.viewer = viewer
 
         self.signal_layer = None

cellfinder/napari/detect/detect.py

@@ -253,8 +253,9 @@ def detect_widget() -> FunctionGui:
         max_cluster_size: int,
         classification_options,
         skip_classification: bool,
-        trained_model: Optional[Path],
         use_pre_trained_weights: bool,
+        trained_model: Optional[Path],
+        batch_size: int,
         misc_options,
         start_plane: int,
         end_plane: int,
@@ -298,6 +299,8 @@ def detect_widget() -> FunctionGui:
             should be attempted
         use_pre_trained_weights : bool
             Select to use pre-trained model weights
+        batch_size : int
+            How many points to classify at one time
         skip_classification : bool
             If selected, the classification step is skipped and all cells from
             the detection stage are added
@@ -372,7 +375,10 @@ def detect_widget() -> FunctionGui:
         if use_pre_trained_weights:
             trained_model = None
         classification_inputs = ClassificationInputs(
-            skip_classification, use_pre_trained_weights, trained_model
+            skip_classification,
+            use_pre_trained_weights,
+            trained_model,
+            batch_size,
         )
 
         if analyse_local:

cellfinder/napari/detect/detect_containers.py

@@ -114,6 +114,7 @@ class ClassificationInputs(InputContainer):
     skip_classification: bool = False
     use_pre_trained_weights: bool = True
     trained_model: Optional[Path] = Path.home()
+    batch_size: int = 64
 
     def as_core_arguments(self) -> dict:
         args = super().as_core_arguments()
@@ -131,6 +132,7 @@ class ClassificationInputs(InputContainer):
             skip_classification=dict(
                 value=cls.defaults()["skip_classification"]
             ),
+            batch_size=dict(value=cls.defaults()["batch_size"]),
         )
 
 

cellfinder/napari/detect/thread_worker.py

@@ -72,10 +72,10 @@ class Worker(WorkerBase):
         def classify_callback(batch: int) -> None:
             self.update_progress_bar.emit(
                 "Classifying cells",
-                # Default cellfinder-core batch size is 32. This seems to give
+                # Default cellfinder-core batch size is 64. This seems to give
                 # a slight underestimate of the number of batches though, so
                 # allow for batch number to go over this
-                max(self.npoints_detected // 32 + 1, batch + 1),
+                max(self.npoints_detected // 64 + 1, batch + 1),
                 batch + 1,
             )
 

{cellfinder-1.2.0.dist-info → cellfinder-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cellfinder
-Version: 1.2.0
+Version: 1.3.0
 Summary: Automated 3D cell detection in large microscopy images
 Author-email: "Adam Tyson, Christian Niedworok, Charly Rousseau" <code@adamltyson.com>
 License: BSD-3-Clause
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
@@ -31,10 +32,10 @@ Requires-Dist: numba
 Requires-Dist: numpy
 Requires-Dist: scikit-image
 Requires-Dist: scikit-learn
+Requires-Dist: keras >=3.0.0
+Requires-Dist: torch >=2.1.0
 Requires-Dist: tifffile
 Requires-Dist: tqdm
-Requires-Dist: tensorflow <2.12.0,>=2.5.0 ; platform_system != "Darwin" or platform_machine != "arm64"
-Requires-Dist: tensorflow-macos <2.12.0,>=2.5.0 ; platform_system == "Darwin" and platform_machine == "arm64"
 Provides-Extra: dev
 Requires-Dist: black ; extra == 'dev'
 Requires-Dist: pre-commit ; extra == 'dev'

{cellfinder-1.2.0.dist-info → cellfinder-1.3.0.dist-info}/RECORD

@@ -1,14 +1,14 @@
-cellfinder/__init__.py,sha256=7H9UKFWCk-TNkFLwk-D4EVU5-rJZsGxfo0qcaFVcfKY,1035
+cellfinder/__init__.py,sha256=S5oQ3EORuyQTMYC4uUuzGKZ23J3Ya6q-1DOBib1KfiA,1166
 cellfinder/cli_migration_warning.py,sha256=gPtNrtnXvWpl5q0k_EGAQZg0DwcpCmuBTgpg56n5NfA,1578
 cellfinder/core/__init__.py,sha256=pRFuQsl78HEK0S6gvhJw70QLbjjSBzP-GFO0AtVaGtk,62
-cellfinder/core/main.py,sha256=Y-4Wt9d9K7o2l3IO7acZbtNuhqZzF5UU-rMJksb6k94,4392
+cellfinder/core/main.py,sha256=t2mkq6iieEytbPckehBB43juwN5E-vhzstLSs620vdM,3625
 cellfinder/core/types.py,sha256=lTqWE4v0SMM0qLAZJdyAzqV1nLgDtobEpglNJcXt160,106
 cellfinder/core/classify/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cellfinder/core/classify/augment.py,sha256=8dMbM7KhimM6NMgdMC53oHoCfYj1CIB-h3Yk8CZAxPw,6321
-cellfinder/core/classify/classify.py,sha256=y4EnHX3IpAUAfhGScoJGv2hLLpy6OtmGvkqxAow3_Kw,3046
-cellfinder/core/classify/cube_generator.py,sha256=4zMaa7bb7vr1QxBFpzqqitnWaNdzgEzZ2gt2xFeN37A,16586
-cellfinder/core/classify/resnet.py,sha256=976Qt0li7ySYx_w2m9QfRMaDAc0rEkA8rO9SYP5a6uw,10057
-cellfinder/core/classify/tools.py,sha256=UqABKS9KnkwR3PIhGRl7xHJbADP0mgVRSxXI5dybpqA,2547
+cellfinder/core/classify/classify.py,sha256=33ZvNDmgVabH_6p4jV9Xi8bLKwDZnVhzlImrL_TlnBk,3269
+cellfinder/core/classify/cube_generator.py,sha256=jC5aogTVy213PHouViSR9CgKkuOks3yk6csQC5kRoOE,17125
+cellfinder/core/classify/resnet.py,sha256=vGa85y_NcQnOXwAt5EtatLx5mrO8IoShCcNKtJ5-EFg,10034
+cellfinder/core/classify/tools.py,sha256=s5PEKAsZVbVuoferZ_nqMUM0f2bV_8WEKsdKe3SXEuE,2560
 cellfinder/core/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cellfinder/core/config/cellfinder.conf,sha256=5i8axif7ekMutKDiVnZRs-LiJrgVQljg_beltidqtNk,56
 cellfinder/core/detect/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -31,30 +31,29 @@ cellfinder/core/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 cellfinder/core/tools/array_operations.py,sha256=LDbqWA_N2YtxeKS877QYdJRL2FCMC1R1ExJtdb6-vEA,3371
 cellfinder/core/tools/geometry.py,sha256=xlEQQmVQ9jcXqRzUqU554P8VxkaWkc5J_YhjkkKlI0Q,1124
 cellfinder/core/tools/image_processing.py,sha256=z27bGjf3Iv3G4Nt1OYzpEnIYQNc4nNomj_QitqvZB78,2269
-cellfinder/core/tools/prep.py,sha256=hGWwXxeP88ahp6ji2bkBu9FOz8BmrZsVwI-CEv4qxQ4,2473
+cellfinder/core/tools/prep.py,sha256=YU4VDyjEsOXn3S3gymEfPUzka-WYiHmR4Q4XdRebhDI,2175
 cellfinder/core/tools/source_files.py,sha256=vvwsIMe1ULKvXg_x22L75iqvCyMjEbUjJsDFQk3GYzY,856
-cellfinder/core/tools/system.py,sha256=6iSOLrz2mtKgOIdwjp_VS0-74XQCuaQYCvWAViJgEus,2232
-cellfinder/core/tools/tf.py,sha256=jNU97-gIB51rLube9EoZ9bREkbqWRk-xPZhDTXG_jE4,1713
+cellfinder/core/tools/system.py,sha256=WvEzPr7v-ohLDnzf4n13TMcN5OYIAXOEkaSmrHzdnwc,2438
 cellfinder/core/tools/tiff.py,sha256=NzIz6wq2GzxmcIhawFMwZADe-uQO2rIG46H7xkpGKLs,2899
 cellfinder/core/tools/tools.py,sha256=G8oDGNRuWkzEJDnnC4r3SNGgpVbqbelCZR5ODk9JRzs,4867
 cellfinder/core/train/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cellfinder/core/train/train_yml.py,sha256=0iiROBdxNoMyMoPGYC8tq6bGM3jzAKlGA2LABgoqd6M,13139
+cellfinder/core/train/train_yml.py,sha256=9QXv2wk24G8hYskMnBij7OngEELUWySK2fH4NFbYWw4,13260
 cellfinder/napari/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cellfinder/napari/curation.py,sha256=jz7GekuuzOdwiEFLarvvO5Bho68k8eL7AjObMdEoPXw,21447
+cellfinder/napari/curation.py,sha256=nbxCwY2bhEPM15Wf3S_Ff8qGdLSWojr4X48mmAJqD3U,21447
 cellfinder/napari/input_container.py,sha256=tkm0dkPt7kSL8Xkvs1fh8M6vKWw57QLIt_wv74HFkGk,2150
 cellfinder/napari/napari.yaml,sha256=WMR1CIAmYIVyQngbdbomTRZLvlDbb6LxsXsvTRClQnE,921
 cellfinder/napari/sample_data.py,sha256=oUST23q09MM8dxHbUCmO0AjtXG6OlR_32LLqP0EU2UA,732
 cellfinder/napari/utils.py,sha256=AwTs76M9azutHhHj2yuaKErDEQ5F6pFbIIakBfzen6M,3824
 cellfinder/napari/detect/__init__.py,sha256=BD9Bg9NTAr6yRTq2A_p58U6j4w5wbY0sdXwhPJ3MSMY,34
-cellfinder/napari/detect/detect.py,sha256=gX10U4HHcni1ff3zYXkXYX6dWMZErYpvzVde56yV80w,13658
-cellfinder/napari/detect/detect_containers.py,sha256=Xx-E2JTbe5Ih_1m5ZhQ0pCYFxfdQxVJMMfGh9EZExms,5377
-cellfinder/napari/detect/thread_worker.py,sha256=3xGUPSthycpErcVxl_yc5soIv_z0Hf0iEaZq0Yx90OM,3078
+cellfinder/napari/detect/detect.py,sha256=VB3SLZvqJhjuypjpaZuV9JscQ-sE8yVHG8RzEsZWfeA,13809
+cellfinder/napari/detect/detect_containers.py,sha256=j9NTsIyyDNrhlI2dc7hvc7QlxvI1NRHlCe137v7fsPg,5467
+cellfinder/napari/detect/thread_worker.py,sha256=PWM3OE-FpK-dpdhaE_Gi-2lD3u8sL-SJ13mp0pMhTyI,3078
 cellfinder/napari/train/__init__.py,sha256=xo4CK-DvSecInGEc2ohcTgQYlH3iylFnGvKTCoq2WkI,35
 cellfinder/napari/train/train.py,sha256=zJY7zKcLqDTDtD76thmbwViEU4tTFCmXZze-zHsTpoo,5941
 cellfinder/napari/train/train_containers.py,sha256=1wZ_GPe7B5XsLYs5XIx4m8GMw5KeVhg6SchhPtXu4V8,4386
-cellfinder-1.2.0.dist-info/LICENSE,sha256=Tw8iMytIDXLSmcIUsbQmRWojstl9yOWsPCx6ZT6dZLY,1564
-cellfinder-1.2.0.dist-info/METADATA,sha256=0_L_i066TdGFnfaRZe8zY-y0XwF7HX88F0mzV4AqLEw,6632
-cellfinder-1.2.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-cellfinder-1.2.0.dist-info/entry_points.txt,sha256=cKKjU8GPiN-TRelG2sT2JCKAcB9XDCjP6g9atE9pSoY,247
-cellfinder-1.2.0.dist-info/top_level.txt,sha256=jyTQzX-tDjbsMr6s-E71Oy0IKQzmHTXSk4ZhpG5EDSE,11
-cellfinder-1.2.0.dist-info/RECORD,,
+cellfinder-1.3.0.dist-info/LICENSE,sha256=Tw8iMytIDXLSmcIUsbQmRWojstl9yOWsPCx6ZT6dZLY,1564
+cellfinder-1.3.0.dist-info/METADATA,sha256=b40Zf7cFJV6nv3HxqXCtJubKshRQmXNpXaPx0vVrFtI,6528
+cellfinder-1.3.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+cellfinder-1.3.0.dist-info/entry_points.txt,sha256=cKKjU8GPiN-TRelG2sT2JCKAcB9XDCjP6g9atE9pSoY,247
+cellfinder-1.3.0.dist-info/top_level.txt,sha256=jyTQzX-tDjbsMr6s-E71Oy0IKQzmHTXSk4ZhpG5EDSE,11
+cellfinder-1.3.0.dist-info/RECORD,,

cellfinder/core/tools/tf.py

@@ -1,46 +0,0 @@
-import tensorflow as tf
-
-from cellfinder.core import logger
-
-
-def allow_gpu_memory_growth():
-    """
-    If a gpu is present, prevent tensorflow from using all the memory straight
-    away. Allows multiple processes to use the GPU (and avoid occasional
-    errors on some systems) at the cost of a slight performance penalty.
-    """
-    gpus = tf.config.experimental.list_physical_devices("GPU")
-    if gpus:
-        logger.debug("Allowing GPU memory growth")
-        try:
-            # Currently, memory growth needs to be the same across GPUs
-            for gpu in gpus:
-                tf.config.experimental.set_memory_growth(gpu, True)
-            logical_gpus = tf.config.experimental.list_logical_devices("GPU")
-            logger.debug(
-                f"{len(gpus)} physical GPUs, {len(logical_gpus)} logical GPUs"
-            )
-        except RuntimeError as e:
-            # Memory growth must be set before GPUs have been initialized
-            print(e)
-    else:
-        logger.debug("No GPUs found, using CPU.")
-
-
-def set_tf_threads(max_threads):
-    """
-    Limit the number of threads that tensorflow uses
-    :param max_threads: Maximum number of threads to use
-    :return:
-    """
-    logger.debug(
-        f"Setting maximum number of threads for tensorflow "
-        f"to: {max_threads}"
-    )
-
-    # If statements are for testing. If tf is initialised, then setting these
-    # parameters throws an error
-    if tf.config.threading.get_inter_op_parallelism_threads() != 0:
-        tf.config.threading.set_inter_op_parallelism_threads(max_threads)
-    if tf.config.threading.get_intra_op_parallelism_threads() != 0:
-        tf.config.threading.set_intra_op_parallelism_threads(max_threads)