initial

README.md

@@ -1,13 +1,183 @@
 ---
 title: SpectraGAN
-emoji: 🔥
+emoji: 🖼️
 colorFrom: blue
-colorTo: indigo
+colorTo: green
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 5.31.0
 app_file: app.py
 pinned: false
-license: mit
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🖼️ SpectraGAN — Multi-Model Upscaler Comparison
+
+A Gradio web app that lets you upscale an image with **multiple SR models simultaneously** and compare results side by side.
+
+Supported models:
+
+| Model | Architecture | Scale |
+|-------|-------------|-------|
+| Real-ESRGAN ×2 | GAN (residual-in-residual dense block) | ×2 |
+| Real-ESRGAN ×4 | GAN (residual-in-residual dense block) | ×4 |
+| SRCNN ×4 | Shallow 3-layer CNN | ×4 |
+| HResNet ×4 | Deep residual network (EDSR-style) | ×4 |
+| SR3 *(stub)* | Diffusion model | ×4 — see note below |
+
+---
+
+## 📑 Table of Contents
+
+1. [Features](#features)
+2. [Project Structure](#project-structure)
+3. [Prerequisites](#prerequisites)
+4. [Installation](#installation)
+5. [Adding Your ONNX Models](#adding-your-onnx-models)
+6. [Running Locally](#running-locally)
+7. [SR3 Integration Guide](#sr3-integration-guide)
+8. [Contributing](#contributing)
+9. [License](#license)
+
+---
+
+## ✨ Features
+
+- **Side-by-side comparison** — run up to 4 models at once, results displayed in a 4-panel grid.
+- **Selective execution** — toggle any model on/off before running; unchecked models are skipped.
+- **×8 post-resize** — optionally apply a bicubic ×2 pass on top of any ×4 result.
+- **Tile-based inference** — large images are split into tiles matching each model's fixed input size, then stitched back together seamlessly.
+- **Per-result download** — each panel has its own PNG download button.
+- **Graceful degradation** — if a model file is missing (e.g. Drive ID not yet set), that panel is skipped without crashing the others.
+
+---
+
+## 📁 Project Structure
+
+```
+spectragan/
+├── model/
+│   ├── Real-ESRGAN_x2plus.onnx   # auto-downloaded
+│   ├── Real-ESRGAN-x4plus.onnx   # auto-downloaded
+│   ├── SRCNN_x4.onnx              # you provide — see below
+│   └── HResNet_x4.onnx            # you provide — see below
+├── app.py
+├── requirements.txt
+└── README.md
+```
+
+---
+
+## ⚙️ Prerequisites
+
+- Python 3.10+
+- `git`
+- A terminal / command prompt
+
+---
+
+## 🔧 Installation
+
+```bash
+git clone https://github.com/ParamAhuja/SpectraGAN.git
+cd SpectraGAN
+python -m venv .venv
+source .venv/bin/activate      # Linux/macOS
+# .venv\Scripts\activate       # Windows
+pip install -r requirements.txt
+```
+
+---
+
+## 🗂️ Adding Your ONNX Models
+
+The Real-ESRGAN weights are downloaded automatically from Google Drive on first run.
+
+For **SRCNN** and **HResNet** you need to:
+
+1. Export your trained PyTorch model to ONNX:
+
+```python
+import torch
+
+# SRCNN example
+from srcnn import SRCNN
+model = SRCNN()
+model.load_state_dict(torch.load("srcnn.pth"))
+model.eval()
+
+dummy = torch.randn(1, 3, 128, 128)
+torch.onnx.export(
+    model, dummy, "SRCNN_x4.onnx",
+    input_names=["input"], output_names=["output"],
+    dynamic_axes={"input": {2: "H", 3: "W"}, "output": {2: "H", 3: "W"}}
+)
+```
+
+2. Upload the `.onnx` file to Google Drive and set **"Anyone with the link can view"**.
+
+3. Copy the file ID from the share URL and update `DRIVE_IDS` in `app.py`:
+
+```python
+DRIVE_IDS = {
+    ...
+    "srcnn_x4":   "YOUR_SRCNN_DRIVE_FILE_ID_HERE",
+    "hresnet_x4": "YOUR_HRESNET_DRIVE_FILE_ID_HERE",
+}
+```
+
+---
+
+## 🚀 Running Locally
+
+```bash
+python app.py
+```
+
+Open `http://127.0.0.1:7860` in your browser.
+
+---
+
+## 🌀 SR3 Integration Guide
+
+SR3 (Super-Resolution via Repeated Refinement) is a **diffusion model** — it cannot be exported to a static ONNX graph because its inference involves a variable-length denoising loop.
+
+To add SR3:
+
+1. Clone the reference implementation:
+   ```bash
+   git clone https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement
+   ```
+
+2. Place your trained checkpoint at `model/sr3_x4.pth`.
+
+3. Add `torch` and `torchvision` to `requirements.txt`.
+
+4. Write a wrapper in `app.py`:
+   ```python
+   def run_sr3(input_img: Image.Image) -> Image.Image:
+       # load config + model, run the denoising loop, return result
+       ...
+   ```
+
+5. Add `"sr3_x4"` to the `PANEL_KEYS` list and wire `run_sr3` into `compare_models`.
+
+---
+
+## 🤝 Contributing
+
+Pull requests welcome. Please open an issue first to discuss significant changes.
+
+---
+
+## 📄 License
+
+Apache 2.0 — see `LICENSE`.
+
+---
+
+## 👤 Author & Credits
+
+- Real-ESRGAN by [xinntao](https://github.com/xinntao/Real-ESRGAN)
+- SRCNN by Dong et al. (2014)
+- HResNet / EDSR by Lim et al. (2017)
+- SR3 by Ho et al. (2022) — [paper](https://arxiv.org/abs/2104.07636)

app.py

@@ -0,0 +1,363 @@
+import os
+import math
+import uuid
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+import gradio as gr
+import tempfile
+import requests
+
+# ---------------------------------------------------------------------------
+# Directory & model paths
+# ---------------------------------------------------------------------------
+MODEL_DIR = "model"
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+MODEL_PATHS = {
+    "esrgan_x2":  os.path.join(MODEL_DIR, "Real-ESRGAN_x2plus.onnx"),
+    "esrgan_x4":  os.path.join(MODEL_DIR, "Real-ESRGAN-x4plus.onnx"),
+    "srcnn_x4":   os.path.join(MODEL_DIR, "SRCNN_x4.onnx"),
+    "hresnet_x4": os.path.join(MODEL_DIR, "HResNet_x4.onnx"),
+}
+
+# ---------------------------------------------------------------------------
+# Google Drive file IDs
+# TODO: Replace the SRCNN / HResNet IDs with your own uploaded ONNX exports.
+#       Steps to export SRCNN to ONNX:
+#         import torch; from srcnn import SRCNN
+#         model = SRCNN(); model.load_state_dict(torch.load("srcnn.pth"))
+#         dummy = torch.randn(1, 3, 128, 128)
+#         torch.onnx.export(model, dummy, "SRCNN_x4.onnx",
+#                           input_names=["input"], output_names=["output"],
+#                           dynamic_axes={"input":{2:"H",3:"W"},"output":{2:"H",3:"W"}})
+#       Same pattern applies for HResNet / EDSR.
+# ---------------------------------------------------------------------------
+DRIVE_IDS = {
+    "esrgan_x2":  "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6",
+    "esrgan_x4":  "1wDBHad9RCJgJDGsPdapLYl3cr8j-PMJ6",
+    "srcnn_x4":   "YOUR_SRCNN_DRIVE_FILE_ID_HERE",    # <-- replace
+    "hresnet_x4": "YOUR_HRESNET_DRIVE_FILE_ID_HERE",  # <-- replace
+}
+
+# Scale factor each model produces
+MODEL_SCALES = {
+    "esrgan_x2":  2,
+    "esrgan_x4":  4,
+    "srcnn_x4":   4,
+    "hresnet_x4": 4,
+}
+
+# Human-readable labels shown in the UI
+MODEL_LABELS = {
+    "esrgan_x2":  "Real-ESRGAN ×2",
+    "esrgan_x4":  "Real-ESRGAN ×4",
+    "srcnn_x4":   "SRCNN ×4",
+    "hresnet_x4": "HResNet ×4",
+}
+
+# ---------------------------------------------------------------------------
+# SR3 NOTE
+# ---------------------------------------------------------------------------
+# SR3 (Super-Resolution via Repeated Refinement, Ho et al. 2022) is a
+# *diffusion model* and cannot be exported to ONNX in the same way as
+# feed-forward CNNs.  To integrate SR3:
+#   1. Clone the official repo: https://github.com/Janspiry/Image-Super-Resolution-via-Iterative-Refinement
+#   2. Place your checkpoint in model/sr3_x4.pth
+#   3. Add a `run_sr3(image: Image) -> Image` function that runs the
+#      denoising loop using PyTorch directly (add `torch` to requirements).
+#   4. Wire `run_sr3` into `upscale_one_model` for key "sr3_x4".
+# SR3 is intentionally omitted from the ONNX pipeline to avoid misleading
+# model shape assumptions.
+# ---------------------------------------------------------------------------
+
+
+# ---------------------------------------------------------------------------
+# Google Drive downloader
+# ---------------------------------------------------------------------------
+def download_from_drive(file_id: str, dest_path: str):
+    URL = "https://drive.google.com/uc?export=download"
+    session = requests.Session()
+    response = session.get(URL, params={"id": file_id}, stream=True)
+    token = None
+    for key, value in response.cookies.items():
+        if key.startswith("download_warning"):
+            token = value
+            break
+    if token:
+        response = session.get(URL, params={"id": file_id, "confirm": token}, stream=True)
+    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+    with open(dest_path, "wb") as f:
+        for chunk in response.iter_content(chunk_size=32768):
+            if chunk:
+                f.write(chunk)
+    print(f"Downloaded → {dest_path}")
+
+
+# ---------------------------------------------------------------------------
+# Download models if missing
+# ---------------------------------------------------------------------------
+for key, path in MODEL_PATHS.items():
+    if not os.path.isfile(path):
+        file_id = DRIVE_IDS[key]
+        if file_id.startswith("YOUR_"):
+            print(f"[WARN] Skipping {key}: Google Drive ID not set. "
+                  "Update DRIVE_IDS in app.py with your ONNX export.")
+        else:
+            print(f"Downloading {MODEL_LABELS[key]} …")
+            download_from_drive(file_id, path)
+
+
+# ---------------------------------------------------------------------------
+# Load ONNX sessions (only for models that have been downloaded)
+# ---------------------------------------------------------------------------
+sess_opts = ort.SessionOptions()
+sess_opts.intra_op_num_threads = 2
+sess_opts.inter_op_num_threads = 2
+
+SESSIONS = {}      # key → ort.InferenceSession
+INPUT_SHAPES = {}  # key → (H_in, W_in)
+
+for key, path in MODEL_PATHS.items():
+    if os.path.isfile(path):
+        try:
+            sess = ort.InferenceSession(
+                path,
+                sess_options=sess_opts,
+                providers=["CPUExecutionProvider"]
+            )
+            meta = sess.get_inputs()[0]
+            shape = tuple(meta.shape)
+            # shape is (1, 3, H, W) for fixed-size models, or (1,3,None,None) for dynamic
+            h = int(shape[2]) if shape[2] is not None and str(shape[2]).isdigit() else 128
+            w = int(shape[3]) if shape[3] is not None and str(shape[3]).isdigit() else 128
+            SESSIONS[key] = (sess, meta)
+            INPUT_SHAPES[key] = (h, w)
+            print(f"Loaded {MODEL_LABELS[key]}  tile={h}×{w}")
+        except Exception as e:
+            print(f"[ERROR] Could not load {key}: {e}")
+    else:
+        print(f"[INFO] {key} not available — will be skipped in comparisons.")
+
+
+# ---------------------------------------------------------------------------
+# Tile-based upscale for a single ONNX model
+# ---------------------------------------------------------------------------
+def run_onnx_tile(sess, meta, tile_np: np.ndarray) -> np.ndarray:
+    """Run one tile through any ONNX session. tile_np is HWC float32 [0,1]."""
+    patch = np.transpose(tile_np, (2, 0, 1))[None, ...]  # NCHW
+    out   = sess.run(None, {meta.name: patch})[0]
+    out   = np.squeeze(out, axis=0)
+    return np.transpose(out, (1, 2, 0))                   # back to HWC
+
+
+def tile_upscale_model(input_img: Image.Image, key: str, max_dim: int = 1024) -> Image.Image:
+    """
+    Upscale *input_img* using the ONNX model identified by *key*.
+    Returns a PIL Image at the upscaled resolution.
+    """
+    if key not in SESSIONS:
+        raise ValueError(f"Model '{key}' is not loaded. Check the Drive ID / path.")
+
+    sess, meta = SESSIONS[key]
+    H_in, W_in = INPUT_SHAPES[key]
+    scale       = MODEL_SCALES[key]
+
+    # Optionally cap input size to avoid OOM on large images
+    w, h = input_img.size
+    if w > max_dim or h > max_dim:
+        factor = max_dim / float(max(w, h))
+        input_img = input_img.resize((int(w * factor), int(h * factor)), Image.LANCZOS)
+
+    arr = np.array(input_img.convert("RGB")).astype(np.float32) / 255.0
+    h_orig, w_orig, _ = arr.shape
+
+    tiles_h = math.ceil(h_orig / H_in)
+    tiles_w = math.ceil(w_orig / W_in)
+    pad_h   = tiles_h * H_in - h_orig
+    pad_w   = tiles_w * W_in - w_orig
+
+    arr_padded = np.pad(arr, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
+    out_arr    = np.zeros((tiles_h * H_in * scale, tiles_w * W_in * scale, 3), dtype=np.float32)
+
+    for i in range(tiles_h):
+        for j in range(tiles_w):
+            y0, x0   = i * H_in, j * W_in
+            tile      = arr_padded[y0:y0 + H_in, x0:x0 + W_in, :]
+            up_tile   = run_onnx_tile(sess, meta, tile)
+            oy0, ox0  = i * H_in * scale, j * W_in * scale
+            out_arr[oy0:oy0 + H_in * scale, ox0:ox0 + W_in * scale, :] = up_tile
+
+    final = np.clip(out_arr[0:h_orig * scale, 0:w_orig * scale, :], 0.0, 1.0)
+    return Image.fromarray((final * 255.0).round().astype(np.uint8))
+
+
+def upscale_8x_from_4x(input_img: Image.Image, key: str) -> Image.Image:
+    """Run ×4 model, then bicubic ×2 to reach ×8."""
+    img_4x = tile_upscale_model(input_img, key)
+    w, h   = input_img.size
+    return img_4x.resize((w * 8, h * 8), Image.LANCZOS)
+
+
+# ---------------------------------------------------------------------------
+# Core comparison function  (called by the Gradio button)
+# ---------------------------------------------------------------------------
+def compare_models(
+    input_img: Image.Image,
+    use_esrgan_x2: bool,
+    use_esrgan_x4: bool,
+    use_srcnn:     bool,
+    use_hresnet:   bool,
+    include_8x:    bool,
+):
+    if input_img is None:
+        return [None] * 8   # 4 preview + 4 download slots
+
+    selection = []
+    if use_esrgan_x2: selection.append("esrgan_x2")
+    if use_esrgan_x4: selection.append("esrgan_x4")
+    if use_srcnn:     selection.append("srcnn_x4")
+    if use_hresnet:   selection.append("hresnet_x4")
+
+    previews   = []
+    downloads  = []
+
+    for key in selection:
+        if key not in SESSIONS:
+            previews.append(None)
+            downloads.append(gr.DownloadButton(label=f"{MODEL_LABELS[key]} – not loaded",
+                                               visible=True, value=None))
+            continue
+
+        try:
+            if include_8x and MODEL_SCALES[key] == 4:
+                result = upscale_8x_from_4x(input_img, key)
+                suffix = "×8"
+            else:
+                result = tile_upscale_model(input_img, key)
+                suffix = f"×{MODEL_SCALES[key]}"
+
+            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+            result.save(tmp.name, format="PNG")
+            tmp.close()
+
+            previews.append(result)
+            downloads.append(tmp.name)
+        except Exception as e:
+            print(f"[ERROR] {key}: {e}")
+            previews.append(None)
+            downloads.append(None)
+
+    # Pad to always return exactly 4 preview + 4 download values
+    while len(previews)  < 4: previews.append(None)
+    while len(downloads) < 4: downloads.append(None)
+
+    return previews + downloads   # 8-element list
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI – side-by-side comparison layout
+# ---------------------------------------------------------------------------
+css = """
+body { font-family: 'Segoe UI', sans-serif; }
+
+.panel-title {
+    text-align: center;
+    font-weight: 700;
+    font-size: 0.85rem;
+    letter-spacing: 0.08em;
+    text-transform: uppercase;
+    margin-bottom: 4px;
+    color: #555;
+}
+
+#run-btn {
+    background: linear-gradient(135deg, #1a1a2e, #16213e) !important;
+    color: #e2e2e2 !important;
+    font-size: 1rem !important;
+    font-weight: 600 !important;
+    border-radius: 8px !important;
+    padding: 12px 28px !important;
+}
+#run-btn:hover {
+    background: linear-gradient(135deg, #0f3460, #533483) !important;
+}
+
+.dl-btn button {
+    background: #f0f4ff !important;
+    border: 1px solid #c5d0f5 !important;
+    color: #333 !important;
+    font-size: 0.78rem !important;
+    border-radius: 6px !important;
+    width: 100%;
+}
+
+.model-toggle label { font-size: 0.9rem; }
+"""
+
+ALL_KEYS    = ["esrgan_x2", "esrgan_x4", "srcnn_x4", "hresnet_x4"]
+PANEL_KEYS  = ALL_KEYS   # order for the 4 comparison panels
+
+with gr.Blocks(css=css, title="SpectraGAN — Multi-Model Comparison") as demo:
+
+    gr.Markdown("""
+    # 🖼️ SpectraGAN — Multi-Model Upscaler Comparison
+    Upload an image, select models, and compare results side by side.
+    """)
+
+    # ── Input row ──────────────────────────────────────────────────────────
+    with gr.Row():
+        inp_image = gr.Image(type="pil", label="Source Image", scale=2)
+
+        with gr.Column(scale=1):
+            gr.Markdown("### Models to compare")
+            chk_esrgan_x2  = gr.Checkbox(label="Real-ESRGAN ×2",  value=True,  elem_classes="model-toggle")
+            chk_esrgan_x4  = gr.Checkbox(label="Real-ESRGAN ×4",  value=True,  elem_classes="model-toggle")
+            chk_srcnn      = gr.Checkbox(label="SRCNN ×4",        value=True,  elem_classes="model-toggle")
+            chk_hresnet    = gr.Checkbox(label="HResNet ×4",      value=True,  elem_classes="model-toggle")
+
+            gr.Markdown("### Options")
+            chk_8x = gr.Checkbox(label="Also apply ×8 post-resize on ×4 models", value=False)
+
+            run_btn = gr.Button("⚡ Run Comparison", elem_id="run-btn")
+
+    # ── Comparison grid ────────────────────────────────────────────────────
+    gr.Markdown("---")
+    gr.Markdown("## Results")
+
+    previews  = []
+    dl_btns   = []
+
+    with gr.Row():
+        for key in PANEL_KEYS:
+            with gr.Column():
+                gr.HTML(f'<div class="panel-title">{MODEL_LABELS[key]}</div>')
+                img_out = gr.Image(
+                    type="pil",
+                    label=MODEL_LABELS[key],
+                    show_label=False,
+                    height=320,
+                )
+                dl_out = gr.DownloadButton(
+                    label="⬇ Download PNG",
+                    elem_classes="dl-btn",
+                    visible=True,
+                )
+                previews.append(img_out)
+                dl_btns.append(dl_out)
+
+    # ── Wire up ─────────────────────────────────────────────────────────────
+    run_btn.click(
+        fn=compare_models,
+        inputs=[
+            inp_image,
+            chk_esrgan_x2,
+            chk_esrgan_x4,
+            chk_srcnn,
+            chk_hresnet,
+            chk_8x,
+        ],
+        outputs=previews + dl_btns,   # 8 outputs: 4 images + 4 download buttons
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -0,0 +1,9 @@
+onnxruntime         # ONNX inference engine (CPU)
+numpy               # Array manipulation
+Pillow              # Image I/O
+gradio>=4.0         # Web UI (4.x needed for DownloadButton stability)
+requests            # Google Drive model downloader
+
+# --- Optional: needed only if you integrate SR3 (PyTorch diffusion model) ---
+# torch              # PyTorch inference for SR3
+# torchvision        # Required by SR3 repo