Reducing memory usage (#17)

* fixed deep learning modules utilities imports

* optimized hierarchy segmentation memory usage

* deleting viewer after usage

* updated README doc section

* fixing plantseg/microsam test skip

* updated examples documentation and dependencies

* updated examples

README.md

@@ -47,11 +47,12 @@ Usage examples can be found [here](examples), including their environment files
 
 ## Documentation
 
-The `ultrack` library relies on a configuration schema, its description is [here](ultrack/config/README.md).
+The official documentation is available [here](https://royerlab.github.io/ultrack/).
 
-The segmentation and tracking data are stored in an SQL database, described [here](ultrack/core/README.md).
+These additional developer documentation are available:
 
-Helper functions to export to the cell tracking challenge and napari formats are available [here](ultrack/core/export).
+- Parameter [configuration schema](ultrack/config/README.md).
+- Intermediate segmentation and tracking SQL database are [here](ultrack/core/README.md).
 
 ## Citing
 

examples/README.md

@@ -10,10 +10,11 @@ conda activate <your new env>
 
 The existing examples are:
 
+- [multi_color_ensemble](./multi_color_ensemble): Multi-colored cytoplasm cell tracking using Cellpose and Watershed segmentation ensemble. Data provided by [The Lammerding Lab ](https://lammerding.wicmb.cornell.edu/).
 - [flow_field_3d](./flow_field_3d): Tracking demo on a cartographic projection of Tribolium Castaneum embryo from the [cell-tracking challenge](http://celltrackingchallenge.net/3d-datasets/), using a flow field estimation to assist tracking of motile cells.
 - [stardist_2d](./stardist_2d): Tracking demo on HeLa GPF nuclei from the [cell-tracking challenge](http://celltrackingchallenge.net/2d-datasets/) using Stardist 2D fluorescence images pre-trained model.
 - [zebrahub](./zebrahub/): Tracking demo on zebrafish tail data from the [zebrahub](https://zebrahub.ds.czbiohub.org/) acquired with [DaXi](https://www.nature.com/articles/s41592-022-01417-2) using Ultrack's image processing helper functions.
-- [neuromast_plantseg](./neuromast_plantseg/): Tracking demo membrane-labeled zebrafish neuromast from the [TODO](TODO) using [PlantSeg's](https://github.com/hci-unihd/plant-seg) membrane detection model.
+- [neuromast_plantseg](./neuromast_plantseg/): Tracking demo membrane-labeled zebrafish neuromast from [Jacobo Group of CZ Biohub](https://www.czbiohub.org/jacobo/) using [PlantSeg's](https://github.com/hci-unihd/plant-seg) membrane detection model.
 - [micro_sam](./micro_sam/): Tracking demo with [MicroSAM](https://github.com/computational-cell-analytics/micro-sam) instance segmentation package.
 
 

examples/multi_color_ensemble/environment_cpu.yml

@@ -18,3 +18,4 @@ dependencies:
     - napari==0.4.18
     - traccuracy
     - napari-arboretum
+    - pyift

examples/multi_color_ensemble/environment_gpu.yml

@@ -18,3 +18,4 @@ dependencies:
     - napari==0.4.18
     - traccuracy
     - napari-arboretum
+    - pyift

examples/refresh_examples.sh

@@ -1,5 +1,8 @@
 #! /bin/bash
 
+# Terminates if any command fails
+set -e
+
 # subject to change depending on your conda setup
 CONDA_SETUP_PATH=$HOME/miniconda3/etc/profile.d/conda.sh
 UPDATE_JUPYTER="jupyter nbconvert --execute --to notebook --inplace"
@@ -12,10 +15,6 @@ function install () {
 	pip install -e ..
 }
 
-# stardist
-install ultrack-stardist stardist_2d
-$UPDATE_JUPYTER stardist_2d/2d_tracking.ipynb
-
 # multi color
 install ultrack-multi-color multi_color_ensemble
 $UPDATE_JUPYTER multi_color_ensemble/multi_color_ensemble.ipynb
@@ -35,3 +34,7 @@ $UPDATE_JUPYTER neuromast_plantseg/neuromast_plantseg.ipynb
 # micro-sam
 install ultrack-micro-sam micro_sam
 $UPDATE_JUPYTER micro_sam/micro_sam_tracking.ipynb
+
+# stardist
+install ultrack-stardist stardist_2d
+$UPDATE_JUPYTER stardist_2d/2d_tracking.ipynb

examples/stardist_2d/environment_cpu.yml

@@ -10,6 +10,7 @@ dependencies:
   - gurobi
   - jupyter
   - pip
+  - pytorch
   - zarr
   - pip:
     - pyqt5

examples/stardist_2d/environment_gpu.yml

@@ -14,6 +14,7 @@ dependencies:
   - coin-or-cbc
   - gurobi
   - jupyter
+  - pytorch
   - pip
   - zarr
   - pip:

ultrack/cli/estimate_params.py

@@ -70,6 +70,8 @@ def estimate_params_cli(
     except (KeyError, ValueError):
         labels = viewer.layers[layer_key].data
 
+    del viewer
+
     df = estimate_parameters_from_labels(labels, is_timelapse=timelapse)
 
     covariables = {"area", "distance"}

ultrack/cli/export.py

@@ -151,6 +151,7 @@ def zarr_napari_cli(
                 layer.data[0] if layer.multiscale else layer.data
                 for layer in viewer.open(image_path, plugin=reader_plugin)
             ]
+            del viewer
 
         tracks_w_measures = tracks_properties(
             segments=segments,

ultrack/cli/flow.py

@@ -35,6 +35,7 @@ def add_flow_cli(
         layer.data[0] if layer.multiscale else layer.data
         for layer in viewer.open(paths, channel_axis=channel_axis, plugin=reader_plugin)
     ]
+    del viewer
 
     add_flow(
         config,

ultrack/cli/labels_to_edges.py

@@ -49,8 +49,11 @@ def labels_to_edges_cli(
     viewer = ViewerModel()
     viewer.open(path=paths, plugin=reader_plugin)
 
+    labels = [layer.data for layer in viewer.layers]
+    del viewer
+
     labels_to_edges(
-        [layer.data for layer in viewer.layers],
+        labels,
         sigma=sigma,
         detection_store_or_path=detection_path,
         edges_store_or_path=edges_path,

ultrack/cli/link.py

@@ -51,6 +51,7 @@ def link_cli(
             layer.data[0] if layer.multiscale else layer.data
             for layer in viewer.open(paths, **kwargs, plugin=reader_plugin)
         ]
+        del viewer
 
     link(
         config,

ultrack/cli/segment.py

@@ -71,6 +71,8 @@ def segmentation_cli(
             {"scale": viewer.layers[edge_layer].scale.tolist()}
         )
 
+    del viewer
+
     segment(
         detection,
         edge,

ultrack/core/segmentation/hierarchy.py

@@ -1,7 +1,8 @@
 import logging
-from typing import List
+from typing import Iterator
 
 import numpy as np
+import scipy.ndimage as ndi
 from numpy.typing import ArrayLike
 from skimage import measure, morphology
 from skimage.measure._regionprops import RegionProperties
@@ -33,7 +34,7 @@ def create_hierarchies(
     binary_detection: ArrayLike,
     edge: ArrayLike,
     **kwargs,
-) -> List[Hierarchy]:
+) -> Iterator[Hierarchy]:
     """Computes a collection of hierarchical watersheds inside `binary_detection` mask.
 
     Parameters
@@ -46,26 +47,25 @@ def create_hierarchies(
 
     Returns
     -------
-    List[Hierarchy]
+    Iterator[Hierarchy]
         List of hierarchical watersheds.
     """
     binary_detection = np.asarray(binary_detection)
-    edge = np.asarray(edge)
 
     assert (
         issubclass(binary_detection.dtype.type, np.integer)
         or binary_detection.dtype == bool
     )
 
     LOG.info("Labeling connected components.")
-    labels, num_labels = measure.label(
-        binary_detection, return_num=True, connectivity=1
-    )
-    labels = labels.astype(np.int32)
+    labels, num_labels = ndi.label(binary_detection, output=np.int32)
+    del binary_detection
 
     if "min_area" in kwargs and num_labels > 1:
         LOG.info("Filtering small connected components.")
-        labels = morphology.remove_small_objects(labels, min_size=kwargs["min_area"])
+        morphology.remove_small_objects(labels, min_size=kwargs["min_area"], out=labels)
+
+    edge = np.asarray(edge)
 
     if "max_area" in kwargs:
         LOG.info("Oversegmenting connected components.")
@@ -77,6 +77,5 @@ def create_hierarchies(
         )
 
     LOG.info("Creating hierarchies (lazy).")
-    return [
-        Hierarchy(c, **kwargs) for c in measure.regionprops(labels, edge, cache=True)
-    ]
+    for c in measure.regionprops(labels, edge, cache=True):
+        yield Hierarchy(c, **kwargs)

ultrack/imgproc/_test/test_plantseg.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-plantseg = pytest.importorskip("ultrack.imgproc.plantseg")
+from ultrack.imgproc.plantseg import PlantSeg
 
 
 @pytest.mark.parametrize(
@@ -19,9 +19,13 @@ def test_plantseg(
 
     image = np.random.rand(100, 100, 100)
 
-    seg_model = plantseg.PlantSeg(
-        model_name="generic_light_sheet_3D_unet",
-        batch_size=1,
-        patch=image.shape,
-    )
+    try:
+        seg_model = PlantSeg(
+            model_name="generic_light_sheet_3D_unet",
+            batch_size=1,
+            patch=image.shape,
+        )
+    except ModuleNotFoundError:
+        pytest.skip("PlantSeg not installed")
+
     seg_model(image, transpose=transpose)

ultrack/imgproc/_test/test_sam.py

@@ -1,12 +1,17 @@
 import numpy as np
 import pytest
 
-sam = pytest.importorskip("ultrack.imgproc.sam")
+from ultrack.imgproc import sam
 
 
 def test_sam() -> None:
     image = np.random.rand(100, 100)
-    seg_model = sam.MicroSAM()
+
+    try:
+        seg_model = sam.MicroSAM()
+    except ModuleNotFoundError:
+        pytest.skip("MicroSAM not installed")
+
     seg_model(image)
 
     # with maxima prompt

ultrack/imgproc/plantseg.py

@@ -4,13 +4,6 @@
 import numpy as np
 import torch as th
 from numpy.typing import ArrayLike
-from plantseg.dataprocessing import fix_input_shape
-from plantseg.predictions.functional.array_predictor import ArrayPredictor
-from plantseg.predictions.functional.utils import (
-    get_array_dataset,
-    get_model_config,
-    get_patch_halo,
-)
 
 from ultrack.utils.cuda import import_module, to_cpu, torch_default_device
 
@@ -67,6 +60,12 @@ def __init__(
         """
         Initialized Plant-Seg model.
         """
+        from plantseg.predictions.functional.array_predictor import ArrayPredictor
+        from plantseg.predictions.functional.utils import (
+            get_model_config,
+            get_patch_halo,
+        )
+
         if device is None:
             device = torch_default_device()
 
@@ -121,6 +120,8 @@ def __call__(
         np.ndarray
             Segmentation boundary probability map as a numpy array.
         """
+        from plantseg.dataprocessing import fix_input_shape
+        from plantseg.predictions.functional.utils import get_array_dataset
 
         if isinstance(image, da.Array):
             # avoiding building a large compute graph

ultrack/imgproc/sam.py

@@ -6,9 +6,7 @@
 import numpy as np
 import scipy.ndimage as ndi
 import torch as th
-from micro_sam import instance_segmentation, util
 from numpy.typing import ArrayLike
-from segment_anything.utils.amg import area_from_rle, rle_to_mask
 from skimage.feature import peak_local_max
 from skimage.segmentation import find_boundaries
 
@@ -92,6 +90,7 @@ def __init__(
         tile_shape: Tuple[int, int] = (512, 512),
         halo_shape: Tuple[int, int] = (128, 128),
     ) -> None:
+        from micro_sam import instance_segmentation, util
 
         if device is None:
             device = torch_default_device()
@@ -148,6 +147,8 @@ def __call__(self, image: ArrayLike) -> np.ndarray:
         np.ndarray
             The processed image with contours derived from the identified masks.
         """
+        from micro_sam.util import precompute_image_embeddings
+        from segment_anything.utils.amg import area_from_rle, rle_to_mask
 
         image = np.asarray(image)
 
@@ -159,7 +160,7 @@ def __call__(self, image: ArrayLike) -> np.ndarray:
         if embedding_path.exists():
             shutil.rmtree(embedding_path)
 
-        embeddings = util.precompute_image_embeddings(
+        embeddings = precompute_image_embeddings(
             self._predictor,
             image,
             save_path=str(embedding_path),