cellfinder 1.3.3__py3-none-any.whl → 1.4.0a0__py3-none-any.whl

cellfinder/core/detect/filters/volume/structure_detection.py

@@ -4,18 +4,20 @@ from typing import Dict, Optional, Tuple, TypeVar, Union
 import numba.typed
 import numpy as np
 import numpy.typing as npt
-from numba import njit, typed
+from numba import njit, objmode, typed
 from numba.core import types
 from numba.experimental import jitclass
 from numba.types import DictType
 
+from cellfinder.core.tools.tools import get_max_possible_int_value
+
 T = TypeVar("T")
 # type used for the domain of the volume - the size of the vol
 vol_np_type = np.int64
 vol_numba_type = types.int64
 # type used for the structure id
-sid_np_type = np.int64
-sid_numba_type = types.int64
+sid_np_type = np.uint64
+sid_numba_type = types.uint64
 
 
 @dataclass
@@ -26,14 +28,17 @@ class Point:
 
 
 @njit
-def get_non_zero_dtype_min(values: np.ndarray) -> int:
+def get_non_zero_dtype_min(values: np.ndarray) -> sid_numba_type:
     """
     Get the minimum of non-zero entries in *values*.
 
     If all entries are zero, returns maximum storeable number
     in the values array.
     """
-    min_val = np.iinfo(values.dtype).max
+    # we don't know how big the int is, so make it as large as possible (64)
+    with objmode(min_val="u8"):
+        min_val = get_max_possible_int_value(values.dtype)
+
     for v in values:
         if v != 0 and v < min_val:
             min_val = v
@@ -97,6 +102,7 @@ list_of_points_type = types.ListType(tuple_point_type)
 spec = [
     ("z", vol_numba_type),
     ("next_structure_id", sid_numba_type),
+    ("soma_centre_value", sid_numba_type),  # as large as possible
     ("shape", types.UniTuple(vol_numba_type, 2)),
     ("obsolete_ids", DictType(sid_numba_type, sid_numba_type)),
     ("coords_maps", DictType(sid_numba_type, list_of_points_type)),
@@ -133,18 +139,25 @@ class CellDetector:
         points.
     """
 
-    def __init__(self, width: int, height: int, start_z: int):
+    def __init__(
+        self,
+        height: int,
+        width: int,
+        start_z: int,
+        soma_centre_value: sid_numba_type,
+    ):
         """
         Parameters
         ----------
-        width, height
+        height, width:
             Shape of the planes input to self.process()
         start_z:
             The z-coordinate of the first processed plane.
         """
-        self.shape = width, height
+        self.shape = height, width
         self.z = start_z
         self.next_structure_id = 1
+        self.soma_centre_value = soma_centre_value
 
         # Mapping from obsolete IDs to the IDs that they have been
         # made obsolete by
@@ -156,11 +169,18 @@ class CellDetector:
             key_type=sid_numba_type, value_type=list_of_points_type
         )
 
+    def _set_soma(self, soma_centre_value: sid_numba_type):
+        # Due to https://github.com/numba/numba/issues/9576. For testing we try
+        # different data types. Because of that issue we cannot pass a uint64
+        # soma_centre_value to constructor after we pass a uint32. This is the
+        # only way for now until numba fixes the issue
+        self.soma_centre_value = soma_centre_value
+
     def process(
         self, plane: np.ndarray, previous_plane: Optional[np.ndarray]
     ) -> np.ndarray:
         """
-        Process a new plane.
+        Process a new plane (should be in Y, X axis order).
         """
         if plane.shape[:2] != self.shape:
             raise ValueError("plane does not have correct shape")
@@ -185,21 +205,21 @@ class CellDetector:
         -------
         plane :
             Plane with pixels either set to zero (no structure) or labelled
-            with their structure ID.
+            with their structure ID. Plane is in Y, X axis order.
         """
-        SOMA_CENTRE_VALUE = np.iinfo(plane.dtype).max
-        for y in range(plane.shape[1]):
-            for x in range(plane.shape[0]):
-                if plane[x, y] == SOMA_CENTRE_VALUE:
+        soma_centre_value = self.soma_centre_value
+        for y in range(plane.shape[0]):
+            for x in range(plane.shape[1]):
+                if plane[y, x] == soma_centre_value:
                     # Labels of structures below, left and behind
                     neighbour_ids = np.zeros(3, dtype=sid_np_type)
                     # If in bounds look at neighbours
-                    if x > 0:
-                        neighbour_ids[0] = plane[x - 1, y]
                     if y > 0:
-                        neighbour_ids[1] = plane[x, y - 1]
+                        neighbour_ids[0] = plane[y - 1, x]
+                    if x > 0:
+                        neighbour_ids[1] = plane[y, x - 1]
                     if previous_plane is not None:
-                        neighbour_ids[2] = previous_plane[x, y]
+                        neighbour_ids[2] = previous_plane[y, x]
 
                     if is_new_structure(neighbour_ids):
                         neighbour_ids[0] = self.next_structure_id
@@ -210,17 +230,20 @@ class CellDetector:
                     # structure in next iterations
                     struct_id = 0
 
-                plane[x, y] = struct_id
+                plane[y, x] = struct_id
 
         return plane
 
     def get_cell_centres(self) -> np.ndarray:
+        """
+        Returns the 2D array of cell centers. It's (N, 3) with X, Y, Z columns.
+        """
         return self.structures_to_cells()
 
     def get_structures(self) -> Dict[int, np.ndarray]:
         """
         Gets the structures as a dict of structure IDs mapped to the 2D array
-        of structure points.
+        of structure points (points vs x, y, z columns).
         """
         d = {}
         for sid, points in self.coords_maps.items():
@@ -228,7 +251,10 @@ class CellDetector:
             # `item = np.array(points, dtype=vol_np_type)` so we need to create
             # array and then fill in the point
             item = np.empty((len(points), 3), dtype=vol_np_type)
-            d[sid] = item
+            # need to cast to int64, otherwise when dict is used we can get
+            # warnings as in numba issue #8829 b/c it assumes it's uint64.
+            # Python uses int(64) as the type
+            d[types.int64(sid)] = item
 
             for i, point in enumerate(points):
                 item[i, :] = point
@@ -239,33 +265,41 @@ class CellDetector:
         self, sid: int, point: Union[tuple, list, np.ndarray]
     ) -> None:
         """
-        Add single 3d *point* to the structure with the given *sid*.
+        Add single 3d (x, y, z) *point* to the structure with the given *sid*.
         """
-        if sid not in self.coords_maps:
-            self.coords_maps[sid] = typed.List.empty_list(tuple_point_type)
+        # cast in case user passes in int64 (default type for int in python)
+        # and numba complains
+        key = sid_numba_type(sid)
+        if key not in self.coords_maps:
+            self.coords_maps[key] = typed.List.empty_list(tuple_point_type)
 
-        self._add_point(sid, (int(point[0]), int(point[1]), int(point[2])))
+        self._add_point(key, (int(point[0]), int(point[1]), int(point[2])))
 
     def add_points(self, sid: int, points: np.ndarray):
         """
         Adds ndarray of *points* to the structure with the given *sid*.
-        Each row is a 3d point.
+        Each row is a 3-column (x, y, z) point.
         """
-        if sid not in self.coords_maps:
-            self.coords_maps[sid] = typed.List.empty_list(tuple_point_type)
+        # cast in case user passes in int64 (default type for int in python)
+        # and numba complains
+        key = sid_numba_type(sid)
+        if key not in self.coords_maps:
+            self.coords_maps[key] = typed.List.empty_list(tuple_point_type)
 
-        append = self.coords_maps[sid].append
+        append = self.coords_maps[key].append
         pts = np.round(points).astype(vol_np_type)
         for point in pts:
             append((point[0], point[1], point[2]))
 
-    def _add_point(self, sid: int, point: Tuple[int, int, int]) -> None:
+    def _add_point(
+        self, sid: sid_numba_type, point: Tuple[int, int, int]
+    ) -> None:
         # sid must exist
         self.coords_maps[sid].append(point)
 
     def add(
         self, x: int, y: int, z: int, neighbour_ids: npt.NDArray[sid_np_type]
-    ) -> int:
+    ) -> sid_numba_type:
         """
         For the current coordinates takes all the neighbours and find the
         minimum structure including obsolete structures mapping to any of
@@ -287,7 +321,9 @@ class CellDetector:
         self._add_point(updated_id, (int(x), int(y), int(z)))
         return updated_id
 
-    def sanitise_ids(self, neighbour_ids: npt.NDArray[sid_np_type]) -> int:
+    def sanitise_ids(
+        self, neighbour_ids: npt.NDArray[sid_np_type]
+    ) -> sid_numba_type:
         """
         Get the smallest ID of all the structures that are connected to IDs
         in `neighbour_ids`.
@@ -300,15 +336,17 @@ class CellDetector:
         """
         for i, neighbour_id in enumerate(neighbour_ids):
             # walk up the chain of obsolescence
-            neighbour_id = int(traverse_dict(self.obsolete_ids, neighbour_id))
+            neighbour_id = traverse_dict(self.obsolete_ids, neighbour_id)
             neighbour_ids[i] = neighbour_id
 
         # Get minimum of all non-obsolete IDs
         updated_id = get_non_zero_dtype_min(neighbour_ids)
-        return int(updated_id)
+        return updated_id
 
     def merge_structures(
-        self, updated_id: int, neighbour_ids: npt.NDArray[sid_np_type]
+        self,
+        updated_id: sid_numba_type,
+        neighbour_ids: npt.NDArray[sid_np_type],
     ) -> None:
         """
         For all the neighbours, reassign all the points of neighbour to

cellfinder/core/detect/filters/volume/structure_splitting.py

@@ -1,9 +1,14 @@
-from typing import List, Tuple
+from typing import List, Tuple, Type
 
 import numpy as np
+import torch
 
 from cellfinder.core import logger
-from cellfinder.core.detect.filters.volume.ball_filter import BallFilter
+from cellfinder.core.detect.filters.setup_filters import DetectionSettings
+from cellfinder.core.detect.filters.volume.ball_filter import (
+    BallFilter,
+    InvalidVolume,
+)
 from cellfinder.core.detect.filters.volume.structure_detection import (
     CellDetector,
     get_structure_centre,
@@ -14,41 +19,62 @@ class StructureSplitException(Exception):
     pass
 
 
-def get_shape(xs: np.ndarray, ys: np.ndarray, zs: np.ndarray) -> List[int]:
+def get_shape(
+    xs: np.ndarray, ys: np.ndarray, zs: np.ndarray
+) -> Tuple[int, int, int]:
+    """
+    Takes a list of x, y, z coordinates and returns a volume size such that
+    all the points will fit into it. With axis order = x, y, z.
+    """
     # +1 because difference. TEST:
-    shape = [int((dim.max() - dim.min()) + 1) for dim in (xs, ys, zs)]
+    shape = tuple(int((dim.max() - dim.min()) + 1) for dim in (xs, ys, zs))
     return shape
 
 
 def coords_to_volume(
-    xs: np.ndarray, ys: np.ndarray, zs: np.ndarray, ball_radius: int = 1
-) -> np.ndarray:
+    xs: np.ndarray,
+    ys: np.ndarray,
+    zs: np.ndarray,
+    volume_shape: Tuple[int, int, int],
+    ball_radius: int,
+    dtype: Type[np.number],
+    threshold_value: int,
+) -> torch.Tensor:
+    """
+    Takes the series of x, y, z points along with the shape of the volume
+    that fully enclose them (also x, y, z order). It than expands the
+    shape by the ball diameter in each axis. Then, each point, shifted
+    by the radius internally is set to the threshold value.
+
+    The volume is then transposed and returned in the Z, Y, X order.
+    """
+    # it's faster doing the work in numpy and then returning as torch array,
+    # than doing the work in torch
     ball_diameter = ball_radius * 2
     # Expanded to ensure the ball fits even at the border
-    expanded_shape = [
-        dim_size + ball_diameter for dim_size in get_shape(xs, ys, zs)
-    ]
-    volume = np.zeros(expanded_shape, dtype=np.uint32)
+    expanded_shape = [dim_size + ball_diameter for dim_size in volume_shape]
+    # volume is now x, y, z order
+    volume = np.zeros(expanded_shape, dtype=dtype)
 
     x_min, y_min, z_min = xs.min(), ys.min(), zs.min()
 
+    # shift the points so any sphere centered on it would not have its
+    # radius expand beyond the volume
     relative_xs = np.array((xs - x_min + ball_radius), dtype=np.int64)
     relative_ys = np.array((ys - y_min + ball_radius), dtype=np.int64)
     relative_zs = np.array((zs - z_min + ball_radius), dtype=np.int64)
 
-    # OPTIMISE: vectorize
+    # set each point as the center with a value of threshold
     for rel_x, rel_y, rel_z in zip(relative_xs, relative_ys, relative_zs):
-        volume[rel_x, rel_y, rel_z] = np.iinfo(volume.dtype).max - 1
-    return volume
+        volume[rel_x, rel_y, rel_z] = threshold_value
+
+    volume = volume.swapaxes(0, 2)
+    return torch.from_numpy(volume)
 
 
 def ball_filter_imgs(
-    volume: np.ndarray,
-    threshold_value: int,
-    soma_centre_value: int,
-    ball_xy_size: int = 3,
-    ball_z_size: int = 3,
-) -> Tuple[np.ndarray, np.ndarray]:
+    volume: torch.Tensor, settings: DetectionSettings
+) -> np.ndarray:
     """
     Apply ball filtering to a 3D volume and detect cell centres.
 
@@ -56,105 +82,118 @@ def ball_filter_imgs(
     and the `CellDetector` class to detect cell centres.
 
     Args:
-        volume (np.ndarray): The 3D volume to be filtered.
-        threshold_value (int): The threshold value for ball filtering.
-        soma_centre_value (int): The value representing the soma centre.
-        ball_xy_size (int, optional):
-            The size of the ball filter in the XY plane. Defaults to 3.
-        ball_z_size (int, optional):
-            The size of the ball filter in the Z plane. Defaults to 3.
+        volume (torch.Tensor): The 3D volume to be filtered (Z, Y, X order).
+        settings (DetectionSettings):
+            The settings to use.
 
     Returns:
-        Tuple[np.ndarray, np.ndarray]:
-            A tuple containing the filtered volume and the cell centres.
+        The 2D array of cell centres (N, 3) - X, Y, Z order.
 
     """
-    # OPTIMISE: reuse ball filter instance
-
-    good_tiles_mask = np.ones((1, 1, volume.shape[2]), dtype=np.bool_)
-
-    plane_width, plane_height = volume.shape[:2]
-    current_z = ball_z_size // 2
-
-    bf = BallFilter(
-        plane_width,
-        plane_height,
-        ball_xy_size,
-        ball_z_size,
-        overlap_fraction=0.8,
-        tile_step_width=plane_width,
-        tile_step_height=plane_height,
-        threshold_value=threshold_value,
-        soma_centre_value=soma_centre_value,
+    detection_convert = settings.detection_data_converter_func
+    batch_size = settings.batch_size
+
+    # make sure volume is not less than kernel etc
+    try:
+        bf = BallFilter(
+            plane_height=settings.plane_height,
+            plane_width=settings.plane_width,
+            ball_xy_size=settings.ball_xy_size,
+            ball_z_size=settings.ball_z_size,
+            overlap_fraction=settings.ball_overlap_fraction,
+            threshold_value=settings.threshold_value,
+            soma_centre_value=settings.soma_centre_value,
+            tile_height=settings.tile_height,
+            tile_width=settings.tile_width,
+            dtype=settings.filtering_dtype.__name__,
+            batch_size=batch_size,
+            torch_device=settings.torch_device,
+            use_mask=False,  # we don't need a mask here
+        )
+    except InvalidVolume:
+        return np.empty((0, 3))
+
+    start_z = bf.first_valid_plane
+    cell_detector = CellDetector(
+        settings.plane_height,
+        settings.plane_width,
+        start_z=start_z,
+        soma_centre_value=settings.detection_soma_centre_value,
     )
-    cell_detector = CellDetector(plane_width, plane_height, start_z=current_z)
 
-    # FIXME: hard coded type
-    ball_filtered_volume = np.zeros(volume.shape, dtype=np.uint32)
     previous_plane = None
-    for z in range(volume.shape[2]):
-        bf.append(volume[:, :, z].astype(np.uint32), good_tiles_mask[:, :, z])
+    for z in range(0, volume.shape[0], batch_size):
+        bf.append(volume[z : z + batch_size, :, :])
+
         if bf.ready:
             bf.walk()
-            middle_plane = bf.get_middle_plane()
 
-            # first valid middle plane is the current_z, not z
-            ball_filtered_volume[:, :, current_z] = middle_plane[:]
-            current_z += 1
+            middle_planes = bf.get_processed_planes()
+            n = middle_planes.shape[0]
 
-            # DEBUG: TEST: transpose
-            previous_plane = cell_detector.process(
-                middle_plane.copy(), previous_plane
+            # we edit volume, but only for planes already processed that won't
+            # be passed to the filter in this run
+            volume[start_z : start_z + n, :, :] = torch.from_numpy(
+                middle_planes
             )
-    return ball_filtered_volume, cell_detector.get_cell_centres()
+            start_z += n
+
+            # convert to type needed for detection
+            middle_planes = detection_convert(middle_planes)
+            for plane in middle_planes:
+                previous_plane = cell_detector.process(plane, previous_plane)
+
+    return cell_detector.get_cell_centres()
 
 
 def iterative_ball_filter(
-    volume: np.ndarray, n_iter: int = 10
+    volume: torch.Tensor, settings: DetectionSettings
 ) -> Tuple[List[int], List[np.ndarray]]:
     """
     Apply iterative ball filtering to the given volume.
     The volume is eroded at each iteration, by subtracting 1 from the volume.
 
     Parameters:
-        volume (np.ndarray): The input volume.
-        n_iter (int): The number of iterations to perform. Default is 10.
+        volume (torch.Tensor): The input volume. It is edited inplace.
+            Of shape Z, Y, X.
+        settings (DetectionSettings): The settings to use.
 
     Returns:
-        Tuple[List[int], List[np.ndarray]]: A tuple containing two lists:
-            The structures found in each iteration.
+        tuple: A tuple containing two lists:
+            The number of structures found in each iteration.
             The cell centres found in each iteration.
     """
     ns = []
     centres = []
 
-    threshold_value = np.iinfo(volume.dtype).max - 1
-    soma_centre_value = np.iinfo(volume.dtype).max
-
-    vol = volume.copy()  # TODO: check if required
-
-    for i in range(n_iter):
-        vol, cell_centres = ball_filter_imgs(
-            vol, threshold_value, soma_centre_value
-        )
-
-        # vol is unsigned, so can't let zeros underflow to max value
-        vol[:, :, :] = np.where(vol != 0, vol - 1, 0)
+    for i in range(settings.n_splitting_iter):
+        cell_centres = ball_filter_imgs(volume, settings)
+        volume.sub_(1)
 
         n_structures = len(cell_centres)
         ns.append(n_structures)
         centres.append(cell_centres)
         if n_structures == 0:
             break
+
     return ns, centres
 
 
 def check_centre_in_cuboid(centre: np.ndarray, max_coords: np.ndarray) -> bool:
     """
-    Checks whether a coordinate is in a cuboid
-    :param centre: x,y,z coordinate
-    :param max_coords: far corner of cuboid
-    :return: True if within cuboid, otherwise False
+    Checks whether a coordinate is in a cuboid.
+
+    Parameters
+    ----------
+
+    centre : np.ndarray
+        x, y, z coordinate.
+    max_coords : np.ndarray
+        Far corner of cuboid.
+
+    Returns
+    -------
+    True if within cuboid, otherwise False.
     """
     relative_coords = centre
     if (relative_coords > max_coords).all():
@@ -168,7 +207,7 @@ def check_centre_in_cuboid(centre: np.ndarray, max_coords: np.ndarray) -> bool:
 
 
 def split_cells(
-    cell_points: np.ndarray, outlier_keep: bool = False
+    cell_points: np.ndarray, settings: DetectionSettings
 ) -> np.ndarray:
     """
     Split the given cell points into individual cell centres.
@@ -177,28 +216,24 @@ def split_cells(
         cell_points (np.ndarray): Array of cell points with shape (N, 3),
             where N is the number of cell points and each point is represented
             by its x, y, and z coordinates.
-        outlier_keep (bool, optional): Flag indicating whether to keep outliers
-            during the splitting process. Defaults to False.
+        settings (DetectionSettings) : The settings to use for splitting. It is
+            modified inplace.
 
     Returns:
         np.ndarray: Array of absolute cell centres with shape (M, 3),
             where M is the number of individual cells and each centre is
             represented by its x, y, and z coordinates.
     """
+    # these points are in x, y, z order columnwise, in absolute pixels
     orig_centre = get_structure_centre(cell_points)
 
     xs = cell_points[:, 0]
     ys = cell_points[:, 1]
     zs = cell_points[:, 2]
 
-    orig_corner = np.array(
-        [
-            orig_centre[0] - (orig_centre[0] - xs.min()),
-            orig_centre[1] - (orig_centre[1] - ys.min()),
-            orig_centre[2] - (orig_centre[2] - zs.min()),
-        ]
-    )
-
+    # corner coordinates in absolute pixels
+    orig_corner = np.array([xs.min(), ys.min(), zs.min()])
+    # volume center relative to corner
     relative_orig_centre = np.array(
         [
             orig_centre[0] - orig_corner[0],
@@ -207,22 +242,51 @@ def split_cells(
         ]
     )
 
+    # total volume enclosing all points
     original_bounding_cuboid_shape = get_shape(xs, ys, zs)
 
-    ball_radius = 1
-    vol = coords_to_volume(xs, ys, zs, ball_radius=ball_radius)
+    ball_radius = settings.ball_xy_size // 2
+    # they should be the same dtype so as to not need a conversion before
+    # passing the input data with marked cells to the filters (we currently
+    # set both to float32)
+    assert settings.filtering_dtype == settings.plane_original_np_dtype
+    # volume will now be z, y, x order
+    vol = coords_to_volume(
+        xs,
+        ys,
+        zs,
+        volume_shape=original_bounding_cuboid_shape,
+        ball_radius=ball_radius,
+        dtype=settings.filtering_dtype,
+        threshold_value=settings.threshold_value,
+    )
+
+    # get an estimate of how much memory processing a single batch of original
+    # input planes takes. For this much smaller volume, our batch will be such
+    # that it uses at most that much memory
+    total_vol_size = (
+        settings.plane_height * settings.plane_width * settings.batch_size
+    )
+    batch_size = total_vol_size // (vol.shape[1] * vol.shape[2])
+    batch_size = min(batch_size, vol.shape[0])
+
+    # update settings with our volume data
+    settings.plane_shape = vol.shape[1:]
+    settings.start_plane = 0
+    settings.end_plane = vol.shape[0]
+    settings.batch_size = batch_size
 
     # centres is a list of arrays of centres (1 array of centres per ball run)
-    ns, centres = iterative_ball_filter(vol)
+    # in x, y, z order
+    ns, centres = iterative_ball_filter(vol, settings)
     ns.insert(0, 1)
     centres.insert(0, np.array([relative_orig_centre]))
 
     best_iteration = ns.index(max(ns))
-
     # TODO: put constraint on minimum centres distance ?
     relative_centres = centres[best_iteration]
 
-    if not outlier_keep:
+    if not settings.outlier_keep:
         # TODO: change to checking whether in original cluster shape
         original_max_coords = np.array(original_bounding_cuboid_shape)
         relative_centres = np.array(
@@ -234,7 +298,7 @@ def split_cells(
         )
 
     absolute_centres = np.empty((len(relative_centres), 3))
-    # FIXME: extract functionality
+    # convert centers to absolute pixels
     absolute_centres[:, 0] = orig_corner[0] + relative_centres[:, 0]
     absolute_centres[:, 1] = orig_corner[1] + relative_centres[:, 1]
     absolute_centres[:, 2] = orig_corner[2] + relative_centres[:, 2]