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