diff --git a/cleopatra/array_glyph.py b/cleopatra/array_glyph.py index ee28109..de11451 100644 --- a/cleopatra/array_glyph.py +++ b/cleopatra/array_glyph.py @@ -548,7 +548,7 @@ def apply_colormap(self, cmap: Union[Colormap, str]) -> np.ndarray: Parameters ---------- - cmap: LinearSegmentedColormap + cmap: LinearSegmentedColormap/str colormap. Returns @@ -597,8 +597,8 @@ def to_image(self) -> Image.Image: """ # This is done to scale the values between 0 and 255 - # TODO: not sure if scaling the array to rgb is necessary here. - return Image.fromarray(self.scale_to_rgb()).convert("RGB") + arr = self.arr if self.arr.dtype == "uint8" else self.scale_to_rgb() + return Image.fromarray(arr).convert("RGB") def scale_to_rgb(self) -> np.ndarray: """Create an RGB image. diff --git a/tests/test_array_glyph.py b/tests/test_array_glyph.py index 519e917..dcac53b 100644 --- a/tests/test_array_glyph.py +++ b/tests/test_array_glyph.py @@ -355,12 +355,20 @@ def test_scale_to_rgb(): assert rgb_arr.dtype == "uint8" -def test_to_image(): - arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]) - my_glyph = ArrayGlyph(arr) +class TestToImage: + def test_int64_arr(self): + arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]) + my_glyph = ArrayGlyph(arr) + + image = my_glyph.to_image() + assert isinstance(image, Image.Image) + + def test_uint_arr(self): + arr = np.random.randint(0, 255, size=(4, 4)).astype(np.uint8) + my_glyph = ArrayGlyph(arr) - image = my_glyph.to_image() - assert isinstance(image, Image.Image) + image = my_glyph.to_image() + assert isinstance(image, Image.Image) def test_adjust_ticks():