fix: multichannel export as tiff (#565)

Author: KRiedmiller

panseg/core/image.py

@@ -509,15 +509,8 @@ def _get_data_channel_layout(
             data = dp.normalize_01(data)
         return data
 
-    def _get_data(self, normalize_01: bool = True) -> np.ndarray:
-        """Get the data if the layout is not multichannel."""
-        data = self._data
-        if normalize_01:
-            data = dp.normalize_01(data)
-        return data
-
     def get_data(
-        self, channel: int | None = None, normalize_01: bool = True
+        self, channel: int | None = None, normalize_01: bool = False
     ) -> np.ndarray:
         """Returns the data of the image.
 
@@ -527,13 +520,15 @@ def get_data(
             normalize_01 (bool): Normalize the data between 0 and 1, if the
                 image is a Label image, the data is not normalized.
         """
+        data = self._data
         if self.image_type == ImageType.LABEL:
-            return self._data
+            return data
 
         if self.channel_axis is not None:
-            return self._get_data_channel_layout(channel, normalize_01)
-
-        return self._get_data(normalize_01)
+            data = self._get_data_channel_layout(channel, normalize_01)
+        elif normalize_01:
+            data = dp.normalize_01(data)
+        return data
 
     @property
     def scale(self) -> tuple[float, ...]:
@@ -728,6 +723,8 @@ def import_image(
             original_voxel_size=voxel_size,
             source_file_name=path.stem,
         )
+        if image_properties.image_type == ImageType.IMAGE:
+            data = dp.normalize_01(data)
 
         return PanSegImage(data=data, properties=image_properties)
 

panseg/io/tiff.py

@@ -198,7 +198,8 @@ def create_tiff(
 
     elif layout == "CZYX":
         assert stack.ndim == 4, "Stack dimensions must be in CZYX order"
-        c, z, y, x = stack.shape
+        stack = np.transpose(stack, (1, 0, 2, 3))
+        z, c, y, x = stack.shape
         stack = stack.reshape(1, z, c, y, x, 1)
 
     else:

panseg/tasks/io_tasks.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Optional
 
+import numpy as np
+
 from panseg.core.image import PanSegImage, import_image, save_image
 from panseg.tasks import task_tracker
 from panseg.tasks.workflow_handler import RunTimeInputSchema, Task_message
@@ -110,9 +112,10 @@ def export_image_task(
 def merge_channels_task(**kwargs) -> PanSegImage:
     """Merge an arbitrary number of PanSegImages
 
-    Pass each image as a named argument, the name doesn't matter.
+    Pass each image as a named argument.
+    The images get sorted according to the key
     """
-    images: list[PanSegImage] = list(kwargs.values())
+    images: list[PanSegImage] = [kwargs[k] for k in sorted(kwargs.keys())]
     image = images[0].derive_new(images[0].get_data(), images[0].name + "_merged")
     for im in images[1:]:
         image = image.merge_with(im)

tests/core/test_image.py

@@ -249,7 +249,7 @@ def test_panseg_image_to_napari_layer_tuple():
 
     assert isinstance(layer_tuple, tuple)
     layer_tuple = tuple(layer_tuple)
-    np.testing.assert_allclose(layer_tuple[0], ps_image.get_data(normalize_01=True))
+    np.testing.assert_allclose(layer_tuple[0], ps_image.get_data(normalize_01=False))
     assert "metadata" in layer_tuple[1]
     assert layer_tuple[2] == ps_image.image_type.value
 

tests/functionals/dataprocessing/test_image_math.py

@@ -9,13 +9,7 @@
     process_images,
     subtract_images,
 )
-
-
-def normalize_01(image: np.ndarray) -> np.ndarray:
-    min_val, max_val = image.min(), image.max()
-    if max_val - min_val == 0:
-        return np.zeros_like(image)
-    return (image - min_val) / (max_val - min_val)
+from panseg.functionals.dataprocessing.dataprocessing import normalize_01
 
 
 @pytest.mark.parametrize(

tests/headless/test_headless.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import yaml
+from numpy.testing import assert_allclose
 
 from panseg.core.image import PanSegImage
 from panseg.headless.headless import run_headless_workflow
@@ -101,6 +102,8 @@ def test_create_workflow_channels(tmp_path):
         **{f"image_{i}": ps for i, ps in enumerate(ps_1s + [ps_2])}
     )
     assert isinstance(ps_3, PanSegImage)
+    np.testing.assert_allclose(ps_1s[0].get_data(), ps_3.get_data()[0])
+    np.testing.assert_allclose(ps_2.get_data(), ps_3.get_data()[4])
     export_image_task(
         image=ps_3,
         export_directory=path_tiff.parent,

tests/tasks/test_io_tasks.py

@@ -7,6 +7,7 @@
     PanSegImage,
     SemanticType,
 )
+from panseg.functionals.dataprocessing.dataprocessing import normalize_01
 from panseg.io.voxelsize import VoxelSize
 from panseg.tasks.io_tasks import (
     export_image_task,
@@ -17,7 +18,79 @@
 
 
 @pytest.mark.parametrize(
-    "shape, layout, export_format",
+    "shape,layout,export_format",
+    [
+        ((2, 64, 64), ImageLayout.CYX, "tiff"),
+        ((2, 64, 64), ImageLayout.CYX, "h5"),
+        ((2, 64, 64), ImageLayout.CYX, "zarr"),
+        ((2, 64, 32, 32), ImageLayout.CZYX, "h5"),
+        ((2, 64, 32, 32), ImageLayout.CZYX, "tiff"),
+        ((2, 64, 32, 32), ImageLayout.CZYX, "zarr"),
+        # ((64, 64), ImageLayout.YX, "tiff"),
+        # ((64, 64), ImageLayout.YX, "h5"),
+        # ((64, 64), ImageLayout.YX, "zarr"),
+    ],
+)
+def test_image_io_round_trip_multichannel(tmp_path, shape, layout, export_format):
+    mock_data = normalize_01(np.random.rand(*shape).astype("float32"))
+
+    property = ImageProperties(
+        name="test",
+        voxel_size=VoxelSize(voxels_size=(1.0, 1.0, 1.0), unit="um"),
+        semantic_type=SemanticType.RAW,
+        image_layout=layout,
+        original_voxel_size=VoxelSize(voxels_size=(1.0, 1.0, 1.0), unit="um"),
+        source_file_name="test",
+    )
+    image = PanSegImage(data=mock_data, properties=property)
+
+    export_image_task(
+        image=image,
+        export_directory=tmp_path,
+        name_pattern="test",
+        key="raw",
+        export_format=export_format,
+        data_type="float32",
+    )
+
+    if export_format == "tiff":
+        file_path = tmp_path / "test.tiff"
+        key = None
+        # tiff alwayes saved as ZCYX
+        if layout == ImageLayout.CZYX:
+            layout = ImageLayout.ZCYX
+
+    elif export_format == "h5":
+        file_path = tmp_path / "test.h5"
+        key = "raw"
+    else:
+        file_path = tmp_path / "test.zarr"
+        key = "raw"
+
+    imported_image = import_image_task(
+        input_path=file_path,
+        key=key,
+        image_name="test_import",
+        semantic_type="raw",
+        stack_layout=layout.name,
+        m_slicing=None,
+    )
+    assert isinstance(imported_image, list)
+
+    for i in [0, 1]:
+        original_data = image.get_data()[i]
+        imported_data = imported_image[i].get_data()
+
+        assert np.allclose(original_data, imported_data)
+        assert original_data.max() <= 1.0  # check if the normalization is applied
+        assert imported_data.max() <= 1.0
+
+        assert image.voxel_size == imported_image[i].voxel_size
+        assert image.semantic_type == imported_image[i].semantic_type
+
+
+@pytest.mark.parametrize(
+    "shape,layout,export_format",
     [
         ((32, 64, 64), ImageLayout.ZYX, "tiff"),
         ((32, 64, 64), ImageLayout.ZYX, "h5"),
@@ -69,7 +142,8 @@ def test_image_io_round_trip(tmp_path, shape, layout, export_format):
     )
     assert isinstance(imported_image, PanSegImage)
 
-    original_data = image.get_data()
+    # would be normalized during import
+    original_data = image.get_data(normalize_01=True)
     imported_data = imported_image.get_data()
 
     assert np.allclose(original_data, imported_data)
@@ -133,7 +207,7 @@ def test_label_io_round_trip(tmp_path, shape, layout, export_format):
     imported_image = import_image_task(
         input_path=file_path,
         key=key,
-        image_name="tesi_import",
+        image_name="test_import",
         semantic_type="segmentation",
         stack_layout=layout.name,
         m_slicing=None,