augment.py

"""The image-augmentation operation set and the policy that composes them.

These are the *actions* the RL search chooses among. Every op takes a float image
in ``[0, 1]`` (grayscale ``(H, W)`` or colour ``(H, W, C)``), a ``magnitude`` in
``[0, 1]`` (strength, not direction — the sign is drawn per image so a single
magnitude covers both ways), and a seeded ``rng`` so a reward evaluation is fully
deterministic.

The op set deliberately includes ``flip_h``, which is *wrong* for digits (it turns
valid digits into invalid ones). A good search must learn to down-weight it — that
the search does so is one of the honest results this repo reports, not a bug.
"""

from __future__ import annotations

from collections.abc import Callable

import numpy as np
from scipy import ndimage

Array = np.ndarray
Op = Callable[[Array, float, np.random.Generator], Array]


def rotate(img: Array, mag: float, rng: np.random.Generator) -> Array:
    angle = rng.uniform(-1.0, 1.0) * mag * 25.0
    return ndimage.rotate(img, angle, axes=(0, 1), reshape=False, order=1, mode="nearest")


def translate_x(img: Array, mag: float, rng: np.random.Generator) -> Array:
    dx = rng.uniform(-1.0, 1.0) * mag * 0.25 * img.shape[1]
    shift = [0.0, dx] + ([0.0] if img.ndim == 3 else [])
    return ndimage.shift(img, shift, order=1, mode="nearest")


def translate_y(img: Array, mag: float, rng: np.random.Generator) -> Array:
    dy = rng.uniform(-1.0, 1.0) * mag * 0.25 * img.shape[0]
    shift = [dy, 0.0] + ([0.0] if img.ndim == 3 else [])
    return ndimage.shift(img, shift, order=1, mode="nearest")


def zoom(img: Array, mag: float, rng: np.random.Generator) -> Array:
    factor = 1.0 + rng.uniform(-1.0, 1.0) * mag * 0.3
    h, w = img.shape[0], img.shape[1]
    zoomed = ndimage.zoom(img, [factor, factor] + ([1.0] if img.ndim == 3 else []), order=1)
    return _center_to(zoomed, h, w)


def shear_x(img: Array, mag: float, rng: np.random.Generator) -> Array:
    s = rng.uniform(-1.0, 1.0) * mag * 0.3
    matrix = np.array([[1.0, s], [0.0, 1.0]])
    return _affine2d(img, matrix)


def flip_h(img: Array, mag: float, rng: np.random.Generator) -> Array:
    return img[:, ::-1].copy()


def brightness(img: Array, mag: float, rng: np.random.Generator) -> Array:
    return np.clip(img * (1.0 + rng.uniform(-1.0, 1.0) * mag * 0.5), 0.0, 1.0)


def contrast(img: Array, mag: float, rng: np.random.Generator) -> Array:
    mean = float(img.mean())
    return np.clip((img - mean) * (1.0 + rng.uniform(-1.0, 1.0) * mag * 0.5) + mean, 0.0, 1.0)


def noise(img: Array, mag: float, rng: np.random.Generator) -> Array:
    return np.clip(img + rng.normal(0.0, mag * 0.15, img.shape), 0.0, 1.0)


def cutout(img: Array, mag: float, rng: np.random.Generator) -> Array:
    h, w = img.shape[0], img.shape[1]
    side = int(mag * 0.5 * min(h, w))
    if side < 1:
        return img
    out = img.copy()
    top = int(rng.integers(0, max(1, h - side)))
    left = int(rng.integers(0, max(1, w - side)))
    out[top : top + side, left : left + side, ...] = 0.0
    return out


# Order is the canonical op index used by the search's parameter vectors.
OPS: dict[str, Op] = {
    "rotate": rotate,
    "translate_x": translate_x,
    "translate_y": translate_y,
    "zoom": zoom,
    "shear_x": shear_x,
    "flip_h": flip_h,
    "brightness": brightness,
    "contrast": contrast,
    "noise": noise,
    "cutout": cutout,
}
OP_NAMES: list[str] = list(OPS)
N_OPS = len(OP_NAMES)


def _affine2d(img: Array, matrix2x2: Array) -> Array:
    """Apply a 2x2 affine about the image centre, channel-aware."""
    h, w = img.shape[0], img.shape[1]
    centre = np.array([h, w]) / 2.0
    offset = centre - matrix2x2 @ centre

    def warp(plane: Array) -> Array:
        return ndimage.affine_transform(plane, matrix2x2, offset=offset, order=1, mode="nearest")

    if img.ndim == 3:
        return np.stack([warp(img[..., c]) for c in range(img.shape[2])], axis=-1)
    return warp(img)


def _center_to(img: Array, h: int, w: int) -> Array:
    """Centre-crop or pad ``img`` back to ``(h, w[, C])`` after a zoom."""
    ch, cw = img.shape[0], img.shape[1]
    if ch >= h:
        top = (ch - h) // 2
        img = img[top : top + h]
    else:
        pad = (h - ch) // 2
        widths = [(pad, h - ch - pad), (0, 0)] + ([(0, 0)] if img.ndim == 3 else [])
        img = np.pad(img, widths, mode="edge")
    if cw >= w:
        left = (cw - w) // 2
        img = img[:, left : left + w]
    else:
        pad = (w - cw) // 2
        widths = [(0, 0), (pad, w - cw - pad)] + ([(0, 0)] if img.ndim == 3 else [])
        img = np.pad(img, widths, mode="edge")
    return img


# Magnitude bins shared across ops: the categorical the search picks from.
MAG_BINS: tuple[float, ...] = (0.2, 0.4, 0.6, 0.8, 1.0)
N_MAG = len(MAG_BINS)


class Policy:
    """A concrete augmentation policy: which ops are on, and at what magnitude.

    ``include[o]`` and ``mag_bin[o]`` are the action the search sampled. Applying
    the policy to an image turns each included op on with per-image probability
    ``p_apply`` (so a single policy still yields varied augmented copies).
    """

    def __init__(self, include: Array, mag_bin: Array, p_apply: float = 0.5) -> None:
        self.include = include.astype(bool)
        self.mag_bin = mag_bin.astype(int)
        self.p_apply = p_apply

    def apply(self, img: Array, rng: np.random.Generator) -> Array:
        out = img
        for o in range(N_OPS):
            if self.include[o] and rng.random() < self.p_apply:
                out = OPS[OP_NAMES[o]](out, MAG_BINS[self.mag_bin[o]], rng)
        return out

    def to_dict(self) -> dict:
        return {
            "ops": [
                {"op": OP_NAMES[o], "magnitude": MAG_BINS[self.mag_bin[o]]}
                for o in range(N_OPS)
                if self.include[o]
            ],
            "p_apply": self.p_apply,
        }

    @classmethod
    def from_dict(cls, d: dict) -> Policy:
        include = np.zeros(N_OPS, dtype=bool)
        mag_bin = np.zeros(N_OPS, dtype=int)
        for entry in d.get("ops", []):
            o = OP_NAMES.index(entry["op"])
            include[o] = True
            mag_bin[o] = int(np.argmin([abs(m - entry["magnitude"]) for m in MAG_BINS]))
        return cls(include, mag_bin, d.get("p_apply", 0.5))