From 0953d5864971b12850638c4ee53a6c3be16505e7 Mon Sep 17 00:00:00 2001 From: Ajey Pai K Date: Wed, 20 Dec 2023 13:22:42 +0100 Subject: [PATCH] Refactoring, Exception handling. --- dlup/writers.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_writers.py | 39 +++++++++++++++++---------------------- 2 files changed, 52 insertions(+), 22 deletions(-) diff --git a/dlup/writers.py b/dlup/writers.py index 42860b39..3b2e38e6 100644 --- a/dlup/writers.py +++ b/dlup/writers.py @@ -13,6 +13,7 @@ import numpy as np import numpy.typing as npt import PIL.Image +import PIL.ImageColor from pyvips.enums import Kernel from tifffile import tifffile @@ -62,6 +63,38 @@ class TiffCompression(str, Enum): } +def _color_dict_to_clut(color_map: dict[int, str]) -> npt.NDArray[np.uint16]: + """ + Convert a color map to a color look-up table (LUT). + + Parameters + ---------- + color_map : dict + Color map to convert. The keys are the indices and the values are the color names. + + Returns + ------- + npt.NDArray[np.uint16] + Color LUT as a 3x256 array. + """ + # Convert color names to RGB values + rgb_color_map = {index: PIL.ImageColor.getrgb(color_name) for index, color_name in color_map.items()} + + # Initialize a 3x256 CLUT (for 8-bit images) + color_lut = np.zeros((3, 256), dtype=np.uint16) + + # Prepare indices and corresponding colors for assignment + indices = np.array(list(rgb_color_map.keys())) + colors = np.array(list(rgb_color_map.values())) * 256 # Scale to 16-bit color depth + + # Assign colors to clut using advanced indexing + color_lut[0, indices] = colors[:, 0] # Red channel + color_lut[1, indices] = colors[:, 1] # Green channel + color_lut[2, indices] = colors[:, 2] # Blue channel + + return color_lut + + class ImageWriter: """Base writer class""" @@ -194,6 +227,8 @@ def from_tiles_iterator(self, iterator: Iterator[npt.NDArray[np.int_]]) -> None: _compression = TIFFFILE_COMPRESSION[self._compression.value] is_rgb = self._size[-1] in (3, 4) + if is_rgb and self._colormap is not None: + raise ValueError("Cannot use a colormap with an RGB image.") with tempfile.TemporaryDirectory() as temp_dir: temp_filename = pathlib.Path(temp_dir) / filename.name tiff_writer = tifffile.TiffWriter(temp_filename, bigtiff=True) diff --git a/tests/test_writers.py b/tests/test_writers.py index d94fae2a..bc37ed8b 100644 --- a/tests/test_writers.py +++ b/tests/test_writers.py @@ -10,34 +10,16 @@ from dlup import SlideImage from dlup.experimental_backends import ImageBackend from dlup.utils.pyvips_utils import vips_to_numpy -from dlup.writers import TiffCompression, TifffileImageWriter - - -def _color_dict_to_clut(color_map: dict[int, str]) -> tuple[np.ndarray, dict[int, tuple[int, int, int]]]: - # Convert color names to RGB values - rgb_color_map = {index: ImageColor.getrgb(color_name) for index, color_name in color_map.items()} - - # Initialize a 3x256 CLUT (for 8-bit images) - clut = np.zeros((3, 256), dtype=np.uint16) - - # Prepare indices and corresponding colors for assignment - indices = np.array(list(rgb_color_map.keys())) - colors = np.array(list(rgb_color_map.values())) * 256 # Scale to 16-bit color depth - - # Assign colors to clut using advanced indexing - clut[0, indices] = colors[:, 0] # Red channel - clut[1, indices] = colors[:, 1] # Green channel - clut[2, indices] = colors[:, 2] # Blue channel - - return clut, rgb_color_map - +from dlup.writers import TiffCompression, TifffileImageWriter, _color_dict_to_clut color_map = { 1: "green", 2: "red", 3: "yellow", } -clut, rgb_color_map = _color_dict_to_clut(color_map) +clut = _color_dict_to_clut(color_map) + +rgb_color_map = {index: ImageColor.getrgb(color_name) for index, color_name in color_map.items()} class TestTiffWriter: @@ -129,3 +111,16 @@ def test_color_map(self, shape, target_mpp): assert np.all(bottom_left == rgb_color_map[2]) bottom_right = np.asarray(thumbnail).astype(np.uint8)[256:512, 256:512] assert np.all(bottom_right == rgb_color_map[3]) + + def test_image_type(self): + # Test to raise a value error if color_map is defined for an RGB image. + random_array = np.random.randint(low=0, high=255, size=(512, 512, 3), dtype=np.uint8) + pil_image = Image.fromarray(random_array, mode="RGB") + size = (*pil_image.size, 3) + + with tempfile.NamedTemporaryFile(suffix=".tiff") as temp_tiff: + writer = TifffileImageWriter( + temp_tiff.name, size=size, mpp=(0.25, 0.25), compression=TiffCompression.NONE, colormap=clut + ) + with pytest.raises(ValueError, match="Cannot use a colormap with an RGB image."): + writer.from_pil(pil_image)