Spaces:
Running on Zero
Running on Zero
File size: 5,881 Bytes
8d928e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
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)
|