Add 2D input support for detection with clear guardrails

cellfinder/core/classify/classify.py

@@ -19,7 +19,7 @@
 from cellfinder.core.train.train_yaml import depth_type, models
 
 
 def main(
     points: List[Cell],
     signal_array: types.array,
     background_array: types.array,
@@ -90,9 +90,19 @@
         A callback function that is called during classification. Called with
         the batch number once that batch has been classified.
     """
-    if signal_array.ndim != 3:
-        raise IOError("Signal data must be 3D")
-    if background_array.ndim != 3:
+    if signal_array.ndim == 2:
+        if background_array.ndim != 2:
+            raise IOError("For 2D signal data, background must also be 2D")
+        if cube_depth != 1:
+            raise ValueError(
+                "For 2D data, cube_depth must be 1. "
+                "Train/use a 2D classifier or set skip_classification."
+            )
+        signal_array = signal_array[None, ...]
+        background_array = background_array[None, ...]
+    elif signal_array.ndim != 3:
+        raise IOError("Signal data must be 2D or 3D")
+    elif background_array.ndim != 3:
         raise IOError("Background data must be 3D")
 
     # Too many workers doesn't increase speed, and uses huge amounts of RAM
@@ -101,6 +111,8 @@
 
     start_time = datetime.now()
 
+    if len(voxel_sizes) == 2:
+        voxel_sizes = (1.0, *voxel_sizes)
     voxel_sizes = list(map(float, voxel_sizes))
     logger.debug("Initialising cube generator")
     dataset = CuboidArrayDataset(

cellfinder/core/detect/detect.py

@@ -29,7 +29,7 @@
 
 
 @inference_wrapper
 def main(
     signal_array: types.array,
     start_plane: int = 0,
     end_plane: int = -1,
@@ -66,6 +66,7 @@
     ----------
     signal_array : numpy.ndarray or dask array
         3D array representing the signal data in z, y, x order.
+        2D arrays are also accepted and treated as a single z-plane.
     start_plane : int
         First plane index to process (inclusive, to process a subset of the
         data).
@@ -176,8 +177,23 @@
             f"{signal_array.dtype}"
         )
 
-    if signal_array.ndim != 3:
-        raise ValueError("Input data must be 3D")
+    is_2d = False
+    if signal_array.ndim == 2:
+        is_2d = True
+        signal_array = signal_array[None, ...]
+    elif signal_array.ndim != 3:
+        raise ValueError("Input data must be 2D or 3D")
+
+    if is_2d:
+        if start_plane != 0:
+            raise ValueError("For 2D data, start_plane must be 0.")
+        if end_plane not in (-1, 1):
+            raise ValueError(
+                "For 2D data, end_plane must be 1 (or -1 for default)."
+            )
+        end_plane = 1
+        ball_z_size = 1
+        split_ball_z_size = 1
 
     if end_plane < 0:
         end_plane = len(signal_array)
@@ -192,6 +208,8 @@
 
     batch_size = max(batch_size, 1)
     # brainmapper can pass them in as str
+    if len(voxel_sizes) == 2:
+        voxel_sizes = (1.0, *voxel_sizes)
     voxel_sizes = list(map(float, voxel_sizes))
 
     settings = DetectionSettings(

cellfinder/core/main.py

@@ -8,7 +8,7 @@
 from cellfinder.core.train.train_yaml import depth_type
 
 
 def main(
     signal_array: types.array,
     background_array: types.array,
     voxel_sizes: Tuple[float, float, float],
@@ -55,10 +55,13 @@
     ----------
     signal_array : numpy.ndarray or dask array
         3D array representing the signal data in z, y, x order.
+        2D arrays are also accepted and treated as a single z-plane.
     background_array : numpy.ndarray or dask array
         3D array representing the signal data in z, y, x order.
+        2D arrays are also accepted and treated as a single z-plane.
     voxel_sizes : 3-tuple of floats
         Size of your voxels in the z, y, and x dimensions (microns).
+        For 2D inputs, you may pass a 2-tuple (y, x); z is assumed to be 1.0.
     start_plane : int
         First plane index to process (inclusive, to process a subset of the
         data).
@@ -188,6 +191,20 @@
     from cellfinder.core.detect import detect
     from cellfinder.core.tools import prep
 
+    if signal_array.ndim == 2:
+        if background_array.ndim != 2:
+            raise ValueError(
+                "For 2D signal inputs, background_array must also be 2D."
+            )
+        if not skip_classification:
+            raise ValueError(
+                "2D classification is not supported in this pipeline yet. "
+                "Set skip_classification=True to run detection only, or "
+                "train a dedicated 2D classifier and wire it in."
+            )
+        if len(voxel_sizes) == 2:
+            voxel_sizes = (1.0, *voxel_sizes)
+
     if not skip_detection:
         logger.info("Detecting cell candidates")