Merge branch 'main' into public-main

.github/workflows/test_pull_request.yml

@@ -3,12 +3,14 @@ name: test PR
 
 on:
   pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
     branches:
       - main
 
 jobs:
   test:
     name: ${{ matrix.platform }} ${{ matrix.python }} ${{ matrix.toxenv || matrix.backend }}
+    if: github.event.pull_request.draft == false
     runs-on: ${{ matrix.platform }}
     strategy:
       fail-fast: false

ultrack/core/database.py

@@ -1,5 +1,6 @@
 import enum
 import logging
+import pickle
 from pathlib import Path
 from typing import Any, List, Union
 
@@ -39,7 +40,11 @@ def bind_processor(self, dialect):
         def _process(value):
             if isinstance(value, (bytes, memoryview)):
                 return value
-            return processor(value)
+            try:
+                return processor(value)
+            except pickle.UnpicklingError:
+                # for some reason, when converting database it has a few extra bytes
+                return processor(bytes.fromhex(value[3:].decode("utf-8")))
 
         return _process
 

ultrack/core/segmentation/hierarchy.py

@@ -63,7 +63,12 @@ def create_hierarchies(
 
     if "min_area" in kwargs and num_labels > 1:
         LOG.info("Filtering small connected components.")
-        morphology.remove_small_objects(labels, min_size=kwargs["min_area"], out=labels)
+        #  To avoid removing small objects, divide by 4 to remove the smallest ones.
+        #  Nodes in hierarchies are still filtered by minimum area.
+        #  This is mainly for lonely cells.
+        morphology.remove_small_objects(
+            labels, min_size=int(kwargs["min_area"] / 4), out=labels
+        )
 
     edge = np.asarray(edge)
 

ultrack/core/solve/solver/_test/test_solvers.py

@@ -39,7 +39,7 @@ def test_solvers_optimize(solver: BaseSolver, config_instance: MainConfig) -> No
 
     1 - 0.5 - 2 - 0.5 - 3 - 0.5 - 4
               |  \\             /      \\ due linting software
-              C   1.0      0.95
+              C   1.0      0.9
               |       \\  /
               5 - 0.5 - 6 - 0.7 - 7
 
@@ -102,7 +102,7 @@ def test_fixed_nodes_constraint_solver(config_instance: MainConfig) -> None:
 
     1 - 0.5 - 2 - 0.5 - 3 - 0.5 - 4
               |  \\             /      \\ due linting software
-              C   1.0      0.95
+              C   1.0      0.9
               |       \\  /
               5 - 0.5 - 6 - 0.7 - 7
                         ^
@@ -113,11 +113,11 @@ def test_fixed_nodes_constraint_solver(config_instance: MainConfig) -> None:
 
     1 - 0.5 - 2                   4
                 \\              /
-                  1.0      0.95
+                  1.0      0.9
                      \\  /
                        6 - 0.7 - 7
 
-    Result: 0.5 + 1.0 + 0.7 + 0.95 - division_weight
+    Result: 0.5 + 1.0 + 0.7 + 0.9 - division_weight
     """
     solver = MIPSolver(config_instance.tracking_config)
 
@@ -170,7 +170,7 @@ def test_fixed_edges_constraint_solver(config_instance: MainConfig) -> None:
 
     1 - 0.5 - 2 - 0.5 - 3 - 0.5 - 4
               |  \\             /      \\ due linting software
-              C   1.0      0.95
+              C   1.0      0.9
               |       \\  /
               5 - 0.5 - 6 - 0.7 - 7
                              ^
@@ -181,11 +181,11 @@ def test_fixed_edges_constraint_solver(config_instance: MainConfig) -> None:
 
     1 - 0.5 - 2                   4
                 \\              /
-                  1.0      0.95
+                  1.0      0.9
                      \\  /
                        6
 
-    Result: 0.5 + 1.0 + 0.95
+    Result: 0.5 + 1.0 + 0.9
     """
     solver = MIPSolver(config_instance.tracking_config)
 

ultrack/imgproc/_test/test_segmentation.py

@@ -1,11 +1,15 @@
 import logging
+from types import ModuleType
 from typing import Optional
 
 import numpy as np
 import pytest
+import scipy.ndimage as ndi
 from skimage.data import cells3d
+from skimage.morphology import reconstruction
 
 from ultrack.imgproc import Cellpose, detect_foreground, inverted_edt, robust_invert
+from ultrack.imgproc.segmentation import reconstruction_by_dilation
 from ultrack.utils.cuda import to_cpu
 
 LOG = logging.getLogger(__name__)
@@ -21,19 +25,20 @@
     LOG.info("cupy not found using numpy.")
 
 
-@pytest.mark.parametrize("channel_axis", [None, 0, 2])
+@pytest.mark.parametrize("np_module,channel_axis", [(xp, None), (xp, 0), (np, 2)])
 def test_foreground_detection(
+    np_module: ModuleType,
     channel_axis: Optional[int],
     request,
 ) -> None:
 
     cells = cells3d()
-    nuclei = xp.asarray(cells[:, 1])
-    membrane = xp.asarray(cells[:, 0])
+    nuclei = np_module.asarray(cells[:, 1])
+    membrane = np_module.asarray(cells[:, 0])
 
     if channel_axis is not None:
-        nuclei = xp.stack([nuclei] * 2, axis=channel_axis)
-        membrane = xp.stack([membrane] * 2, axis=channel_axis)
+        nuclei = np_module.stack([nuclei] * 2, axis=channel_axis)
+        membrane = np_module.stack([membrane] * 2, axis=channel_axis)
 
     contours = robust_invert(membrane, [1, 1, 1], channel_axis=channel_axis)
     foreground = detect_foreground(
@@ -57,6 +62,30 @@ def test_foreground_detection(
         napari.run()
 
 
+def test_reconstruction_by_dilation(request) -> None:
+    cells = cells3d()
+    membrane = cells[cells.shape[0] // 2, 0]
+
+    seed = ndi.gaussian_filter(membrane, 5)
+
+    iterative_bkg = reconstruction_by_dilation(seed, membrane, 250)
+    exact_bkg = reconstruction(seed, membrane, method="dilation")
+
+    if request.config.getoption("--show-napari-viewer"):
+        import napari
+
+        viewer = napari.Viewer()
+
+        viewer.add_image(membrane, blending="additive", visible=False)
+        viewer.add_image(seed, blending="additive", visible=False)
+        viewer.add_image(iterative_bkg, blending="additive", colormap="red")
+        viewer.add_image(exact_bkg, blending="additive", colormap="green")
+
+        napari.run()
+
+    np.testing.assert_array_almost_equal(iterative_bkg, exact_bkg)
+
+
 def test_inverted_edt() -> None:
     mask = np.array([[0, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=bool)
     expected_output = np.array(

ultrack/imgproc/intensity.py

@@ -111,7 +111,7 @@ def robust_invert(
             _image = np.take(image, indices=ch, axis=channel_axis)
 
         _image = xp.asarray(_image)
-        ndi.gaussian_filter(_image, sigma=sigmas, output=_image)
+        _image = ndi.gaussian_filter(_image, sigma=sigmas)
 
         flat_small_img = ndi.zoom(_image, (0.25,) * _image.ndim, order=1).ravel()
 

ultrack/imgproc/segmentation.py

@@ -4,10 +4,11 @@
 import edt
 import numpy as np
 from numpy.typing import ArrayLike
+from skimage.morphology import reconstruction
 
 from ultrack.imgproc.utils import _channel_iterator, _parse_voxel_size
 from ultrack.utils.constants import ULTRACK_DEBUG
-from ultrack.utils.cuda import import_module, to_cpu
+from ultrack.utils.cuda import import_module, is_cupy_array, to_cpu
 
 LOG = logging.getLogger(__name__)
 
@@ -43,7 +44,7 @@ def reconstruction_by_dilation(
     -------
         Image reconstructed by dilation.
     """
-    ndi = import_module("scipy", "ndimage")
+    ndi = import_module("scipy", "ndimage", arr=mask)
 
     seed = np.minimum(seed, mask, out=seed)  # just making sure
 
@@ -110,7 +111,15 @@ def detect_foreground(
         _image = xp.asarray(_image)
 
         seed = ndi.gaussian_filter(_image, sigma=sigmas)
-        background = reconstruction_by_dilation(seed, _image, 100)
+        if is_cupy_array(_image):
+            background = reconstruction_by_dilation(seed, _image, 100)
+        else:
+            if seed.ndim > 2:
+                LOG.warning(
+                    "Using CPU background reconstruction, this could take a while, consider using GPU for 3D images."
+                )
+            np.minimum(seed, _image, out=seed)  # required condition for reconstruction
+            background = reconstruction(seed, _image, method="dilation")
         del seed
 
         foreground = _image - background

ultrack/tracks/_test/test_tracks_stats.py

@@ -80,6 +80,32 @@ def test_tracks_df_movement():
     pd.testing.assert_frame_equal(result, expected)
 
 
+def test_tracks_df_movement_2d():
+    # Sample test data
+    df = pd.DataFrame(
+        {
+            "track_id": [1, 1, 2, 2],
+            "t": [1, 2, 1, 2],
+            "y": [1, 2, 1, 2],
+            "x": [2, 3, 2, 2],
+        }
+    )
+
+    # Call the function
+    result = tracks_df_movement(df)
+
+    # Expected result
+    expected = pd.DataFrame(
+        {
+            "y": [0.0, 1.0, 0.0, 1.0],
+            "x": [0.0, 1.0, 0.0, 0.0],
+        }
+    )
+
+    # Assert that the result matches the expected dataframe
+    pd.testing.assert_frame_equal(result, expected)
+
+
 def test_tracks_profile_matrix_one_track_one_attribute():
     tracks_df = pd.DataFrame(
         {"track_id": [1, 1, 1], "t": [0, 1, 2], "attribute_1": [10, 20, 30]}

ultrack/tracks/stats.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 import pandas as pd
@@ -13,7 +13,7 @@
 def tracks_df_movement(
     tracks_df: pd.DataFrame,
     lag: int = 1,
-    cols: tuple[str, ...] = ("z", "y", "x"),
+    cols: Optional[tuple[str, ...]] = None,
 ) -> pd.DataFrame:
     """
     Compute the displacement for track data across given time lags.
@@ -33,7 +33,8 @@ def tracks_df_movement(
         Number of periods to compute the difference over. Default is 1.
 
     cols : tuple[str, ...], optional
-        Columns to compute the displacement for. Default is ("z", "y", "x").
+        Columns to compute the displacement for. If not provided, it will try to
+        find any of ["z", "y", "x"] columns in the dataframe and use them.
 
     Returns
     -------
@@ -62,7 +63,13 @@ def tracks_df_movement(
 
     tracks_df.sort_values(by=["track_id", "t"], inplace=True)
 
-    cols = list(cols)
+    if cols is None:
+        cols = []
+        for c in ["z", "y", "x"]:
+            if c in tracks_df.columns:
+                cols.append(c)
+    else:
+        cols = list(cols)
 
     out = tracks_df.groupby("track_id", as_index=False)[cols].diff(periods=lag)
     out.fillna(0, inplace=True)

ultrack/utils/cuda.py

@@ -35,6 +35,18 @@
 }
 
 
+def is_cupy_array(arr: ArrayLike) -> bool:
+    """
+    Checks if array is a cupy array.
+
+    Parameters
+    ----------
+    arr : ArrayLike
+        Array to be checked.
+    """
+    return cp is not None and isinstance(arr, cp.ndarray)
+
+
 @contextmanager
 def unified_memory() -> Generator:
     """

ultrack/widgets/ultrackwidget/data_forms.py

@@ -3,7 +3,7 @@
 from typing import Any, Callable, Dict, Optional
 
 import qtawesome as qta
-from napari.layers import Image
+from napari.layers import Image, Layer
 from qtpy.QtWidgets import (
     QCheckBox,
     QGroupBox,
@@ -414,9 +414,9 @@ def setup_additional_options(self, workflow_choice: WorkflowChoice) -> None:
         # if no additional options are visible, hide the tab
         self._tab.setTabVisible(0, any_visible)
 
-    def notify_image_update(self, image: Image) -> None:
+    def notify_image_update(self, image: Layer) -> None:
         channel = ""
-        if image and image.rgb:
+        if image and isinstance(image, Image) and image.rgb:
             channel = image.data.ndim - 2  # last dimension is channel
 
         for widget in self._update_channel_axis_bindings: