feat: Add saving color maps in OME tiffs

package/PartSegImage/image.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import re
 import sys
 import typing
@@ -918,7 +919,33 @@ def get_imagej_colors(self):
                             ]
                         )
                     )
-                res.append(color)
+                res.append(color.astype(np.uint8))
+        return res
+
+    def get_ome_colors(self) -> list[int]:
+        """The ome stores colors as single integer encoding RGB value
+
+        :returns: list of integers representing colors
+        """
+
+        res = []
+        default_colors = ["red", "blue", "green", "yellow", "magenta", "cyan"]
+        for i, color in enumerate(self.default_coloring):
+            if isinstance(color, str):
+                if color.startswith("#"):
+                    color_array = _hex_to_rgb(color)
+                else:
+                    color_array = _name_to_rgb(color)
+                res.append(_rgb_to_signed_int(color_array))
+            elif color.ndim == 1:
+                # treat as RGB
+                res.append(_rgb_to_signed_int(tuple(color)))
+            else:
+                logging.warning(
+                    "Do not support custom colormap in ome colors. Use %s", default_colors[i % len(default_colors)]
+                )
+                color_array = _name_to_rgb(default_colors[i % len(default_colors)])
+                res.append(_rgb_to_signed_int(color_array))
         return res
 
     def get_colors(self) -> list[str | list[int]]:
@@ -1001,6 +1028,12 @@ def _name_to_rgb(name: str) -> tuple[int, int, int]:
     return _hex_to_rgb(_NAMED_COLORS[name])
 
 
+def _rgb_to_signed_int(rgb: tuple[int, int, int]) -> int:
+    """Convert an RGB tuple to a signed integer representation."""
+    r, g, b = rgb[:3]
+    return np.uint32((r << 24) | (g << 16) | (b << 8) | 255).view(np.int32).item()
+
+
 try:
     from vispy.color import get_color_dict
 except ImportError:  # pragma: no cover

package/PartSegImage/image_writer.py

@@ -76,6 +76,7 @@ def save(cls, image: Image, save_path: typing.Union[str, BytesIO, Path], compres
         metadata["Channel"] = {
             "Name": image.channel_names,
             "axes": "TZYXC",
+            "Color": image.get_ome_colors(),
         }
         cls._save(data, save_path, metadata, compression)
 

package/tests/test_PartSegImage/test_image_writer.py

@@ -97,8 +97,9 @@ def test_imagej_write_all_metadata(tmp_path, data_test_dir):
     npt.assert_array_equal(image2.default_coloring, image.get_imagej_colors())
 
 
-def test_imagej_save_color(tmp_path):
-    data = np.zeros((4, 20, 20), dtype=np.uint8)
+@pytest.fixture
+def data_to_save():
+    data = np.zeros((5, 20, 20), dtype=np.uint8)
     data[:, 2:-2, 2:-2] = 20
     img = Image(
         data,
@@ -109,18 +110,37 @@ def test_imagej_save_color(tmp_path):
             ChannelInfo(name="ch2", color_map="#FFAA00", contrast_limits=(0, 30)),
             ChannelInfo(name="ch3", color_map="#FB1", contrast_limits=(0, 25)),
             ChannelInfo(name="ch4", color_map=(0, 180, 0), contrast_limits=(0, 22)),
+            ChannelInfo(
+                name="ch5", color_map=np.linspace((0, 0, 0), (128, 255, 0), num=256).T, contrast_limits=(0, 20)
+            ),
         ],
     )
     assert img.get_colors()[:3] == ["blue", "#FFAA00", "#FB1"]
     assert tuple(img.get_colors()[3]) == (0, 180, 0)
-    IMAGEJImageWriter.save(img, tmp_path / "image.tif")
+    return img
+
+
+def test_imagej_save_color(tmp_path, data_to_save):
+    IMAGEJImageWriter.save(data_to_save, tmp_path / "image.tif")
     image2 = TiffImageReader.read_image(tmp_path / "image.tif")
-    assert image2.channel_names == ["ch1", "ch2", "ch3", "ch4"]
-    assert image2.ranges == [(0, 20), (0, 30), (0, 25), (0, 22)]
+    assert image2.channel_names == ["ch1", "ch2", "ch3", "ch4", "ch5"]
+    npt.assert_array_equal(image2.ranges, [(0, 20), (0, 30), (0, 25), (0, 22), (0, 20)])
     assert tuple(image2.default_coloring[0][:, -1]) == (0, 0, 255)
     assert tuple(image2.default_coloring[1][:, -1]) == (255, 170, 0)
     assert tuple(image2.default_coloring[2][:, -1]) == (255, 187, 17)
     assert tuple(image2.default_coloring[3][:, -1]) == (0, 180, 0)
+    assert tuple(image2.default_coloring[4][:, -1]) == (128, 255, 0)
+
+
+def test_ome_save_color(tmp_path, data_to_save):
+    ImageWriter.save(data_to_save, tmp_path / "image.tif")
+    image2 = TiffImageReader.read_image(tmp_path / "image.tif")
+    assert image2.channel_names == ["ch1", "ch2", "ch3", "ch4", "ch5"]
+    assert tuple(image2.default_coloring[0]) == (0, 0, 255)
+    assert tuple(image2.default_coloring[1]) == (255, 170, 0)
+    assert tuple(image2.default_coloring[2]) == (255, 187, 17)
+    assert tuple(image2.default_coloring[3]) == (0, 180, 0)
+    assert tuple(image2.default_coloring[4]) == (255, 0, 255)  # fallback to magenta
 
 
 def test_save_mask_imagej(tmp_path):

pyproject.toml

@@ -1,6 +1,6 @@
 [build-system]
 # Minimum requirements for the build system to execute.
-requires = ["setuptools>=61.2.0", "setuptools_scm[toml]>=8"]  # PEP 508 specifications.
+requires = ["setuptools>=77.0.0", "setuptools_scm[toml]>=8"]  # PEP 508 specifications.
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools_scm]