astrophotography/session-scripts/compose.py

"""Pass 4: turn the four masters into the finished LRGB image.

The order of operations matters and is the usual one for a linear stack:

1. Crop the registration border, where not every frame contributed.
2. Remove the sky gradient. A 46% moon was up about 30 degrees away, so each
   channel carries a smooth ramp; a plane (not a higher-order surface) is fitted
   to tiles OUTSIDE a generous ellipse around the galaxy, because Centaurus A's
   halo fills much of this field and a flexible model would happily eat it.
3. Colour-calibrate on stars. Aperture photometry of a few hundred field stars
   in R, G and B is scaled so their average colour is neutral. This is the
   "average field star is grey" assumption, which is the standard cheap
   substitute for a full photometric calibration and is well behaved here
   because the field is rich.
4. Stretch. A midtone transfer function moves the sky background to a chosen
   level while keeping the highlights unclipped: gentler on the core than a
   plain gamma, and reversible arithmetic rather than a curve drawn by hand.
5. LRGB assembly. Colour comes from the 20-minute-per-channel RGB, detail and
   noise from the 60-minute luminance: the RGB is scaled pixel-by-pixel to the
   luminance's brightness, which is why the colour data being four times
   shallower does not matter much.
"""
import os

import numpy as np
import sep
import tifffile
from astropy.io import fits
from PIL import Image
from scipy.ndimage import gaussian_filter, median_filter
from skimage.restoration import denoise_tv_chambolle

import layout

OUT = layout.SESSION
CROP = 48                     # registration border, in pixels
GALAXY_MASK = (1250, 1000)    # semi-axes of the halo exclusion ellipse, px
BG_TARGET = 0.10              # where the sky sits in the stretched image
SATURATION = 1.35
CHANNELS = ("Luminance", "Red", "Green", "Blue")


def load():
    data, hdr = {}, None
    for name in CHANNELS:
        with fits.open(layout.path(f"master-{name}.fit")) as hd:
            arr = hd[0].data.astype(np.float32)
            if hdr is None:
                hdr = hd[0].header.copy()
        data[name] = arr[CROP:-CROP, CROP:-CROP]
    return data, hdr


def galaxy_mask(shape):
    ny, nx = shape
    yy, xx = np.mgrid[0:ny, 0:nx]
    cy, cx = ny / 2.0, nx / 2.0
    a, b = GALAXY_MASK
    return ((xx - cx) / a) ** 2 + ((yy - cy) / b) ** 2 < 1.0


def remove_gradient(img, mask):
    """Subtract a least-squares plane fitted to tile medians outside `mask`."""
    ny, nx = img.shape
    step = 96
    xs, ys, zs = [], [], []
    for y0 in range(0, ny - step, step):
        for x0 in range(0, nx - step, step):
            tile = img[y0:y0 + step, x0:x0 + step]
            if mask[y0:y0 + step, x0:x0 + step].any():
                continue
            # The median of a tile is dominated by sky even with stars in it.
            zs.append(np.median(tile))
            xs.append(x0 + step / 2.0)
            ys.append(y0 + step / 2.0)
    xs, ys, zs = map(np.asarray, (xs, ys, zs))
    keep = np.ones(len(zs), bool)
    for _ in range(3):                       # clip tiles containing companions
        A = np.column_stack([xs[keep], ys[keep], np.ones(keep.sum())])
        coef, *_ = np.linalg.lstsq(A, zs[keep], rcond=None)
        model = coef[0] * xs + coef[1] * ys + coef[2]
        resid = zs - model
        s = 1.4826 * np.median(np.abs(resid - np.median(resid)))
        keep = np.abs(resid - np.median(resid)) < 2.5 * s
    yy, xx = np.mgrid[0:ny, 0:nx]
    plane = (coef[0] * xx + coef[1] * yy + coef[2]).astype(np.float32)
    return img - plane, coef, int(keep.sum()), len(zs)


def star_photometry(lum, channels):
    """Aperture flux in each colour at the position of every luminance star."""
    bkg = sep.Background(lum, bw=64, bh=64, fw=3, fh=3)
    sub = lum - bkg.back()
    o = sep.extract(sub, 12.0, err=bkg.globalrms, minarea=9,
                    deblend_cont=0.005)
    o = o[(o["flag"] == 0) & (o["npix"] > 12) & (o["npix"] < 800)]
    ny, nx = lum.shape
    cy, cx = ny / 2.0, nx / 2.0
    a, b = GALAXY_MASK
    outside = ((o["x"] - cx) / a) ** 2 + ((o["y"] - cy) / b) ** 2 > 1.0
    o = o[outside]                      # keep the galaxy out of the white balance
    o = o[np.argsort(o["flux"])[::-1][:500]]
    flux = {}
    for name in ("Red", "Green", "Blue"):
        img = np.ascontiguousarray(channels[name])
        f, _, _ = sep.sum_circle(img, o["x"], o["y"], 6.0, subpix=5)
        flux[name] = f
    return o, flux


def mtf(x, midtone):
    """PixInsight-style midtone transfer function on data already in [0, 1]."""
    x = np.clip(x, 0.0, 1.0)
    return ((midtone - 1.0) * x) / ((2.0 * midtone - 1.0) * x - midtone)


def autostretch(img, target=BG_TARGET, shadow_sigma=2.8):
    """Black-point just below the sky, then an MTF that puts sky at `target`."""
    sky = np.median(img)
    mad = 1.4826 * np.median(np.abs(img - sky))
    black = sky - shadow_sigma * mad
    white = np.percentile(img, 99.995)
    norm = np.clip((img - black) / (white - black), 0.0, 1.0)
    sky_norm = (sky - black) / (white - black)
    # Solve the MTF midtone that maps sky_norm exactly onto target.
    m = ((target - 1.0) * sky_norm) / (2.0 * target * sky_norm - target -
                                       sky_norm)
    return mtf(norm, m), dict(black=float(black), white=float(white),
                              midtone=float(m), sky=float(sky), mad=float(mad))


def main():
    data, hdr = load()
    shape = data["Luminance"].shape
    print(f"working frame {shape[1]} x {shape[0]} px after {CROP} px crop")

    mask = galaxy_mask(shape)
    print(f"halo exclusion covers {mask.mean():.1%} of the frame")
    for name in CHANNELS:
        data[name], coef, kept, total = remove_gradient(data[name], mask)
        print(f"  {name:10s} plane dz/dx={coef[0]*1e3:+.3f} "
              f"dz/dy={coef[1]*1e3:+.3f} ADU/kpx, offset {coef[2]:8.2f}, "
              f"{kept}/{total} sky tiles used")

    o, flux = star_photometry(data["Luminance"], data)
    good = (flux["Red"] > 0) & (flux["Green"] > 0) & (flux["Blue"] > 0)
    r, g, b = (flux[k][good] for k in ("Red", "Green", "Blue"))
    print(f"colour calibration on {good.sum()} field stars")
    # Neutral point: the median star should come out white.
    gr = np.median(g / r)
    gb = np.median(g / b)
    print(f"  raw median star colour G/R={1/gr:.3f} G/B={1/gb:.3f}")
    data["Red"] *= gr
    data["Blue"] *= gb

    lum_lin = data["Luminance"]
    rgb_lin = np.dstack([data["Red"], data["Green"], data["Blue"]])
    del data

    lum, params = autostretch(lum_lin)
    print(f"luminance stretch: black={params['black']:.2f} "
          f"white={params['white']:.1f} midtone={params['midtone']:.4f} "
          f"(sky {params['sky']:.2f} +/- {params['mad']:.2f})")

    # The colour channels get their own black point but SHARE the luminance
    # midtone, so the colour balance set above survives the stretch.
    rgb = np.empty_like(rgb_lin)
    for i in range(3):
        ch = rgb_lin[:, :, i]
        sky = np.median(ch)
        mad = 1.4826 * np.median(np.abs(ch - sky))
        black = sky - 2.8 * mad
        white = np.percentile(ch, 99.995)
        rgb[:, :, i] = mtf(np.clip((ch - black) / (white - black), 0, 1),
                           params["midtone"])
    del rgb_lin

    # Neutralise the sky. Calibrating on stars makes STARS grey but leaves the
    # residual sky tinted, because moonlight is blue-ish and each channel kept
    # its own black point. Forcing the three sky medians together is what stops
    # the empty parts of the frame reading brown.
    sky_med = [float(np.median(rgb[:, :, i][~mask])) for i in range(3)]
    target = float(np.mean(sky_med))
    print(f"  sky medians after stretch R/G/B "
          f"{sky_med[0]:.4f}/{sky_med[1]:.4f}/{sky_med[2]:.4f} -> {target:.4f}")
    for i in range(3):
        rgb[:, :, i] = np.clip(rgb[:, :, i] - (sky_med[i] - target), 0.0, 1.0)

    # Chroma noise is the ugliest part of a 20-minute colour stack. Blurring
    # colour alone is invisible at this scale because the eye takes structure
    # from the luminance, which is untouched.
    rgb_lum = rgb.mean(axis=2, keepdims=True)
    chroma = rgb - rgb_lum
    for i in range(3):
        chroma[:, :, i] = median_filter(chroma[:, :, i], size=3)
        chroma[:, :, i] = gaussian_filter(chroma[:, :, i], 1.5)
    rgb = np.clip(rgb_lum + chroma * SATURATION, 0.0, 1.0)
    del chroma, rgb_lum

    # Gentle local contrast on the luminance to lift the dust lane, held back
    # in the noise floor so the sky does not get grainier.
    detail = lum - gaussian_filter(lum, 2.0)
    weight = np.clip((lum - BG_TARGET) * 4.0, 0.0, 1.0)
    lum = np.clip(lum + 0.35 * detail * weight, 0.0, 1.0)

    # Edge-preserving smoothing, applied ONLY where there is nothing but sky
    # (the same weight, inverted). Structure and the galaxy halo keep their
    # full resolution; the empty 70% of the frame loses its grain.
    smooth = denoise_tv_chambolle(lum, weight=0.012)
    lum = np.clip(lum * weight + smooth * (1.0 - weight), 0.0, 1.0)
    del detail, weight, smooth

    # LRGB: keep the RGB hue, take the brightness from the deep luminance.
    ratio = lum / np.maximum(rgb.mean(axis=2), 1e-5)
    out = np.clip(rgb * ratio[:, :, None], 0.0, 1.0)
    del ratio, rgb

    # Nothing in this field is genuinely green: no astronomical source between
    # the H-alpha reds and the OIII blues emits there, so any green excess is
    # noise or residual moonlight. Pulling green down to the neutral average
    # wherever it exceeds it (the standard SCNR operation) is safe and takes
    # the last of the cast out of the sky.
    neutral = 0.5 * (out[:, :, 0] + out[:, :, 2])
    amount = 0.85
    green = out[:, :, 1]
    out[:, :, 1] = np.where(green > neutral,
                            green * (1.0 - amount) + neutral * amount, green)

    out16 = (out * 65535.0 + 0.5).astype(np.uint16)
    tif = layout.path("NGC5128-LRGB.tif")
    tifffile.imwrite(tif, out16, photometric="rgb")
    print("wrote", tif)

    png = layout.path("NGC5128-LRGB.png")
    Image.fromarray((out * 255 + 0.5).astype(np.uint8)).save(png)
    print("wrote", png)

    prev = Image.fromarray((out * 255 + 0.5).astype(np.uint8))
    prev.thumbnail((2000, 2000), Image.LANCZOS)
    prevpath = layout.path("NGC5128-LRGB-preview.jpg")
    prev.save(prevpath, quality=92)
    print("wrote", prevpath, prev.size)

    hdr["CRPIX1"] = hdr.get("CRPIX1", 0) - CROP
    hdr["CRPIX2"] = hdr.get("CRPIX2", 0) - CROP
    hdr["NCOMBINE"] = (24, "frames across all filters")
    hdr["EXPTOTAL"] = (7200.0, "[s] total integration, all filters")
    hdr["COMMENT"] = "LRGB composite: L 12x300s, R/G/B 4x300s each"
    cube = np.moveaxis((out * 65535).astype(np.uint16), 2, 0)
    fitsout = layout.path("NGC5128-LRGB.fit")
    fits.PrimaryHDU(cube, hdr).writeto(fitsout, overwrite=True)
    print("wrote", fitsout)


if __name__ == "__main__":
    main()