""" 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