inference/pipeline.py

from pathlib import Path
from PIL import Image

from .tiling import create_tiles
from .stitching import stitch_tiles

from models.hatsat.hatsat_inference import HATSATInference
from models.esrgan.esrgan_inference import ESRGANInference


# ==========================================================
# Project paths
# ==========================================================

# satellite_sr_deploy/ is the parent of inference/
ROOT = Path(__file__).resolve().parent.parent

HATSAT_CHECKPOINT = (
    ROOT /
    "weights/hatsat/net_g_150000.pth"
)

ESRGAN_CHECKPOINT = (
    ROOT /
    "weights/esrgan/RRDB_ESRGAN_x4.pth"
)


# ==========================================================
# Model cache
# ==========================================================

_MODEL_CACHE = {}


def get_model(model_name, device=None):
    """
    Load a model once and reuse it.

    Parameters
    ----------
    model_name : str
        'hatsat' or 'esrgan'

    device : optional
        'cuda' or 'cpu'

    Returns
    -------
    Loaded inference model.
    """

    model_name = model_name.lower().strip()

    if model_name not in ("hatsat", "esrgan"):
        raise ValueError(
            "model_name must be 'hatsat' or 'esrgan'."
        )

    cache_key = (
        model_name,
        str(device) if device is not None else "auto"
    )

    if cache_key in _MODEL_CACHE:
        return _MODEL_CACHE[cache_key]

    # ------------------------------------------------------
    # HATSAT
    # ------------------------------------------------------

    if model_name == "hatsat":

        model = HATSATInference(
            checkpoint_path=HATSAT_CHECKPOINT,
            device=device
        )

    # ------------------------------------------------------
    # ESRGAN
    # ------------------------------------------------------

    else:

        model = ESRGANInference(
            checkpoint_path=ESRGAN_CHECKPOINT,
            device=device
        )

    _MODEL_CACHE[cache_key] = model

    return model


# ==========================================================
# Main pipeline
# ==========================================================

def super_resolve(
    image,
    model_name="hatsat",
    tile_size=128,
    overlap=32,
    device=None
):
    """
    Complete satellite super-resolution pipeline.

    Pipeline:

        Input image
            ↓
        Tiling
            ↓
        Selected model
            ↓
        4x SR
            ↓
        Overlap stitching
            ↓
        Final SR image

    Parameters
    ----------
    image : PIL.Image.Image

    model_name : str
        'hatsat' or 'esrgan'

    tile_size : int
        Input tile size.

    overlap : int
        Input overlap.

    device : optional
        'cuda' or 'cpu'

    Returns
    -------
    PIL.Image.Image
        Final 4x super-resolved image.
    """

    # ------------------------------------------------------
    # Validate image
    # ------------------------------------------------------

    if not isinstance(image, Image.Image):
        raise TypeError(
            "image must be a PIL.Image.Image"
        )

    image = image.convert("RGB")

    original_size = image.size

    print(
        f"Pipeline input: "
        f"{original_size[0]} × {original_size[1]}"
    )

    # ------------------------------------------------------
    # Get selected model
    # ------------------------------------------------------

    model = get_model(
        model_name,
        device=device
    )

    print(
        f"Selected model: "
        f"{model_name.upper()}"
    )

    # ------------------------------------------------------
    # Create tiles
    # ------------------------------------------------------

    tiles = create_tiles(
        image,
        tile_size=tile_size,
        overlap=overlap
    )

    print(
        f"Number of tiles: {len(tiles)}"
    )

    # ------------------------------------------------------
    # Run SR on each tile
    # ------------------------------------------------------

    sr_tiles = []

    for tile_info in tiles:

        print(
            f"Processing tile "
            f"{tile_info['index'] + 1}/"
            f"{len(tiles)} "
            f"(x={tile_info['x']}, "
            f"y={tile_info['y']})"
        )

        sr_image = model.predict(
            tile_info["image"]
        )

        sr_tiles.append(
            {
                "image": sr_image,
                "x": tile_info["x"],
                "y": tile_info["y"]
            }
        )

    # ------------------------------------------------------
    # Stitch SR tiles
    # ------------------------------------------------------

    result = stitch_tiles(
        sr_tiles,
        scale=4,
        original_size=original_size,
        overlap=overlap
    )

    print(
        f"Pipeline output: "
        f"{result.size[0]} × "
        f"{result.size[1]}"
    )

    return result


# ==========================================================
# Convenience functions
# ==========================================================

def run_hatsat_pipeline(
    image,
    tile_size=128,
    overlap=32,
    device=None
):
    return super_resolve(
        image=image,
        model_name="hatsat",
        tile_size=tile_size,
        overlap=overlap,
        device=device
    )


def run_esrgan_pipeline(
    image,
    tile_size=128,
    overlap=32,
    device=None
):
    return super_resolve(
        image=image,
        model_name="esrgan",
        tile_size=tile_size,
        overlap=overlap,
        device=device
    )


# ==========================================================
# COMPLETE SATELLITE PROCESSING
# ==========================================================

from .geotiff import (
    sentinel_rgb_from_bands,
    save_sr_geotiff
)


def process_satellite(
    b02_path,
    b03_path,
    b04_path,
    model_name="hatsat",
    output_dir=None,
    tile_size=128,
    overlap=32,
    device=None
):
    """
    Complete Sentinel-2 satellite super-resolution workflow.

    Input:
        B02, B03 and B04 Sentinel-2 GeoTIFF files.

    Processing:
        B02/B03/B04
            ↓
        RGB conversion
            ↓
        Tiling
            ↓
        HATSAT / ESRGAN
            ↓
        Stitching
            ↓
        4x SR

    Output:
        PNG preview
        Georeferenced RGB GeoTIFF

    Returns
    -------
    dict
        {
            "model": model_name,
            "input_size": (...),
            "output_size": (...),
            "png": Path,
            "geotiff": Path
        }
    """

    # ------------------------------------------------------
    # Paths
    # ------------------------------------------------------

    b02_path = Path(b02_path)
    b03_path = Path(b03_path)
    b04_path = Path(b04_path)

    if output_dir is None:
        output_dir = (
            ROOT /
            "outputs" /
            model_name.lower() /
            "results"
        )

    output_dir = Path(output_dir)
    output_dir.mkdir(
        parents=True,
        exist_ok=True
    )

    # ------------------------------------------------------
    # Validate inputs
    # ------------------------------------------------------

    for path in [
        b02_path,
        b03_path,
        b04_path
    ]:

        if not path.exists():
            raise FileNotFoundError(
                f"Input file not found: {path}"
            )

    # ------------------------------------------------------
    # Read Sentinel-2 bands
    # ------------------------------------------------------

    print("\n========================================")
    print("SATELLITE SUPER-RESOLUTION PIPELINE")
    print("========================================")

    print("\nReading Sentinel-2 bands...")

    image, source_profile = sentinel_rgb_from_bands(
        b02_path,
        b03_path,
        b04_path
    )

    print(
        "Input RGB size:",
        image.size
    )

    print(
        "Source CRS:",
        source_profile["crs"]
    )

    # ------------------------------------------------------
    # Run existing SR pipeline
    # ------------------------------------------------------

    print(
        "\nRunning model:",
        model_name.upper()
    )

    sr_image = super_resolve(
        image=image,
        model_name=model_name,
        tile_size=tile_size,
        overlap=overlap,
        device=device
    )

    print(
        "\nSR image size:",
        sr_image.size
    )

    # ------------------------------------------------------
    # Output filenames
    # ------------------------------------------------------

    model_tag = model_name.lower()

    png_path = (
        output_dir /
        f"sentinel2_{model_tag}_4x.png"
    )

    geotiff_path = (
        output_dir /
        f"sentinel2_{model_tag}_4x_georef.tif"
    )

    # ------------------------------------------------------
    # Save PNG preview
    # ------------------------------------------------------

    sr_image.save(
        png_path,
        format="PNG"
    )

    print(
        "\nPNG saved:",
        png_path
    )

    # ------------------------------------------------------
    # Save georeferenced GeoTIFF
    # ------------------------------------------------------

    save_sr_geotiff(
        sr_image=sr_image,
        source_profile=source_profile,
        output_path=geotiff_path,
        scale=4
    )

    print(
        "GeoTIFF saved:",
        geotiff_path
    )

    # ------------------------------------------------------
    # Final summary
    # ------------------------------------------------------

    print("\n========================================")
    print("PROCESSING COMPLETE")
    print("========================================")

    print(
        "Model:",
        model_name.upper()
    )

    print(
        "Input:",
        image.size
    )

    print(
        "Output:",
        sr_image.size
    )

    print(
        "Scale: 4x"
    )

    print(
        "PNG:",
        png_path
    )

    print(
        "GeoTIFF:",
        geotiff_path
    )

    return {
        "model": model_name.lower(),
        "input_size": image.size,
        "output_size": sr_image.size,
        "png": png_path,
        "geotiff": geotiff_path
    }