From c38399a34deb848214c9cddf1306f1acb0a6ae6e Mon Sep 17 00:00:00 2001
From: Adam Tyson <code@adamltyson.com>
Date: Fri, 5 Jan 2024 09:28:49 +0000
Subject: [PATCH] Prep merge (#349)

* Remove unnecessary code duplication

* Remove redundant parameter

* Make brainmapper check optional in this repo for releases and status checks on PRs

---------

Co-authored-by: willGraham01 <1willgraham@gmail.com>
---
 .github/workflows/test_and_deploy.yml |  2 +-
 benchmarks/benchmarks/tools/prep.py   | 16 ++--------------
 cellfinder/core/main.py               |  4 ++--
 cellfinder/core/tools/prep.py         | 23 ++---------------------
 cellfinder/core/train/train_yml.py    | 10 +++-------
 5 files changed, 10 insertions(+), 45 deletions(-)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index e620eb5b..79e765b4 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -127,7 +127,7 @@ jobs:
 
   build_sdist_wheel:
     name: Build source distribution and wheel
-    needs: [test, test_brainmapper_cli, test_numba_disabled]
+    needs: [test, test_numba_disabled]
     if: github.event_name == 'push' && github.ref_type == 'tag'
     runs-on: ubuntu-latest
     steps:
diff --git a/benchmarks/benchmarks/tools/prep.py b/benchmarks/benchmarks/tools/prep.py
index 27d9620b..0e2bbce4 100644
--- a/benchmarks/benchmarks/tools/prep.py
+++ b/benchmarks/benchmarks/tools/prep.py
@@ -4,10 +4,9 @@
 from brainglobe_utils.general.system import get_num_processes
 
 from cellfinder.core.tools.prep import (
-    prep_classification,
+    prep_model_weights,
     prep_models,
     prep_tensorflow,
-    prep_training,
 )
 
 
@@ -43,30 +42,19 @@ def teardown(self, model_name):
 
     def time_prep_models(self, model_name):
         prep_models(
-            self.trained_model,
             self.model_weights,
             self.install_path,
             model_name,
         )
 
     def time_prep_classification(self, model_name):
-        prep_classification(
-            self.trained_model,
+        prep_model_weights(
             self.model_weights,
             self.install_path,
             model_name,
             self.n_free_cpus,
         )
 
-    def time_prep_training(self, model_name):
-        prep_training(
-            self.n_free_cpus,
-            self.trained_model,
-            self.model_weights,
-            self.install_path,
-            model_name,
-        )
-
 
 class PrepTF:
     def setup(self):
diff --git a/cellfinder/core/main.py b/cellfinder/core/main.py
index 4ad395e2..bc2f0cf3 100644
--- a/cellfinder/core/main.py
+++ b/cellfinder/core/main.py
@@ -87,8 +87,8 @@ def main(
         detect_finished_callback(points)
 
     install_path = None
-    model_weights = prep.prep_classification(
-        trained_model, model_weights, install_path, model, n_free_cpus
+    model_weights = prep.prep_model_weights(
+        model_weights, install_path, model, n_free_cpus
     )
     if len(points) > 0:
         logger.info("Running classification")
diff --git a/cellfinder/core/tools/prep.py b/cellfinder/core/tools/prep.py
index dc2eda54..0e7e7217 100644
--- a/cellfinder/core/tools/prep.py
+++ b/cellfinder/core/tools/prep.py
@@ -20,8 +20,7 @@
 DEFAULT_INSTALL_PATH = home / ".cellfinder"
 
 
-def prep_classification(
-    trained_model: Optional[os.PathLike],
+def prep_model_weights(
     model_weights: Optional[os.PathLike],
     install_path: Optional[os.PathLike],
     model_name: model_download.model_type,
@@ -29,35 +28,17 @@ def prep_classification(
 ) -> Path:
     n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
     prep_tensorflow(n_processes)
-    model_weights = prep_models(
-        trained_model, model_weights, install_path, model_name
-    )
+    model_weights = prep_models(model_weights, install_path, model_name)
 
     return model_weights
 
 
-def prep_training(
-    n_free_cpus: int,
-    trained_model: Optional[os.PathLike],
-    model_weights: Optional[os.PathLike],
-    install_path: Optional[os.PathLike],
-    model_name: model_download.model_type,
-) -> Path:
-    n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
-    prep_tensorflow(n_processes)
-    model_weights = prep_models(
-        trained_model, 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(
-    trained_model_path: Optional[os.PathLike],
     model_weights_path: Optional[os.PathLike],
     install_path: Optional[os.PathLike],
     model_name: model_download.model_type,
diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index 5596dcff..63e97e00 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -332,17 +332,13 @@ def run(
 
     from cellfinder.core.classify.cube_generator import CubeGeneratorFromDisk
     from cellfinder.core.classify.tools import get_model, make_lists
-    from cellfinder.core.tools.prep import prep_training
+    from cellfinder.core.tools.prep import prep_model_weights
 
     start_time = datetime.now()
 
     ensure_directory_exists(output_dir)
-    model_weights = prep_training(
-        n_free_cpus,
-        trained_model,
-        install_path,
-        model_weights,
-        model,
+    model_weights = prep_model_weights(
+        install_path, model_weights, model, n_free_cpus
     )
 
     yaml_contents = parse_yaml(yaml_file)