Add Super-Resolution API: MewZoom 2X/4X + comparison + metrics

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.10-slim
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libgl1 libglib2.0-0 git curl \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+WORKDIR /app
+
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY --chown=user app.py .
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,195 @@
+import json
+import logging
+import time
+from io import BytesIO
+from pathlib import Path
+from contextlib import asynccontextmanager
+from typing import Literal
+
+import numpy as np
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from scipy import ndimage
+
+from fastapi import FastAPI, File, UploadFile, Query, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
+
+from torchvision.io import decode_image, ImageReadMode
+from torchvision.transforms.v2 import ToDtype, ToPILImage
+
+from mewzoom.model import MewZoom
+
+MODELS_CONFIG = {"2x": "andrewdalpino/MewZoom-V1-2X-Unet", "4x": "andrewdalpino/MewZoom-V1-4X-Unet"}
+MAX_DIM = {"2x": 2048, "4x": 1024}
+CACHE_DIR = Path("models")
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+_models: dict[str, MewZoom] = {}
+_image_to_tensor = ToDtype(torch.float32, scale=True)
+_tensor_to_pil = ToPILImage()
+_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _load_model(scale: str) -> MewZoom:
+    if scale in _models:
+        return _models[scale]
+    model_id = MODELS_CONFIG[scale]
+    logger.info("Loading %s (%s) on %s ...", scale, model_id, _DEVICE)
+    CACHE_DIR.mkdir(exist_ok=True)
+    model = MewZoom.from_pretrained(model_id, cache_dir=str(CACHE_DIR))
+    model.to(_DEVICE).eval()
+    _models[scale] = model
+    logger.info("%s loaded (%s params)", scale, f"{sum(p.numel() for p in model.parameters()):,}")
+    return model
+
+
+def _resize_if_needed(img: Image.Image, scale: str) -> tuple[Image.Image, bool]:
+    max_dim = MAX_DIM[scale]
+    w, h = img.size
+    if max(w, h) <= max_dim:
+        return img, False
+    ratio = max_dim / max(w, h)
+    return img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS), True
+
+
+def _pil_to_tensor(img: Image.Image) -> torch.Tensor:
+    arr = np.array(img, dtype=np.float32) / 255.0
+    return torch.from_numpy(arr).permute(2, 0, 1)
+
+
+def upscale_image(image_bytes: bytes, scale: str) -> tuple[bytes, dict]:
+    model = _load_model(scale)
+    factor = int(scale[0])
+    try:
+        pil = Image.open(BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(400, f"Bad image: {e}")
+    orig = (pil.width, pil.height)
+    pil, resized = _resize_if_needed(pil, scale)
+    out_mp = pil.width * factor * pil.height * factor / 1e6
+    if out_mp > 64:
+        raise HTTPException(400, f"Output too large ({out_mp:.0f}MP). Use smaller image.")
+    x = _pil_to_tensor(pil).unsqueeze(0).to(_DEVICE)
+    with torch.inference_mode():
+        y = model.upscale(x)
+    result = _tensor_to_pil(y.squeeze(0).cpu())
+    buf = BytesIO()
+    result.save(buf, format="PNG")
+    buf.seek(0)
+    return buf.getvalue(), {"scale": scale, "input": f"{orig[0]}x{orig[1]}", "output": f"{result.width}x{result.height}", "resized": resized}
+
+
+def _laplacian_variance(img: Image.Image) -> float:
+    lap = ndimage.laplace(np.array(img.convert("L"), dtype=np.float64))
+    return float(lap.var())
+
+
+def _entropy(img: Image.Image) -> float:
+    hist = np.histogram(np.array(img.convert("L")), bins=256, range=(0, 256))[0]
+    hist = hist[hist > 0] / hist.sum()
+    return float(-np.sum(hist * np.log2(hist)))
+
+
+def _edge_density(img: Image.Image) -> float:
+    arr = np.array(img.convert("L"), dtype=np.float64)
+    mag = np.hypot(ndimage.sobel(arr, axis=0), ndimage.sobel(arr, axis=1))
+    return float(np.mean(mag > mag.mean() + mag.std()))
+
+
+def compute_metrics(img: Image.Image) -> dict:
+    return {"size": f"{img.width}x{img.height}", "sharpness": round(_laplacian_variance(img), 4), "entropy": round(_entropy(img), 4), "edge_density": round(_edge_density(img), 4), "contrast_std": round(float(np.array(img).std()), 2)}
+
+
+def generate_comparison(image_bytes: bytes) -> tuple[bytes, dict]:
+    original = Image.open(BytesIO(image_bytes)).convert("RGB")
+    metrics = {"original": compute_metrics(original)}
+    upscaled = {}
+    for scale in MODELS_CONFIG:
+        t0 = time.perf_counter()
+        result_bytes, info = upscale_image(image_bytes, scale)
+        elapsed = time.perf_counter() - t0
+        img = Image.open(BytesIO(result_bytes)).convert("RGB")
+        upscaled[scale] = img
+        metrics[scale] = {**compute_metrics(img), "time_s": round(elapsed, 3), **info}
+    orig_r = original.resize(upscaled["2x"].size, Image.LANCZOS)
+    images = [orig_r, upscaled["2x"], upscaled["4x"]]
+    labels = ["Original", "MewZoom 2X", "MewZoom 4X"]
+    label_h, gap = 30, 8
+    max_h = max(i.height for i in images)
+    total_w = sum(i.width for i in images) + gap * (len(images) - 1)
+    canvas = Image.new("RGB", (total_w, max_h + label_h), (30, 30, 30))
+    draw = ImageDraw.Draw(canvas)
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
+    except Exception:
+        font = ImageFont.load_default()
+    x = 0
+    for img, lbl in zip(images, labels):
+        canvas.paste(img, (x, label_h))
+        bbox = draw.textbbox((0, 0), lbl, font=font)
+        tw = bbox[2] - bbox[0]
+        draw.text((x + (img.width - tw) // 2, 6), lbl, fill=(255, 255, 255), font=font)
+        x += img.width + gap
+    buf = BytesIO()
+    canvas.save(buf, format="PNG")
+    buf.seek(0)
+    return buf.getvalue(), metrics
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logger.info("Starting on %s, loading models...", _DEVICE)
+    for scale in MODELS_CONFIG:
+        _load_model(scale)
+    yield
+
+
+app = FastAPI(
+    title="Super-Resolution API",
+    description="MewZoom 2X/4X upscaling + comparison + quality metrics. InvSR requires GPU (not on free tier).",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+
+@app.get("/")
+@app.get("/health")
+async def health():
+    return JSONResponse({"status": "healthy", "device": str(_DEVICE), "models": list(MODELS_CONFIG.keys()), "gpu": torch.cuda.is_available()})
+
+
+@app.post("/upscale/2x")
+async def route_2x(file: UploadFile = File(...)):
+    result, info = upscale_image(await file.read(), "2x")
+    return StreamingResponse(BytesIO(result), media_type="image/png", headers={"X-Info": json.dumps(info)})
+
+
+@app.post("/upscale/4x")
+async def route_4x(file: UploadFile = File(...)):
+    result, info = upscale_image(await file.read(), "4x")
+    return StreamingResponse(BytesIO(result), media_type="image/png", headers={"X-Info": json.dumps(info)})
+
+
+@app.post("/upscale/compare")
+async def route_compare(file: UploadFile = File(...), format: Literal["image", "json", "both"] = Query("both")):
+    img, metrics = generate_comparison(await file.read())
+    if format == "json":
+        return JSONResponse(metrics)
+    if format == "image":
+        return StreamingResponse(BytesIO(img), media_type="image/png")
+    return StreamingResponse(BytesIO(img), media_type="image/png", headers={"X-Metrics": json.dumps(metrics)})
+
+
+@app.post("/upscale/metrics")
+async def route_metrics(file: UploadFile = File(...)):
+    _, metrics = generate_comparison(await file.read())
+    return JSONResponse(metrics)
+
+
+@app.post("/upscale/invsr")
+async def route_invsr(file: UploadFile = File(...)):
+    raise HTTPException(400, detail="InvSR (diffusion 4X) needs GPU. This Space is CPU. Use /upscale/2x or /upscale/4x.")

requirements.txt

@@ -0,0 +1,9 @@
+fastapi>=0.110.0
+uvicorn[standard]>=0.29.0
+python-multipart>=0.0.9
+mewzoom~=1.0.0
+torch>=2.0.0
+torchvision>=0.15.0
+Pillow>=10.0.0
+scipy>=1.10.0
+numpy>=1.23.0