Spaces:
Running on Zero
Running on Zero
| 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 | |
| # ------------------------------------------------------------------ | |
| 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) | |