Merge pull request #16 from fmi-faim/copilot/add-channel-list-parameter

Add channel_list parameter to sanitize_pixels/sanitize_image

Author: imagejan

src/ome2xarray/companion.py

@@ -11,7 +11,7 @@
 import xarray as xr
 
 
-def sanitize_pixels(image: Image, include_sg: bool = False) -> Pixels:
+def sanitize_pixels(image: Image, *, include_sg: bool = False, channel_list: list[str] | None = None) -> Pixels:
     """
     Sanitize incomplete/corrupted pixels by regenerating tiff_data_blocks and planes.
     
@@ -27,6 +27,10 @@ def sanitize_pixels(image: Image, include_sg: bool = False) -> Pixels:
     include_sg : bool, optional
         If True, include stage group suffix (_sg1, _sg2, etc.) in file names.
         Default is False.
+    channel_list : list[str] | None, optional
+        List of channel names in the correct order. If provided, this list will be
+        used instead of pixels.channels to determine channel names.
+        Default is None (use pixels.channels).
         
     Returns:
     --------
@@ -50,11 +54,8 @@ def sanitize_pixels(image: Image, include_sg: bool = False) -> Pixels:
         # Parse stage_label.name like "0:Number1_sg:0" or "4:Position5:0"
         # Extract the first number after splitting by ":"
         parts = image.stage_label.name.split(':')
-        try:
-            stage_pos = int(parts[0])
-            stage_suffix = f"_s{stage_pos + 1}"
-        except (ValueError, IndexError):
-            pass
+        stage_pos = int(parts[0])
+        stage_suffix = f"_s{stage_pos + 1}"
 
         # Extract stage group if include_sg is True
         if include_sg:
@@ -96,24 +97,24 @@ def sanitize_pixels(image: Image, include_sg: bool = False) -> Pixels:
 
     # Track UUIDs per file name to ensure consistency
     file_uuids = {}
-    
-    for c in range(pixels.size_c):
-        channel = pixels.channels[c]
-        channel_name = channel.name or f"Channel{c}"
-        
+
+    if channel_list is None:
+        channel_list = [channel.name for channel in pixels.channels]
+
+    for c, channel_name in enumerate(channel_list):
         # Generate file name for this channel
         # Format: {base_name}_w{c+1}{channel_name}{sg_suffix}{stage_suffix}{time_suffix}.ome.tif
         file_base = f"{base_name}_w{c+1}{channel_name}{sg_suffix}{stage_suffix}"
-        
+
         for t in range(pixels.size_t):
             # Add time suffix only if there are multiple timepoints
             time_suffix = f"_t{t+1}" if pixels.size_t > 1 else ""
             file_name = f"{file_base}{time_suffix}.ome.tif"
-            
+
             # Generate a unique UUID for this file if not already done
             if file_name not in file_uuids:
                 file_uuids[file_name] = f"urn:uuid:{uuid4()}"
-            
+
             for z in range(pixels.size_z):
                 # Create TiffData block
                 tiff_data = TiffData(
@@ -213,7 +214,7 @@ def get_ome_metadata(self) -> OME:
         """
         return self._ome
 
-    def sanitize_image(self, image_index: int, include_sg: bool = False) -> None:
+    def sanitize_image(self, image_index: int, include_sg: bool = False, channel_list: list[str] | None = None) -> None:
         """
         Sanitize an image's pixels by regenerating tiff_data_blocks and planes.
         
@@ -227,6 +228,10 @@ def sanitize_image(self, image_index: int, include_sg: bool = False) -> None:
         include_sg : bool, optional
             If True, include stage group suffix (_sg1, _sg2, etc.) in file names.
             Default is False.
+        channel_list : list[str] | None, optional
+            List of channel names in the correct order. If provided, this list will be
+            used instead of pixels.channels to determine channel names. The length must
+            match pixels.size_c. Default is None (use pixels.channels).
         """
         if image_index < 0 or image_index >= len(self._ome.images):
             raise IndexError(
@@ -236,7 +241,7 @@ def sanitize_image(self, image_index: int, include_sg: bool = False) -> None:
         image = self._ome.images[image_index]
 
         # Sanitize the pixels
-        sanitized_pixels = sanitize_pixels(image, include_sg=include_sg)
+        sanitized_pixels = sanitize_pixels(image, include_sg=include_sg, channel_list=channel_list)
 
         # Replace the image's pixels with sanitized version
         # We need to create a new image with the sanitized pixels

tests/test_sanitize.py

@@ -289,29 +289,6 @@ def test_companion_file_sanitize_image_invalid_index():
         companion_file.sanitize_image(-1)
 
 
-def test_sanitize_pixels_channel_without_name():
-    """Test sanitization with a channel that has no name."""
-    pixels = Pixels(
-        id="Pixels:0",
-        dimension_order="XYZCT",
-        type="uint16",
-        size_x=512,
-        size_y=512,
-        size_z=1,
-        size_c=1,
-        size_t=1,
-        channels=[Channel(id="Channel:0", name=None)],  # No name
-        tiff_data_blocks=[],
-        planes=[]
-    )
-
-    image = Image(id="Image:0", name="test", pixels=pixels)
-    sanitized = sanitize_pixels(image)
-
-    # Should generate a default channel name
-    assert sanitized.tiff_data_blocks[0].uuid.file_name == "test_w1Channel0.ome.tif"
-
-
 def test_sanitize_pixels_removes_stage_label_from_image_name():
     """Test that stage label is removed from image.name to form basename."""
     pixels = Pixels(
@@ -379,13 +356,17 @@ def test_sanitize_pixels_with_include_sg():
     assert sanitized.tiff_data_blocks[0].uuid.file_name == "test_w1DAPI_sg1_s1.ome.tif"
 
 
-def test_companion_file_filenames_with_and_without_sanitize_sg1():
+@pytest.mark.parametrize("filename", [
+    "20250910_test4ch_2roi_3z_1_sg1.companion.ome",
+    "20250910_test4ch_2roi_3z_1_sg2.companion.ome"
+])
+def test_companion_file_filenames_with_and_without_sanitize_sg1(filename):
     """Test that filenames match between original and sanitized (with include_sg=True) for sg1 dataset."""
     companion_file_path = (
         Path(__file__).parent
         / "resources"
         / "20250910_VV7-0-0-6-ScanSlide"
-        / "20250910_test4ch_2roi_3z_1_sg1.companion.ome"
+        / filename
     )
 
     companion_file = CompanionFile(companion_file_path)
@@ -409,13 +390,13 @@ def test_companion_file_filenames_with_and_without_sanitize_sg1():
     )
 
 
-def test_companion_file_filenames_with_and_without_sanitize_sg2():
-    """Test that filenames match between original and sanitized (with include_sg=True) for sg2 dataset."""
+def test_companion_file_filenames_without_include_sg():
+    """Test that filenames are different when sanitizing without include_sg for sg1/sg2 datasets."""
     companion_file_path = (
         Path(__file__).parent
         / "resources"
         / "20250910_VV7-0-0-6-ScanSlide"
-        / "20250910_test4ch_2roi_3z_1_sg2.companion.ome"
+        / "20250910_test4ch_2roi_3z_1_sg1.companion.ome"
     )
 
     companion_file = CompanionFile(companion_file_path)
@@ -425,22 +406,99 @@ def test_companion_file_filenames_with_and_without_sanitize_sg2():
     original_image = metadata.images[0]
     original_filenames = {block.uuid.file_name for block in original_image.pixels.tiff_data_blocks}
 
-    # Sanitize with include_sg=True
-    companion_file.sanitize_image(0, include_sg=True)
+    # Sanitize without include_sg (default behavior)
+    companion_file.sanitize_image(0, include_sg=False)
     sanitized_metadata = companion_file.get_ome_metadata()
     sanitized_image = sanitized_metadata.images[0]
     sanitized_filenames = {block.uuid.file_name for block in sanitized_image.pixels.tiff_data_blocks}
 
-    # The filenames should match
-    assert original_filenames == sanitized_filenames, (
-        f"Filenames don't match.\n"
-        f"Original: {sorted(original_filenames)}\n"
-        f"Sanitized: {sorted(sanitized_filenames)}"
+    # The filenames should NOT match (original has _sg1, sanitized without include_sg doesn't)
+    assert original_filenames != sanitized_filenames, (
+        "Filenames should be different when include_sg=False"
     )
+    
+    # Verify that original has _sg1 and sanitized doesn't
+    for orig_fn in original_filenames:
+        assert "_sg1_" in orig_fn, f"Original filename should contain _sg1_: {orig_fn}"
+    
+    for san_fn in sanitized_filenames:
+        assert "_sg" not in san_fn, f"Sanitized filename should not contain _sg: {san_fn}"
 
 
-def test_companion_file_filenames_without_include_sg():
-    """Test that filenames are different when sanitizing without include_sg for sg1/sg2 datasets."""
+def test_sanitize_pixels_with_channel_list():
+    """Test sanitization with custom channel_list parameter."""
+    pixels = Pixels(
+        id="Pixels:0",
+        dimension_order="XYZCT",
+        type="uint16",
+        size_x=512,
+        size_y=512,
+        size_z=1,
+        size_c=3,
+        size_t=1,
+        channels=[
+            Channel(id="Channel:0", name="Red"),
+            Channel(id="Channel:1", name="Green"),
+            Channel(id="Channel:2", name="Blue")
+        ],
+        tiff_data_blocks=[],
+        planes=[
+            Plane(the_c=0, the_t=0, the_z=0, position_x=0.0, position_y=0.0, position_z=0.0),
+        ]
+    )
+
+    image = Image(
+        id="Image:0",
+        name="test_image",
+        pixels=pixels
+    )
+
+    # Sanitize with custom channel list (different order)
+    channel_list = ["DAPI", "GFP"]
+    sanitized = sanitize_pixels(image, channel_list=channel_list)
+
+    # Check that the custom channel names are used
+    assert len(sanitized.tiff_data_blocks) == 2
+    assert sanitized.tiff_data_blocks[0].uuid.file_name == "test_image_w1DAPI.ome.tif"
+    assert sanitized.tiff_data_blocks[1].uuid.file_name == "test_image_w2GFP.ome.tif"
+
+
+def test_sanitize_pixels_channel_list_none_uses_default():
+    """Test that passing channel_list=None uses the default behavior."""
+    pixels = Pixels(
+        id="Pixels:0",
+        dimension_order="XYZCT",
+        type="uint16",
+        size_x=512,
+        size_y=512,
+        size_z=1,
+        size_c=2,
+        size_t=1,
+        channels=[
+            Channel(id="Channel:0", name="OriginalA"),
+            Channel(id="Channel:1", name="OriginalB")
+        ],
+        tiff_data_blocks=[],
+        planes=[]
+    )
+
+    image = Image(
+        id="Image:0",
+        name="test_image",
+        pixels=pixels
+    )
+
+    # Sanitize with channel_list=None (explicit)
+    sanitized = sanitize_pixels(image, channel_list=None)
+
+    # Check that the original channel names are used
+    assert len(sanitized.tiff_data_blocks) == 2
+    assert sanitized.tiff_data_blocks[0].uuid.file_name == "test_image_w1OriginalA.ome.tif"
+    assert sanitized.tiff_data_blocks[1].uuid.file_name == "test_image_w2OriginalB.ome.tif"
+
+
+def test_companion_file_sanitize_image_with_channel_list():
+    """Test CompanionFile.sanitize_image with channel_list parameter."""
     companion_file_path = (
         Path(__file__).parent
         / "resources"
@@ -451,24 +509,23 @@ def test_companion_file_filenames_without_include_sg():
     companion_file = CompanionFile(companion_file_path)
     metadata = companion_file.get_ome_metadata()
 
-    # Get the first image
+    # Get original channel names for the first image
     original_image = metadata.images[0]
-    original_filenames = {block.uuid.file_name for block in original_image.pixels.tiff_data_blocks}
+    original_channels = [ch.name for ch in original_image.pixels.channels]
 
-    # Sanitize without include_sg (default behavior)
-    companion_file.sanitize_image(0, include_sg=False)
-    sanitized_metadata = companion_file.get_ome_metadata()
-    sanitized_image = sanitized_metadata.images[0]
-    sanitized_filenames = {block.uuid.file_name for block in sanitized_image.pixels.tiff_data_blocks}
+    # Create a custom channel list with different names
+    custom_channel_list = [f"Custom{i}" for i in range(len(original_channels))]
 
-    # The filenames should NOT match (original has _sg1, sanitized without include_sg doesn't)
-    assert original_filenames != sanitized_filenames, (
-        "Filenames should be different when include_sg=False"
-    )
+    # Sanitize with custom channel list
+    companion_file.sanitize_image(0, channel_list=custom_channel_list)
 
-    # Verify that original has _sg1 and sanitized doesn't
-    for orig_fn in original_filenames:
-        assert "_sg1_" in orig_fn, f"Original filename should contain _sg1_: {orig_fn}"
+    # Get sanitized metadata
+    sanitized_metadata = companion_file.get_ome_metadata()
+    sanitized_image = sanitized_metadata.images[0]
 
-    for san_fn in sanitized_filenames:
-        assert "_sg" not in san_fn, f"Sanitized filename should not contain _sg: {san_fn}"
+    # Check that custom channel names are used in filenames
+    for block in sanitized_image.pixels.tiff_data_blocks:
+        # The filename should contain one of the custom channel names
+        assert any(custom_ch in block.uuid.file_name for custom_ch in custom_channel_list), (
+            f"Filename {block.uuid.file_name} doesn't contain any custom channel name"
+        )