Use numpy/dask array type

brainglobe · Mar 28, 2023 · d46f865 · d46f865
1 parent 2261036
commit d46f865
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/cellfinder_core/classify/classify.py b/src/cellfinder_core/classify/classify.py
@@ -4,16 +4,16 @@
 from imlib.general.system import get_num_processes
 from tensorflow import keras
 
-from cellfinder_core import logger
+from cellfinder_core import logger, types
 from cellfinder_core.classify.cube_generator import CubeGeneratorFromFile
 from cellfinder_core.classify.tools import get_model
 from cellfinder_core.train.train_yml import models
 
 
 def main(
     points,
-    signal_array: np.ndarray,
-    background_array: np.ndarray,
+    signal_array: types.array,
+    background_array: types.array,
     n_free_cpus: int,
     voxel_sizes,
     network_voxel_sizes,

diff --git a/src/cellfinder_core/types.py b/src/cellfinder_core/types.py
@@ -0,0 +1,6 @@
+from typing import Union
+
+import dask.array as da
+import numpy as np
+
+array = Union[da.Array, np.ndarray]