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

cellfinder/core/main.py

@@ -1,22 +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.general.logging import suppress_specific_logs
+from brainglobe_utils.cells.cells import Cell
 
 from cellfinder.core import logger
-from cellfinder.core.download.models import model_type
+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,
@@ -27,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,
@@ -42,6 +33,9 @@ def main(
     cube_height: int = 50,
     cube_depth: int = 20,
     network_depth: depth_type = "50",
+    skip_detection: bool = False,
+    skip_classification: bool = False,
+    detected_cells: List[Cell] = None,
     *,
     detect_callback: Optional[Callable[[int], None]] = None,
     classify_callback: Optional[Callable[[int], None]] = None,
@@ -54,73 +48,64 @@ 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
 
-    logger.info("Detecting cell candidates")
+    if not skip_detection:
+        logger.info("Detecting cell candidates")
 
-    points = detect.main(
-        signal_array,
-        start_plane,
-        end_plane,
-        voxel_sizes,
-        soma_diameter,
-        max_cluster_size,
-        ball_xy_size,
-        ball_z_size,
-        ball_overlap_fraction,
-        soma_spread_factor,
-        n_free_cpus,
-        log_sigma_size,
-        n_sds_above_mean_thresh,
-        callback=detect_callback,
-    )
-
-    if detect_finished_callback is not None:
-        detect_finished_callback(points)
-
-    install_path = None
-    model_weights = prep.prep_model_weights(
-        model_weights, install_path, model, n_free_cpus
-    )
-    if len(points) > 0:
-        logger.info("Running classification")
-        points = classify.main(
-            points,
+        points = detect.main(
             signal_array,
-            background_array,
-            n_free_cpus,
+            start_plane,
+            end_plane,
             voxel_sizes,
-            network_voxel_sizes,
-            batch_size,
-            cube_height,
-            cube_width,
-            cube_depth,
-            trained_model,
-            model_weights,
-            network_depth,
-            callback=classify_callback,
+            soma_diameter,
+            max_cluster_size,
+            ball_xy_size,
+            ball_z_size,
+            ball_overlap_fraction,
+            soma_spread_factor,
+            n_free_cpus,
+            log_sigma_size,
+            n_sds_above_mean_thresh,
+            callback=detect_callback,
         )
-    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"
+        if detect_finished_callback is not None:
+            detect_finished_callback(points)
+    else:
+        points = detected_cells or []  # if None
+        detect_finished_callback(points)
 
-    for message in tf_suppress_log_messages:
-        suppress_specific_logs("tensorflow", message)
+    if not skip_classification:
+        install_path = None
+        model_weights = prep.prep_model_weights(
+            model_weights, install_path, model
+        )
+        if len(points) > 0:
+            logger.info("Running classification")
+            points = classify.main(
+                points,
+                signal_array,
+                background_array,
+                n_free_cpus,
+                voxel_sizes,
+                network_voxel_sizes,
+                batch_size,
+                cube_height,
+                cube_width,
+                cube_depth,
+                trained_model,
+                model_weights,
+                network_depth,
+                callback=classify_callback,
+            )
+        else:
+            logger.info("No candidates, skipping classification")
+    return points

cellfinder/core/tools/prep.py

@@ -9,42 +9,34 @@ 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 import models as model_download
-from cellfinder.core.download.download import amend_user_configuration
+from cellfinder.core.download.download import (
+    DEFAULT_DOWNLOAD_DIRECTORY,
+    amend_user_configuration,
+    download_models,
+    model_type,
+)
 from cellfinder.core.tools.source_files import user_specific_configuration_path
 
-home = Path.home()
-DEFAULT_INSTALL_PATH = home / ".cellfinder"
-
 
 def prep_model_weights(
     model_weights: Optional[os.PathLike],
     install_path: Optional[os.PathLike],
-    model_name: model_download.model_type,
-    n_free_cpus: int,
+    model_name: model_type,
 ) -> 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],
-    model_name: model_download.model_type,
+    model_name: model_type,
 ) -> Path:
-    install_path = install_path or DEFAULT_INSTALL_PATH
+    install_path = install_path or DEFAULT_DOWNLOAD_DIRECTORY
     # if no model or weights, set default weights
     if model_weights_path is None:
         logger.debug("No model supplied, so using the default")
@@ -53,13 +45,13 @@ def prep_models(
 
         if not Path(config_file).exists():
             logger.debug("Custom config does not exist, downloading models")
-            model_path = model_download.main(model_name, install_path)
+            model_path = download_models(model_name, install_path)
             amend_user_configuration(new_model_path=model_path)
 
         model_weights = get_model_weights(config_file)
         if not model_weights.exists():
             logger.debug("Model weights do not exist, downloading")
-            model_path = model_download.main(model_name, install_path)
+            model_path = download_models(model_name, install_path)
             amend_user_configuration(new_model_path=model_path)
             model_weights = get_model_weights(config_file)
     else:

cellfinder/core/tools/source_files.py

@@ -1,5 +1,7 @@
 from pathlib import Path
 
+from cellfinder import DEFAULT_CELLFINDER_DIRECTORY
+
 
 def default_configuration_path():
     """
@@ -17,11 +19,11 @@ def user_specific_configuration_path():
 
     This function returns the path to the user-specific configuration file
     for cellfinder. The user-specific configuration file is located in the
-    user's home directory under the ".cellfinder" folder and is named
-    "cellfinder.conf.custom".
+    user's home directory under the ".brainglobe/cellfinder" folder
+    and is named "cellfinder.conf.custom".
 
     Returns:
         Path: The path to the custom configuration file.
 
     """
-    return Path.home() / ".cellfinder" / "cellfinder.conf.custom"
+    return DEFAULT_CELLFINDER_DIRECTORY / "cellfinder.conf.custom"

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
@@ -31,12 +34,7 @@ from sklearn.model_selection import train_test_split
 import cellfinder.core as program_for_log
 from cellfinder.core import logger
 from cellfinder.core.classify.resnet import layer_type
-from cellfinder.core.tools.prep import DEFAULT_INSTALL_PATH
-
-tf_suppress_log_messages = [
-    "sample_weight modes were coerced from",
-    "multiprocessing can interact badly with TensorFlow",
-]
+from cellfinder.core.download.download import DEFAULT_DOWNLOAD_DIRECTORY
 
 depth_type = Literal["18", "34", "50", "101", "152"]
 
@@ -112,8 +110,7 @@ def misc_parse(parser):
 
 def training_parse():
     from cellfinder.core.download.cli import (
-        download_directory_parser,
-        model_parser,
+        download_parser,
     )
 
     training_parser = ArgumentParser(
@@ -223,8 +220,7 @@ def training_parse():
     )
 
     training_parser = misc_parse(training_parser)
-    training_parser = model_parser(training_parser)
-    training_parser = download_directory_parser(training_parser)
+    training_parser = download_parser(training_parser)
     args = training_parser.parse_args()
 
     return args
@@ -306,7 +302,7 @@ def run(
     n_free_cpus=2,
     trained_model=None,
     model_weights=None,
-    install_path=DEFAULT_INSTALL_PATH,
+    install_path=DEFAULT_DOWNLOAD_DIRECTORY,
     model="resnet50_tv",
     network_depth="50",
     learning_rate=0.0001,
@@ -320,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,
@@ -341,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)
@@ -363,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")
         (
@@ -389,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,
@@ -407,6 +401,8 @@ def run(
         shuffle=True,
         train=True,
         augment=not no_augment,
+        use_multiprocessing=False,
+        workers=n_processes,
     )
     callbacks = []
 
@@ -423,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,
@@ -434,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