Deploy ZeroGPU compatible code with Sen2SR

.gitattributes

@@ -33,3 +33,16 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+outputs/local_test/test_esrgan_output.png filter=lfs diff=lfs merge=lfs -text
+outputs/local_test/test_hatsat_output.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/esrgan_100x80.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/esrgan_198x139.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/esrgan_256x256.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/esrgan_64x64.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_100x80.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_128x128_mod16.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_129x129.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_198x139.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_255x255.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_256x256_mod16.png filter=lfs diff=lfs merge=lfs -text
+outputs/robustness_test/hatsat_64x64_mod16.png filter=lfs diff=lfs merge=lfs -text

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.huggingfaceignore

@@ -0,0 +1,12 @@
+.venv/
+__pycache__/
+.gradio/
+outputs/
+test_*.py
+test*.jpg
+test*.png
+test*.tif
+local_test.py
+robustness_test.py
+validation.py
+geotiff_test.py

geotiff_test.py

@@ -0,0 +1,390 @@
+"""
+Phase 4 -- GeoTIFF Geospatial Verification Test
+satellite_sr_deploy
+
+Tests:
+  1. save_sr_geotiff() writes a valid GeoTIFF from a PIL image + profile
+  2. CRS is preserved exactly in the output
+  3. Geographic bounds are preserved (same extent as input)
+  4. Output pixel dimensions are exactly 4x the input
+  5. Affine transform is correctly scaled (pixel size = input/4)
+  6. Band ordering is R=1, G=2, B=3
+  7. Pixel values are preserved correctly (uint8 RGB)
+  8. Round-trip: create synthetic GeoTIFF -> read -> SR -> save -> verify
+  9. Verify against an existing real GeoTIFF (from Colab outputs)
+ 10. normalize_rgb() clamps properly and returns uint8
+
+Run from satellite_sr_deploy/:
+    $env:PYTHONIOENCODING="utf-8"; .venv\\Scripts\\python.exe geotiff_test.py
+"""
+
+import sys
+import math
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import rasterio
+from rasterio.transform import Affine
+from rasterio.crs import CRS
+
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+# Import only geotiff.py (no model loading needed for most tests)
+from inference.geotiff import save_sr_geotiff, normalize_rgb
+
+results = []
+
+OUTDIR = ROOT / "outputs" / "geotiff_test"
+OUTDIR.mkdir(parents=True, exist_ok=True)
+
+PASS_SYM = "[PASS]"
+FAIL_SYM = "[FAIL]"
+
+
+def record(name, passed, note=""):
+    tag = PASS_SYM if passed else FAIL_SYM
+    msg = f"{tag}  {name}"
+    if note:
+        msg += f"  -- {note}"
+    print(msg)
+    results.append((name, passed, note))
+
+
+def make_synthetic_geotiff(path, width=50, height=40, crs_epsg=4326):
+    """
+    Create a small synthetic single-band GeoTIFF with known transform/CRS.
+    Returns the profile used to write it.
+    """
+    # Mimic a Sentinel-2 scene over India (approximate EPSG:4326 coords)
+    west, north = 77.0, 13.0      # top-left corner
+    pixel_size_x =  0.0001        # degrees/pixel (~10 m)
+    pixel_size_y = -0.0001        # negative (north-to-south)
+
+    transform = Affine.translation(west, north) * Affine.scale(pixel_size_x, pixel_size_y)
+
+    profile = {
+        "driver": "GTiff",
+        "height": height,
+        "width": width,
+        "count": 1,
+        "dtype": "uint16",
+        "crs": CRS.from_epsg(crs_epsg),
+        "transform": transform,
+    }
+
+    data = np.random.randint(0, 3000, (height, width), dtype=np.uint16)
+
+    with rasterio.open(path, "w", **profile) as dst:
+        dst.write(data, 1)
+
+    return profile
+
+
+# ================================================================
+# GROUP 1 -- save_sr_geotiff: basic output verification
+# ================================================================
+print("=" * 65)
+print("PHASE 4 -- GEOTIFF GEOSPATIAL VERIFICATION")
+print("=" * 65)
+print()
+print("-" * 65)
+print("GROUP 1 -- save_sr_geotiff: CRS, transform, dimensions, bands")
+print("-" * 65)
+
+# Build a synthetic source profile (50x40, EPSG:4326)
+SRC_W, SRC_H = 50, 40
+SCALE = 4
+EXPECTED_W = SRC_W * SCALE
+EXPECTED_H = SRC_H * SCALE
+
+west, north = 77.0, 13.0
+px_x, px_y  = 0.0001, -0.0001
+src_transform = Affine.translation(west, north) * Affine.scale(px_x, px_y)
+
+source_profile = {
+    "driver":    "GTiff",
+    "height":    SRC_H,
+    "width":     SRC_W,
+    "count":     1,
+    "dtype":     "uint16",
+    "crs":       CRS.from_epsg(4326),
+    "transform": src_transform,
+}
+
+# Synthetic SR image (4x size, gradient RGB)
+sr_arr = np.zeros((EXPECTED_H, EXPECTED_W, 3), dtype=np.uint8)
+sr_arr[:, :, 0] = np.linspace(0, 255, EXPECTED_W, dtype=np.uint8)   # R
+sr_arr[:, :, 1] = np.linspace(0, 200, EXPECTED_H, dtype=np.uint8).reshape(-1, 1)  # G
+sr_arr[:, :, 2] = 128  # B flat
+sr_pil = Image.fromarray(sr_arr, mode="RGB")
+
+output_tif = OUTDIR / "synthetic_sr_4x.tif"
+
+try:
+    save_sr_geotiff(sr_pil, source_profile, output_tif, scale=SCALE)
+    saved_ok = True
+except Exception as e:
+    saved_ok = False
+    record("save_sr_geotiff() runs without error", False, str(e))
+
+if saved_ok:
+    record("save_sr_geotiff() runs without error", True)
+
+    with rasterio.open(output_tif) as out:
+        out_profile = out.profile
+        out_transform = out.transform
+        out_crs  = out.crs
+        out_w    = out.width
+        out_h    = out.height
+        out_count = out.count
+        out_dtype = out.dtypes[0]
+        band_r = out.read(1)
+        band_g = out.read(2)
+        band_b = out.read(3)
+
+    # 1a. Dimensions
+    record(
+        "Output dimensions are 4x input",
+        out_w == EXPECTED_W and out_h == EXPECTED_H,
+        f"expected {EXPECTED_W}x{EXPECTED_H}, got {out_w}x{out_h}"
+    )
+
+    # 1b. CRS preserved
+    record(
+        "CRS preserved (EPSG:4326)",
+        out_crs.to_epsg() == 4326,
+        f"got CRS: {out_crs}"
+    )
+
+    # 1c. Transform: pixel size scaled by 1/4
+    expected_px_x = px_x / SCALE
+    expected_px_y = px_y / SCALE
+    tol = 1e-12
+    px_x_ok = abs(out_transform.a - expected_px_x) < tol
+    px_y_ok = abs(out_transform.e - expected_px_y) < tol
+    record(
+        "Pixel size scaled to 1/4 of source (x)",
+        px_x_ok,
+        f"expected {expected_px_x:.8f}, got {out_transform.a:.8f}"
+    )
+    record(
+        "Pixel size scaled to 1/4 of source (y)",
+        px_y_ok,
+        f"expected {expected_px_y:.8f}, got {out_transform.e:.8f}"
+    )
+
+    # 1d. Origin (top-left corner) unchanged
+    origin_ok = (
+        abs(out_transform.c - west) < tol and
+        abs(out_transform.f - north) < tol
+    )
+    record(
+        "Geographic origin (top-left) unchanged",
+        origin_ok,
+        f"expected ({west}, {north}), got ({out_transform.c:.6f}, {out_transform.f:.6f})"
+    )
+
+    # 1e. Geographic extent / bounds preserved
+    src_right  = west  + SRC_W * px_x
+    src_bottom = north + SRC_H * px_y
+    out_right  = out_transform.c + out_w * out_transform.a
+    out_bottom = out_transform.f + out_h * out_transform.e
+    bounds_ok = (
+        abs(out_right  - src_right)  < tol and
+        abs(out_bottom - src_bottom) < tol
+    )
+    record(
+        "Geographic bounds (extent) preserved",
+        bounds_ok,
+        f"src right={src_right:.6f}, out right={out_right:.6f} | "
+        f"src bottom={src_bottom:.6f}, out bottom={out_bottom:.6f}"
+    )
+
+    # 1f. Band count
+    record("Band count is 3 (RGB)", out_count == 3, f"got {out_count}")
+
+    # 1g. dtype is uint8
+    record("Output dtype is uint8", out_dtype == "uint8", f"got {out_dtype}")
+
+    # 1h. Band ordering R=1, G=2, B=3
+    # Check representative pixels in the gradient
+    r_corner = int(sr_arr[0, EXPECTED_W - 1, 0])   # max-R corner
+    g_corner = int(sr_arr[EXPECTED_H - 1, 0, 1])   # max-G corner
+    b_mid    = int(sr_arr[EXPECTED_H // 2, EXPECTED_W // 2, 2])  # flat B=128
+
+    r_read   = int(band_r[0, EXPECTED_W - 1])
+    g_read   = int(band_g[EXPECTED_H - 1, 0])
+    b_read   = int(band_b[EXPECTED_H // 2, EXPECTED_W // 2])
+
+    record(
+        "Band 1 = Red channel",
+        r_read == r_corner,
+        f"expected R={r_corner}, got R={r_read}"
+    )
+    record(
+        "Band 2 = Green channel",
+        g_read == g_corner,
+        f"expected G={g_corner}, got G={g_read}"
+    )
+    record(
+        "Band 3 = Blue (flat 128)",
+        b_read == 128,
+        f"expected B=128, got B={b_read}"
+    )
+
+print()
+
+# ================================================================
+# GROUP 2 -- normalize_rgb: clipping and dtype
+# ================================================================
+print("-" * 65)
+print("GROUP 2 -- normalize_rgb: output range and dtype")
+print("-" * 65)
+
+# Typical Sentinel-2 DN range (0-10000+ reflectance)
+test_raw = np.random.randint(0, 10000, (40, 50, 3), dtype=np.uint16).astype(np.float32)
+
+try:
+    norm = normalize_rgb(test_raw.copy())
+
+    record(
+        "normalize_rgb returns uint8",
+        norm.dtype == np.uint8,
+        f"got dtype={norm.dtype}"
+    )
+    record(
+        "normalize_rgb output range [0, 255]",
+        norm.min() >= 0 and norm.max() <= 255,
+        f"min={norm.min()}, max={norm.max()}"
+    )
+    record(
+        "normalize_rgb shape preserved (H, W, 3)",
+        norm.shape == test_raw.shape,
+        f"expected {test_raw.shape}, got {norm.shape}"
+    )
+except Exception as e:
+    record("normalize_rgb runs without error", False, str(e))
+
+# Edge case: single-value band (all pixels same value)
+flat = np.full((10, 10, 3), 5000, dtype=np.float32)
+try:
+    norm_flat = normalize_rgb(flat)
+    record(
+        "normalize_rgb handles flat (constant) band",
+        norm_flat.dtype == np.uint8,
+        f"min={norm_flat.min()}, max={norm_flat.max()}"
+    )
+except Exception as e:
+    record("normalize_rgb handles flat band", False, str(e))
+
+print()
+
+# ================================================================
+# GROUP 3 -- save_sr_geotiff with numpy array input (not PIL)
+# ================================================================
+print("-" * 65)
+print("GROUP 3 -- save_sr_geotiff accepts numpy array input")
+print("-" * 65)
+
+np_input_tif = OUTDIR / "numpy_input_sr.tif"
+try:
+    save_sr_geotiff(sr_arr, source_profile, np_input_tif, scale=SCALE)
+    with rasterio.open(np_input_tif) as f:
+        dims_ok = f.width == EXPECTED_W and f.height == EXPECTED_H
+    record("save_sr_geotiff accepts numpy ndarray", dims_ok,
+           f"{EXPECTED_W}x{EXPECTED_H} ok={dims_ok}")
+except Exception as e:
+    record("save_sr_geotiff accepts numpy ndarray", False, str(e))
+
+print()
+
+# ================================================================
+# GROUP 4 -- Verify existing real GeoTIFF from Colab outputs
+# ================================================================
+print("-" * 65)
+print("GROUP 4 -- Verify existing Colab-produced GeoTIFFs")
+print("-" * 65)
+
+REAL_TIFS = [
+    Path(r"D:\downloads\satellite_resolution_increaser\satellite_resolution_increaser\outputs\hatsat\results\sentinel2_hatsat_4x.tif"),
+    Path(r"D:\downloads\satellite_resolution_increaser\satellite_resolution_increaser\outputs\hatsat\results\sentinel2_hatsat_4x_georef.tif"),
+    Path(r"D:\downloads\satellite_resolution_increaser\satellite_resolution_increaser\outputs\esrgan\results\sentinel2_esrgan_4x_georef.tif"),
+]
+
+for tif_path in REAL_TIFS:
+    label = tif_path.name
+    if not tif_path.exists():
+        record(f"{label} -- file exists", False, "not found on disk")
+        continue
+
+    try:
+        with rasterio.open(tif_path) as src:
+            crs_epsg = src.crs.to_epsg() if src.crs else None
+            band_count = src.count
+            dtype = src.dtypes[0]
+            w, h = src.width, src.height
+            t = src.transform
+            bounds = src.bounds
+
+        record(
+            f"{label} -- opens without error",
+            True,
+            f"{w}x{h}, {band_count} bands, {dtype}, EPSG:{crs_epsg}"
+        )
+        record(
+            f"{label} -- has CRS",
+            crs_epsg is not None,
+            f"EPSG:{crs_epsg}"
+        )
+        record(
+            f"{label} -- has 3 bands (RGB)",
+            band_count == 3,
+            f"bands={band_count}"
+        )
+        record(
+            f"{label} -- dimensions multiple of 4 (4x SR)",
+            w % 4 == 0 and h % 4 == 0,
+            f"{w}x{h}"
+        )
+        # Check pixel size is positive x / negative y (north-up)
+        north_up = t.a > 0 and t.e < 0
+        record(
+            f"{label} -- north-up orientation",
+            north_up,
+            f"x_res={t.a:.8f}, y_res={t.e:.8f}"
+        )
+
+    except Exception as e:
+        record(f"{label} -- opens without error", False, str(e))
+
+print()
+
+# ================================================================
+# SUMMARY
+# ================================================================
+print("=" * 65)
+print("PHASE 4 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 "x"
+    print(f"  [{tag}]  {name:<50}  {note}")
+
+print()
+print(f"  Result: {passed}/{total} passed", end="")
+if failed:
+    print(f"  --  {failed} FAILED")
+else:
+    print("  --  ALL PASSED")
+print("=" * 65)
+print(f"  Output GeoTIFFs: {OUTDIR}")
+print("=" * 65)
+
+sys.exit(0 if failed == 0 else 1)

inference/geotiff.py

@@ -88,19 +88,55 @@ def normalize_rgb(
 ):
     """
     Convert Sentinel-2 reflectance/DN values to uint8 RGB.
-    Global percentile normalization preserves color balance.
+
+    Percentile normalization is performed independently
+    for each channel.
     """
+
     rgb = rgb.astype(np.float32)
-    
-    # Calculate global percentiles across all channels to preserve color balance
-    low = np.percentile(rgb, lower_percentile)
-    high = np.percentile(rgb, upper_percentile)
 
-    if high <= low:
-        result = np.clip(rgb, 0, 255).astype(np.uint8)
-    else:
-        normalized = (rgb - low) / (high - low)
-        result = (np.clip(normalized, 0, 1) * 255).astype(np.uint8)
+    result = np.zeros_like(
+        rgb,
+        dtype=np.uint8
+    )
+
+    for channel in range(3):
+
+        band = rgb[:, :, channel]
+
+        low = np.percentile(
+            band,
+            lower_percentile
+        )
+
+        high = np.percentile(
+            band,
+            upper_percentile
+        )
+
+        if high <= low:
+            result[:, :, channel] = np.clip(
+                band,
+                0,
+                255
+            ).astype(np.uint8)
+
+        else:
+
+            normalized = (
+                (band - low) /
+                (high - low)
+            )
+
+            normalized = np.clip(
+                normalized,
+                0,
+                1
+            )
+
+            result[:, :, channel] = (
+                normalized * 255
+            ).astype(np.uint8)
 
     return result
 

inference/pipeline.py

@@ -1,504 +1,504 @@
-
-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
-    }
-
+
+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
+    }
+

local_test.py

@@ -0,0 +1,219 @@
+"""
+Local test script for satellite_sr_deploy.
+Tests:
+  1. Import checks for all modules
+  2. HATSAT inference on a synthetic 198x139 image
+  3. ESRGAN inference on a synthetic 198x139 image
+  4. Router switch (HATSAT → ESRGAN)
+"""
+
+import sys
+import time
+from pathlib import Path
+
+# -------------------------------------------------------
+# Make sure we run from the deploy root
+# -------------------------------------------------------
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+print("=" * 60)
+print("STEP 1 — IMPORT CHECK")
+print("=" * 60)
+
+try:
+    import torch
+    print(f"  torch        : {torch.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] torch : {e}")
+    sys.exit(1)
+
+try:
+    import numpy as np
+    print(f"  numpy        : {np.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] numpy : {e}")
+    sys.exit(1)
+
+try:
+    from PIL import Image
+    import PIL
+    print(f"  Pillow       : {PIL.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] Pillow : {e}")
+    sys.exit(1)
+
+try:
+    import einops
+    print(f"  einops       : {einops.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] einops : {e}")
+    sys.exit(1)
+
+try:
+    import gradio as gr
+    print(f"  gradio       : {gr.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] gradio : {e}")
+    sys.exit(1)
+
+try:
+    import rasterio
+    print(f"  rasterio     : {rasterio.__version__}")
+except ImportError as e:
+    print(f"  [FAIL] rasterio : {e}")
+    sys.exit(1)
+
+print("\nAll imports OK.\n")
+
+# -------------------------------------------------------
+# Check weight files
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 2 — WEIGHT FILE CHECK")
+print("=" * 60)
+
+hatsat_ckpt = ROOT / "weights" / "hatsat" / "net_g_150000.pth"
+esrgan_ckpt = ROOT / "weights" / "esrgan" / "RRDB_ESRGAN_x4.pth"
+
+if hatsat_ckpt.exists():
+    size_mb = hatsat_ckpt.stat().st_size / (1024 * 1024)
+    print(f"  HATSAT  : {hatsat_ckpt}  ({size_mb:.1f} MB)  OK")
+else:
+    print(f"  [FAIL] HATSAT checkpoint not found: {hatsat_ckpt}")
+    sys.exit(1)
+
+if esrgan_ckpt.exists():
+    size_mb = esrgan_ckpt.stat().st_size / (1024 * 1024)
+    print(f"  ESRGAN  : {esrgan_ckpt}  ({size_mb:.1f} MB)  OK")
+else:
+    print(f"  [FAIL] ESRGAN checkpoint not found: {esrgan_ckpt}")
+    sys.exit(1)
+
+print()
+
+# -------------------------------------------------------
+# Model router import
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 3 — MODEL ROUTER IMPORT")
+print("=" * 60)
+
+try:
+    from model_router import SatelliteSRRouter
+    print("  SatelliteSRRouter imported OK")
+except Exception as e:
+    print(f"  [FAIL] model_router import: {e}")
+    sys.exit(1)
+
+print()
+
+# -------------------------------------------------------
+# Create a synthetic test image (198 x 139, RGB)
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 4 — CREATE SYNTHETIC TEST IMAGE (198x139 RGB)")
+print("=" * 60)
+
+import numpy as np
+test_arr = np.random.randint(0, 255, (139, 198, 3), dtype=np.uint8)
+test_image = Image.fromarray(test_arr, mode="RGB")
+print(f"  Input size: {test_image.size}  (W x H)")
+print()
+
+# -------------------------------------------------------
+# Instantiate router (lazy — no model loaded yet)
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 5 — INIT ROUTER")
+print("=" * 60)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"  Device: {device}")
+
+router = SatelliteSRRouter(device=device)
+print()
+
+# -------------------------------------------------------
+# HATSAT inference
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 6 — HATSAT INFERENCE")
+print("=" * 60)
+
+t0 = time.time()
+try:
+    hatsat_output = router.predict(test_image, model_name="hatsat")
+    elapsed = time.time() - t0
+    print(f"  Input  : {test_image.size}")
+    print(f"  Output : {hatsat_output.size}")
+    expected = (test_image.width * 4, test_image.height * 4)
+    if hatsat_output.size == expected:
+        print(f"  Scale check: OK  (expected {expected})")
+    else:
+        print(f"  [WARN] Scale mismatch: expected {expected}, got {hatsat_output.size}")
+    print(f"  Time   : {elapsed:.2f}s")
+except Exception as e:
+    print(f"  [FAIL] HATSAT inference: {e}")
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+
+print()
+
+# -------------------------------------------------------
+# ESRGAN inference (router should unload HATSAT first)
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 7 — ESRGAN INFERENCE (router switches model)")
+print("=" * 60)
+
+t0 = time.time()
+try:
+    esrgan_output = router.predict(test_image, model_name="esrgan")
+    elapsed = time.time() - t0
+    print(f"  Input  : {test_image.size}")
+    print(f"  Output : {esrgan_output.size}")
+    expected = (test_image.width * 4, test_image.height * 4)
+    if esrgan_output.size == expected:
+        print(f"  Scale check: OK  (expected {expected})")
+    else:
+        print(f"  [WARN] Scale mismatch: expected {expected}, got {esrgan_output.size}")
+    print(f"  Time   : {elapsed:.2f}s")
+except Exception as e:
+    print(f"  [FAIL] ESRGAN inference: {e}")
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+
+print()
+
+# -------------------------------------------------------
+# Save outputs
+# -------------------------------------------------------
+print("=" * 60)
+print("STEP 8 — SAVE TEST OUTPUTS")
+print("=" * 60)
+
+out_dir = ROOT / "outputs" / "local_test"
+out_dir.mkdir(parents=True, exist_ok=True)
+
+hatsat_path = out_dir / "test_hatsat_output.png"
+esrgan_path = out_dir / "test_esrgan_output.png"
+
+hatsat_output.save(hatsat_path)
+esrgan_output.save(esrgan_path)
+
+print(f"  HATSAT output saved : {hatsat_path}")
+print(f"  ESRGAN output saved : {esrgan_path}")
+print()
+
+# -------------------------------------------------------
+# Summary
+# -------------------------------------------------------
+print("=" * 60)
+print("LOCAL TEST COMPLETE")
+print("=" * 60)
+print("  All tests passed.")
+print(f"  HATSAT : {test_image.size} -> {hatsat_output.size}")
+print(f"  ESRGAN : {test_image.size} -> {esrgan_output.size}")
+print("=" * 60)

model_router.py

@@ -10,9 +10,11 @@ class SatelliteSRRouter:
 
         from models.hatsat.hatsat_inference import HATSATInference
         from models.esrgan.esrgan_inference import ESRGANInference
+        from models.sen2sr.sen2sr_inference import Sen2SRInference
 
         self.HATSATInference = HATSATInference
         self.ESRGANInference = ESRGANInference
+        self.Sen2SRInference = Sen2SRInference
 
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -23,6 +25,7 @@ class SatelliteSRRouter:
         # Only the selected model is loaded.
         self.hatsat = None
         self.esrgan = None
+        self.sen2sr = None
 
         print("=" * 55)
         print("SATELLITE SR ROUTER")
@@ -71,6 +74,22 @@ class SatelliteSRRouter:
 
             self._clear_gpu()
 
+    def _unload_sen2sr(self):
+
+        if self.sen2sr is not None:
+
+            print("Unloading Sen2SR...")
+
+            try:
+                if hasattr(self.sen2sr, "model") and hasattr(self.sen2sr.model, "model"):
+                    self.sen2sr.model.model.cpu()
+            except Exception:
+                pass
+
+            self.sen2sr = None
+
+            self._clear_gpu()
+
     def _load_hatsat(self):
 
         if self.hatsat is None:
@@ -103,6 +122,18 @@ class SatelliteSRRouter:
 
             print("ESRGAN loaded successfully.")
 
+    def _load_sen2sr(self):
+
+        if self.sen2sr is None:
+
+            print("Loading Sen2SR...")
+
+            self.sen2sr = self.Sen2SRInference(
+                device=self.device
+            )
+
+            print("Sen2SR loaded successfully.")
+
     def available_models(self):
 
         return {
@@ -116,6 +147,12 @@ class SatelliteSRRouter:
                 "name": "ESRGAN",
                 "description": "General-purpose super-resolution baseline",
                 "scale": 4
+            },
+
+            "sen2sr": {
+                "name": "Sen2SR",
+                "description": "WEO-SAS Sentinel-2 CNN super-resolution (HuggingFace)",
+                "scale": 4
             }
         }
 
@@ -137,8 +174,9 @@ class SatelliteSRRouter:
 
         if model_name == "hatsat":
 
-            # Free GPU memory used by ESRGAN
+            # Free GPU memory used by other models
             self._unload_esrgan()
+            self._unload_sen2sr()
 
             # Load HATSAT only when required
             self._load_hatsat()
@@ -157,8 +195,9 @@ class SatelliteSRRouter:
 
         elif model_name == "esrgan":
 
-            # Free GPU memory used by HATSAT
+            # Free GPU memory used by HATSAT and Sen2SR
             self._unload_hatsat()
+            self._unload_sen2sr()
 
             # Load ESRGAN only when required
             self._load_esrgan()
@@ -171,9 +210,30 @@ class SatelliteSRRouter:
 
             return result
 
+        # ==========================================
+        # SEN2SR
+        # ==========================================
+
+        elif model_name == "sen2sr":
+
+            # Free GPU memory used by other models
+            self._unload_hatsat()
+            self._unload_esrgan()
+
+            # Load Sen2SR only when required
+            self._load_sen2sr()
+
+            print("Running Sen2SR...")
+
+            result = self.sen2sr.predict(image)
+
+            print("Sen2SR inference complete.")
+
+            return result
+
         else:
 
             raise ValueError(
                 f"Unknown model '{model_name}'. "
-                f"Choose 'hatsat' or 'esrgan'."
+                f"Choose 'hatsat', 'esrgan', or 'sen2sr'."
             )

models/sen2sr/__init__.py

@@ -0,0 +1 @@
+# Sen2SR model package

models/sen2sr/sen2sr_inference.py

@@ -0,0 +1,162 @@
+
+import sys
+import os
+import numpy as np
+import torch
+from pathlib import Path
+from PIL import Image
+
+
+class Sen2SRInference:
+    """
+    Inference wrapper for WEO-SAS/sen2sr — a Sentinel-2 super-resolution model
+    hosted on Hugging Face (https://huggingface.co/WEO-SAS/sen2sr).
+
+    The model's `main` branch expects 4-channel RGBN input (B04, B03, B02, B08)
+    as a float32 NumPy array in [0, 1] with shape (C, H, W) and returns a
+    4× upscaled output of the same channel order.
+
+    This wrapper accepts standard 3-channel RGB PIL images and adapts them
+    by treating the R channel as a proxy for NIR (B08), since single-image
+    workflows won't generally have a separate NIR band.
+    """
+
+    MODEL_REPO = "WEO-SAS/sen2sr"
+
+    def __init__(self, device=None, local_dir=None):
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.device = str(device)
+        self.local_dir = local_dir
+        self.model = None
+
+        self._load_model()
+
+    # ------------------------------------------------------------------
+    # Model loading
+    # ------------------------------------------------------------------
+
+    def _load_model(self):
+        """Download and load the sen2sr model from Hugging Face Hub."""
+
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            raise ImportError(
+                "huggingface_hub is required for Sen2SR. "
+                "Install it with: pip install huggingface-hub"
+            )
+
+        try:
+            from safetensors import safe_open  # noqa: F401 — ensure available
+        except ImportError:
+            raise ImportError(
+                "safetensors is required for Sen2SR. "
+                "Install it with: pip install safetensors"
+            )
+
+        print("Sen2SR: downloading / verifying model from Hugging Face Hub …")
+
+        if self.local_dir is None:
+            self.local_dir = snapshot_download(
+                repo_id=self.MODEL_REPO,
+                revision="main",
+            )
+
+        print(f"Sen2SR: model files at {self.local_dir}")
+
+        # The model repo exposes a `model.py` with a `Model` class.
+        if str(self.local_dir) not in sys.path:
+            sys.path.insert(0, str(self.local_dir))
+
+        from model import Model  # type: ignore  # loaded dynamically
+
+        self.model = Model(local_dir=str(self.local_dir))
+
+        print("Sen2SR: model loaded successfully.")
+
+    # ------------------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(self, image) -> Image.Image:
+        """
+        Run 4× super-resolution on a PIL RGB image.
+
+        Parameters
+        ----------
+        image : PIL.Image.Image | np.ndarray
+            Input satellite image (RGB, any uint8 size).
+
+        Returns
+        -------
+        PIL.Image.Image
+            4× upscaled RGB image.
+        """
+
+        if self.model is None:
+            raise RuntimeError("Sen2SR model is not loaded.")
+
+        # ── Normalise input ────────────────────────────────────────────
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        image = image.convert("RGB")
+
+        img_np = np.array(image).astype(np.float32) / 255.0  # (H, W, 3)
+
+        # ── Build RGBN array: (4, H, W) ───────────────────────────────
+        # Model expects: B04(R), B03(G), B02(B), B08(NIR).
+        # Approximate NIR with the Red channel as the best available proxy.
+        r = img_np[:, :, 0]
+        g = img_np[:, :, 1]
+        b = img_np[:, :, 2]
+        nir_proxy = r
+
+        rgbn = np.stack([r, g, b, nir_proxy], axis=0)  # (4, H, W)
+
+        # ── Pad input to a multiple of patch_size=128 ─────────────────
+        # The HardConstraint FFT low-pass mask is baked at patch_size=128
+        # (output size 512). Both the direct path (image ≤ 128) and the
+        # tiled predict_large path need H and W to be exact multiples of
+        # 128 so every tile lands on a 128×128 boundary.
+        PATCH = 128
+        _, orig_h, orig_w = rgbn.shape
+
+        pad_h = (PATCH - orig_h % PATCH) % PATCH
+        pad_w = (PATCH - orig_w % PATCH) % PATCH
+
+        if pad_h > 0 or pad_w > 0:
+            rgbn = np.pad(
+                rgbn,
+                ((0, 0), (0, pad_h), (0, pad_w)),
+                mode="reflect",
+            )
+
+        # ── Run model ─────────────────────────────────────────────────
+        sr_rgbn = self.model.predict(rgbn)  # (4, H_padded*4, W_padded*4)
+
+        # ── Crop back to the true 4× output size ──────────────────────
+        out_h = orig_h * 4
+        out_w = orig_w * 4
+        sr_rgbn = sr_rgbn[:, :out_h, :out_w]
+
+        # ── Convert to uint8 RGB PIL image ────────────────────────────
+        sr_rgb = sr_rgbn[:3, :, :]           # (3, H*4, W*4)
+        sr_rgb = np.clip(sr_rgb, 0.0, 1.0)
+        sr_rgb = (sr_rgb * 255.0).round().astype(np.uint8)
+        sr_rgb = sr_rgb.transpose(1, 2, 0)   # (H*4, W*4, 3)
+
+        return Image.fromarray(sr_rgb, mode="RGB")
+
+
+# -----------------------------------------------------------------------
+# Convenience function
+# -----------------------------------------------------------------------
+
+def run_sen2sr(image, device=None, local_dir=None):
+    model = Sen2SRInference(device=device, local_dir=local_dir)
+    return model.predict(image)