Merge pull request #72 from imbasimba/integer-fits

Support for 32 bit integer FITS

toasty/image.py

@@ -93,6 +93,8 @@ def _array_to_mode(array):
             return ImageMode.F64
         elif array.dtype.kind == "u" and array.dtype.itemsize == 1:
             return ImageMode.U8
+        elif array.dtype.kind == "i" and array.dtype.itemsize == 4:
+            return ImageMode.I32
     elif array.ndim == 3:
         if array.shape[2] == 3:
             if array.dtype.kind == "f" and array.itemsize == 2:
@@ -142,15 +144,23 @@ class ImageMode(Enum):
 
     U8 = "U8"
 
+    I32 = "I32"
+    "32-bit integer data."
+
     @classmethod
     def from_array_info(cls, shape, dtype):
+        # Make sure we have an actual dtype instance to work with:
+        dtype = np.dtype(dtype)
+
         if len(shape) == 2:
             if dtype.kind == "f" and dtype.itemsize == 4:
                 return cls.F32
             elif dtype.kind == "f" and dtype.itemsize == 8:
                 return cls.F64
             elif dtype.kind == "u" and dtype.itemsize == 1:
                 return cls.U8
+            elif dtype.kind == "i" and dtype.itemsize == 4:
+                return cls.I32
         elif len(shape) == 3:
             if shape[2] == 3:
                 if dtype.kind == "f" and dtype.itemsize == 2:
@@ -202,6 +212,8 @@ def make_maskable_buffer(self, buf_height, buf_width):
             arr = np.empty((buf_height, buf_width, 3), dtype=np.float16)
         elif self == ImageMode.U8:
             arr = np.empty((buf_height, buf_width), dtype=np.uint8)
+        elif self == ImageMode.I32:
+            arr = np.empty((buf_height, buf_width), dtype=np.int32)
         else:
             raise Exception("unhandled mode in make_maskable_buffer()")
 
@@ -891,6 +903,7 @@ def default_format(self):
                 ImageMode.F64,
                 ImageMode.F16x3,
                 ImageMode.U8,
+                ImageMode.I32,
             ):
                 return "npy"
         else:
@@ -1046,7 +1059,7 @@ def fill_into_maskable_buffer(self, buffer, iy_idx, ix_idx, by_idx, bx_idx):
             b.fill(0)
             b[by_idx, bx_idx, :3] = i[iy_idx, ix_idx]
             b[by_idx, bx_idx, 3] = 255
-        elif self.mode in (ImageMode.RGBA, ImageMode.U8):
+        elif self.mode in (ImageMode.RGBA, ImageMode.U8, ImageMode.I32):
             b.fill(0)
             b[by_idx, bx_idx] = i[iy_idx, ix_idx]
         elif self.mode in (ImageMode.F32, ImageMode.F64, ImageMode.F16x3):
@@ -1103,7 +1116,7 @@ def update_into_maskable_buffer(self, buffer, iy_idx, ix_idx, by_idx, bx_idx):
             valid = ~np.any(np.isnan(sub_i), axis=2)
             valid = np.broadcast_to(valid[..., None], sub_i.shape)
             np.putmask(sub_b, valid, sub_i)
-        elif self.mode == ImageMode.U8:
+        elif self.mode in (ImageMode.U8, ImageMode.I32):
             # zero is our maskval, so here's a convenient way to get pretty good
             # update semantics. It will behave unusually if two buffers overlap
             # and disagree on their non-zero pixel values: instead of the second
@@ -1193,7 +1206,7 @@ def make_thumbnail_bitmap(self):
         be saved in JPEG format.
 
         """
-        if self.mode in (ImageMode.F32, ImageMode.F64, ImageMode.F16x3):
+        if self.mode in (ImageMode.I32, ImageMode.F32, ImageMode.F64, ImageMode.F16x3):
             raise Exception("cannot thumbnail-ify non-RGB Image")
 
         THUMB_SHAPE = (96, 45)
@@ -1242,7 +1255,7 @@ def clear(self):
         with NaNs.
 
         """
-        if self._mode in (ImageMode.RGB, ImageMode.RGBA, ImageMode.U8):
+        if self._mode in (ImageMode.RGB, ImageMode.RGBA, ImageMode.U8, ImageMode.I32):
             self.asarray().fill(0)
         elif self._mode in (ImageMode.F32, ImageMode.F64, ImageMode.F16x3):
             self.asarray().fill(np.nan)

toasty/tests/test_modes.py

@@ -0,0 +1,62 @@
+# -*- mode: python; coding: utf-8 -*-
+# Copyright 2022 the AAS WorldWide Telescope project
+# Licensed under the MIT License.
+
+import numpy as np
+import pytest
+
+from ..image import Image, ImageMode
+
+
+class ModeInfo(object):
+    def __init__(self, mode, dtype, demo_shape, default_format):
+        self.mode = mode
+        self.dtype = dtype
+        self.demo_shape = demo_shape
+        self.default_format = default_format
+
+
+MODE_INFO = [
+    ModeInfo(ImageMode.RGB, np.uint8, (10, 10, 3), "png"),
+    ModeInfo(ImageMode.RGBA, np.uint8, (10, 10, 4), "png"),
+    ModeInfo(ImageMode.F32, np.float32, (10, 10), "npy"),
+    ModeInfo(ImageMode.F64, np.float64, (10, 10), "npy"),
+    ModeInfo(ImageMode.F16x3, np.float16, (10, 10, 3), "npy"),
+    ModeInfo(ImageMode.U8, np.uint8, (10, 10), "npy"),
+    ModeInfo(ImageMode.I32, np.int32, (10, 10), "npy"),
+]
+
+
+def test_from_array_info():
+    for mi in MODE_INFO:
+        assert ImageMode.from_array_info(mi.demo_shape, mi.dtype) == mi.mode
+
+    with pytest.raises(ValueError):
+        s = (10, 10)
+        ImageMode.from_array_info(s, np.complex64)
+
+        s = (10, 10, 3)
+        ImageMode.from_array_info(s, np.float32)
+
+        s = (10, 10, 4)
+        ImageMode.from_array_info(s, np.int8)
+
+
+def test_image_interface():
+    for mi in MODE_INFO:
+        arr = np.zeros(mi.demo_shape, dtype=mi.dtype)
+        img = Image.from_array(arr)
+        img._default_format = None
+        assert img.default_format == mi.default_format
+        mb = mi.mode.make_maskable_buffer(8, 8)
+        img.fill_into_maskable_buffer(
+            mb, slice(1, 5), slice(1, 5), slice(4, 8), slice(4, 8)
+        )
+
+        mb.asarray().fill(1)
+        img.update_into_maskable_buffer(
+            mb, slice(1, 5), slice(1, 5), slice(4, 8), slice(4, 8)
+        )
+        assert mb.asarray().flat[0] == 1
+
+        img.clear()