app.py

"""TimEE: End-to-end Time Series Classification via In-Context Learning.

A Gradio demo that lets users provide labeled support time series and unlabeled
query series, then classifies the queries in a single forward pass using the
pretrained TimEE foundation model (4.5M params).
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch
import torch
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.ticker import MultipleLocator
import io
import csv
import math

from timee import TimeeClassifier

# ---------------------------------------------------------------------------
# Load model at module scope (ZeroGPU rule: eager .to("cuda"))
# ---------------------------------------------------------------------------
clf = TimeeClassifier.from_pretrained(
    "liamsbhoo/timee",
    device=torch.device("cuda"),
    use_ensemble=True,
)

PATCH_SIZE = 16
COLORS = ["#08519C", "#DD8452", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf"]

# ---------------------------------------------------------------------------
# CSV parsing helpers
# ---------------------------------------------------------------------------

def _parse_csv_text(text):
    """Parse CSV text into list of rows (each row = list of floats)."""
    if not text or not text.strip():
        return []
    reader = csv.reader(io.StringIO(text))
    rows = []
    for row in reader:
        vals = []
        for cell in row:
            cell = cell.strip()
            if cell:
                try:
                    vals.append(float(cell))
                except ValueError:
                    pass
        if vals:
            rows.append(vals)
    return rows


def _parse_series_csv(csv_text, has_label=False):
    """Parse CSV where each row is one time series.

    If has_label=True, the first column is the integer/string label and
    the rest are the series values.

    Returns:
        X: np.ndarray (n, 1, seq_len) float32
        y: np.ndarray (n,) — labels (only if has_label=True, else None)
    """
    rows = _parse_csv_text(csv_text)
    if not rows:
        return np.zeros((0, 1, 1), dtype=np.float32), np.array([]) if has_label else None

    if has_label:
        labels = [r[0] for r in rows]
        series = [r[1:] for r in rows]
    else:
        labels = None
        series = [r for r in rows]

    # Pad to same length
    max_len = max(len(s) for s in series)
    padded = []
    for s in series:
        if len(s) < max_len:
            s = list(s) + [0.0] * (max_len - len(s))
        padded.append(s)

    X = np.array(padded, dtype=np.float32)
    if X.ndim == 1:
        X = X.reshape(1, -1)
    X = X[:, np.newaxis, :]  # (n, 1, seq_len)

    y = np.array(labels) if has_label else None
    return X, y


# ---------------------------------------------------------------------------
# Synthetic data generators for examples
# ---------------------------------------------------------------------------

def _generate_sine_data(
    T=160,
    k_shot=5,
    n_query=4,
    f0=5,
    f1=13,
    f2=14,
    noise_std=0.3,
    seed=42,
):
    """Generate a synthetic two-class time series benchmark."""
    rng = np.random.default_rng(seed)
    t = np.linspace(0, 1, T, endpoint=False)

    def _make_signal(freq_extra, phi):
        s = np.sin(2 * np.pi * f0 * t + phi)
        if freq_extra is not None:
            s = s + np.sin(2 * np.pi * freq_extra * t + phi)
        return s

    def _noisy(freq_extra, n):
        phases = rng.uniform(0, 2 * np.pi, n)
        return np.array([_make_signal(freq_extra, p) + rng.normal(0, noise_std, T) for p in phases])

    X0_ctx = _noisy(f1, k_shot)
    X1_ctx = _noisy(f2, k_shot)
    X0_qry = _noisy(f1, n_query)
    X1_qry = _noisy(f2, n_query)

    X_train = np.concatenate([X0_ctx, X1_ctx])[:, np.newaxis, :]
    y_train = np.array([0] * k_shot + [1] * k_shot)
    X_test = np.concatenate([X0_qry, X1_qry])[:, np.newaxis, :]
    y_test = np.array([0] * n_query + [1] * n_query)
    return X_train, y_train, X_test, y_test


def _generate_ecg_data(
    n_support=8,
    n_query=8,
    seed=42,
):
    """Generate synthetic ECG-like signals (two classes)."""
    rng = np.random.default_rng(seed)
    T = 140

    def _ecg_normal(phi, hr=72):
        t = np.linspace(0, 1, T, endpoint=False)
        # PQRST complex approximation
        ecg = np.zeros(T)
        for i, ti in enumerate(t):
            # P wave
            ecg[i] += 0.15 * math.exp(-((ti - 0.2 + phi) % 1 - 0.2) ** 2 / 0.002)
            # QRS complex
            ecg[i] -= 0.1 * math.exp(-((ti - 0.35 + phi) % 1 - 0.35) ** 2 / 0.001)
            ecg[i] += 1.0 * math.exp(-((ti - 0.4 + phi) % 1 - 0.4) ** 2 / 0.0015)
            ecg[i] -= 0.3 * math.exp(-((ti - 0.45 + phi) % 1 - 0.45) ** 2 / 0.002)
            # T wave
            ecg[i] += 0.35 * math.exp(-((ti - 0.65 + phi) % 1 - 0.65) ** 2 / 0.005)
        return ecg

    def _ecg_abnormal(phi, hr=72):
        t = np.linspace(0, 1, T, endpoint=False)
        ecg = np.zeros(T)
        for i, ti in enumerate(t):
            ecg[i] += 0.1 * math.exp(-((ti - 0.2 + phi) % 1 - 0.2) ** 2 / 0.002)
            ecg[i] -= 0.05 * math.exp(-((ti - 0.35 + phi) % 1 - 0.35) ** 2 / 0.001)
            # elevated ST segment (the abnormality)
            ecg[i] += 0.3 * math.exp(-((ti - 0.4 + phi) % 1 - 0.4) ** 2 / 0.01)
            ecg[i] += 0.4 * math.exp(-((ti - 0.55 + phi) % 1 - 0.55) ** 2 / 0.01)
            ecg[i] += 0.2 * math.exp(-((ti - 0.65 + phi) % 1 - 0.65) ** 2 / 0.005)
        return ecg

    def _noisy(fn, n):
        return np.array([fn(rng.uniform(0, 0.5)) + rng.normal(0, 0.05, T) for _ in range(n)])

    X0_ctx = _noisy(_ecg_normal, n_support)
    X1_ctx = _noisy(_ecg_abnormal, n_support)
    X0_qry = _noisy(_ecg_normal, n_query)
    X1_qry = _noisy(_ecg_abnormal, n_query)

    X_train = np.concatenate([X0_ctx, X1_ctx])[:, np.newaxis, :]
    y_train = np.array([0] * n_support + [1] * n_support)
    X_test = np.concatenate([X0_qry, X1_qry])[:, np.newaxis, :]
    y_test = np.array([0] * n_query + [1] * n_query)
    return X_train, y_train, X_test, y_test


def _generate_trend_data(
    n_support=5,
    n_query=5,
    T=128,
    seed=42,
):
    """Generate synthetic trend data (two classes: up vs down)."""
    rng = np.random.default_rng(seed)
    t = np.linspace(0, 1, T)

    def _up(i):
        slope = rng.uniform(1.5, 3.0)
        return slope * t + rng.normal(0, 0.3, T) + rng.uniform(-1, 1)

    def _down(i):
        slope = rng.uniform(-3.0, -1.5)
        return slope * t + rng.normal(0, 0.3, T) + rng.uniform(-1, 1)

    X0_ctx = np.array([_up(i) for i in range(n_support)])
    X1_ctx = np.array([_down(i) for i in range(n_support)])
    X0_qry = np.array([_up(i + 100) for i in range(n_query)])
    X1_qry = np.array([_down(i + 100) for i in range(n_query)])

    X_train = np.concatenate([X0_ctx, X1_ctx])[:, np.newaxis, :]
    y_train = np.array([0] * n_support + [1] * n_support)
    X_test = np.concatenate([X0_qry, X1_qry])[:, np.newaxis, :]
    y_test = np.array([0] * n_query + [1] * n_query)
    return X_train, y_train, X_test, y_test


def _series_to_csv(X, y=None):
    """Convert numpy (n, 1, T) to CSV text. If y provided, label is first column."""
    lines = []
    for i in range(X.shape[0]):
        vals = X[i, 0, :].tolist()
        if y is not None:
            lines.append(f"{y[i]}," + ",".join(f"{v:.6f}" for v in vals))
        else:
            lines.append(",".join(f"{v:.6f}" for v in vals))
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Pre-generate example CSV text
# ---------------------------------------------------------------------------

_X_tr_sine, _y_tr_sine, _X_te_sine, _y_te_sine = _generate_sine_data()
_SINE_TRAIN_CSV = _series_to_csv(_X_tr_sine, _y_tr_sine)
_SINE_TEST_CSV = _series_to_csv(_X_te_sine)

_X_tr_ecg, _y_tr_ecg, _X_te_ecg, _y_te_ecg = _generate_ecg_data()
_ECG_TRAIN_CSV = _series_to_csv(_X_tr_ecg, _y_tr_ecg)
_ECG_TEST_CSV = _series_to_csv(_X_te_ecg)

_X_tr_trend, _y_tr_trend, _X_te_trend, _y_te_trend = _generate_trend_data()
_TREND_TRAIN_CSV = _series_to_csv(_X_tr_trend, _y_tr_trend)
_TREND_TEST_CSV = _series_to_csv(_X_te_trend)

# Also store the ground truth for examples (for visualization)
_EXAMPLES_DATA = {
    "sine": (_X_tr_sine, _y_tr_sine, _X_te_sine, _y_te_sine),
    "ecg": (_X_tr_ecg, _y_tr_ecg, _X_te_ecg, _y_te_ecg),
    "trend": (_X_tr_trend, _y_tr_trend, _X_te_trend, _y_te_trend),
}


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------

def _plot_episode(
    X_train, y_train, X_test, preds, probs, class_names=None,
    title="TimEE Classification Results",
):
    """Plot support series, query series, and predictions."""
    n_train = len(X_train)
    n_test = len(X_test)
    classes = np.unique(y_train)
    n_classes = len(classes)

    if class_names is None:
        class_names = [f"Class {c}" for c in classes]

    colors = COLORS[:n_classes]
    T_len = X_train.shape[-1]
    patch_boundaries = np.arange(0, T_len + 1, PATCH_SIZE)

    # Build list of (series, cls_idx, role, pred, prob)
    series_list = []
    for i in range(n_train):
        ci = int(np.where(classes == y_train[i])[0][0])
        series_list.append((X_train[i, 0], ci, "Support", None, None))
    for i in range(n_test):
        series_list.append((X_test[i, 0], None, "Query", int(preds[i]), probs[i]))

    n_rows = len(series_list)
    all_s = np.concatenate([X_train[:, 0], X_test[:, 0]])
    y_lo, y_hi = all_s.min(), all_s.max()
    y_pad = (y_hi - y_lo) * 0.12

    fig = plt.figure(figsize=(10, n_rows * 0.48))
    gs = gridspec.GridSpec(n_rows, 1, hspace=0.05, figure=fig)
    first_ax = None

    for row, (series, cls_idx, role, pred, prob) in enumerate(series_list):
        kw = {"sharex": first_ax, "sharey": first_ax} if first_ax else {}
        ax = fig.add_subplot(gs[row, 0], **kw)
        if first_ax is None:
            first_ax = ax

        is_query = role == "Query"
        if is_query:
            ax.set_facecolor("#f5f5f5")
            c = colors[pred] if pred is not None else "#888888"
        else:
            c = colors[cls_idx]

        ax.plot(np.arange(T_len), series, color=c, lw=1.8, alpha=0.85)
        for xb in patch_boundaries:
            ax.axvline(xb, color="gray", lw=0.8, ls=":")

        if is_query:
            annotation = f"Query {row - n_train + 1}"
            text_color = "#888888"
        else:
            annotation = f"{class_names[cls_idx]}\nSupport {row + 1}"
            text_color = c

        ax.text(
            -0.02, 0.5, annotation,
            transform=ax.transAxes, fontsize=8, color=text_color,
            ha="right", va="center", multialignment="right",
        )

        if is_query and pred is not None and prob is not None:
            pred_name = class_names[pred] if pred < len(class_names) else f"C{pred}"
            conf = prob[pred]
            ax.text(
                1.01, 0.5,
                f"→ {pred_name}\n({conf:.0%})",
                transform=ax.transAxes, fontsize=8, color=c,
                ha="left", va="center", multialignment="left",
            )

        ax.xaxis.set_major_locator(MultipleLocator(PATCH_SIZE))
        ax.set_ylim(y_lo - y_pad, y_hi + y_pad)
        ax.yaxis.set_visible(False)
        plt.setp(ax.get_xticklabels(), visible=(row == n_rows - 1), rotation=45, ha="right", fontsize=7)
        for spine in ["top", "right", "left"]:
            ax.spines[spine].set_visible(False)

    fig.suptitle(title, fontsize=11, y=1.01)
    fig.subplots_adjust(left=0.15, right=0.92, top=0.97, bottom=0.08)
    return fig


def _plot_probabilities(probs, class_names, n_test):
    """Plot the predicted class distribution as grouped bar chart."""
    n_classes = len(class_names)
    fig, ax = plt.subplots(figsize=(8, max(3, min(6, n_test * 0.8))))

    x = np.arange(n_test)
    width = 0.8 / max(n_classes, 1)
    colors = COLORS[:n_classes]

    for c_idx in range(n_classes):
        offset = c_idx - (n_classes - 1) / 2
        ax.barh(x + offset * width * 0.0, probs[:, c_idx], height=width * 0.9,
                label=class_names[c_idx], color=colors[c_idx], alpha=0.85)

    ax.set_yticks(x)
    ax.set_yticklabels([f"Query {i+1}" for i in range(n_test)])
    ax.set_xlabel("Probability")
    ax.set_title("Predicted Class Probabilities")
    ax.set_xlim(0, 1)
    ax.legend(loc="upper right", fontsize=8)
    ax.invert_yaxis()
    plt.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
def classify(
    support_csv: str,
    query_csv: str,
    use_ensemble: bool = True,
):
    """Classify query time series using labeled support examples via TimEE in-context learning.

    Args:
        support_csv: CSV text where each row is a labeled time series. First column is the
            class label (int or string), remaining columns are the time series values.
        query_csv: CSV text where each row is an unlabeled time series. All columns are values.
        use_ensemble: Whether to use the 4-member preprocessing ensemble (interpolate x {256,512} x
            {raw, first_difference}) for more robust predictions.

    Returns:
        A matplotlib figure showing the support/query series with predictions, and
        a bar chart of the predicted class probabilities.
    """
    if not support_csv or not support_csv.strip():
        raise gr.Error("Please provide labeled support time series (CSV format).")
    if not query_csv or not query_csv.strip():
        raise gr.Error("Please provide unlabeled query time series (CSV format).")

    X_train, y_train = _parse_series_csv(support_csv, has_label=True)
    X_test, _ = _parse_series_csv(query_csv, has_label=False)

    if X_train.shape[0] == 0:
        raise gr.Error("No valid support series found. Each row needs: label,val1,val2,...")
    if X_test.shape[0] == 0:
        raise gr.Error("No valid query series found. Each row needs: val1,val2,...")
    if len(np.unique(y_train)) < 2:
        raise gr.Error("Support series must contain at least 2 distinct classes.")

    # Use the pre-loaded classifier with the requested ensemble setting
    clf.use_ensemble = use_ensemble
    if use_ensemble and clf.transforms is None:
        from timee.transforms import default_ensemble_transforms
        clf.transforms = default_ensemble_transforms()
    elif not use_ensemble:
        clf.transforms = None

    predictions, probabilities = clf.predict(X_train, y_train, X_test)

    # Get class names from labels
    classes = np.unique(y_train)
    class_names = [str(c) for c in classes]

    # Map predictions back to original label space
    pred_labels = predictions

    # Generate plots
    fig_episode = _plot_episode(
        X_train, y_train, X_test, pred_labels, probabilities,
        class_names=class_names,
        title="TimEE: In-Context Classification",
    )

    fig_probs = _plot_probabilities(probabilities, class_names, len(X_test))

    # Build text summary
    summary_lines = []
    summary_lines.append(f"**{len(X_train)} support series** | **{len(X_test)} query series** | **{len(classes)} classes**")
    summary_lines.append("")
    for i in range(len(X_test)):
        pred_c = int(np.argmax(probabilities[i]))
        conf = probabilities[i][pred_c]
        summary_lines.append(f"- Query {i+1}: → **{class_names[pred_c]}** (confidence: {conf:.1%})")

    summary = "\n".join(summary_lines)

    plt.close("all")
    return fig_episode, fig_probs, summary


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

DESCRIPTION = """
# TimEE: End-to-end Time Series Classification via In-Context Learning

TimEE is a 4.5M-parameter foundation model that classifies time series **in a single forward pass**
using a few labeled examples — no training or fine-tuning required.

**How to use:**
1. **Support CSV**: Each row is one labeled time series. First column = class label, remaining columns = values.
2. **Query CSV**: Each row is one unlabeled time series. All columns = values.
3. Click **Classify** to see predictions and class probability distributions.

📄 [Paper](https://arxiv.org/abs/2607.07500) | 🐙 [GitHub](https://github.com/automl/timee) | 🤗 [Model](https://huggingface.co/liamsbhoo/timee)
"""

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Support (Labeled) Series")
            gr.Markdown("Each row: `label, value1, value2, value3, ...`")
            support_input = gr.Textbox(
                label="Support CSV",
                value=_SINE_TRAIN_CSV,
                lines=12,
                placeholder="0, 0.1, 0.2, ...\n1, 0.3, 0.4, ...",
            )
        with gr.Column(scale=1):
            gr.Markdown("### 🔍 Query (Unlabeled) Series")
            gr.Markdown("Each row: `value1, value2, value3, ...`")
            query_input = gr.Textbox(
                label="Query CSV",
                value=_SINE_TEST_CSV,
                lines=12,
                placeholder="0.1, 0.2, 0.3, ...",
            )

    with gr.Row():
        use_ensemble = gr.Checkbox(label="Use 4-member preprocessing ensemble", value=True)
        classify_btn = gr.Button("Classify", variant="primary", scale=2)

    with gr.Row():
        episode_plot = gr.Plot(label="Series & Predictions")
        prob_plot = gr.Plot(label="Class Probability Distribution")

    summary_output = gr.Markdown(label="Summary")

    with gr.Accordion("Examples", open=False):
        gr.Markdown("Click an example to load pre-generated data, then click **Classify**.")
        gr.Examples(
            examples=[
                [_SINE_TRAIN_CSV, _SINE_TEST_CSV, True],
                [_ECG_TRAIN_CSV, _ECG_TEST_CSV, True],
                [_TREND_TRAIN_CSV, _TREND_TEST_CSV, True],
            ],
            inputs=[support_input, query_input, use_ensemble],
            outputs=[episode_plot, prob_plot, summary_output],
            fn=classify,
            cache_examples=True,
            cache_mode="lazy",
        )

    classify_btn.click(
        fn=classify,
        inputs=[support_input, query_input, use_ensemble],
        outputs=[episode_plot, prob_plot, summary_output],
        api_name="classify",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)