merged ultrack-array into ultrackastra/node-gt-matching-teun

ultrack/config/config.py

@@ -72,3 +72,9 @@ def load_config(path: Union[str, Path]) -> MainConfig:
         data = toml.load(f)
         LOG.info(data)
         return MainConfig.parse_obj(data)
+
+
+def save_config(config, path: Union[str, Path]):
+    """Saved MainConfig to TOML file."""
+    with open(path, mode="w") as f:
+        toml.dump(config.dict(by_alias=True), f)

ultrack/napari.yaml

@@ -17,6 +17,10 @@ contributions:
       python_name: ultrack.widgets:UltrackWidget
       title: Ultrack
 
+    - id: ultrack.hierarchy_viz_widget
+      python_name: ultrack.widgets:HierarchyVizWidget
+      title: Hierarchy visualization
+
 ###### DEPRECATED & WIP WIDGETS  #####
 #     - id: ultrack.labels_to_edges_widget
 #       python_name: ultrack.widgets:LabelsToContoursWidget
@@ -67,5 +71,8 @@ contributions:
     # - command: ultrack.division_annotation_widget
     #   display_name: Division annotation
 
-    # - command: ultrack.track_inspection
-    #   display_name: Track inspection
+    - command: ultrack.track_inspection
+      display_name: Track inspection
+
+    - command: ultrack.hierarchy_viz_widget
+      display_name: Hierarchy visualization

ultrack/utils/_test/test_utils_array.py

@@ -1,8 +1,10 @@
 import numpy as np
 import pytest
 
+from typing import Tuple
 from ultrack.utils.array import array_apply
-
+from ultrack.config import MainConfig
+from ultrack.utils.array import UltrackArray
 
 @pytest.mark.parametrize("axis", [0, 1])
 def test_array_apply_parametrized(axis):
@@ -17,3 +19,31 @@ def sample_func(arr_1, arr_2):
     array_apply(in_data, in_data, out_array=out_data, func=sample_func, axis=axis)
     other_axes_length = in_data.shape[1 - axis]
     assert np.array_equal(out_data, 2 * in_data + other_axes_length)
+
+@pytest.mark.parametrize(
+    "key,timelapse_mock_data",
+    [
+        (1,{'n_dim':3}),
+        (1,{'n_dim':2}),
+        ((slice(None), 1),{'n_dim':3}),
+        ((slice(None), 1),{'n_dim':2}),
+        ((0, [1, 2]),{'n_dim':3}),
+        ((0, [1, 2]),{'n_dim':2}),
+        # ((-1, np.asarray([0, 3])),{'n_dim':3}),       #does testing negative time make sense?
+        # ((-1, np.asarray([0, 3])),{'n_dim':2}),
+        ((slice(1), -2),{'n_dim':3}),
+        ((slice(1), -2),{'n_dim':2}),
+        ((np.asarray(0),),{'n_dim':3}),
+        ((np.asarray(0),),{'n_dim':2}),
+        ((0, 0, slice(32)),{'n_dim':3}),
+        ((0, 0, slice(32)),{'n_dim':2}),
+    ],
+    indirect=["timelapse_mock_data",],
+)
+def test_ultrack_array(
+        segmentation_database_mock_data: MainConfig,
+        key: Tuple,
+    ):
+    ua = UltrackArray(segmentation_database_mock_data)
+    ua_numpy = ua[slice(None)]
+    np.testing.assert_equal(ua_numpy[key], ua[key])

ultrack/utils/array.py

@@ -6,11 +6,16 @@
 from typing import Any, Callable, Dict, Literal, Optional, Tuple, Type, Union
 
 import numpy as np
+import sqlalchemy as sqla
 import zarr
 from numpy.typing import ArrayLike
+from sqlalchemy.orm import Session
 from tqdm import tqdm
 from zarr.storage import Store
 
+from ultrack.core.database import NodeDB
+from ultrack.config import MainConfig
+
 LOG = logging.getLogger(__name__)
 
 
@@ -213,6 +218,180 @@ def create_zarr(
     return zarr.zeros(shape, dtype=dtype, store=store, chunks=chunks, **kwargs)
 
 
+class UltrackArray:
+    def __init__(
+        self,
+        config: MainConfig,
+        dtype: np.dtype = np.int32,
+    ):
+        """Create an array that directly visualizes the segments in the ultrack database.
+
+        Parameters
+        ----------
+        config : MainConfig
+            Configuration file of Ultrack.
+        dtype : np.dtype
+            Data type of the array.
+        """
+
+        self.config = config
+        self.shape = tuple(config.data_config.metadata["shape"])  # (t,(z),y,x)
+        self.dtype = dtype
+        self.t_max = self.shape[0]
+        self.ndim = len(self.shape)
+        self.array = np.zeros(self.shape[1:], dtype=self.dtype)
+
+        self.database_path = config.data_config.database_path
+        self.minmax = self.find_min_max_volume_entire_dataset()
+        self.volume = self.minmax.mean().astype(int)
+
+    def __getitem__(self, 
+                    indexing: Union[Tuple[Union[int, slice]], int, slice],
+        ) -> np.ndarray:
+        """Indexing the ultrack-array
+
+        Parameters
+        ----------
+        indexing : Tuple or Array
+
+        Returns
+        -------
+        array : numpy array
+            array with painted segments
+        """
+        # print('indexing in getitem:',indexing)
+
+        if isinstance(indexing, tuple):
+            time, volume_slicing = indexing[0], indexing[1:]
+        else:       #if only 1 (time) is provided
+            time = indexing
+            volume_slicing = tuple()
+
+        if isinstance(time, slice): #if all time points are requested
+            return np.stack([
+                self.__getitem__((t,) + volume_slicing)
+                for t in range(*time.indices(self.shape[0]))
+            ])
+        else:
+            try:
+                time = time.item()  # convert from numpy.int to int
+            except AttributeError:
+                time = time
+
+        self.fill_array(
+            time=time,
+        )
+
+        return self.array[volume_slicing]
+
+    def fill_array(
+        self,
+        time: int,
+    ) -> None:
+        """Paint all segments of specific time point which volume is bigger than self.volume
+        Parameters
+        ----------
+        time : int
+            time point to paint the segments
+        """
+
+        engine = sqla.create_engine(self.database_path)
+        self.array.fill(0)
+
+        with Session(engine) as session:
+            query = list(
+                session.query(NodeDB.id, NodeDB.pickle, NodeDB.hier_parent_id).where(
+                    NodeDB.t == time
+                )
+            )
+
+            idx_to_plot = []
+
+            for idx, q in enumerate(query):
+                if q[1].area <= self.volume:
+                    idx_to_plot.append(idx)
+
+            id_to_plot = [q[0] for idx, q in enumerate(query) if idx in idx_to_plot]
+            label_list = np.arange(1, len(query) + 1, dtype=int)
+
+            to_remove = []
+            for idx in idx_to_plot:
+                if query[idx][2] in id_to_plot:  # if parent is also printed
+                    to_remove.append(idx)
+
+            for idx in to_remove:
+                idx_to_plot.remove(idx)
+
+            if len(query) == 0:
+                print("query is empty!")
+
+            for idx in idx_to_plot:
+                query[idx][1].paint_buffer(
+                    self.array, value=label_list[idx], include_time=False
+                )
+
+    def get_tp_num_pixels(
+            self, 
+            timeStart:int,
+            timeStop:int,
+    ) -> list:
+        """Gets a list of number of pixels of all segments range of time points (timeStart to timeStop)
+        Parameters
+        ----------
+        timeStart : int
+        timeStop : int
+        Returns
+        -------
+        num_pix_list : list
+            list with all num_pixels for timeStart to timeStop
+        """
+        engine = sqla.create_engine(self.database_path)
+        num_pix_list = []
+        with Session(engine) as session:
+            query = list(session.query(NodeDB.area).where(NodeDB.t >= timeStart).where(NodeDB.t <= timeStop))
+            for num_pix in query:
+                num_pix_list.append(int(np.array(num_pix)))
+        return num_pix_list
+
+    def get_tp_num_pixels_minmax(
+        self,
+        time: int,
+    ) -> np.ndarray:
+        """Find minimum and maximum segment volume for single time point
+
+        Parameters
+        ----------
+        time : int
+
+        Returns
+        -------
+        num_pix_list : list
+            array with two elements: [min_volume, max_volume]
+        """
+        num_pix_list = self.get_tp_num_pixels(time,time)
+        return (min(num_pix_list), max(num_pix_list))
+
+
+    def find_min_max_volume_entire_dataset(self):
+        """Find minimum and maximum segment volume for ALL time point
+
+        Returns
+        -------
+        np.array : np.array
+            array with two elements: [min_volume, max_volume]
+        """
+        min_vol = np.inf
+        max_vol = 0
+        for t in range(self.t_max): #range(self.shape[0]):
+            minmax = self.get_tp_num_pixels_minmax(t)
+            if minmax[0] < min_vol:
+                min_vol = minmax[0]
+            if minmax[1] > max_vol:
+                max_vol = minmax[1]
+
+        return np.array([min_vol, max_vol], dtype=int)
+
+
 def assert_same_length(**kwargs) -> None:
     """Validates if key-word arguments have the same length."""
     for k1, v1 in kwargs.items():

ultrack/utils/segmentation.py

@@ -266,6 +266,6 @@ def copy_segments(
         )  # not sure why this is necessary in large datasets
     else:
         for t in tqdm(range(segments.shape[0]), "Copying segments"):
-            out_segments[t] = segments[t]
+            out_segments[t] = np.asarray(segments[t])
 
     return out_segments

ultrack/widgets/__init__.py

@@ -1,5 +1,6 @@
 from ultrack.widgets.division_annotation_widget import DivisionAnnotationWidget
 from ultrack.widgets.hypotheses_viz_widget import HypothesesVizWidget
+from ultrack.widgets.hierarchy_viz_widget import HierarchyVizWidget
 from ultrack.widgets.labels_to_edges_widget import LabelsToContoursWidget
 from ultrack.widgets.node_annotation_widget import NodeAnnotationWidget
 from ultrack.widgets.track_inspection_widget import TrackInspectionWidget

ultrack/widgets/_test/test_hierarchy_viz_widget.py

@@ -0,0 +1,95 @@
+from typing import Callable, Tuple
+
+import napari
+import numpy as np
+import pytest
+import zarr
+
+from ultrack.config import MainConfig
+from ultrack.widgets.ultrackwidget import UltrackWidget
+from ultrack.widgets.hierarchy_viz_widget import HierarchyVizWidget
+from ultrack.widgets.ultrackwidget.workflows import WorkflowChoice
+from ultrack.widgets.ultrackwidget.utils import UltrackInput
+
+
+
+
+
+def test_hierarchy_viz_widget(
+        make_napari_viewer: Callable[[],napari.Viewer],
+        segmentation_database_mock_data: MainConfig,
+        timelapse_mock_data: Tuple[zarr.Array, zarr.Array, zarr.Array],
+        request,
+    ):
+
+    ####################################################################################
+    #OPTION 1: run widget using config    
+    ####################################################################################
+    config = segmentation_database_mock_data
+    viewer = make_napari_viewer()
+    widget = HierarchyVizWidget(viewer,config)
+    viewer.window.add_dock_widget(widget)
+
+    assert "hierarchy" in viewer.layers
+
+    #test moving sliders:
+    widget._slider_update(0.75)
+    widget._slider_update(0.25)
+
+    #test is shape of layer.data has same shape as the data shape reported in config:
+    assert tuple(config.data_config.metadata["shape"]) == viewer.layers['hierarchy'].data.shape     #metadata["shape"] is a list, data.shape is a tuple
+
+
+    ####################################################################################
+    #OPTION 2: run widget by taking config from Ultrack-widget
+    ####################################################################################
+    #make napari viewer
+    viewer2 = make_napari_viewer()
+
+    #get mock segmentation data + add to viewer
+    segments = timelapse_mock_data[2]
+    print('segments shape',segments.shape)
+    viewer2.add_labels(segments,name='segments')
+
+    #open ultrack widget
+    widget_ultrack = UltrackWidget(viewer2)
+    viewer2.window.add_dock_widget(widget_ultrack)
+
+    #setup ultrack widget for 'Labels' input
+    layers = viewer2.layers
+    workflow = WorkflowChoice.AUTO_FROM_LABELS
+    workflow_idx = widget_ultrack._cb_workflow.findData(workflow)
+    widget_ultrack._cb_workflow.setCurrentIndex(workflow_idx)
+    widget_ultrack._cb_workflow.currentIndexChanged.emit(workflow_idx)
+    # setting combobox choices manually, because they were not working automatically
+    widget_ultrack._cb_images[UltrackInput.LABELS].choices = layers
+    # # selecting layers
+    widget_ultrack._cb_images[UltrackInput.LABELS].value = layers['segments']
+
+    #load config
+    widget_ultrack._data_forms.load_config(config)
+
+
+    widget_hier = HierarchyVizWidget(viewer2)
+    viewer2.window.add_dock_widget(widget_hier)
+
+    assert "hierarchy" in viewer.layers
+
+    #test moving sliders:
+    widget._slider_update(0.75)
+    widget._slider_update(0.25)
+
+    # test is shape of layer.data has same shape as the data shape reported in config:
+    assert tuple(config.data_config.metadata["shape"]) == viewer2.layers['hierarchy'].data.shape     #metadata["shape"] is a list, data.shape in layer is a tuple
+
+    ####################################################################################
+    #TO DO: 
+    #   - test other datatypes than labels (contours, detection, image, etc.), but shouldn't make a difference
+    #   - how to test that the 'hierarchy' layer actually shows data? (apart from the layer having a .data property)
+    ####################################################################################
+
+
+    if request.config.getoption("--show-napari-viewer"):
+        napari.run()
+
+