Add serving-oriented tiled inference and jittered training support

models/wildfire_fm/tiled_inference.py

@@ -0,0 +1,112 @@
+"""Overlap-tiled inference helpers for WildFIRE-FM probability maps."""
+
+from __future__ import annotations
+
+from typing import Iterable, Tuple
+
+import torch
+import torch.nn.functional as F
+
+
+def _starts(length: int, tile_size: int, stride: int) -> list[int]:
+    if length <= tile_size:
+        return [0]
+    out = list(range(0, max(length - tile_size + 1, 1), stride))
+    last = length - tile_size
+    if out[-1] != last:
+        out.append(last)
+    return out
+
+
+def _crop_slices(top: int, left: int, tile_size: int, height: int, width: int, halo: int) -> Tuple[slice, slice, slice, slice]:
+    y0 = 0 if top == 0 else halo
+    x0 = 0 if left == 0 else halo
+    y1 = tile_size if top + tile_size >= height else tile_size - halo
+    x1 = tile_size if left + tile_size >= width else tile_size - halo
+    return slice(y0, y1), slice(x0, x1), slice(top + y0, top + y1), slice(left + x0, left + x1)
+
+
+def predict_probability_tiled(
+    model: torch.nn.Module,
+    x: torch.Tensor,
+    tile_size: int = 32,
+    stride: int = 16,
+    halo: int = 8,
+    device: torch.device | str | None = None,
+    batch_size: int = 16,
+) -> torch.Tensor:
+    """Predict a full probability map from an input tensor using overlap tiles.
+
+    Parameters
+    ----------
+    model:
+        WildFIRE-FM model returning logits or ``(logits, aux_logits)``.
+    x:
+        Input tensor in ``[C, H, W]`` or ``[1, C, H, W]`` order.
+    tile_size:
+        Spatial tile size used for model calls.
+    stride:
+        Distance between tile origins. Use a value smaller than ``tile_size``
+        for overlap.
+    halo:
+        Number of pixels cropped away from interior tile borders before
+        stitching. Border tiles keep the image edge.
+    device:
+        Device for inference. Defaults to the model parameter device.
+    batch_size:
+        Number of tiles evaluated per model call.
+
+    Returns
+    -------
+    torch.Tensor
+        Probability map in ``[H, W]`` order on CPU.
+    """
+    if x.ndim == 3:
+        x = x.unsqueeze(0)
+    if x.ndim != 4 or x.shape[0] != 1:
+        raise ValueError("x must have shape [C, H, W] or [1, C, H, W].")
+    if tile_size <= 0 or stride <= 0:
+        raise ValueError("tile_size and stride must be positive.")
+    if halo < 0 or halo * 2 >= tile_size:
+        raise ValueError("halo must be non-negative and smaller than tile_size / 2.")
+
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+    device = torch.device(device)
+    model.eval()
+
+    _, channels, height, width = x.shape
+    pad_h = max(tile_size - height, 0)
+    pad_w = max(tile_size - width, 0)
+    if pad_h or pad_w:
+        x_work = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate")
+    else:
+        x_work = x
+    _, _, work_h, work_w = x_work.shape
+
+    output = torch.zeros((work_h, work_w), dtype=torch.float32)
+    weight = torch.zeros((work_h, work_w), dtype=torch.float32)
+    coords = [(top, left) for top in _starts(work_h, tile_size, stride) for left in _starts(work_w, tile_size, stride)]
+
+    with torch.no_grad():
+        for start in range(0, len(coords), batch_size):
+            batch_coords = coords[start : start + batch_size]
+            tiles = torch.cat(
+                [x_work[:, :, top : top + tile_size, left : left + tile_size] for top, left in batch_coords],
+                dim=0,
+            ).to(device)
+            pred = model(tiles)
+            logits = pred[0] if isinstance(pred, tuple) else pred
+            probs = torch.sigmoid(logits.float()).detach().cpu()[:, 0]
+            for prob, (top, left) in zip(probs, batch_coords):
+                sy, sx, dy, dx = _crop_slices(top, left, tile_size, work_h, work_w, halo)
+                output[dy, dx] += prob[sy, sx]
+                weight[dy, dx] += 1.0
+    output = output / weight.clamp_min(1.0)
+    return output[:height, :width].contiguous()
+
+
+__all__ = ["predict_probability_tiled"]