Source code for pyobs.images.image

from __future__ import annotations

import copy
import io
from typing import Any, TypeVar, cast

import numpy as np
import numpy.typing as npt
from astropy.io import fits
from astropy.io.fits import ImageHDU, table_to_hdu
from astropy.nddata import CCDData, StdDevUncertainty
from astropy.table import Table
from numpy.typing import NDArray

import pyobs.utils.exceptions as exc
from pyobs.utils.fits import FilenameFormatter

MetaClass = TypeVar("MetaClass")


class Image:
    """
    A container class for astronomical image data and associated metadata.

    This class represents a two-dimensional astronomical image, typically loaded
    from or saved to a FITS file. It provides unified access to image data, mask,
    uncertainty, catalogs, raw calibration frames, and meta information. The
    `Image` class serves as the fundamental data structure within the `pyobs`
    imaging pipeline, enabling reading, writing, and conversion between FITS,
    `astropy.CCDData`, and in-memory representations.

    The image may optionally contain:
      - **data**: 2D array of pixel values.
      - **mask**: Boolean or integer mask indicating invalid or excluded pixels.
      - **uncertainty**: Per-pixel uncertainty values.
      - **catalog**: Source catalog (as an `astropy.table.Table`).
      - **raw**: Reference to the unprocessed raw image data.
      - **meta**: Arbitrary metadata dictionary, typically used to store runtime
        or class-based contextual information (not preserved in FITS I/O).

    The class supports deep/shallow copying, FITS serialization, arithmetic
    operations (e.g., division), and format conversions (e.g., to JPEG or
    `astropy.CCDData`). It is designed to be interoperable with the `pyobs.images`
    processing framework, and compatible with source detection and photometry
    tools such as SEP or Source Extractor.

    Parameters
    ----------
    data : numpy.ndarray[float], optional
        2D array containing the image pixel data.
    header : astropy.io.fits.Header, optional
        FITS header containing image metadata. If omitted, a new empty header is created.
    mask : numpy.ndarray[float] or numpy.ndarray[bool], optional
        Image mask. Pixels marked True or non-zero are ignored during processing.
    uncertainty : numpy.ndarray[float], optional
        Per-pixel uncertainty array.
    catalog : astropy.table.Table, optional
        Source catalog table associated with the image.
    raw : numpy.ndarray[float], optional
        Raw image data before calibration.
    meta : dict, optional
        Dictionary of additional metadata. Not serialized in I/O operations.
    *args, **kwargs : Any
        Additional positional or keyword arguments for subclass compatibility.

    Behavior
    --------
    - If `data` is provided, the FITS header keywords ``NAXIS1`` and ``NAXIS2`` are
      automatically set to match the array shape.
    - The `Image` can be created from FITS files (`from_file`), byte arrays
      (`from_bytes`), or `astropy.CCDData` objects (`from_ccddata`).
    - The `writeto()` method saves the image to a FITS file, including associated
      mask, uncertainty, raw, and catalog extensions when present.
    - The `to_bytes()` method serializes the image into an in-memory FITS byte stream.
    - Metadata entries can be managed via `set_meta()`, `get_meta()`, and
      `get_meta_safe()` for safe retrieval of class-based meta objects.
    - Supports lightweight arithmetic, such as division via `__truediv__`.

    Input/Output
    ------------
    - **Input**:
        - FITS file or HDU list
        - In-memory bytes (FITS format)
        - `astropy.CCDData` object
    - **Output**:
        - `Image` instance
        - FITS file or bytes
        - `astropy.CCDData` object
        - JPEG bytes (grayscale representation)

    Examples
    --------
    Load an image from a FITS file and convert it to JPEG:

    .. code-block:: python

        from pyobs.images import Image

        image = Image.from_file("science_frame.fits")
        jpeg_data = image.to_jpeg()

    Create an image from CCDData and write it back to disk:

    .. code-block:: python

        from astropy.nddata import CCDData

        ccd = CCDData.read("flat_field.fits")
        image = Image.from_ccddata(ccd)
        image.writeto("flat_field_copy.fits")

    Access image data safely:

    .. code-block:: python

        if image.safe_data is not None:
            print(f"Mean value: {image.data.mean():.3f}")

    Notes
    -----
    - Header keywords such as ``BUNIT`` and ``CD1_1`` / ``CDELT1`` are used to
      infer pixel units and scale (in arcsec/pixel).
    - The `meta` dictionary is **not persisted** in FITS I/O; it is intended for
      runtime information storage.
    - JPEG conversion (`to_jpeg`) automatically determines display scaling
      based on 5th–95th percentile clipping unless explicit limits are provided.
    - For color images (3×H×W arrays), `to_grayscale()` converts to a single
      channel using ITU-R BT.709 luminance coefficients by default.

    See Also
    --------
    astropy.io.fits : For FITS file I/O.
    astropy.nddata.CCDData : For conversion to/from standard CCD image structures.
    pyobs.images.processors.detection.SepSourceDetection : For source extraction
        using SEP on `Image` objects.
    """

    __module__ = "pyobs.images"

    def __init__(
        self,
        data: npt.NDArray[np.floating[Any]] | None = None,
        header: fits.Header | None = None,
        mask: npt.NDArray[np.floating[Any]] | None = None,
        uncertainty: npt.NDArray[np.floating[Any]] | None = None,
        catalog: Table | None = None,
        raw: npt.NDArray[np.floating[Any]] | None = None,
        meta: dict[Any, Any] | None = None,
        *args: Any,
        **kwargs: Any,
    ):
        """Init a new image.

        Args:
            data: Numpy array containing data for image.
            header: Header for the new image.
            mask: Mask for the image.
            uncertainty: Uncertainty image.
            catalog: Catalog table.
            raw: If image is calibrated, this should be the raw image.
            meta: Dictionary with meta information (note: not preserved in I/O operations!).
        """

        # store
        self._data: npt.NDArray[np.floating[Any]] | None = data
        self._header = fits.Header() if header is None else header.copy()
        self._mask: npt.NDArray[np.floating[Any]] | None = None if mask is None else mask.copy()
        self._uncertainty: npt.NDArray[np.floating[Any]] | None = None if uncertainty is None else uncertainty.copy()
        self._catalog = None if catalog is None else catalog.copy()
        self._raw: npt.NDArray[np.floating[Any]] | None = None if raw is None else raw.copy()
        self._meta = {} if meta is None else copy.deepcopy(meta)

        # add basic header stuff
        if data is not None and self._header is not None:
            self.header["NAXIS1"] = data.shape[1]
            self.header["NAXIS2"] = data.shape[0]

[docs] @classmethod def from_bytes(cls, data: bytes) -> Image: """Create Image from a bytes array containing a FITS file. Args: data: Bytes array to create image from. Returns: The new image. """ # create hdu with io.BytesIO(data) as bio: # read whole file d = fits.open(bio, memmap=False, lazy_load_hdus=False) # load image return cls._from_hdu_list(d)
[docs] @classmethod def from_file(cls, filename: str) -> Image: """Create image from FITS file. Args: filename: Name of file to load image from. Returns: New image. """ with fits.open(filename, memmap=False, lazy_load_hdus=False) as data: return cls._from_hdu_list(data)
[docs] @classmethod def from_ccddata(cls, data: CCDData) -> Image: """Create image from astropy.CCDData. Args: data: CCDData to create image from. Returns: New image. """ # create image and assign data image = Image( data=data.data.astype(np.float32), header=data.header, mask=data.mask, uncertainty=None if data.uncertainty is None else data.uncertainty.array.astype(np.float32), ) return image
@classmethod def _from_hdu_list(cls, data: fits.HDUList) -> Image: """Load Image from HDU list. Args: data: HDU list. Returns: Image. """ # create image image = cls() # find HDU with image data for hdu in data: if ( isinstance(hdu, fits.PrimaryHDU) and hdu.header["NAXIS"] > 0 or isinstance(hdu, fits.ImageHDU) and hdu.name == "SCI" or isinstance(hdu, fits.CompImageHDU) ): # found image HDU image_hdu = hdu break else: raise ValueError("Could not find HDU with main image.") # get data image._data = image_hdu.data image._header = image_hdu.header # mask if "MASK" in data: image._mask = data["MASK"].data # uncertainties if "UNCERT" in data: image._uncertainty = data["UNCERT"].data # catalog if "CAT" in data: image._catalog = Table(data["CAT"].data) # raw if "RAW" in data: image._raw = data["RAW"].data # finished return image @property def unit(self) -> str: """Returns units of pixels in image.""" if "BUNIT" in self.header: return str(self.header["BUNIT"]).lower() else: return "adu" def __deepcopy__(self) -> Image: """Returns a shallow copy of this image.""" return self.copy()
[docs] def copy(self) -> Image: """Returns a copy of this image.""" return Image( data=self._data, header=self._header, mask=self._mask, uncertainty=self._uncertainty, catalog=self._catalog, raw=self._raw, meta=self._meta, )
def __truediv__(self, other: Image) -> Image: """Divides this image by other.""" img = self.copy() if img._data is not None and other._data is not None: img._data = img._data / other._data return img else: raise ValueError("One image in division is None.")
[docs] def writeto(self, f: Any, *args: Any, **kwargs: Any) -> None: """Write image as FITS to given file object. Args: f: File object to write to. """ # create HDU list hdu_list = fits.HDUList([]) # create image HDU hdu = fits.PrimaryHDU(self.data, header=self.header) hdu_list.append(hdu) # catalog? if self._catalog is not None: hdu = table_to_hdu(self._catalog) hdu.name = "CAT" hdu_list.append(hdu) # mask? if self._mask is not None: hdu = ImageHDU(self._mask.astype(np.uint8)) hdu.name = "MASK" hdu_list.append(hdu) # errors? if self._uncertainty is not None: hdu = ImageHDU(self._uncertainty.data) hdu.name = "UNCERT" hdu_list.append(hdu) # raw? if self._raw is not None: hdu = ImageHDU(self._raw.data) hdu.name = "RAW" hdu_list.append(hdu) # write it hdu_list.writeto(f, *args, **kwargs)
[docs] def to_bytes(self) -> bytes: """Write to a bytes array and return it.""" with io.BytesIO() as bio: self.writeto(bio) return bio.getvalue()
[docs] def write_catalog(self, f: Any, *args: Any, **kwargs: Any) -> None: """Write catalog to file object.""" if self._catalog is None: return hdu = table_to_hdu(self._catalog) hdu.writeto(f, *args, **kwargs)
[docs] def to_ccddata(self) -> CCDData: """Convert Image to CCDData""" return CCDData( data=self._data, meta=self._header, mask=self._mask, uncertainty=None if self._uncertainty is None else StdDevUncertainty(self._uncertainty), unit="adu", )
[docs] def format_filename(self, formatter: FilenameFormatter) -> str: """Format filename with given formatter.""" self._header["FNAME"] = formatter(self._header) return str(self._header["FNAME"])
@property def pixel_scale(self) -> float | None: """Returns pixel scale in arcsec/pixel.""" if "CD1_1" in self._header: return abs(float(self._header["CD1_1"])) * 3600.0 elif "CDELT1" in self._header: return abs(float(self._header["CDELT1"])) * 3600.0 else: return None
[docs] def to_jpeg(self, vmin: float | None = None, vmax: float | None = None) -> bytes: """Returns a JPEG image created from this image. Returns: The image. """ # import PIL Image import PIL.Image # copy data data: NDArray[Any] = np.copy(self._data) # type: ignore if data is None: raise ValueError("No data in image.") # no vmin/vmax? if vmin is None or vmax is None: flattened = sorted(data.flatten()) vmin = flattened[int(0.05 * len(flattened))] vmax = flattened[int(0.95 * len(flattened))] if vmin is None or vmax is None: raise ValueError("Could not determine vmin/vmax.") # Clip data to brightness limits data[data > vmax] = vmax data[data < vmin] = vmin # Scale data to range [0, 1] data = (data - vmin) / (vmax - vmin) # Convert to 8-bit integer data = (255 * data).astype(np.uint8) # Invert y axis data = data[::-1, :] # create image from data array image = PIL.Image.fromarray(data, "L") with io.BytesIO() as bio: image.save(bio, format="jpeg") return bio.getvalue()
[docs] def set_meta(self, meta: Any) -> None: """Sets meta information, storing it under it class. Note that it is possible to store, e.g., strings, but they would be stored as img.meta[str] and be overwritten with every new string, which is probably not what you want. Use the img.meta dict directly for this and set_meta/get_meta only for class-based data. Args: meta: Meta information to store. """ # store it self._meta[meta.__class__] = meta
[docs] def has_meta(self, meta_class: type[MetaClass]) -> bool: """Whether meta exists.""" return meta_class in self._meta
[docs] def get_meta(self, meta_class: type[MetaClass]) -> MetaClass: """Returns meta information, assuming that it is stored under the class of the object. Args: meta_class: Class to return meta information for. Returns: Meta information of the given class. """ # return default? if meta_class not in self._meta: raise ValueError("Meta value not found.") # correct type? if not isinstance(self._meta[meta_class], meta_class): raise ValueError("Stored meta information is of wrong type.") # return it return cast(MetaClass, self._meta[meta_class])
[docs] def get_meta_safe(self, meta_class: type[MetaClass], default: MetaClass | None = None) -> MetaClass | None: """Calls get_meta in a safe way and returns default value in case of an exception.""" try: return self.get_meta(meta_class) except Exception: return default
@property def safe_data(self) -> npt.NDArray[np.floating[Any]] | None: return self._data @property def data(self) -> npt.NDArray[np.floating[Any]]: if self._data is None: raise exc.ImageError("No data found in image.") return self._data @data.setter def data(self, val: npt.NDArray[np.floating[Any]] | None) -> None: self._data = val @property def safe_header(self) -> fits.Header | None: return self._header @property def header(self) -> fits.Header: if self._header is None: raise exc.ImageError("No header found in image.") return self._header @header.setter def header(self, val: fits.Header | None) -> None: self._header = val @property def safe_mask(self) -> npt.NDArray[np.floating[Any]] | None: return self._mask @property def mask(self) -> npt.NDArray[np.floating[Any]]: if self._mask is not None: return self._mask else: raise exc.ImageError("No mask found in image.") @mask.setter def mask(self, val: npt.NDArray[np.floating[Any]] | None) -> None: self._mask = val @property def safe_uncertainty(self) -> npt.NDArray[np.floating[Any]] | None: return self._uncertainty @property def uncertainty(self) -> npt.NDArray[np.floating[Any]]: if self._uncertainty is None: raise exc.ImageError("No uncertainties found in image.") return self._uncertainty @uncertainty.setter def uncertainty(self, val: npt.NDArray[np.floating[Any]] | None) -> None: self._uncertainty = val @property def safe_catalog(self) -> Table | None: return self._catalog @property def catalog(self) -> Table: if self._catalog is None: raise exc.ImageError("No catalog found in image.") return self._catalog @catalog.setter def catalog(self, val: Table | None) -> None: self._catalog = val @property def safe_raw(self) -> npt.NDArray[np.floating[Any]] | None: return self._raw @property def raw(self) -> npt.NDArray[np.floating[Any]]: if self._raw is None: raise exc.ImageError("No raw data found in image.") return self._raw @raw.setter def raw(self, val: npt.NDArray[np.floating[Any]] | None) -> None: self._raw = val @property def meta(self) -> dict[Any, Any]: if self._meta is None: self._meta = {} return self._meta @meta.setter def meta(self, val: dict[Any, Any] | None) -> None: self._meta = {} if val is None else val @property def is_color(self) -> bool: return len(self.data.shape) == 3 and self.data.shape[0] == 3
[docs] def to_grayscale(self, r: float = 0.2126, g: float = 0.7152, b: float = 0.0722, **kwargs: Any) -> Image: """Convert RGB image to grayscale. Args: r: Weight for red. g: Weight for green. b: Weight for blue. """ if not self.is_color: return self image = self.copy() data = image.data image.data = (r * data[0, :, :] + g * data[1, :, :] + b * data[2, :, :]).astype(np.float32) return image
__all__ = ["Image"]