Implement Cellpose distributed module and fix bugs in processing/trackmate

Co-authored-by: lguerard <6182107+lguerard@users.noreply.github.com>
Agent-Logs-Url: https://github.com/imcf/python-imcflibs/sessions/52e63697-e516-4bbc-9f13-96067017fda2

Author: Copilot

conftest.py

@@ -1,3 +1,15 @@
 """Pytest configuration."""
 
+import sys
+from unittest.mock import MagicMock
+
 collect_ignore = ["tests/interactive-imagej"]
+
+# Add mock stubs for BIOP Cellpose wrappers not yet covered by imcf-fiji-mocks.
+for _mod in [
+    "ch.epfl.biop.wrappers",
+    "ch.epfl.biop.wrappers.cellpose",
+    "ch.epfl.biop.wrappers.cellpose.ij2commands",
+]:
+    if _mod not in sys.modules:
+        sys.modules[_mod] = MagicMock()

pyproject.toml

@@ -29,6 +29,9 @@ sjlogging = ">=0.5.2"
 pytest = "^8.0.1"
 pytest-cov = "^6.0.0"
 
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+
 [build-system]
 build-backend = "poetry.core.masonry.api"
 requires = ["poetry-core"]

src/imcflibs/imagej/__init__.py

@@ -6,6 +6,7 @@
 """
 
 from . import bioformats
+from . import cellpose
 from . import misc
 from . import prefs
 from . import projections

src/imcflibs/imagej/cellpose.py

@@ -0,0 +1,226 @@
+"""Functions for running [Cellpose] segmentation in Fiji.
+
+This module provides wrappers around the [BIOP Cellpose plugin] for Fiji,
+including support for distributed (batch) processing of multiple images.
+
+[Cellpose]: https://cellpose.readthedocs.io/
+[BIOP Cellpose plugin]: https://github.com/BIOP/ijl-utilities-wrappers
+"""
+
+import os
+
+from ch.epfl.biop.wrappers.cellpose import CellposeTaskSettings
+from ch.epfl.biop.wrappers.cellpose.ij2commands import Cellpose_SegmentImgPlusAdvanced
+from ij import IJ
+
+from ..log import LOG as log
+
+
+VALID_MODELS = ["cyto", "cyto2", "cyto3", "nuclei"]
+
+DIMENSION_MODE_2D = "2D"
+DIMENSION_MODE_3D = "3D"
+
+
+def build_cellpose_settings(
+    model,
+    diameter,
+    cyto_channel,
+    nuclei_channel=0,
+    cellproba_threshold=0.0,
+    flow_threshold=0.4,
+    use_gpu=True,
+    dimension_mode=DIMENSION_MODE_2D,
+    stitch_threshold=-1.0,
+    omni=False,
+    cluster=False,
+    additional_flags="",
+):
+    """Build a Cellpose settings object for use with the BIOP Cellpose plugin.
+
+    Parameters
+    ----------
+    model : str
+        Name of the Cellpose model to use. Must be one of: cyto, cyto2, cyto3,
+        nuclei.
+    diameter : float
+        Expected diameter of the objects to segment in pixels. Use 0 to let
+        Cellpose estimate the diameter automatically.
+    cyto_channel : int
+        Index of the cytoplasm channel (1-based). Use 0 if not applicable.
+    nuclei_channel : int, optional
+        Index of the nuclei channel (1-based). Use 0 to disable the nuclei
+        channel, by default 0.
+    cellproba_threshold : float, optional
+        Cell probability threshold. Lower values result in more detected cells,
+        by default 0.0.
+    flow_threshold : float, optional
+        Flow error threshold for discarding masks. Higher values allow more
+        errors, by default 0.4.
+    use_gpu : bool, optional
+        Whether to use GPU acceleration, by default True.
+    dimension_mode : str, optional
+        Dimension mode for segmentation, either ``"2D"`` or ``"3D"``, by
+        default ``"2D"``.
+    stitch_threshold : float, optional
+        Threshold for stitching masks across z-planes. A value of -1.0
+        disables stitching, by default -1.0.
+    omni : bool, optional
+        Whether to use OmniPose segmentation, by default False.
+    cluster : bool, optional
+        Whether to use DBSCAN clustering for cell detection, by default False.
+    additional_flags : str, optional
+        Additional command-line flags passed to Cellpose, by default ``""``.
+
+    Returns
+    -------
+    ch.epfl.biop.wrappers.cellpose.CellposeTaskSettings
+        A configured ``CellposeTaskSettings`` object.
+
+    Raises
+    ------
+    ValueError
+        If ``model`` is not one of the valid model names.
+    ValueError
+        If ``dimension_mode`` is not ``"2D"`` or ``"3D"``.
+    """
+    if model.lower() not in VALID_MODELS:
+        raise ValueError(
+            "model '%s' is not valid, must be one of: %s"
+            % (model, ", ".join(VALID_MODELS))
+        )
+
+    if dimension_mode not in (DIMENSION_MODE_2D, DIMENSION_MODE_3D):
+        raise ValueError(
+            "dimension_mode must be '%s' or '%s', got '%s'"
+            % (DIMENSION_MODE_2D, DIMENSION_MODE_3D, dimension_mode)
+        )
+
+    log.debug(
+        "Building Cellpose settings: model=%s, diameter=%s, cyto_ch=%s, "
+        "nuclei_ch=%s, dimension_mode=%s"
+        % (model, diameter, cyto_channel, nuclei_channel, dimension_mode)
+    )
+
+    settings = CellposeTaskSettings()
+    settings.setModel(model.lower())
+    settings.setDiameter(diameter)
+    settings.setCytoCh(cyto_channel)
+    settings.setNucleiCh(nuclei_channel)
+    settings.setCellProbaTreshold(cellproba_threshold)
+    settings.setFlowThreshold(flow_threshold)
+    settings.useGpu(use_gpu)
+    settings.setDimensionMode(dimension_mode)
+    settings.setStitchThreshold(stitch_threshold)
+    settings.setOmni(omni)
+    settings.setCluster(cluster)
+    settings.setAdditionalFlags(additional_flags)
+
+    return settings
+
+
+def run_cellpose(imp, settings):
+    """Run Cellpose segmentation on an ImagePlus.
+
+    Parameters
+    ----------
+    imp : ij.ImagePlus
+        Input ImagePlus to segment.
+    settings : ch.epfl.biop.wrappers.cellpose.CellposeTaskSettings
+        Configured ``CellposeTaskSettings`` object, see
+        :func:`build_cellpose_settings`.
+
+    Returns
+    -------
+    ij.ImagePlus
+        Label image produced by Cellpose, or ``None`` if segmentation fails.
+    """
+    log.info("Running Cellpose on image '%s'" % imp.getTitle())
+
+    command = Cellpose_SegmentImgPlusAdvanced()
+    command.imp = imp
+    command.cellposeSettings = settings
+
+    try:
+        command.run()
+    except Exception as err:
+        log.error("Cellpose segmentation failed: %s" % err)
+        return None
+
+    result = command.output_imp
+    log.info(
+        "Cellpose segmentation completed, result image: '%s'" % result.getTitle()
+    )
+
+    return result
+
+
+def run_cellpose_distributed(imp_list, settings, show_progress=True):
+    """Run Cellpose segmentation on a list of ImagePlus objects.
+
+    Processes each image sequentially while reporting progress, allowing
+    Cellpose to be applied across a large batch of images (distributed
+    processing).
+
+    Parameters
+    ----------
+    imp_list : list(ij.ImagePlus)
+        List of ImagePlus objects to segment.
+    settings : ch.epfl.biop.wrappers.cellpose.CellposeTaskSettings
+        Configured ``CellposeTaskSettings`` object shared across all images,
+        see :func:`build_cellpose_settings`.
+    show_progress : bool, optional
+        Whether to show progress in the ImageJ progress bar, by default True.
+
+    Returns
+    -------
+    list(ij.ImagePlus)
+        List of label images produced by Cellpose, with ``None`` entries for
+        images where segmentation failed.
+
+    Example
+    -------
+    >>> settings = build_cellpose_settings(
+    ...     model="cyto2",
+    ...     diameter=30.0,
+    ...     cyto_channel=1,
+    ...     nuclei_channel=2,
+    ... )
+    >>> results = run_cellpose_distributed(image_list, settings)
+    """
+    total = len(imp_list)
+    log.info("Starting distributed Cellpose run on %d images" % total)
+
+    results = []
+
+    for idx, imp in enumerate(imp_list):
+        log.info(
+            "Processing image %d / %d: '%s'" % (idx + 1, total, imp.getTitle())
+        )
+
+        if show_progress:
+            IJ.showProgress(idx, total)
+            IJ.showStatus(
+                "Cellpose: processing image %d / %d" % (idx + 1, total)
+            )
+
+        result = run_cellpose(imp, settings)
+        results.append(result)
+
+    if show_progress:
+        IJ.showProgress(total, total)
+        IJ.showStatus("Cellpose: done processing %d images" % total)
+
+    failed = sum(1 for r in results if r is None)
+    if failed:
+        log.warning(
+            "Cellpose distributed run finished: %d / %d images failed"
+            % (failed, total)
+        )
+    else:
+        log.info(
+            "Cellpose distributed run finished successfully: %d images processed"
+            % total
+        )
+
+    return results

src/imcflibs/imagej/processing.py

@@ -53,7 +53,7 @@ def apply_filter(imp, filter_method, filter_radius, do_3d=False):
         filter = filter_method + "..."
 
     options = (
-        "sigma="
+        "sigma=" + str(filter_radius)
         if filter_method == "Gaussian Blur"
         else "radius=" + str(filter_radius) + " stack"
     )
@@ -85,12 +85,12 @@ def apply_rollingball_bg_subtraction(imp, rolling_ball_radius, do_3d=False):
     """
     log.info("Applying rolling ball with radius %d" % rolling_ball_radius)
 
-    options = "rolling=" + str(rolling_ball_radius) + " stack" if do_3d else ""
+    options = "rolling=" + str(rolling_ball_radius) + " stack" if do_3d else "rolling=" + str(rolling_ball_radius)
 
     log.debug("Background subtraction options: %s" % options)
 
     imageplus = imp.duplicate()
-    IJ.run(imageplus, "Substract Background...", options)
+    IJ.run(imageplus, "Subtract Background...", options)
 
     return imageplus
 
@@ -118,7 +118,7 @@ def apply_threshold(imp, threshold_method, do_3d=True):
     imageplus = imp.duplicate()
 
     auto_threshold_options = (
-        threshold_method + " " + "dark" + " " + "stack" if do_3D else ""
+        threshold_method + " " + "dark" + " " + "stack" if do_3d else threshold_method + " " + "dark"
     )
 
     log.debug("Auto threshold options: %s" % auto_threshold_options)

src/imcflibs/imagej/trackmate.py

@@ -77,10 +77,11 @@ def cellpose_detector(
     settings.detectorSettings["OPTIONAL_CHANNEL_2"] = optional_channel
 
     settings.detectorSettings["CELLPOSE_PYTHON_FILEPATH"] = pathtools.join2(
-        cellpose_env_path, "python.exe"
+        cellpose_env_path, "python.exe" if os.name == "nt" else "python"
     )
+    home_dir = os.environ.get("USERPROFILE") or os.environ.get("HOME", "")
     settings.detectorSettings["CELLPOSE_MODEL_FILEPATH"] = os.path.join(
-        os.environ["USERPROFILE"], ".cellpose", "models"
+        home_dir, ".cellpose", "models"
     )
     input_to_model = {
         "nuclei": PretrainedModel.NUCLEI,
@@ -90,7 +91,9 @@ def cellpose_detector(
     if model_to_use.lower() in input_to_model:
         selected_model = input_to_model[model_to_use.lower()]
     else:
-        print("Selected Model Does Not Exist")
+        print("Selected model '%s' does not exist, valid options are: %s" % (
+            model_to_use, ", ".join(input_to_model.keys())
+        ))
         return
 
     settings.detectorSettings["CELLPOSE_MODEL"] = selected_model

tests/test_cellpose.py

@@ -0,0 +1,76 @@
+"""Tests for `imcflibs.imagej.cellpose`."""
+
+import pytest
+
+from imcflibs.imagej.cellpose import (
+    DIMENSION_MODE_2D,
+    DIMENSION_MODE_3D,
+    VALID_MODELS,
+    build_cellpose_settings,
+)
+
+
+def test_valid_models_list():
+    """Test that VALID_MODELS contains the expected model names."""
+    assert "cyto" in VALID_MODELS
+    assert "cyto2" in VALID_MODELS
+    assert "cyto3" in VALID_MODELS
+    assert "nuclei" in VALID_MODELS
+
+
+def test_build_cellpose_settings_invalid_model():
+    """Test that build_cellpose_settings raises ValueError for invalid model."""
+    with pytest.raises(ValueError, match="not valid"):
+        build_cellpose_settings(
+            model="invalid_model",
+            diameter=30.0,
+            cyto_channel=1,
+        )
+
+
+def test_build_cellpose_settings_invalid_dimension_mode():
+    """Test that build_cellpose_settings raises ValueError for invalid dimension_mode."""
+    with pytest.raises(ValueError, match="dimension_mode"):
+        build_cellpose_settings(
+            model="cyto2",
+            diameter=30.0,
+            cyto_channel=1,
+            dimension_mode="4D",
+        )
+
+
+def test_build_cellpose_settings_case_insensitive_model():
+    """Test that build_cellpose_settings accepts model names case-insensitively."""
+    # Should not raise
+    build_cellpose_settings(model="CYTO2", diameter=30.0, cyto_channel=1)
+    build_cellpose_settings(model="Nuclei", diameter=30.0, cyto_channel=1)
+    build_cellpose_settings(model="CYTO3", diameter=30.0, cyto_channel=1)
+
+
+def test_build_cellpose_settings_valid_2d():
+    """Test that build_cellpose_settings succeeds for 2D mode."""
+    settings = build_cellpose_settings(
+        model="cyto2",
+        diameter=30.0,
+        cyto_channel=1,
+        nuclei_channel=2,
+        dimension_mode=DIMENSION_MODE_2D,
+    )
+    assert settings is not None
+
+
+def test_build_cellpose_settings_valid_3d():
+    """Test that build_cellpose_settings succeeds for 3D mode."""
+    settings = build_cellpose_settings(
+        model="nuclei",
+        diameter=0.0,
+        cyto_channel=1,
+        dimension_mode=DIMENSION_MODE_3D,
+    )
+    assert settings is not None
+
+
+def test_dimension_mode_constants():
+    """Test that dimension mode constants have expected values."""
+    assert DIMENSION_MODE_2D == "2D"
+    assert DIMENSION_MODE_3D == "3D"