import io
import time
import sys
import tempfile
from pathlib import Path
import gradio as gr
from PIL import Image
import spaces
import numpy as np
import rasterio
from inference.geotiff import normalize_rgb, save_sr_geotiff
# ============================================================
# PATHS
# ============================================================
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
# ============================================================
# ZERO-GPU MODEL LOADING (CPU initially)
# ============================================================
from models.hatsat.hatsat_inference import HATSATInference
from models.esrgan.esrgan_inference import ESRGANInference
from models.sen2sr.sen2sr_inference import Sen2SRInference
print("Loading models to CPU for ZeroGPU deployment...")
try:
hatsat_model = HATSATInference(device="cpu")
esrgan_model = ESRGANInference(
checkpoint_path=str(ROOT / "weights" / "esrgan" / "RRDB_ESRGAN_x4.pth"),
device="cpu"
)
sen2sr_model = Sen2SRInference(device="cpu")
print("Models loaded successfully on CPU.")
except Exception as e:
print(f"Error loading models: {e}")
hatsat_model = None
esrgan_model = None
sen2sr_model = None
# ============================================================
# ZERO-GPU INFERENCE WRAPPERS
# ============================================================
@spaces.GPU(duration=120)
def infer_hatsat(image):
if hatsat_model is None:
raise RuntimeError("HATSAT model is not loaded.")
return hatsat_model.predict(image)
@spaces.GPU(duration=120)
def infer_esrgan(image):
if esrgan_model is None:
raise RuntimeError("ESRGAN model is not loaded.")
return esrgan_model.predict(image)
@spaces.GPU(duration=120)
def infer_sen2sr(image):
if sen2sr_model is None:
raise RuntimeError("Sen2SR model is not loaded.")
return sen2sr_model.predict(image)
# ============================================================
# CUSTOM CSS
# ============================================================
CSS = """
/* Google Font */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
body, .gradio-container {
font-family: 'Inter', sans-serif !important;
background-color: #f8fafc !important;
}
/* Hide default Gradio padding/footers */
footer { display: none !important; }
.gradio-container { max-width: 100% !important; padding: 0 !important; margin: 0 !important; }
/* Top Navbar */
.top-nav {
background: linear-gradient(135deg, #03182b 0%, #07355c 100%);
padding: 20px 40px;
display: flex;
justify-content: space-between;
align-items: center;
color: white;
}
.nav-left { display: flex; align-items: center; gap: 16px; }
.nav-logo-icon { width: 36px; height: 36px; fill: #38bdf8; }
.nav-title { display: flex; flex-direction: column; }
.nav-logo-text { font-size: 1.6rem; font-weight: 700; color: #38bdf8; margin: 0; line-height: 1.1; }
.nav-logo-text span { color: white; }
.nav-subtitle { font-size: 0.8rem; color: #cbd5e1; margin: 0; }
.nav-links { display: flex; gap: 30px; font-size: 0.95rem; align-items: center; }
.nav-links a { color: #cbd5e1; text-decoration: none; font-weight: 500; transition: color 0.2s; }
.nav-links a:hover { color: white; }
.nav-links a.active { color: white; border-bottom: 2px solid white; padding-bottom: 4px; }
.nav-btn { background: #34d399; color: #064e3b !important; padding: 10px 20px; border-radius: 6px; font-weight: 600; text-decoration: none; transition: background 0.2s; }
.nav-btn:hover { background: #10b981; }
/* Sidebar */
.sidebar-container { background: #f8fafc; padding: 30px 20px; height: 100%; border-right: 1px solid #e2e8f0; display: flex; flex-direction: column; justify-content: space-between; }
.sidebar-links { display: flex; flex-direction: column; gap: 8px; }
.sidebar-link { display: flex; align-items: center; padding: 14px 20px; border-radius: 10px; color: #475569; text-decoration: none; font-weight: 500; font-size: 0.95rem; gap: 16px; transition: background 0.2s; cursor: pointer; }
.sidebar-link:hover { background: #f1f5f9; }
.sidebar-link.active { background: #d1fae5; color: #065f46; font-weight: 600; }
.sidebar-link svg { width: 20px; height: 20px; }
.sidebar-banner { background: #e0f2fe; padding: 24px 20px; border-radius: 16px; text-align: center; color: #0369a1; display: flex; flex-direction: column; align-items: center; gap: 12px; }
.sidebar-banner svg { width: 40px; height: 40px; color: #0284c7; }
.sidebar-banner div { font-weight: 600; font-size: 0.95rem; line-height: 1.4; }
/* Main Content Area */
.main-wrapper { padding: 30px 40px; }
.page-title { font-size: 2.2rem; font-weight: 700; color: #0f172a; margin: 0 0 8px; letter-spacing: -0.5px; }
.page-title span { color: #0369a1; }
.page-subtitle { color: #64748b; font-size: 1.05rem; margin: 0 0 30px; }
/* Upload box */
.upload-card { background: white !important; border-radius: 16px !important; padding: 0 !important; border: 1px dashed #cbd5e1 !important; box-shadow: 0 4px 20px rgba(0,0,0,0.03) !important; overflow: hidden; }
/* Model Select & Enhance Button */
.model-card { background: white; border-radius: 12px; padding: 20px; border: 1px solid #e2e8f0; box-shadow: 0 4px 6px rgba(0,0,0,0.02); margin-bottom: 20px; }
.model-card-title { font-size: 1rem; font-weight: 600; color: #1e293b; margin-bottom: 12px; }
.model-info-box { background: #f0f9ff; color: #0369a1; padding: 14px; border-radius: 8px; font-size: 0.85rem; margin-top: 12px; display: flex; gap: 12px; align-items: flex-start; line-height: 1.5; }
#enhance-btn { background: #0f9d58 !important; color: white !important; font-weight: 600 !important; font-size: 1.1rem !important; padding: 16px !important; border-radius: 8px !important; border: none !important; width: 100%; transition: all 0.2s; }
#enhance-btn:hover { background: #0b8043 !important; transform: translateY(-1px); }
/* Features List */
.features-card { background: #f8fafc; border-radius: 12px; padding: 24px; border: 1px solid #e2e8f0; height: 100%; display: flex; flex-direction: column; justify-content: center; }
.features-title { font-size: 1.1rem; font-weight: 600; color: #0f172a; margin-bottom: 20px; }
.feature-item { display: flex; align-items: center; gap: 14px; margin-bottom: 16px; color: #334155; font-size: 0.95rem; font-weight: 500; }
.feature-icon { width: 24px; height: 24px; color: #0f9d58; flex-shrink: 0; }
/* Results */
.results-header { font-size: 1.4rem; font-weight: 700; color: #0f172a; margin: 40px 0 20px; }
.image-card { border-radius: 12px !important; overflow: hidden; border: 1px solid #e2e8f0; background: white; padding: 0 !important; box-shadow: 0 4px 12px rgba(0,0,0,0.04); }
/* Buttons */
.download-col { display: flex; flex-direction: column; gap: 12px; padding-top: 36px; }
.btn-primary { background: #0f9d58 !important; color: white !important; padding: 12px 16px !important; border-radius: 8px !important; font-weight: 600 !important; text-align: center; border: none !important; }
.btn-outline { background: white !important; color: #334155 !important; border: 1px solid #cbd5e1 !important; padding: 12px 16px !important; border-radius: 8px !important; font-weight: 500 !important; text-align: center; box-shadow: 0 1px 2px rgba(0,0,0,0.02) !important; }
/* Footer */
.footer-strip { background: #03182b; color: white; padding: 24px 40px; display: flex; justify-content: space-between; align-items: center; font-size: 0.9rem; margin-top: 40px; }
.footer-left { display: flex; gap: 10px; align-items: center; color: #cbd5e1; }
.footer-left strong { color: white; }
.footer-links { display: flex; gap: 24px; color: #94a3b8; align-items: center; }
.footer-links a { color: #94a3b8; text-decoration: none; transition: color 0.2s; }
.footer-links a:hover { color: white; }
"""
# ============================================================
# INFERENCE FUNCTION
# ============================================================
def run_super_resolution(image_path, model_choice, progress=gr.Progress()):
if image_path is None:
raise gr.Error("⚠️ Please upload a satellite image before running.")
model_map = {
"HAT-SAT (Recommended)": "hatsat",
"ESRGAN (Baseline)": "esrgan",
"Sen2SR — WEO-SAS (Sentinel-2 CNN)": "sen2sr",
}
selected_model = model_map.get(model_choice)
if selected_model is None:
raise gr.Error("⚠️ Please select a valid model.")
if not isinstance(image_path, list):
image_path = [image_path]
is_tiff = False
if len(image_path) == 3:
paths = [str(p) for p in image_path]
if not all(p.lower().endswith(('.tif', '.tiff')) for p in paths):
raise gr.Error("⚠️ When uploading multiple files, they must all be TIFF files.")
red_path = next((p for p in paths if "B04" in p.upper()), None)
green_path = next((p for p in paths if "B03" in p.upper()), None)
blue_path = next((p for p in paths if "B02" in p.upper()), None)
if not (red_path and green_path and blue_path):
paths = sorted(paths)
blue_path, green_path, red_path = paths[0], paths[1], paths[2]
try:
with rasterio.open(red_path) as src_r, rasterio.open(green_path) as src_g, rasterio.open(blue_path) as src_b:
profile = src_r.profile.copy()
profile.update(count=3)
if not (src_r.width == src_g.width == src_b.width and src_r.height == src_g.height == src_b.height):
raise gr.Error("⚠️ The uploaded bands have mismatched dimensions.")
r_data = src_r.read(1)
g_data = src_g.read(1)
b_data = src_b.read(1)
data = np.stack([r_data, g_data, b_data], axis=-1)
input_w, input_h = src_r.width, src_r.height
is_tiff = True
except Exception as e:
raise gr.Error(f"⚠️ Error reading 3-band TIFFs: {str(e)}")
if data.dtype != np.uint8:
data = normalize_rgb(data)
original_image = Image.fromarray(data, mode="RGB")
elif len(image_path) == 1:
file_path = str(image_path[0])
is_tiff = file_path.lower().endswith(('.tif', '.tiff'))
if is_tiff:
try:
with rasterio.open(file_path) as src:
profile = src.profile.copy()
count = src.count
if count not in [1, 3, 4]:
raise gr.Error(f"⚠️ Unsupported band count: {count}. Expected 1, 3, or 4 bands.")
if count == 3:
data = src.read([1, 2, 3])
data = np.transpose(data, (1, 2, 0))
elif count == 4:
data = src.read([1, 2, 3])
data = np.transpose(data, (1, 2, 0))
else:
data = src.read(1)
data = np.stack([data, data, data], axis=-1)
input_w, input_h = src.width, src.height
except rasterio.errors.RasterioIOError:
raise gr.Error("⚠️ Corrupted TIFF or unsupported format.")
except Exception as e:
raise gr.Error(f"⚠️ Error reading TIFF: {str(e)}")
if data.dtype == np.uint8:
pass
else:
data = normalize_rgb(data)
original_image = Image.fromarray(data, mode="RGB")
else:
try:
original_image = Image.open(file_path).convert("RGB")
input_w, input_h = original_image.size
except Exception:
raise gr.Error("⚠️ Could not read image.")
else:
raise gr.Error("⚠️ Please upload either 1 file (multiband image) or 3 files (B04, B03, B02 GeoTIFFs).")
if input_w < 16 or input_h < 16:
raise gr.Error(f"⚠️ Image too small ({input_w}x{input_h}). Please upload an image at least 16x16 pixels.")
if input_w > 1024 or input_h > 1024:
raise gr.Error(f"⚠️ Image too large ({input_w}x{input_h}). Please use images up to 1024x1024 pixels.")
start_time = time.time()
try:
progress(0.5, desc=f"Running {selected_model} 4x super-resolution (ZeroGPU) 🚀")
if selected_model == "hatsat":
result = infer_hatsat(original_image)
elif selected_model == "esrgan":
result = infer_esrgan(original_image)
else:
result = infer_sen2sr(original_image)
except Exception as e:
raise gr.Error(f"{selected_model} inference failed: {e}")
progress(0.9, desc="Finalising output 🚀")
# Generate PNG download
tmp_png = tempfile.NamedTemporaryFile(suffix=f"_sr_{selected_model}_4x.png", delete=False)
result.save(tmp_png.name, format="PNG")
tmp_png.close()
out_png = tmp_png.name
# Generate TIFF download if input was TIFF
if is_tiff:
tmp_tiff = tempfile.NamedTemporaryFile(suffix=f"_sr_{selected_model}_4x_georef.tif", delete=False)
tmp_tiff.close()
save_sr_geotiff(result, profile, tmp_tiff.name, scale=4)
out_tiff = tmp_tiff.name
else:
out_tiff = None
progress(1.0, desc="Done.")
# Return: original preview, enhanced preview, png download, tiff download
return original_image, result, gr.update(value=out_png, interactive=True), gr.update(value=out_tiff, interactive=True)
# ============================================================
# HTML COMPONENTS
# ============================================================
TOP_NAV = """
SentiEnhance
AI-Powered Super-Resolution for Sentinel-2 Imagery
"""
SIDEBAR = """
"""
WHY_USE_US = """
Why Use Our Platform?
AI-powered super-resolution
Supports Sentinel-2 imagery
Preserves geospatial metadata
Easy to use web interface
Free for research & education
"""
FOOTER = """
"""
# ============================================================
# APP LAYOUT
# ============================================================
with gr.Blocks(theme=gr.themes.Base(), css=CSS) as demo:
gr.HTML(TOP_NAV)
with gr.Row():
with gr.Column(scale=1, min_width=240, elem_classes="sidebar"):
gr.HTML(SIDEBAR)
with gr.Column(scale=5, elem_classes="main-wrapper"):
gr.HTML("""
Enhance Sentinel-2 Satellite Images
Upload your Sentinel-2 image (GeoTIFF, TIFF, JP2, or PNG) and use AI to generate a high-resolution version.
""")
with gr.Row():
with gr.Column(scale=4, elem_classes="upload-card"):
input_image = gr.File(
label="Drag & drop your satellite image here",
file_count="multiple",
file_types=["image", ".tif", ".tiff", ".jp2"]
)
with gr.Column(scale=3):
with gr.Group(elem_classes="model-card"):
gr.HTML("Model Selection
")
model_selector = gr.Dropdown(
choices=[
"HAT-SAT (Recommended)",
"ESRGAN (Baseline)",
"Sen2SR — WEO-SAS (Sentinel-2 CNN)",
],
value="HAT-SAT (Recommended)",
show_label=False,
container=False
)
gr.HTML("""
HAT-SAT: State-of-the-art transformer model. ESRGAN: General-purpose RRDB baseline. Sen2SR: WEO-SAS Sentinel-2 CNN — HuggingFace ↗ .
""")
enhance_btn = gr.Button("✨ Enhance Image", elem_id="enhance-btn")
gr.HTML("Processing may take a few moments...
")
with gr.Column(scale=3, elem_classes="features-card"):
gr.HTML(WHY_USE_US)
# Results Section
gr.HTML("")
with gr.Row():
with gr.Column(scale=4, elem_classes="image-card"):
original_image = gr.Image(label="Original (Sentinel-2, 10m)", type="pil", interactive=False)
with gr.Column(scale=4, elem_classes="image-card"):
enhanced_image = gr.Image(label="Enhanced (HAT-SAT, 2x)", type="pil", interactive=False)
with gr.Column(scale=3, elem_classes="download-col"):
gr.HTML("Download & Share
")
download_png = gr.DownloadButton("📥 Download Enhanced Image", elem_classes="btn-primary")
download_tiff = gr.DownloadButton("📄 Download as GeoTIFF", elem_classes="btn-outline")
compare_btn = gr.Button("🔍 Compare Images", elem_classes="btn-outline")
share_btn = gr.Button("🔗 Share Result", elem_classes="btn-outline")
gr.HTML(FOOTER)
# ============================================================
# EVENTS
# ============================================================
enhance_btn.click(
fn=run_super_resolution,
inputs=[input_image, model_selector],
outputs=[original_image, enhanced_image, download_png, download_tiff],
show_progress="full",
)
# Interactive elements that don't do anything functionally but can show alerts
def show_alert():
gr.Info("This feature is coming soon!")
compare_btn.click(fn=show_alert)
share_btn.click(fn=show_alert)
# ============================================================
# LAUNCH
# ============================================================
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True
)