From d46f865318b6e7c4b65148a5bb072a22de004324 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Tue, 28 Mar 2023 14:39:55 +0100
Subject: [PATCH] Use numpy/dask array type

---
 src/cellfinder_core/classify/classify.py | 6 +++---
 src/cellfinder_core/types.py             | 6 ++++++
 2 files changed, 9 insertions(+), 3 deletions(-)
 create mode 100644 src/cellfinder_core/types.py

diff --git a/src/cellfinder_core/classify/classify.py b/src/cellfinder_core/classify/classify.py
index 4697fed2..044b4152 100644
--- a/src/cellfinder_core/classify/classify.py
+++ b/src/cellfinder_core/classify/classify.py
@@ -4,7 +4,7 @@
 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
@@ -12,8 +12,8 @@
 
 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
new file mode 100644
index 00000000..dc72b2d6
--- /dev/null
+++ 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]