Merge pull request #7 from rhornb/single_channel

Single channel input model fix

Author: rhornb

.github/workflows/build_and_test.yml

@@ -17,10 +17,10 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-22.04, macos-latest]
-        python-version: ["3.10"]
+        python-version: ["3.11"]
     name: "Core, Python ${{ matrix.python-version }}, ${{ matrix.os }}"
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 10
+    timeout-minutes: 30
 
     steps:
       - uses: actions/checkout@v4
@@ -29,7 +29,7 @@ jobs:
 
       - uses: prefix-dev/[email protected]
         with:
-          pixi-version: v0.44.0
+          pixi-version: v0.50.2
 
       - name: Check if manifest has changed
         run: pixi run fractal-manifest check --package ilastik-tasks

pyproject.toml

@@ -118,4 +118,4 @@ docstring_parser = "==0.15"
 ilastik-core = ">=1.4.2a1,<2"
 vigra = ">=1.12.1,<2"
 fractal-tasks-core = "==1.5.3"
-fractal-task-tools = "==0.0.12"
+fractal-task-tools = "==0.0.12"

src/ilastik_tasks/ilastik_pixel_classification_segmentation.py

@@ -34,6 +34,7 @@
 from ilastik_tasks.ilastik_utils import (
     IlastikChannel1InputModel,
     IlastikChannel2InputModel,
+    get_expected_number_of_channels,
 )
 from fractal_tasks_core.labels import prepare_label_group
 from fractal_tasks_core.masked_loading import masked_loading_wrapper
@@ -251,18 +252,25 @@ def ilastik_pixel_classification_segmentation(
 
     # Setup Ilastik headless shell
     shell = setup_ilastik(ilastik_model)
-
-    # Check model channel requirements
-    expected_channels = check_ilastik_model_channels(shell)
-    if expected_channels == 2 and channel2 is None:
+    
+    # Check if channel input fits expected number of channels of model
+    expected_num_channels = get_expected_number_of_channels(shell)
+    
+    if expected_num_channels == 2 and not channel2.is_set():
         raise ValueError(
             "Ilastik model expects two channels as "
             "input but only one channel was provided"
         )
-    elif expected_channels == 1 and channel2 is not None:
+    elif expected_num_channels == 1 and channel2.is_set():
         raise ValueError(
             "Ilastik model expects 1 channel as " "input but two channels were provided"
         )
+        
+    elif expected_num_channels > 2:
+        raise NotImplementedError(
+            f"Expected {expected_num_channels} channels, "
+            "but a maximum of channels are currently supported."
+        )
 
     # Find channel index
     omero_channel = channel.get_omero_channel(zarr_url)
@@ -290,7 +298,7 @@ def ilastik_pixel_classification_segmentation(
     # Load ZYX data
     data_zyx = da.from_zarr(f"{zarr_url}/{level}")[ind_channel]
     logger.info(f"{data_zyx.shape=}")
-    if channel2:
+    if channel2.is_set():
         data_zyx_c2 = da.from_zarr(f"{zarr_url}/{level}")[ind_channel_c2]
         logger.info(f"Second channel: {data_zyx_c2.shape=}")
 
@@ -437,7 +445,7 @@ def ilastik_pixel_classification_segmentation(
         logger.info(f"Now processing ROI {i_ROI+1}/{num_ROIs}")
 
         # Prepare single-channel or dual-channel input for Ilastik
-        if channel2:
+        if channel2.is_set():
             # Dual channel mode
             img_1 = load_region(
                 data_zyx,
@@ -556,24 +564,6 @@ def ilastik_pixel_classification_segmentation(
         )
 
 
-def check_ilastik_model_channels(shell) -> int:
-    """Check number of input channels expected by Ilastik model.
-
-    Args:
-        shell: Initialized Ilastik shell with loaded model
-
-    Returns:
-        int: Number of expected input channels
-    """
-    # Get dataSelection applet from workflow
-    data_selection = shell.workflow.dataSelectionApplet
-
-    # Get slot info containing expected channels
-    slot_info = data_selection.topLevelOperator.DatasetRoles.value
-
-    # Return number of expected channels
-    return len(slot_info)
-
 
 if __name__ == "__main__":
     from fractal_task_tools.task_wrapper import run_fractal_task

src/ilastik_tasks/ilastik_utils.py

@@ -111,3 +111,20 @@ def get_omero_channel(self, zarr_url) -> OmeroChannel:
                 f"Original error: {str(e)}"
             )
             return None
+
+
+def get_expected_number_of_channels(shell) -> int:
+    """
+    Get the expected number of channels from the trained ilastik model
+    """
+    opPixelClassification = shell.workflow.pcApplet.topLevelOperator
+    len_input_images = len(opPixelClassification.InputImages)
+    channel_number = []
+    for i in range(len_input_images):
+        channel = opPixelClassification.InputImages[i].meta.getTaggedShape()["c"]
+        channel_number.append(channel)
+    
+    if len(set(channel_number)) != 1:
+        raise ValueError("Inconsistent number of channels across input images.")
+    
+    return channel_number[0]

tests/data/pixel_classifier_3D_single_channel.ilp

@@ -3,16 +3,18 @@
 
 import pytest
 from devtools import debug
+
+from ilastik_tasks.ilastik_pixel_classification_segmentation import (
+    ilastik_pixel_classification_segmentation,
+)
 from ilastik_tasks.ilastik_utils import (
     IlastikChannel1InputModel,
     IlastikChannel2InputModel,
 )
-from ilastik_tasks.ilastik_pixel_classification_segmentation import (
-    ilastik_pixel_classification_segmentation,
-)
 
 # TODO: add 2D testdata
 
+
 @pytest.fixture(scope="function")
 def test_data_dir_3d(tmp_path: Path, zenodo_zarr_3d: list) -> str:
     """
@@ -26,7 +28,9 @@ def test_data_dir_3d(tmp_path: Path, zenodo_zarr_3d: list) -> str:
     return dest_dir
 
 
-def test_ilastik_pixel_classification_segmentation_task_3D(test_data_dir_3d):
+def test_ilastik_pixel_classification_segmentation_task_3D_dual_channel(
+    test_data_dir_3d,
+):
     """
     Test the 3D ilastik_pixel_classification_segmentation task with dual channel input.
     """
@@ -52,9 +56,44 @@ def test_ilastik_pixel_classification_segmentation_task_3D(test_data_dir_3d):
             zarr_url=zarr_url,
             level=4,
             channel=IlastikChannel1InputModel(label="DAPI_2"),
-            channel2=None,
+            channel2=IlastikChannel2InputModel(label=None),
             ilastik_model=str(ilastik_model),
-            output_label_name="test_label_single_channel",
+            output_label_name="test_label",
+            relabeling=True,
         )
 
 
+def test_ilastik_pixel_classification_segmentation_task_3D_single_channel(
+    test_data_dir_3d,
+):
+    """
+    Test the 3D ilastik_pixel_classification_segmentation task
+    with single channel input.
+    """
+    ilastik_model = (
+        Path(__file__).parent / "data/pixel_classifier_3D_single_channel.ilp"
+    ).as_posix()
+    zarr_url = f"{test_data_dir_3d}/B/03/0"
+
+    ilastik_pixel_classification_segmentation(
+        zarr_url=zarr_url,
+        level=4,
+        channel=IlastikChannel1InputModel(label="DAPI_2"),
+        channel2=IlastikChannel2InputModel(label=None),
+        ilastik_model=str(ilastik_model),
+        output_label_name="test_label",
+        relabeling=True,
+    )
+
+    # Test failing of task if model was trained with one channel
+    # but two are provided
+    with pytest.raises(ValueError):
+        ilastik_pixel_classification_segmentation(
+            zarr_url=zarr_url,
+            level=4,
+            channel=IlastikChannel1InputModel(label="DAPI_2"),
+            channel2=IlastikChannel2InputModel(label="ECadherin_2"),
+            ilastik_model=str(ilastik_model),
+            output_label_name="test_label",
+            relabeling=True,
+        )