toasty/collection.py: start sketching out an image-collection framework

docs/api.rst

@@ -14,6 +14,10 @@ Python API Reference
    :no-inheritance-diagram:
    :no-inherited-members:
 
+.. automodapi:: toasty.collection
+   :no-inheritance-diagram:
+   :no-inherited-members:
+
 .. automodapi:: toasty.image
    :no-inheritance-diagram:
    :no-inherited-members:

docs/api/toasty.collection.ImageCollection.rst

@@ -0,0 +1,19 @@
+ImageCollection
+===============
+
+.. currentmodule:: toasty.collection
+
+.. autoclass:: ImageCollection
+   :show-inheritance:
+
+   .. rubric:: Methods Summary
+
+   .. autosummary::
+
+      ~ImageCollection.descriptions
+      ~ImageCollection.images
+
+   .. rubric:: Methods Documentation
+
+   .. automethod:: descriptions
+   .. automethod:: images

docs/api/toasty.collection.SimpleFitsCollection.rst

@@ -0,0 +1,19 @@
+SimpleFitsCollection
+====================
+
+.. currentmodule:: toasty.collection
+
+.. autoclass:: SimpleFitsCollection
+   :show-inheritance:
+
+   .. rubric:: Methods Summary
+
+   .. autosummary::
+
+      ~SimpleFitsCollection.descriptions
+      ~SimpleFitsCollection.images
+
+   .. rubric:: Methods Documentation
+
+   .. automethod:: descriptions
+   .. automethod:: images

toasty/collection.py

@@ -0,0 +1,94 @@
+# -*- mode: python; coding: utf-8 -*-
+# Copyright 2021 the AAS WorldWide Telescope project
+# Licensed under the MIT License.
+
+"""
+Collections of related input images.
+
+Some Toasty processing tasks operate over collections of input images that are
+related in some way. This module provides standardized mechanisms for
+manipulating such collections. In particular, it provides a framework for
+scanning over image "descriptions" without reading in the complete image data.
+This can be useful because many collection-related operations want to do an
+initial pass over the collection to gather some global information, then a
+second pass with the actual data processing.
+"""
+
+__all__ = '''
+ImageCollection
+SimpleFitsCollection
+'''.split()
+
+from abc import ABC
+import numpy as np
+import warnings
+
+from .image import Image, ImageDescription, ImageMode
+from .study import StudyTiling
+
+
+class ImageCollection(ABC):
+    def descriptions(self):
+        """
+        Generate a sequence of :class:`toasty.image.ImageDescription` items
+        associated with this collection.
+
+        Each description will have an added string attribute ``collection_id``
+        that gives a unique textual identifer for the item in the collection.
+
+        Unlike :meth:`ImageCollection.images`, this function does cause the full
+        data for each image to be loaded.
+        """
+        raise NotImplementedError()
+
+    def images(self):
+        """
+        Generate a sequence of :class:`toasty.image.Image` items associated with
+        this collection.
+
+        Each image will have an added string attribute ``collection_id`` that
+        gives a unique textual identifer for the item in the collection.
+        """
+        raise NotImplementedError()
+
+
+class SimpleFitsCollection(ImageCollection):
+    def __init__(self, paths, hdu_index=None):
+        self._paths = list(paths)
+        self._hdu_index = hdu_index
+
+    def _load(self, actually_load_data):
+        from astropy.io import fits
+        from astropy.wcs import WCS
+
+        for fits_path in self._paths:
+            with fits.open(fits_path) as hdul:
+                if self._hdu_index is not None:
+                    hdu = hdul[self._hdu_index]
+                else:
+                    for hdu in hdul:
+                        if len(hdu.shape) > 1:
+                            break
+
+                wcs = WCS(hdu.header)
+
+                if actually_load_data:
+                    result = Image.from_array(hdu.data, wcs=wcs, default_format='fits')
+                else:
+                    shape = hdu.shape
+                    
+                    if hasattr(hdu, 'dtype'):
+                        mode = ImageMode.from_array_info(shape, hdu.dtype)
+                    else:
+                        mode = None  # CompImageHDU doesn't have dtype
+
+                    result = ImageDescription(mode=mode, shape=shape, wcs=wcs)
+
+                result.collection_id = fits_path
+                yield result
+
+    def descriptions(self):
+        return self._load(False)
+
+    def images(self):
+        return self._load(True)

toasty/image.py

@@ -1,5 +1,5 @@
 # -*- mode: python; coding: utf-8 -*-
-# Copyright 2020 the AAS WorldWide Telescope project
+# Copyright 2020-2021 the AAS WorldWide Telescope project
 # Licensed under the MIT License.
 
 """
@@ -72,10 +72,45 @@ class ImageMode(Enum):
 
     """
     RGB = 'RGB'
+    "24-bit color with three uint8 channels for red, green, and blue."
+
     RGBA = 'RGBA'
+    "32-bit color with four uint8 channels for red, green, blue, and alpha (transparency)."
+
     F32 = 'F'
+    "32-bit floating-point scalar data."
+
     F64 = 'D'
+    "64-bit floating-point scalar data."
+
     F16x3 = 'F16x3'
+    """
+    48-bit color with three 16-bit floating-point channels for red, green, and
+    blue.
+
+    This mode is useful for high-dynamic-range image processing and can be
+    stored in the OpenEXR file format.
+    """
+
+    @classmethod
+    def from_array_info(cls, shape, dtype):
+        if len(shape) == 2 and dtype.kind == 'f':
+            if dtype.itemsize == 4:
+                return cls.F32
+            elif dtype.itemsize == 8:
+                return cls.F64
+        elif len(shape) == 3:
+            if shape[2] == 3:
+                if dtype.kind == 'f' and dtype.itemsize == 2:
+                    return cls.F16x3
+                elif dtype.kind == 'u' and dtype.itemsize == 1:
+                    return cls.RGB
+            elif shape[2] == 4:
+                if dtype.kind == 'u' and dtype.itemsize == 1:
+                    return cls.RGBA
+
+        raise ValueError('Could not determine mode for array with dtype {0} and shape {1}'.format(dtype, shape))
+
 
     def make_maskable_buffer(self, buf_height, buf_width):
         """
@@ -437,12 +472,13 @@ class ImageLoader(object):
 
 
 class Image(object):
-    """A 2D data array stored in memory.
+    """
+    A 2D data array stored in memory, potential with spatial positioning information.
 
     This class primarily exists to help us abstract between the cases where we
     have "bitmap" RGB(A) images and "science" floating-point images.
-
     """
+
     _pil = None
     _array = None
     _mode = None