Band Alignment and Stacking for Multispectral Sets

import cv2
import numpy as np


def match_bands(reference: np.ndarray, target: np.ndarray,
                *, max_features: int = 4000) -> tuple:
    """Feature correspondences between two bands of the same frame.

    Normalising each band independently before detection is what makes this
    work across wavelengths: the absolute levels differ by a factor of
    several between red and near-infrared, and the gradients do not.
    """
    def prep(img):
        f = img.astype(np.float32)
        f = (f - np.percentile(f, 2)) / max(np.percentile(f, 98) - np.percentile(f, 2), 1e-6)
        return np.clip(f * 255, 0, 255).astype(np.uint8)

    a, b = prep(reference), prep(target)
    detector = cv2.AKAZE_create()
    ka, da = detector.detectAndCompute(a, None)
    kb, db = detector.detectAndCompute(b, None)
    if da is None or db is None or len(ka) < 20 or len(kb) < 20:
        raise ValueError("too few features to align these bands")

    matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
    raw = matcher.knnMatch(da, db, k=2)
    good = [m for m, n in raw if m.distance < 0.75 * n.distance]
    if len(good) < 20:
        raise ValueError(f"only {len(good)} good matches — bands too dissimilar")

    src = np.float32([ka[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
    dst = np.float32([kb[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
    return src, dst, good

import cv2
import numpy as np


def estimate_alignment(src: np.ndarray, dst: np.ndarray,
                       *, ransac_px: float = 2.0) -> dict:
    """Projective transform from target band to reference, with a residual.

    The residual is the number that matters. A homography can always be
    fitted; whether it describes the data is a separate question, and the
    median reprojection error of the inliers answers it.
    """
    H, inliers = cv2.findHomography(dst, src, cv2.RANSAC, ransac_px,
                                    maxIters=5000, confidence=0.999)
    if H is None:
        raise ValueError("homography estimation failed")

    mask = inliers.ravel().astype(bool)
    projected = cv2.perspectiveTransform(dst[mask], H)
    residuals = np.linalg.norm(projected - src[mask], axis=2).ravel()

    return {"H": H,
            "inliers": int(mask.sum()),
            "inlier_fraction": float(mask.mean()),
            "median_residual_px": float(np.median(residuals)),
            "p95_residual_px": float(np.percentile(residuals, 95))}


def alignment_acceptable(result: dict, *, max_median_px: float = 0.5,
                         min_inliers: int = 40,
                         min_inlier_fraction: float = 0.4) -> list[str]:
    """Reasons this alignment should not be trusted, if any."""
    problems = []
    if result["median_residual_px"] > max_median_px:
        problems.append(f"median residual {result['median_residual_px']:.2f} px")
    if result["inliers"] < min_inliers:
        problems.append(f"only {result['inliers']} inliers")
    if result["inlier_fraction"] < min_inlier_fraction:
        problems.append(f"inlier fraction {result['inlier_fraction']:.2f}")
    return problems

import cv2
import numpy as np


def warp_band(band: np.ndarray, H: np.ndarray, shape: tuple[int, int],
              *, interpolation: int = cv2.INTER_CUBIC) -> np.ndarray:
    """Apply an alignment homography to one band.

    Cubic rather than nearest: the values are continuous reflectance, and
    nearest-neighbour would preserve the quantisation of the sensor at the
    cost of a half-pixel positional error that is precisely what this step
    exists to remove.
    """
    h, w = shape
    return cv2.warpPerspective(band.astype(np.float32), H, (w, h),
                               flags=interpolation,
                               borderMode=cv2.BORDER_CONSTANT,
                               borderValue=float("nan"))


def compose(H_align: np.ndarray, H_ortho: np.ndarray) -> np.ndarray:
    """One transform instead of two, so the band is resampled once."""
    return H_ortho @ H_align

import numpy as np
import rasterio


def write_band_stack(bands: dict[str, np.ndarray], out_path: str, *,
                     profile: dict, wavelengths: dict[str, float],
                     order: list[str] | None = None) -> dict:
    """Write a band stack with descriptions, wavelengths and NaN NoData.

    The band order is explicit rather than dictionary order, because a
    downstream consumer indexing by position must be able to rely on it.
    """
    order = order or sorted(bands, key=lambda b: wavelengths[b])
    stack = np.stack([bands[b] for b in order]).astype("float32")

    profile = dict(profile)
    profile.update(count=len(order), dtype="float32", nodata=np.nan,
                   compress="deflate", predictor=3, tiled=True)

    with rasterio.open(out_path, "w", **profile) as dst:
        dst.write(stack)
        for i, band in enumerate(order, start=1):
            dst.set_band_description(i, band)
            dst.update_tags(i, WAVELENGTH_NM=f"{wavelengths[band]:.1f}",
                            UNITS="reflectance")
        dst.update_tags(BAND_ORDER=",".join(order))
    return {"bands": order, "shape": stack.shape}


def read_band(path: str, name: str) -> np.ndarray:
    """Read a band by name rather than by index, and fail if it is absent."""
    with rasterio.open(path) as src:
        descriptions = [d or "" for d in src.descriptions]
        if name not in descriptions:
            raise KeyError(f"band {name!r} not in {descriptions}")
        return src.read(descriptions.index(name) + 1, masked=True).filled(np.nan)

import numpy as np
from skimage.registration import phase_cross_correlation


def translation_fallback(reference: np.ndarray, target: np.ndarray,
                         *, upsample: int = 10) -> dict:
    """Sub-pixel translation between two bands, using the whole frame.

    Normalising both to zero mean and unit variance first is what makes this
    work across wavelengths: phase correlation is sensitive to structure and
    indifferent to absolute level, but only once the level is removed.
    """
    def norm(img):
        a = np.nan_to_num(img.astype(np.float32))
        return (a - a.mean()) / max(a.std(), 1e-6)

    shift, error, _ = phase_cross_correlation(norm(reference), norm(target),
                                              upsample_factor=upsample)
    H = np.array([[1.0, 0.0, -shift[1]],
                  [0.0, 1.0, -shift[0]],
                  [0.0, 0.0, 1.0]])
    return {"H": H, "shift_px": shift.tolist(), "error": float(error),
            "method": "phase_correlation"}

def align_with_fallbacks(reference, target, previous_H=None) -> dict:
    """Try features, then phase correlation, then the neighbour, then reject."""
    try:
        src, dst, _ = match_bands(reference, target)
        result = estimate_alignment(src, dst)
        problems = alignment_acceptable(result)
        if not problems:
            return {**result, "method": "features"}
    except ValueError:
        problems = ["feature matching failed"]

    fallback = translation_fallback(reference, target)
    if fallback["error"] < 0.4:
        return {**fallback, "note": f"features rejected: {problems}"}

    if previous_H is not None:
        return {"H": previous_H, "method": "carried_forward",
                "note": "no alignment found; using the previous frame's transform"}

    return {"H": None, "method": "rejected",
            "note": "no usable alignment; exclude this frame"}

import numpy as np
from scipy import ndimage


def edge_agreement(a: np.ndarray, b: np.ndarray, *, percentile: float = 95.0) -> dict:
    """Do the two bands' strong edges fall in the same places?

    Gradient magnitude is comparable between bands even when the contrast
    reverses, so the overlap of their strongest edges is a direct measure of
    alignment that needs no features or transforms.
    """
    def edges(img):
        gy, gx = np.gradient(np.nan_to_num(img.astype(np.float32)))
        g = np.hypot(gx, gy)
        return g > np.nanpercentile(g, percentile)

    ea, eb = edges(a), edges(b)
    intersection = np.count_nonzero(ea & eb)
    union = np.count_nonzero(ea | eb)
    dilated = ndimage.binary_dilation(ea, iterations=1)
    within_one_px = np.count_nonzero(dilated & eb) / max(np.count_nonzero(eb), 1)

    return {"jaccard": intersection / max(union, 1),
            "within_one_px": float(within_one_px),
            "ok": within_one_px > 0.8}