robustness_test.py

"""
Phase 3 — Robustness Test Suite
satellite_sr_deploy

Covers:
  1. Various image sizes (square, non-square)
  2. Non-multiple-of-16 dimensions (HATSAT padding/cropping)
  3. RGB PNG and JPG inputs
  4. Larger images
  5. Invalid / bad inputs (should raise errors, not crash)
  6. Repeated model switching (HATSAT → ESRGAN → HATSAT)
  7. Edge-case dimensions (e.g. exactly 16, 17, 128, 129, 255, 256)

Run from satellite_sr_deploy/:
    .venv\\Scripts\\python.exe robustness_test.py
"""

import sys
import os
import io
import time
import traceback
from pathlib import Path

import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from model_router import SatelliteSRRouter

# ----------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------

PASS  = "  ✓ PASS"
FAIL  = "  ✗ FAIL"
SKIP  = "  - SKIP"

results = []   # (test_name, passed, note)


def record(name, passed, note=""):
    tag = PASS if passed else FAIL
    print(f"{tag}  {name}", f"— {note}" if note else "")
    results.append((name, passed, note))


def make_rgb_image(w, h, mode="random"):
    """Create a synthetic PIL RGB image."""
    if mode == "random":
        arr = np.random.randint(0, 255, (h, w, 3), dtype=np.uint8)
    elif mode == "gradient":
        arr = np.zeros((h, w, 3), dtype=np.uint8)
        arr[:, :, 0] = np.linspace(0, 255, w, dtype=np.uint8)
        arr[:, :, 1] = np.linspace(0, 255, h, dtype=np.uint8).reshape(-1, 1)
        arr[:, :, 2] = 128
    return Image.fromarray(arr, mode="RGB")


def check_output_size(inp: Image.Image, out: Image.Image, scale=4):
    ew = inp.width * scale
    eh = inp.height * scale
    ok = (out.width == ew) and (out.height == eh)
    return ok, f"expected {ew}×{eh}, got {out.width}×{out.height}"


OUTDIR = ROOT / "outputs" / "robustness_test"
OUTDIR.mkdir(parents=True, exist_ok=True)

# ----------------------------------------------------------------
# Initialise router once
# ----------------------------------------------------------------
print("=" * 65)
print("PHASE 3 — ROBUSTNESS TEST SUITE")
print("=" * 65)
print(f"Output dir: {OUTDIR}\n")

print("Initialising SatelliteSRRouter …")
router = SatelliteSRRouter(device="cpu")
print("Router ready.\n")


# ================================================================
# GROUP 1 — Non-multiple-of-16 dimensions (HATSAT critical)
# ================================================================
print("─" * 65)
print("GROUP 1 — Non-multiple-of-16 dimensions (HATSAT padding/crop)")
print("─" * 65)

nmod16_cases = [
    ("17x17",   17,  17),
    ("33x33",   33,  33),
    ("100x80",  100, 80),
    ("129x129", 129, 129),
    ("198x139", 198, 139),   # original test scene size
    ("255x255", 255, 255),
]

for label, w, h in nmod16_cases:
    img = make_rgb_image(w, h)
    try:
        t0 = time.time()
        out = router.predict(img, model_name="hatsat")
        elapsed = time.time() - t0
        ok, size_note = check_output_size(img, out)
        record(
            f"HATSAT {label}",
            ok,
            f"{size_note}  ({elapsed:.1f}s)"
        )
        if ok:
            out.save(OUTDIR / f"hatsat_{label}.png")
    except Exception as e:
        record(f"HATSAT {label}", False, str(e))
        traceback.print_exc()

print()

# ================================================================
# GROUP 2 — Exact multiples of 16 (should need no padding)
# ================================================================
print("─" * 65)
print("GROUP 2 — Exact multiples of 16 (no padding needed)")
print("─" * 65)

mod16_cases = [
    ("16x16",   16,  16),
    ("32x32",   32,  32),
    ("64x64",   64,  64),
    ("128x128", 128, 128),
    ("256x256", 256, 256),
]

for label, w, h in mod16_cases:
    img = make_rgb_image(w, h)
    try:
        t0 = time.time()
        out = router.predict(img, model_name="hatsat")
        elapsed = time.time() - t0
        ok, size_note = check_output_size(img, out)
        record(
            f"HATSAT {label} (exact mod16)",
            ok,
            f"{size_note}  ({elapsed:.1f}s)"
        )
        if ok:
            out.save(OUTDIR / f"hatsat_{label}_mod16.png")
    except Exception as e:
        record(f"HATSAT {label} (exact mod16)", False, str(e))
        traceback.print_exc()

print()

# ================================================================
# GROUP 3 — ESRGAN on same variety of sizes
# ================================================================
print("─" * 65)
print("GROUP 3 — ESRGAN: various sizes")
print("─" * 65)

esrgan_cases = [
    ("17x17",   17,  17),
    ("64x64",   64,  64),
    ("100x80",  100, 80),
    ("198x139", 198, 139),
    ("256x256", 256, 256),
]

for label, w, h in esrgan_cases:
    img = make_rgb_image(w, h)
    try:
        t0 = time.time()
        out = router.predict(img, model_name="esrgan")
        elapsed = time.time() - t0
        ok, size_note = check_output_size(img, out)
        record(
            f"ESRGAN {label}",
            ok,
            f"{size_note}  ({elapsed:.1f}s)"
        )
        if ok:
            out.save(OUTDIR / f"esrgan_{label}.png")
    except Exception as e:
        record(f"ESRGAN {label}", False, str(e))
        traceback.print_exc()

print()

# ================================================================
# GROUP 4 — Input format: JPG bytes (simulating file upload)
# ================================================================
print("─" * 65)
print("GROUP 4 — Input format: JPG round-trip")
print("─" * 65)

for model in ["hatsat", "esrgan"]:
    img = make_rgb_image(80, 60, mode="gradient")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    buf.seek(0)
    jpg_img = Image.open(buf).convert("RGB")

    try:
        out = router.predict(jpg_img, model_name=model)
        ok, size_note = check_output_size(jpg_img, out)
        record(f"{model.upper()} JPG input", ok, size_note)
    except Exception as e:
        record(f"{model.upper()} JPG input", False, str(e))
        traceback.print_exc()

print()

# ================================================================
# GROUP 5 — Model switching (HATSAT → ESRGAN → HATSAT → ESRGAN)
# ================================================================
print("─" * 65)
print("GROUP 5 — Repeated model switching")
print("─" * 65)

switch_seq = [
    ("hatsat", 50, 40),
    ("esrgan", 50, 40),
    ("hatsat", 50, 40),
    ("esrgan", 50, 40),
]

for i, (model, w, h) in enumerate(switch_seq, 1):
    img = make_rgb_image(w, h)
    try:
        out = router.predict(img, model_name=model)
        ok, size_note = check_output_size(img, out)
        record(f"Switch step {i}: {model.upper()}", ok, size_note)
    except Exception as e:
        record(f"Switch step {i}: {model.upper()}", False, str(e))
        traceback.print_exc()

print()

# ================================================================
# GROUP 6 — Invalid / bad inputs (must raise, not hard-crash)
# ================================================================
print("─" * 65)
print("GROUP 6 — Invalid inputs (expected to raise ValueError/Exception)")
print("─" * 65)


def expect_error(name, fn):
    try:
        fn()
        record(name, False, "Expected an error but got none")
    except (ValueError, RuntimeError, TypeError, AttributeError, Exception) as e:
        record(name, True, f"Correctly raised {type(e).__name__}: {str(e)[:60]}")


# None image
expect_error(
    "None image to router",
    lambda: router.predict(None, model_name="hatsat")
)

# Unknown model name
expect_error(
    "Unknown model name 'supermodel'",
    lambda: router.predict(make_rgb_image(32, 32), model_name="supermodel")
)

# Grayscale image — model should either handle it or raise cleanly
# (both models call .convert("RGB") internally — this should succeed)
gray = Image.fromarray(np.random.randint(0, 255, (32, 32), dtype=np.uint8), mode="L")
try:
    out = router.predict(gray, model_name="hatsat")
    # If it succeeds, that's fine — it means convert("RGB") was applied internally
    ok, size_note = check_output_size(gray, out)
    record("Grayscale input (handled by convert RGB)", ok, size_note)
except Exception as e:
    record("Grayscale input (handled by convert RGB)", False, str(e))

print()

# ================================================================
# SUMMARY
# ================================================================
print("=" * 65)
print("ROBUSTNESS TEST SUMMARY")
print("=" * 65)

total  = len(results)
passed = sum(1 for _, p, _ in results if p)
failed = total - passed

for name, p, note in results:
    tag = "✓" if p else "✗"
    print(f"  {tag}  {name:<45}  {note}")

print()
print(f"  Result: {passed}/{total} passed", end="")
if failed:
    print(f"  —  {failed} FAILED  ⚠️")
else:
    print("  —  ALL PASSED  🎉")
print("=" * 65)
print(f"  Output images saved to: {OUTDIR}")
print("=" * 65)

sys.exit(0 if failed == 0 else 1)