"""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)