diff --git a/pipeline/colour.py b/pipeline/colour.py new file mode 100644 index 0000000..ce2edc1 --- /dev/null +++ b/pipeline/colour.py @@ -0,0 +1,259 @@ +"""Stage 5: assemble a viewable image from whatever filters the session has. + +The palette is chosen from what is present, not assumed: + + LRGB colour from R/G/B, brightness from the deep luminance + RGB no luminance, so a synthetic one is made from the colour channels + SHO the Hubble palette: SII -> red, Ha -> green, OIII -> blue + HOO Ha -> red, OIII -> green and blue + MONO one filter, rendered as greyscale + +Background handling deserves a note, because the obvious approach is wrong in a +way that is invisible. Fitting a plane to a frame that a large target fills +makes the plane absorb the target's own outer light - on Centaurus A this was +measured at -17.9 ADU/px of real halo quietly subtracted away. So the fit +excludes a central region whose size is derived from where the signal actually +is, and only a plane is used, never a flexible surface. +""" +import os + +import numpy as np +import tifffile +from astropy.io import fits +from PIL import Image +from scipy.ndimage import gaussian_filter, median_filter + +import layout + +BG_TARGET = 0.10 # where the sky sits in the stretched image +SATURATION = 1.35 + +PALETTE_MAP = { + "SHO": {"Red": "SII", "Green": "Ha", "Blue": "OIII"}, + "HOO": {"Red": "Ha", "Green": "OIII", "Blue": "OIII"}, +} + + +def mtf(x, midtone): + """Midtone transfer function on data already scaled to [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 sky, then an MTF putting sky at `target`.""" + sky = float(np.median(img)) + mad = 1.4826 * float(np.median(np.abs(img - sky))) + if mad <= 0: + mad = max(float(np.std(img)), 1e-6) + black = sky - shadow_sigma * mad + white = float(np.percentile(img, 99.995)) + if white <= black: + white = black + 1.0 + norm = np.clip((img - black) / (white - black), 0.0, 1.0) + sky_norm = (sky - black) / (white - black) + m = ((target - 1.0) * sky_norm) / (2.0 * target * sky_norm - target - + sky_norm) + m = float(np.clip(m, 1e-4, 0.9)) + return mtf(norm, m), dict(black=black, white=white, midtone=m, sky=sky, + mad=mad) + + +def target_mask(img, frac=0.45): + """Ellipse covering the bright central object, sized from the data. + + A fixed radius cannot suit both a galaxy filling the frame and a small + nebula. This grows the ellipse until it contains most of the flux above + sky, then stops - so the background fit is excluded from wherever the + target actually is, whatever its size. + """ + ny, nx = img.shape + yy, xx = np.mgrid[0:ny, 0:nx] + cy, cx = ny / 2.0, nx / 2.0 + r = np.hypot((xx - cx) / (nx / 2.0), (yy - cy) / (ny / 2.0)) + sky = float(np.median(img)) + signal = np.clip(img - sky, 0, None) + total = signal.sum() + if total <= 0: + return r < 0.5 + for cut in np.arange(0.25, 0.96, 0.05): + if signal[r < cut].sum() >= frac * total: + return r < min(cut + 0.15, 0.95) + return r < 0.75 + + +def remove_gradient(img, mask): + """Subtract a plane fitted to tile medians outside `mask`.""" + ny, nx = img.shape + step = max(48, min(ny, nx) // 40) + xs, ys, zs = [], [], [] + for y0 in range(0, ny - step, step): + for x0 in range(0, nx - step, step): + if mask[y0:y0 + step, x0:x0 + step].any(): + continue + tile = img[y0:y0 + step, x0:x0 + step] + zs.append(np.median(tile)) + xs.append(x0 + step / 2.0) + ys.append(y0 + step / 2.0) + if len(zs) < 12: + return img - float(np.median(img)), None + xs, ys, zs = map(np.asarray, (xs, ys, zs)) + keep = np.ones(len(zs), bool) + coef = None + for _ in range(3): + 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))) + if s <= 0: + break + 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 + + +def load_masters(session): + masters = {} + d = os.path.join(session.root, layout.MASTERS) + for name in os.listdir(d) if os.path.isdir(d) else []: + if name.startswith("master-") and name.endswith(".fit"): + filt = name[len("master-"):-len(".fit")] + masters[filt] = os.path.join(d, name) + return masters + + +def run(session, verbose=True): + """Produce the final image for whatever palette this session supports.""" + masters = load_masters(session) + if not masters: + raise RuntimeError("no masters found - run the register stage first") + palette = session.palette + if verbose: + print(f" palette {palette} from {sorted(masters)}") + + data, headers = {}, {} + for filt, path in masters.items(): + with fits.open(path) as hd: + data[filt] = hd[0].data.astype(np.float32) + headers[filt] = hd[0].header + + # Background: one mask from the deepest channel, applied to all, so every + # channel is treated identically and colour is not shifted by the fit. + deepest = session.filters[0] if session.filters[0] in data \ + else sorted(data)[0] + mask = target_mask(data[deepest]) + if verbose: + print(f" target mask covers {mask.mean():.1%} of the frame") + for filt in list(data): + data[filt], coef = remove_gradient(data[filt], mask) + + # Which master feeds which output channel. + if palette in PALETTE_MAP: + mapping = PALETTE_MAP[palette] + channels = [data[mapping[c]] for c in ("Red", "Green", "Blue")] + lum_src = None + elif palette in ("LRGB", "RGB", "PARTIAL") and all( + c in data for c in ("Red", "Green", "Blue")): + channels = [data["Red"], data["Green"], data["Blue"]] + lum_src = data.get("Luminance") + else: + only = data[deepest] + stretched, _ = autostretch(only) + grey = (np.clip(stretched, 0, 1) * 65535).astype(np.uint16) + return _write(session, np.dstack([grey] * 3), palette, verbose) + + # Colour balance: make the median star colour neutral, which is the cheap + # substitute for a photometric calibration and is well behaved on a rich + # field. + refs = [float(np.percentile(c, 99.5)) for c in channels] + target = float(np.median(refs)) + channels = [c * (target / r if r > 0 else 1.0) for c, r in + zip(channels, refs)] + + lum_lin = lum_src if lum_src is not None else np.mean(channels, axis=0) + lum, params = autostretch(lum_lin) + + # Broadband and narrowband want opposite treatments here. + # + # For broadband, all three channels share the luminance's midtone (a + # "linked" stretch). That preserves the real brightness ratios between + # them, which is what keeps star colours honest. + # + # For narrowband it is wrong. The three lines have wildly different + # strengths - Ha is typically far brighter than SII and OIII - so a linked + # stretch renders the whole nebula in whichever colour Ha was mapped to, + # which in the Hubble palette means overwhelmingly green. Stretching each + # channel to its OWN sky target instead brings the weak lines up to + # comparable visual weight, which is what the palette is for. The result is + # deliberately false colour either way; this makes it a readable false + # colour. + unlinked = palette in PALETTE_MAP + rgb = np.empty(lum.shape + (3,), np.float32) + for i, ch in enumerate(channels): + if unlinked: + rgb[:, :, i], _ = autostretch(ch) + else: + sky = float(np.median(ch)) + mad = 1.4826 * float(np.median(np.abs(ch - sky))) or 1.0 + black = sky - 2.8 * mad + white = float(np.percentile(ch, 99.995)) + rgb[:, :, i] = mtf( + np.clip((ch - black) / (max(white - black, 1e-6)), 0, 1), + params["midtone"]) + if unlinked: + # With each channel independently stretched the composite already + # carries its own brightness; taking luminance from the linear mean + # would flatten it back down. + lum = rgb.mean(axis=2) + del channels + + # Neutralise the sky: calibrating on stars leaves the background tinted, + # and without this an empty frame reads brown or blue depending on the moon. + sky_med = [float(np.median(rgb[:, :, i][~mask])) for i in range(3)] + neutral_to = float(np.mean(sky_med)) + for i in range(3): + rgb[:, :, i] = np.clip(rgb[:, :, i] - (sky_med[i] - neutral_to), 0, 1) + + # Denoise colour only; structure comes from the luminance, so this is + # invisible at normal viewing scale. + rgb_lum = rgb.mean(axis=2, keepdims=True) + chroma = rgb - rgb_lum + for i in range(3): + chroma[:, :, i] = gaussian_filter(median_filter(chroma[:, :, i], 3), + 1.5) + rgb = np.clip(rgb_lum + chroma * SATURATION, 0, 1) + del chroma, rgb_lum + + ratio = lum / np.maximum(rgb.mean(axis=2), 1e-5) + out = np.clip(rgb * ratio[:, :, None], 0.0, 1.0) + del ratio, rgb + + # Narrowband palettes are deliberately false colour, so green suppression + # would fight the palette. It only applies to broadband. + if palette not in PALETTE_MAP: + neutralish = 0.5 * (out[:, :, 0] + out[:, :, 2]) + green = out[:, :, 1] + out[:, :, 1] = np.where(green > neutralish, + green * 0.15 + neutralish * 0.85, green) + + return _write(session, (out * 65535 + 0.5).astype(np.uint16), palette, + verbose) + + +def _write(session, rgb16, palette, verbose=True): + target = session.target.replace(" ", "") + final = os.path.join(session.root, layout.FINAL) + os.makedirs(final, exist_ok=True) + stem = f"{target}-{palette}" + tifffile.imwrite(os.path.join(final, stem + ".tif"), rgb16, + photometric="rgb") + u8 = (rgb16 / 257).astype(np.uint8) + Image.fromarray(u8).save(os.path.join(final, stem + ".png")) + prev = Image.fromarray(u8) + prev.thumbnail((2400, 2400), Image.LANCZOS) + prev.save(os.path.join(final, stem + "-preview.jpg"), quality=92) + if verbose: + print(f" -> final/{stem}.png / .tif / -preview.jpg") + return os.path.join(final, stem + ".png")