app.py

"""Live ring/no-ring classifier for Legacy Survey objects.

Give it an object name (e.g. "NGC 7796"), an "ra,dec" pair, or a Legacy Survey
viewer URL. It resolves the coordinates, downloads the grz cutout from Legacy
Survey, builds the exact Lupton RGB the model was trained on, runs the model,
and shows the enhanced image with the fitted galaxy ellipse and inner-ring
candidate (same measurement as TamanoGalaxia.ipynb).

The whole input pipeline reproduces training:
    bands g,r,z -> centroid on r -> crop 320 -> make_lupton_rgb(z,r,g, Q=8, stretch=0.2)
    -> /255 -> ImageNet normalize -> ConvNeXt-Base (Zoobot) -> softmax[ring]
"""

from __future__ import annotations

import io
import os
import re
import urllib.parse
from pathlib import Path

import gradio as gr
import matplotlib
import numpy as np
import requests
import timm
import torch

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.visualization import make_lupton_rgb
from matplotlib.patches import Ellipse

import enhance
from galaxy_ellipse import measure_ellipse

IMG_SIZE = 320
STRETCH, Q = 0.2, 8  # Lupton parameters used to build the training cache
ENCODER = "hf_hub:mwalmsley/zoobot-encoder-convnext_base"
THR = 0.735  # validation-tuned decision threshold
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
LABEL_NAMES = {0: "sin anillo", 1: "con anillo"}

IMEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
ISTD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

CUTOUT_URL = "https://www.legacysurvey.org/viewer/fits-cutout"


# --- Model (loaded once, lazily) ---

_MODEL = None


def _weights_path() -> Path:
    """Locate the checkpoint: env var, the Space dir, or the local outputs dir."""
    candidates = [
        os.environ.get("MODEL_PATH"),
        Path(__file__).resolve().parent / "zoobot_base_natural.pt",
        Path(__file__).resolve().parent.parent / "outputs" / "zoobot_base_natural.pt",
    ]
    for c in candidates:
        if c and Path(c).exists():
            return Path(c)
    raise FileNotFoundError(
        "No encuentro zoobot_base_natural.pt. Copialo al directorio del Space "
        "o definí la variable de entorno MODEL_PATH."
    )


def get_model() -> torch.nn.Module:
    global _MODEL
    if _MODEL is None:
        model = timm.create_model(ENCODER, pretrained=False, num_classes=2)
        model.load_state_dict(torch.load(_weights_path(), map_location=DEVICE))
        _MODEL = model.eval().to(DEVICE)
    return _MODEL


# --- Input resolution ---

def resolve_coords(text: str) -> tuple[float, float, str]:
    """Turn user input into (ra_deg, dec_deg, label).

    Accepts: an "ra,dec" pair, a Legacy Survey viewer URL (?ra=&dec= or #Name),
    or a plain object name resolved through CDS Sesame.
    """
    text = (text or "").strip()
    if not text:
        raise ValueError("Escribí un nombre de objeto, 'ra,dec' o una URL de Legacy Survey.")

    # Legacy Survey viewer URL.
    if "legacysurvey.org" in text or text.startswith("http"):
        parsed = urllib.parse.urlparse(text)
        qs = urllib.parse.parse_qs(parsed.query)
        if "ra" in qs and "dec" in qs:
            return float(qs["ra"][0]), float(qs["dec"][0]), f"ra,dec de la URL"
        if parsed.fragment:
            name = urllib.parse.unquote(parsed.fragment).strip()
            if name:
                c = SkyCoord.from_name(name)
                return float(c.ra.deg), float(c.dec.deg), name
        raise ValueError("No pude extraer coordenadas ni nombre de esa URL.")

    # Bare "ra,dec" pair.
    m = re.match(r"^\s*([-+]?\d+\.?\d*)\s*[, ]\s*([-+]?\d+\.?\d*)\s*$", text)
    if m:
        return float(m.group(1)), float(m.group(2)), "ra,dec"

    # Object name via CDS Sesame (SIMBAD/NED/VizieR).
    c = SkyCoord.from_name(text)
    return float(c.ra.deg), float(c.dec.deg), text


# --- Image pipeline (reproduces training exactly) ---

def _centroid(img, sf=0.4):
    hh, ww = img.shape
    dy, dx = int(hh * sf / 2), int(ww * sf / 2)
    y0, y1, x0, x1 = hh // 2 - dy, hh // 2 + dy, ww // 2 - dx, ww // 2 + dx
    s = img[y0:y1, x0:x1]
    ly, lx = np.unravel_index(np.argmax(s), s.shape)
    return y0 + int(ly), x0 + int(lx)


def _crop(im, yc, xc, sz=IMG_SIZE):
    hh, ww = im.shape
    h2 = sz // 2
    y0, y1, x0, x1 = yc - h2, yc + h2, xc - h2, xc + h2
    py0, py1 = max(0, -y0), max(0, y1 - hh)
    px0, px1 = max(0, -x0), max(0, x1 - ww)
    im = np.pad(im, ((py0, py1), (px0, px1)))
    y0 += py0; y1 += py0; x0 += px0; x1 += px0
    return im[y0:y1, x0:x1]


def fetch_rgb(ra: float, dec: float, layer: str, pixscale: float) -> np.ndarray:
    """Download the grz cutout and build the training-style Lupton RGB."""
    # Pull a slightly larger field so the centroid + 320 crop stays inside.
    url = (
        f"{CUTOUT_URL}?ra={ra}&dec={dec}&layer={layer}"
        f"&pixscale={pixscale}&size={IMG_SIZE + 64}&bands=grz"
    )
    resp = requests.get(url, timeout=60)
    resp.raise_for_status()
    if not resp.content or resp.headers.get("content-type", "").startswith("text"):
        raise ValueError("Legacy Survey no devolvió imagen para esas coordenadas (¿fuera de cobertura?).")

    with fits.open(io.BytesIO(resp.content), memmap=False) as h:
        dd = h[0].data
    if dd is None or dd.ndim != 3 or dd.shape[0] < 3:
        raise ValueError("El cutout no trae las 3 bandas g,r,z.")

    g, r, z = [np.nan_to_num(np.asarray(dd[k], dtype=np.float64)) for k in range(3)]
    if not np.any(np.isfinite(r)):
        raise ValueError("La banda r vino vacía (objeto fuera de cobertura).")

    yc, xc = _centroid(r)
    g, r, z = [_crop(x, yc, xc) for x in (g, r, z)]
    rgb = make_lupton_rgb(z, r, g, minimum=0, stretch=STRETCH, Q=Q)  # z->R, r->G, g->B
    return rgb.astype(np.uint8)


@torch.no_grad()
def predict_prob(rgb: np.ndarray, use_tta: bool) -> float:
    """Probability of class 1 (ring), optional 6-view TTA (matches training)."""
    t = torch.from_numpy(rgb.astype(np.float32) / 255.0).permute(2, 0, 1)
    x = ((t - IMEAN) / ISTD).unsqueeze(0).to(DEVICE)
    views = (
        [x, torch.flip(x, [-1]), torch.flip(x, [-2]),
         torch.rot90(x, 1, [-2, -1]), torch.rot90(x, 2, [-2, -1]), torch.rot90(x, 3, [-2, -1])]
        if use_tta else [x]
    )
    model = get_model()
    acc = sum(torch.softmax(model(v).float(), 1)[:, 1] for v in views)
    return float((acc / len(views)).cpu().item())


# --- Figure ---

def render(rgb: np.ndarray, geo: dict, mode: str = "Original", gray: bool = False):
    """Two-panel figure: enhanced RGB | enhanced RGB + ellipse overlay.

    The enhancement is purely visual; the ellipse geometry was measured on the
    original Lupton RGB, so it still lines up after any realce.
    """
    shown = enhance.apply(rgb, mode, gray)
    title = "Imagen (Lupton grz)" if mode == "Original" else f"Imagen — {mode}"
    fig, axes = plt.subplots(1, 2, figsize=(8, 4.2))
    axes[0].imshow(shown)
    axes[0].set_title(title)
    axes[0].axis("off")

    axes[1].imshow(shown)
    axes[1].set_title("Imagen + elipse")
    axes[1].axis("off")
    if geo.get("status") == "ok":
        cx, cy = geo["cx"], geo["cy"]
        theta_deg = np.degrees(geo["theta"])
        axes[1].add_patch(Ellipse(
            (cx, cy), width=2 * geo["radius_major"], height=2 * geo["radius_minor"],
            angle=theta_deg, fill=False, edgecolor="#3fb950", linewidth=2,
        ))
        if geo.get("ring_radius") is not None:
            axes[1].add_patch(Ellipse(
                (cx, cy), width=2 * geo["ring_radius"], height=2 * geo["ring_radius"] * geo["q"],
                angle=theta_deg, fill=False, edgecolor="#f0883e", linewidth=2, linestyle="--",
            ))
        axes[1].plot([cx], [cy], marker="+", color="white", markersize=9)

    fig.tight_layout()
    return fig


# --- Gradio callback ---

def classify(text: str, layer: str, pixscale: float, use_tta: bool, mode: str, gray: bool):
    try:
        ra, dec, label = resolve_coords(text)
    except Exception as exc:
        return None, f"No pude resolver la entrada: {exc}", None, None

    try:
        rgb = fetch_rgb(ra, dec, layer, float(pixscale))
    except Exception as exc:
        return None, f"No pude bajar el cutout: {exc}", None, None

    # Prediction and geometry use the ORIGINAL Lupton RGB; the realce is only
    # for display and is re-applied (without re-running the model) on change.
    prob = predict_prob(rgb, use_tta)
    pred = int(prob >= THR)
    geo = measure_ellipse(rgb)
    fig = render(rgb, geo, mode, gray)

    ring_txt = (
        f"{geo['ring_radius']:.0f} px (candidato exploratorio)"
        if geo.get("status") == "ok" and geo.get("ring_radius") is not None
        else "no detectado"
    )
    geom_md = ""
    if geo.get("status") == "ok":
        geom_md = (
            f"- **Eje mayor:** {geo['radius_major']:.1f} px\n"
            f"- **Razón axial q:** {geo['q']:.2f}\n"
            f"- **Anillo interno:** {ring_txt}\n"
        )

    summary = (
        f"### {label}\n"
        f"- **Coordenadas:** ra={ra:.5f}, dec={dec:.5f}\n"
        f"- **Predicción:** **{LABEL_NAMES[pred].upper()}**\n"
        f"- **Probabilidad de anillo:** {prob:.3f}  (umbral {THR})\n"
        f"{geom_md}\n"
        f"_La predicción es una sugerencia de apoyo, no la verdad. La elipse verde es el "
        f"ajuste de la galaxia; el anillo naranja punteado es un candidato exploratorio._"
    )
    return fig, summary, rgb, geo


def rerender(rgb, geo, mode: str, gray: bool):
    """Re-apply a realce to the already-fetched image (no model re-run)."""
    if rgb is None or geo is None:
        return None
    return render(rgb, geo, mode, gray)


EXAMPLES = [
    ["NGC 7796"],
    ["NGC 1398"],
    ["https://www.legacysurvey.org/viewer#NGC 7796"],
    ["10.6848, 41.2691"],
]

with gr.Blocks(title="Anillos en galaxias — Legacy Survey") as demo:
    gr.Markdown(
        "# Clasificador de anillos sobre Legacy Survey\n"
        "Escribí un **nombre de objeto** (ej. `NGC 7796`), un par **`ra,dec`**, o una "
        "**URL del visor** de Legacy Survey. Bajo el cutout grz, lo proceso igual que en "
        "el entrenamiento y corro el modelo (ConvNeXt-Base / Zoobot)."
    )
    st_rgb = gr.State(None)
    st_geo = gr.State(None)
    with gr.Row():
        with gr.Column(scale=2):
            inp = gr.Textbox(label="Objeto / ra,dec / URL", placeholder="NGC 7796")
            with gr.Row():
                layer = gr.Dropdown(
                    ["ls-dr10", "ls-dr9"], value="ls-dr10", label="Capa (layer)"
                )
                pixscale = gr.Number(value=0.262, label="Pixscale (arcsec/px)")
            tta = gr.Checkbox(value=False, label="TTA (6 vistas, más lento y estable)")
            btn = gr.Button("Clasificar", variant="primary")
            mode = gr.Radio(
                enhance.MODES, value="Original", label="Realce (visual, no afecta la predicción)"
            )
            gray = gr.Checkbox(value=False, label="Escala de grises (B/N)")
            gr.Examples(EXAMPLES, inputs=[inp])
        with gr.Column(scale=3):
            out_img = gr.Plot(label="Imagen + elipse")
            out_md = gr.Markdown()

    inputs = [inp, layer, pixscale, tta, mode, gray]
    outputs = [out_img, out_md, st_rgb, st_geo]
    btn.click(classify, inputs, outputs)
    inp.submit(classify, inputs, outputs)

    # Changing the realce re-renders the already-fetched image instantly.
    mode.change(rerender, [st_rgb, st_geo, mode, gray], out_img)
    gray.change(rerender, [st_rgb, st_geo, mode, gray], out_img)


if __name__ == "__main__":
    demo.launch()