scripts/infer_whole_scene.py

"""Sliding-window inference for full-scene seaweed segmentation rasters."""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np
import rasterio
import torch
from rasterio.windows import Window
from torchvision import transforms
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus


NORMALIZE_3CH = transforms.Normalize(mean=(0.430, 0.411, 0.296), std=(0.213, 0.156, 0.143))
NORMALIZE_4CH = transforms.Normalize(mean=(0.430, 0.411, 0.296, 0.350), std=(0.213, 0.156, 0.143, 0.180))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", required=True, help="Input multispectral whole-scene raster path.")
    parser.add_argument("--pan", default=None, help="Optional panchromatic raster path for on-the-fly pan sharpening.")
    parser.add_argument("--checkpoint", required=True, help="Model checkpoint .pth path.")
    parser.add_argument("--output-dir", default="outputs/scene_inference", help="Directory for prediction rasters.")
    parser.add_argument("--tile-size", type=int, default=256, help="Inference tile size.")
    parser.add_argument("--overlap", type=int, default=32, help="Tile overlap in output pixels for weighted blending.")
    parser.add_argument("--stripe-height", type=int, default=1024, help="Rows to blend/write at a time.")
    parser.add_argument("--batch-size", type=int, default=4, help="Number of tiles per forward pass.")
    parser.add_argument("--threshold", type=float, default=0.5, help="Foreground probability threshold.")
    parser.add_argument("--max-tiles", type=int, default=0, help="Optional smoke-test tile limit; 0 means full scene.")
    parser.add_argument("--max-stripes", type=int, default=0, help="Optional smoke-test stripe limit; 0 means all stripes.")
    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"], help="Inference device.")
    parser.add_argument("--write-probability", action="store_true", help="Write the foreground probability raster.")
    parser.add_argument("--no-probability", action="store_true", help="Deprecated; probability output is disabled by default.")
    return parser.parse_args()


def get_checkpoint_state(path: Path, device: torch.device) -> dict:
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    if not isinstance(checkpoint, dict):
        raise ValueError(f"Unsupported checkpoint format: {path}")
    return checkpoint


def build_model(checkpoint: dict, device: torch.device) -> DinoV3DeepLabV3Plus:
    config = checkpoint.get("config") or {}
    model = DinoV3DeepLabV3Plus(
        num_classes=int(config.get("num_classes", 2)),
        backbone_name=config.get("backbone_name", "dinov3_vitl16"),
        pretrained=False,
        weights=config.get("backbone_weights", "SAT493M"),
        use_4channel=bool(config.get("use_4channel", True)),
        freeze_backbone=False,
    ).to(device)

    state_dict = checkpoint.get("model_state_dict") or checkpoint.get("state_dict")
    if state_dict is None:
        raise KeyError("Checkpoint does not contain model_state_dict/state_dict.")
    cleaned = {k.removeprefix("module."): v for k, v in state_dict.items()}
    model.load_state_dict(cleaned, strict=True)
    model.eval()
    return model


def axis_starts(length: int, tile_size: int, overlap: int) -> list[int]:
    if length <= tile_size:
        return [0]
    stride = tile_size - overlap
    if stride <= 0:
        raise ValueError("--overlap must be smaller than --tile-size.")
    starts = list(range(0, length - tile_size + 1, stride))
    last = length - tile_size
    if starts[-1] != last:
        starts.append(last)
    return starts


def tile_grid(width: int, height: int, tile_size: int, overlap: int) -> list[tuple[int, int, int, int]]:
    tiles: list[tuple[int, int, int, int]] = []
    for y in axis_starts(height, tile_size, overlap):
        for x in axis_starts(width, tile_size, overlap):
            w = min(tile_size, width - x)
            h = min(tile_size, height - y)
            tiles.append((x, y, w, h))
    return tiles


def blend_weight(tile_size: int, overlap: int) -> np.ndarray:
    if overlap <= 0:
        return np.ones((tile_size, tile_size), dtype=np.float32)
    ramp = np.minimum(np.arange(tile_size, dtype=np.float32) + 1, tile_size - np.arange(tile_size, dtype=np.float32))
    ramp = np.clip(ramp / float(overlap), 1.0 / float(overlap), 1.0)
    return np.minimum(ramp[:, None], ramp[None, :]).astype(np.float32, copy=False)


def pad_chw(tile: np.ndarray, tile_size: int) -> np.ndarray:
    bands, height, width = tile.shape
    padded = np.zeros((bands, tile_size, tile_size), dtype=np.float32)
    padded[:, :height, :width] = tile.astype(np.float32, copy=False)
    return padded


def match_pan_to_intensity(pan: np.ndarray, intensity: np.ndarray) -> np.ndarray:
    pan = pan.astype(np.float32, copy=False)
    intensity = intensity.astype(np.float32, copy=False)
    pan_std = float(np.std(pan))
    intensity_std = float(np.std(intensity))
    if pan_std < 1e-6 or intensity_std < 1e-6:
        return pan
    return (pan - float(np.mean(pan))) * (intensity_std / pan_std) + float(np.mean(intensity))


def pan_sharpen_tile(ms_tile: np.ndarray, pan_tile: np.ndarray, tile_size: int) -> np.ndarray:
    """Lightweight additive component substitution for one tile.

    The result keeps the multispectral band count and PAN spatial resolution.
    It is intended for streaming inference, not radiometric product generation.
    """
    ms_padded = pad_chw(ms_tile, tile_size)
    pan_padded = pad_chw(pan_tile[:1], tile_size)[0]
    bands = ms_padded[:4] if ms_padded.shape[0] >= 4 else ms_padded
    intensity = np.mean(bands, axis=0)
    matched_pan = match_pan_to_intensity(pan_padded, intensity)
    fused = bands + (matched_pan - intensity)[None, :, :]
    return np.clip(fused, 0, 65535).astype(np.float32, copy=False)


def to_model_tensor(tile: np.ndarray, tile_size: int, use_4channel: bool) -> torch.Tensor:
    # rasterio returns C,H,W. Pad edge tiles to the training tile size.
    padded = pad_chw(tile, tile_size)

    if use_4channel:
        if padded.shape[0] < 4:
            padded = np.pad(padded, ((0, 4 - padded.shape[0]), (0, 0), (0, 0)), mode="edge")
        data = padded[:4]
        normalizer = NORMALIZE_4CH
    else:
        if padded.shape[0] >= 4:
            data = padded[[3, 2, 1]]
        else:
            data = padded[: min(3, padded.shape[0])]
            while data.shape[0] < 3:
                data = np.concatenate([data, data[-1:]], axis=0)
        normalizer = NORMALIZE_3CH

    tensor = torch.from_numpy(data)
    if float(tensor.max()) > 1.0:
        tensor = tensor / 65535.0
    return normalizer(tensor)


def read_fused_tile(
    ms_src: rasterio.DatasetReader,
    pan_src: rasterio.DatasetReader,
    x: int,
    y: int,
    w: int,
    h: int,
    tile_size: int,
) -> np.ndarray:
    scale_x = pan_src.width / ms_src.width
    scale_y = pan_src.height / ms_src.height
    ms_window = Window(x / scale_x, y / scale_y, w / scale_x, h / scale_y)
    ms_tile = ms_src.read(
        window=ms_window,
        out_shape=(ms_src.count, h, w),
        resampling=rasterio.enums.Resampling.bilinear,
        boundless=True,
        fill_value=0,
    )
    pan_tile = pan_src.read(1, window=Window(x, y, w, h), boundless=True, fill_value=0)[None, :, :]
    return pan_sharpen_tile(ms_tile, pan_tile, tile_size)


def run_batch(model: torch.nn.Module, batch: list[torch.Tensor], device: torch.device) -> np.ndarray:
    inputs = torch.stack(batch, dim=0).to(device, non_blocking=True)
    with torch.inference_mode():
        with torch.autocast(device_type="cuda", enabled=device.type == "cuda"):
            output = model(inputs)
            logits = output["out"] if isinstance(output, dict) else output
            probs = torch.softmax(logits.float(), dim=1)[:, 1]
    return probs.detach().cpu().numpy()


def add_probs_to_stripe(
    probs: np.ndarray,
    windows: list[tuple[int, int, int, int]],
    stripe_y: int,
    stripe_prob_sum: np.ndarray,
    stripe_weight_sum: np.ndarray,
    weight: np.ndarray,
) -> None:
    for prob, (wx, wy, ww, wh) in zip(probs, windows):
        out_y0 = max(wy, stripe_y)
        out_y1 = min(wy + wh, stripe_y + stripe_prob_sum.shape[0])
        if out_y0 >= out_y1:
            continue
        prob_y0 = out_y0 - wy
        prob_y1 = out_y1 - wy
        stripe_local_y0 = out_y0 - stripe_y
        stripe_local_y1 = out_y1 - stripe_y

        cropped = prob[prob_y0:prob_y1, :ww]
        cropped_weight = weight[prob_y0:prob_y1, :ww]
        stripe_prob_sum[stripe_local_y0:stripe_local_y1, wx : wx + ww] += (
            cropped.astype(np.float32, copy=False) * cropped_weight
        )
        stripe_weight_sum[stripe_local_y0:stripe_local_y1, wx : wx + ww] += cropped_weight


def write_stripe(
    mask_dst: rasterio.DatasetWriter,
    prob_dst: rasterio.DatasetWriter | None,
    stripe_y: int,
    stripe_prob_sum: np.ndarray,
    stripe_weight_sum: np.ndarray,
    threshold: float,
) -> None:
    probs = np.divide(
        stripe_prob_sum,
        stripe_weight_sum,
        out=np.zeros_like(stripe_prob_sum, dtype=np.float32),
        where=stripe_weight_sum > 0,
    )
    mask = (probs >= threshold).astype(np.uint8) * 255
    window = Window(0, stripe_y, probs.shape[1], probs.shape[0])
    mask_dst.write(mask, 1, window=window)
    if prob_dst is not None:
        prob_dst.write(probs.astype(np.float32, copy=False), 1, window=window)


def main() -> None:
    args = parse_args()
    write_probability = bool(args.write_probability) and not bool(args.no_probability)
    image_path = Path(args.image)
    pan_path = Path(args.pan) if args.pan else None
    checkpoint_path = Path(args.checkpoint)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if args.device == "cuda" and not torch.cuda.is_available():
        print("CUDA is not available; falling back to CPU.")
        device = torch.device("cpu")
    else:
        device = torch.device(args.device)

    checkpoint = get_checkpoint_state(checkpoint_path, device)
    config = checkpoint.get("config") or {}
    use_4channel = bool(config.get("use_4channel", True))
    model = build_model(checkpoint, device)

    start = time.time()
    with rasterio.open(image_path) as ms_src:
        if use_4channel and ms_src.count < 4:
            raise ValueError(f"Model expects 4 channels, but image has {ms_src.count}: {image_path}")

        pan_src = rasterio.open(pan_path) if pan_path else None
        ref_src = pan_src or ms_src
        all_tiles = tile_grid(ref_src.width, ref_src.height, args.tile_size, args.overlap)
        tiles = all_tiles[: args.max_tiles] if args.max_tiles > 0 else all_tiles
        stem = image_path.stem
        if pan_src is not None:
            stem = f"{stem}_pansharpened"
        suffix = "smoke" if args.max_tiles > 0 else "full"
        mask_path = output_dir / f"{stem}_{suffix}_mask.tif"
        prob_path = output_dir / f"{stem}_{suffix}_prob.tif"

        profile = ref_src.profile.copy()
        mask_profile = profile.copy()
        mask_profile.update(count=1, dtype="uint8", compress="lzw", nodata=0)
        prob_profile = profile.copy()
        prob_profile.update(count=1, dtype="float32", compress="lzw", nodata=0.0)
        weight = blend_weight(args.tile_size, args.overlap)

        try:
            mode = "Pan-sharpen inference" if pan_src is not None else "Inference"
            prob_dst = None
            with rasterio.open(mask_path, "w", **mask_profile) as mask_dst:
                if write_probability:
                    prob_dst = rasterio.open(prob_path, "w", **prob_profile)
                try:
                    stripe_starts = list(range(0, ref_src.height, args.stripe_height))
                    if args.max_stripes > 0:
                        stripe_starts = stripe_starts[: args.max_stripes]
                    for stripe_y in tqdm(
                        stripe_starts,
                        desc=f"{mode} stripes {image_path.name}",
                        unit="stripe",
                    ):
                        stripe_h = min(args.stripe_height, ref_src.height - stripe_y)
                        stripe_prob_sum = np.zeros((stripe_h, ref_src.width), dtype=np.float32)
                        stripe_weight_sum = np.zeros((stripe_h, ref_src.width), dtype=np.float32)
                        stripe_tiles = [
                            tile
                            for tile in tiles
                            if tile[1] < stripe_y + stripe_h and tile[1] + tile[3] > stripe_y
                        ]

                        batch: list[torch.Tensor] = []
                        windows: list[tuple[int, int, int, int]] = []
                        for x, y, w, h in stripe_tiles:
                            if pan_src is None:
                                tile = ms_src.read(window=Window(x, y, w, h))
                            else:
                                tile = read_fused_tile(ms_src, pan_src, x, y, w, h, args.tile_size)
                            batch.append(to_model_tensor(tile, args.tile_size, use_4channel))
                            windows.append((x, y, w, h))

                            if len(batch) == args.batch_size:
                                probs = run_batch(model, batch, device)
                                add_probs_to_stripe(probs, windows, stripe_y, stripe_prob_sum, stripe_weight_sum, weight)
                                batch.clear()
                                windows.clear()

                        if batch:
                            probs = run_batch(model, batch, device)
                            add_probs_to_stripe(probs, windows, stripe_y, stripe_prob_sum, stripe_weight_sum, weight)

                        write_stripe(mask_dst, prob_dst, stripe_y, stripe_prob_sum, stripe_weight_sum, args.threshold)
                finally:
                    if prob_dst is not None:
                        prob_dst.close()
        finally:
            if pan_src is not None:
                pan_src.close()

    summary = {
        "image": str(image_path),
        "pan": str(pan_path) if pan_path else None,
        "checkpoint": str(checkpoint_path),
        "device": str(device),
        "tile_size": args.tile_size,
        "overlap": args.overlap,
        "batch_size": args.batch_size,
        "tiles_processed": len(tiles),
        "tiles_total": len(all_tiles),
        "max_stripes": args.max_stripes,
        "mask": str(mask_path),
        "probability": str(prob_path) if write_probability else None,
        "seconds": round(time.time() - start, 2),
        "checkpoint_epoch": checkpoint.get("epoch"),
        "checkpoint_best_val_iou": checkpoint.get("best_val_iou"),
    }
    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()