s2sr_pipe/model/inference.py

"""
Inference engine for SEN2SR on large Sentinel-2 tiles.

Strategy
--------
The official `sen2sr.predict_large()` function handles the full tiling loop
internally (128×128 LR patches, configurable overlap, weighted stitching).
We wrap it here to:
  - add GPU/AMP support,
  - accept numpy arrays instead of raw tensors,
  - return numpy arrays consistent with the rest of the pipeline.

For tiles > 10 000 × 10 000 px the function processes patches sequentially
so VRAM usage stays bounded (one batch at a time).

Output size
-----------
For the default variant (×4), a 10 980 × 10 980 px S2 tile produces a
43 920 × 43 920 px output (≈ 7.3 GB float32 × 10 bands).
Consider writing directly to disk in COG/tiled GeoTIFF rather than
materializing the full array in RAM — see pipeline.py for the streaming path.
"""
from __future__ import annotations

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from s2sr_pipe.utils.logging_utils import get_logger

logger = get_logger("inference")


def _predict_large_batched(
    model: nn.Module,
    X: torch.Tensor,
    overlap: int,
    batch_size: int,
) -> torch.Tensor:
    """
    Batched replacement for sen2sr.predict_large().

    Processes patches in groups of batch_size instead of one at a time,
    giving a ~batch_size× speedup on the forward pass (the main bottleneck).
    Results are bitwise-identical to predict_large since each patch is
    independent and model.eval() freezes BatchNorm statistics.
    """
    from sen2sr.utils import define_iteration
    from tqdm import tqdm

    nruns = define_iteration(
        dimension=(X.shape[1], X.shape[2]),
        chunk_size=128,
        overlap=overlap,
    )

    res_n: int | None = None
    output: torch.Tensor | None = None
    skip: int = 0

    total_batches = (len(nruns) + batch_size - 1) // batch_size

    for batch_start in tqdm(range(0, len(nruns), batch_size),
                            total=total_batches,
                            desc="SR inference", unit="batch"):
        batch_points = nruns[batch_start : batch_start + batch_size]

        # define_iteration((H, W)) returns (y, x) = (row_start, col_start).
        # p[0] = row_start, p[1] = col_start.
        # Border patches may be smaller than 128×128; pad to uniform size before stacking.
        patches = [X[:, p[0] : p[0] + 128, p[1] : p[1] + 128] for p in batch_points]
        patches = [
            F.pad(p, (0, 128 - p.shape[2], 0, 128 - p.shape[1])) if p.shape[1] < 128 or p.shape[2] < 128 else p
            for p in patches
        ]
        batch_tensor = torch.stack(patches, dim=0)   # (B, C, 128, 128)

        with torch.no_grad():
            batch_result = model(batch_tensor)        # (B, C, 128*res_n, 128*res_n)

        # Initialise output on first batch
        if res_n is None:
            res_n = batch_result.shape[2] // 128
            skip  = overlap * res_n // 2
            output = torch.zeros(
                (X.shape[0], X.shape[1] * res_n, X.shape[2] * res_n),
                dtype=torch.float32, device="cpu",
            )

        batch_result = batch_result.detach().cpu()

        for i, point in enumerate(batch_points):
            result_patch = batch_result[i]   # (C, 128*res_n, 128*res_n)

            # SR-space offsets (row = point[0], col = point[1])
            offset_row = 0 if point[0] == 0 else point[0] * res_n + skip
            offset_col = 0 if point[1] == 0 else point[1] * res_n + skip

            # Crop columns: skip leading overlap on non-first patches; keep full width on last patch
            if offset_col == 0:
                length_col = 128 * res_n - skip
                result_patch = result_patch[:, :, :length_col]
            elif point[1] + 128 == X.shape[2]:
                length_col = 128 * res_n
                result_patch = result_patch[:, :, :length_col]
            else:
                length_col = 128 * res_n - skip
                result_patch = result_patch[:, :, skip:]

            # Crop rows: skip leading overlap on non-first patches; keep full height on last patch
            if offset_row == 0:
                length_row = 128 * res_n - skip
                result_patch = result_patch[:, :length_row, :]
            elif point[0] + 128 == X.shape[1]:
                length_row = 128 * res_n
                result_patch = result_patch[:, :length_row, :]
            else:
                length_row = 128 * res_n - skip
                result_patch = result_patch[:, skip:, :]

            oy_end = min(offset_row + length_row, output.shape[1])
            ox_end = min(offset_col + length_col, output.shape[2])
            h_valid = max(0, oy_end - offset_row)
            w_valid = max(0, ox_end - offset_col)
            if h_valid == 0 or w_valid == 0:
                continue
            output[
                :,
                offset_row : oy_end,
                offset_col : ox_end,
            ] = result_patch[:, :h_valid, :w_valid]

    return output if output is not None else torch.zeros(
        (X.shape[0], X.shape[1], X.shape[2]), dtype=torch.float32
    )


def infer_large(
    model: nn.Module,
    array: np.ndarray,
    device: torch.device | str,
    overlap: int = 32,
    use_amp: bool = True,
    batch_size: int = 16,
) -> np.ndarray:
    """
    Super-resolve a full (C, H, W) array using sen2sr.predict_large().

    Parameters
    ----------
    model   : Official SEN2SR model (from architecture.build_model()).
    array   : (C, H, W) float32 in [0, 1], already normalised.
    device  : Target device.
    overlap : Overlap in LR pixels between adjacent patches (default 32).
              Larger overlap → smoother seams but slower inference.
    use_amp : Use automatic mixed precision (CUDA only).

    Returns
    -------
    np.ndarray : (C, H*scale, W*scale) float32, values in [0, 1].

    Raises
    ------
    ImportError : if the `sen2sr` package is not installed.
    """
    try:
        import sen2sr.utils  # noqa: F401 — verifies the package is installed
    except ImportError as exc:
        raise ImportError(
            "The official sen2sr package is required.\n"
            "Install with:  pip install sen2sr"
        ) from exc

    device = torch.device(device) if isinstance(device, str) else device

    if use_amp:
        logger.warning(
            "use_amp=True ne peut pas etre applique : le modele SEN2SR utilise "
            "torch.fft.fftn() qui n'accepte que float32. L'inference sera en float32."
        )

    C, H, W = array.shape

    # AMP is disabled unconditionally: the SEN2SR hard_constraint layer runs
    # torch.fft.fftn() which only supports float32 — float16 and bfloat16 both crash.
    logger.info(
        "Starting large-image SR inference | input=(%d,%d,%d) | overlap=%d | batch_size=%d | float32",
        C, H, W, overlap, batch_size,
    )

    # ── Build nodata mask BEFORE replacing zeros ──────────────────────────────
    # Sentinel-2 border pixels are exactly 0 in ALL bands (not a valid reflectance).
    # We record this mask so we can restore DN=0 on the SR output instead of
    # letting the network hallucinate values into the black margin.
    # Shape: (H, W) boolean, True where ALL bands == 0 (nodata pixel).
    nodata_mask_lr = (array == 0).all(axis=0)   # (H, W)
    nodata_fraction = float(nodata_mask_lr.mean()) * 100
    logger.info("Nodata border pixels: %.1f %% of the LR tile", nodata_fraction)

    # (C, H, W) → torch tensor on device
    X = torch.from_numpy(array).float().to(device)
    # Replace any NaN / Inf introduced by the SAFE reader
    X = torch.nan_to_num(X, nan=0.0, posinf=1.0, neginf=0.0)

    model.eval()
    superX: torch.Tensor = _predict_large_batched(
        model=model,
        X=X,
        overlap=overlap,
        batch_size=batch_size,
    )

    # ── Post-process output tensor ────────────────────────────────────────────
    # 1. Replace any NaN/Inf produced by the network (can appear on nodata patches)
    superX = torch.nan_to_num(superX, nan=0.0, posinf=1.0, neginf=0.0)
    # 2. Clamp to valid range
    superX = superX.clamp(0.0, 1.0)

    result = superX.cpu().numpy()   # (C, H*scale, W*scale)

    # 3. Zero-out nodata border on the HR output
    #    The LR nodata mask must be upscaled to match the HR dimensions.
    if nodata_fraction > 0:
        scale_factor = result.shape[1] // H   # infer actual scale from shapes
        if scale_factor > 1:
            # Nearest-neighbour upscale of boolean mask: repeat each pixel
            nodata_mask_hr = np.repeat(
                np.repeat(nodata_mask_lr, scale_factor, axis=0),
                scale_factor,
                axis=1,
            )
        else:
            nodata_mask_hr = nodata_mask_lr
        result[:, nodata_mask_hr] = 0.0
        logger.info("Nodata mask applied to HR output.")
    logger.info(
        "Inference complete | output shape=%s | range=[%.4f, %.4f]",
        result.shape,
        float(result.min()),
        float(result.max()),
    )
    return result