example_interface.py

import logging
import os
from typing import Union

import numpy as np
from astropy.io import fits
from csst_common.data_manager import CsstMsDataManager
from csst_common.logger import get_logger
from csst_common.status import CsstStatus


def read_image(filename_input: str) -> np.ndarray:
    """ Read image. """
    return fits.getdata(filename_input)


def process_data(data: np.ndarray) -> np.ndarray:
    """ Process data. """
    return np.fliplr(np.flipud(data))


def check_results(dm: CsstMsDataManager, logger: logging.Logger) -> bool:
    """ Check whether processed data are generated. """
    existence = [
        os.path.exists(dm.l1_detector(detector=detector, post="L1_processed.fits"))
        for detector in dm.target_detectors
    ]
    if all(existence):
        return True
    else:
        return False


# process a single image
def process_single_image(
        filename_input: str,
        filename_output: str,
        logger: Union[None, logging.Logger] = None
) -> CsstStatus:
    """
    Flip a single image.

    Flip a single image with ``numpy.fliplr`` and ``numpy.flipud``.

    Parameters
    ----------
    filename_input : str
        The input filename.
    filename_output : str
        The output filename.
    logger : logging.Logger
        The logger.

    Returns
    -------
    CsstStatus
        The final status.

    Examples
    --------
    >>> process_single_image(
    >>>     filename_input="input_image.fits",
    >>>     filename_output="output_image.fits",
    >>>     logger=None
    >>> )
    """
    # set default logger
    if logger is None:
        logger = get_logger()

    # process data
    try:
        # this will NOT be written into the log file
        logger.debug("Reading the image {}".format(filename_input))
        # start processing
        data = read_image(filename_input)
        data_processed = process_data(data)
        np.save(filename_output, data_processed)
        # this will be written into the log file
        logger.info("Processed image saved to {}".format(filename_output))
        return CsstStatus.PERFECT
    except DeprecationWarning:
        # this will be written into the log file
        logger.warning("Suffered DeprecationWarning!")
        return CsstStatus.WARNING
    except IOError:
        # this will be written into the log file
        logger.error("Suffered IOError!")
        return CsstStatus.ERROR


# process an exposure (MBI or SLS)
def process_multiple_images(
        dm: CsstMsDataManager,
        logger: Union[None, logging.Logger] = None
) -> CsstStatus:
    """
    Flip all images.

    Flip all images in an exposure in a for-loop (not in parallel).

    Parameters
    ----------
    dm : CsstMsDataManager
        The data manager of the specified exposure.
    logger : {None, logging.Logger}
        The logger. If None, use the default logger.

    Returns
    -------
    CsstStatus
        The final status.

    Examples
    --------
    >>> dm = CsstMsDataManager.quickstart(
    >>>     ver_sim="C5.2", dir_l1=".", datatype="sls", exposure_id=100)
    >>> logger = get_logger()
    >>> process_multiple_images(dm, logger)
    """

    # set default logger
    if logger is None:
        logger = get_logger()

    # process data
    try:
        # dm.target_detectors is a list of detector number that should be processed
        # start processing
        for detector in dm.target_detectors:
            # this will NOT be written into the log file
            logger.debug("Processing for detector {}".format(detector))
            data = read_image(dm.l0_detector(detector=detector))
            data_processed = process_data(data)
            np.save(
                dm.l1_detector(detector=detector, post="L1_processed.fits"),
                data_processed
            )
        # check results
        if check_results(dm=dm, logger=logger):
            # this will be written into the log file
            logger.info("Processed all images")
            return CsstStatus.PERFECT
        else:
            # not all images are properly processed
            return CsstStatus.WARNING
    except DeprecationWarning:
        # this will be written into the log file
        logger.warning("Suffered DeprecationWarning!")
        return CsstStatus.WARNING
    except IOError:
        # this will be written into the log file
        logger.error("Suffered IOError!")
        return CsstStatus.ERROR