retain the other layers of the input L1 images in the injected fits file

import numpy as np

import os

from pathlib import Path

import shutil

import tempfile

import galsim

import traceback

from astropy import wcs

__all__ = ["SingleEpochImage"]

def _find_science_extension(hdul):

    """Return the science-image extension, preferring the ``IMAGE`` name."""

    try:

        return hdul.index_of("IMAGE")

    except KeyError:

        if len(hdul) > 1 and isinstance(hdul[1], fits.ImageHDU):

            return 1

        raise ValueError(

            "Input FITS file has no IMAGE extension or image HDU at index 1."

        )

class SingleEpochImage(object):

    def __init__(self, config, input_img_path, output_dir):

    def read_initial_image(self, filepath):

        data_dir = os.path.dirname(filepath)

        img_name = os.path.basename(filepath)

        hdu = fits.open(filepath)

        self.header0 = hdu[0].header

        self.header_img = hdu[1].header

        self.image = hdu[1].data

        # flag_img = hdu[3].data

        hdu.close()

        self.input_img_path = Path(filepath).resolve()

        with fits.open(self.input_img_path, memmap=False) as hdul:

            science_extension = _find_science_extension(hdul)

            science_hdu = hdul[science_extension]

            self.science_extension = science_hdu.name or science_extension

            self.header0 = hdul[0].header.copy()

            self.header_img = science_hdu.header.copy()

            self.image = np.array(science_hdu.data, copy=True)

            # flag_img = hdul[3].data

        # Determine which CCD

        # self.chip_ID = int(self.header0['DETECTOR'][-2:])

    def inject_objects(self, pos, cat):

        nobj = len(pos)

        self.n_injected = 0

        # Make sure we have enough objects to inject

        assert nobj <= len(cat.objs)

                    # TODO: add up stats

                    # self.chip_output.Log_info("updating output catalog...")

                    self.chip_output.cat_add_obj(obj, pos_img, pos_shear)

                    self.n_injected += 1

                    pass

                else:

                    self.chip_output.Log_info("object omitted")

        pass

    def save_injected_img(self):

        self.chip.img = galsim.Image(self.chip.img.array, dtype=np.float32)

        # [TODO] TEST

        self.chip.img *= self.chip.gain

        self.chip.img /= self.exp_time

        # self.chip.img /= self.chip.gain

        hdu1 = fits.PrimaryHDU(header=self.header0)

        hdu2 = fits.ImageHDU(self.chip.img.array, header=self.header_img)

        hdu1 = fits.HDUList([hdu1, hdu2])

        fname = self.output_img_path

        hdu1.writeto(fname, output_verify='ignore', overwrite=True)

        """Write an injected image while retaining every input FITS extension."""

        output_path = Path(self.output_img_path)

        output_path.parent.mkdir(parents=True, exist_ok=True)

        injected_image = np.asarray(

            self.chip.img.array * self.chip.gain / self.exp_time,

            dtype=np.float32,

        )

        temporary_file = tempfile.NamedTemporaryFile(

            prefix=f".{output_path.stem}.",

            suffix=".fits",

            dir=output_path.parent,

            delete=False,

        )

        temporary_path = Path(temporary_file.name)

        temporary_file.close()

        try:

            # Start with an exact copy so IVAR, FLAG, WCS, calibration, sky,

            # PSF, and any future extensions are retained without rebuilding

            # their headers or data arrays.

            shutil.copyfile(self.input_img_path, temporary_path)

            with fits.open(temporary_path, mode="update", memmap=False) as hdul:

                science_hdu = hdul[self.science_extension]

                if science_hdu.data.shape != injected_image.shape:

                    raise ValueError(

                        "Injected image shape "

                        f"{injected_image.shape} does not match science extension "

                        f"shape {science_hdu.data.shape}."

                    )

                science_hdu.data[...] = injected_image

                science_hdu.header["INJECTED"] = (

                    True,

                    "Synthetic sources injected",

                )

                science_hdu.header["NINJECT"] = (

                    getattr(self, "n_injected", 0),

                    "Number of injected objects",

                )

                run_name = self.config.get("run_name")

                if run_name:

                    science_hdu.header["INJRUN"] = (

                        str(run_name),

                        "Injection run name",

                    )

                science_hdu.add_checksum(override_datasum=True)

                hdul.flush(output_verify="exception")

            os.replace(temporary_path, output_path)

        finally:

            temporary_path.unlink(missing_ok=True)

import argparse

def parse_args():

    '''

    Parse command line arguments. Many of the following