diff --git a/benchmarks/cellfinder_core.py b/benchmarks/cellfinder_core.py index 41821d94..8f94824b 100644 --- a/benchmarks/cellfinder_core.py +++ b/benchmarks/cellfinder_core.py @@ -2,7 +2,6 @@ import shutil from pathlib import Path -import pooch from brainglobe_utils.IO.cells import save_cells from cellfinder.core.main import main as cellfinder_run from cellfinder.core.tools.IO import read_with_dask @@ -81,18 +80,18 @@ class TimeBenchmarkPrepGIN: # Custom attributes input_config_path = str(DEFAULT_JSON_CONFIG_PATH_CELLFINDER) - def setup_cache( - self, - ): + def setup_cache(self): """ Download the input data from the GIN repository to the local - directory specified in the default_config.json + directory specified in the default_config.json. Notes ----- The `setup_cache` method only performs the computations once per benchmark round and then caches the result to disk [1]_. It cannot - be parametrised [2]_. + be parametrised [2]_. Therefore, if we sweep across different input + JSON files, we need to ensure all data for all configs is made + available with this setup function. [1] https://asv.readthedocs.io/en/latest/writing_benchmarks.html#setup-and-teardown-functions @@ -108,15 +107,6 @@ def setup_cache( config_dict = json.load(cfg) config = CellfinderConfig(**config_dict) - # Download data with pooch - _ = pooch.retrieve( - url=config.data_url, - known_hash=config.data_hash, - path=config._install_path, - progressbar=True, - processor=pooch.Unzip(extract_dir=config.data_dir_relative), - ) - # Check paths to input data should now exist in config assert Path(config._signal_dir_path).exists() assert Path(config._background_dir_path).exists() @@ -129,12 +119,7 @@ def setup(self): """ # Run setup - cfg = setup_cellfinder_workflow( - [ - "--config", - self.input_config_path, - ] - ) + cfg = setup_cellfinder_workflow(self.input_config_path) # Save configuration as attribute self.cfg = cfg @@ -162,7 +147,7 @@ class TimeFullWorkflow(TimeBenchmarkPrepGIN): A base class for timing benchmarks for the cellfinder workflow. """ - def time_workflow_from_cellfinder_run(self): + def time_workflow(self): run_workflow_from_cellfinder_run(self.cfg) @@ -177,10 +162,10 @@ class TimeReadInputDask(TimeBenchmarkPrepGIN): """ def time_read_signal_with_dask(self): - read_with_dask(self.cfg._signal_dir_path) + read_with_dask(str(self.cfg._signal_dir_path)) def time_read_background_with_dask(self): - read_with_dask(self.cfg._background_dir_path) + read_with_dask(str(self.cfg._background_dir_path)) class TimeDetectCells(TimeBenchmarkPrepGIN): @@ -198,13 +183,37 @@ def setup(self): # basic setup TimeBenchmarkPrepGIN.setup(self) - # add input data as arrays to config - self.signal_array = read_with_dask(self.cfg._signal_dir_path) - self.background_array = read_with_dask(self.cfg._background_dir_path) + # add input data as arrays to the config + self.signal_array = read_with_dask(str(self.cfg._signal_dir_path)) + self.background_array = read_with_dask( + str(self.cfg._background_dir_path) + ) def time_cellfinder_run(self): cellfinder_run( - self.signal_array, self.background_array, self.cfg.voxel_sizes + self.signal_array, + self.background_array, + self.cfg.voxel_sizes, + self.cfg.start_plane, + self.cfg.end_plane, + self.cfg.trained_model, + self.cfg.model_weights, + self.cfg.model, + self.cfg.batch_size, + self.cfg.n_free_cpus, + self.cfg.network_voxel_sizes, + self.cfg.soma_diameter, + self.cfg.ball_xy_size, + self.cfg.ball_z_size, + self.cfg.ball_overlap_fraction, + self.cfg.log_sigma_size, + self.cfg.n_sds_above_mean_thresh, + self.cfg.soma_spread_factor, + self.cfg.max_cluster_size, + self.cfg.cube_width, + self.cfg.cube_height, + self.cfg.cube_depth, + self.cfg.network_depth, ) @@ -215,12 +224,36 @@ def setup(self): TimeBenchmarkPrepGIN.setup(self) # add input data as arrays to config - self.signal_array = read_with_dask(self.cfg._signal_dir_path) - self.background_array = read_with_dask(self.cfg._background_dir_path) + self.signal_array = read_with_dask(str(self.cfg._signal_dir_path)) + self.background_array = read_with_dask( + str(self.cfg._background_dir_path) + ) # detect cells self.detected_cells = cellfinder_run( - self.signal_array, self.background_array, self.cfg.voxel_sizes + self.signal_array, + self.background_array, + self.cfg.voxel_sizes, + self.cfg.start_plane, + self.cfg.end_plane, + self.cfg.trained_model, + self.cfg.model_weights, + self.cfg.model, + self.cfg.batch_size, + self.cfg.n_free_cpus, + self.cfg.network_voxel_sizes, + self.cfg.soma_diameter, + self.cfg.ball_xy_size, + self.cfg.ball_z_size, + self.cfg.ball_overlap_fraction, + self.cfg.log_sigma_size, + self.cfg.n_sds_above_mean_thresh, + self.cfg.soma_spread_factor, + self.cfg.max_cluster_size, + self.cfg.cube_width, + self.cfg.cube_height, + self.cfg.cube_depth, + self.cfg.network_depth, ) def time_save_cells(self):