first

README.md

@@ -12,3 +12,61 @@ short_description: Interpretable Mesomorphic Neural Networks for 12-Lead ECG
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## MesomorphicECG XAI Space
+
+This Space hosts an interactive Gradio app for the **mesomorphicECG** models in  
+`SEARCH-IHI/mesomorphicECG` (`https://huggingface.co/SEARCH-IHI/mesomorphicECG`).
+
+The app:
+- Loads **IMN** checkpoints (categorical and single-linear) from the model repo.
+- Lets you choose sampling rate (100 / 500 Hz) and task:
+  `norm_vs_cd`, `norm_vs_hyp`, `norm_vs_mi`, `norm_vs_sttc`.
+- Uses pre-packaged PTB-XL examples stored as binary `.npz` files in this Space.
+- Visualizes intrinsic IMN feature attributions (Impact = w·x) as a lead × segment heatmap
+  together with per-lead ECG traces.
+
+### Files
+
+- `app.py` – main Gradio application.
+- `single_linear_imn_core.py` – core single-linear IMN model for inference.
+- `categorical_imn_core.py` – core categorical IMN model for inference.
+- `requirements.txt` – Python dependencies for this Space.
+
+### Required data binaries
+
+For each combination of **sampling rate** and **task**, the app expects a `.npz` file:
+
+- 100 Hz:
+  - `data/ptbxl_100hz_norm_vs_cd_test.npz`
+  - `data/ptbxl_100hz_norm_vs_hyp_test.npz`
+  - `data/ptbxl_100hz_norm_vs_mi_test.npz`
+  - `data/ptbxl_100hz_norm_vs_sttc_test.npz`
+- 500 Hz:
+  - `data/ptbxl_500hz_norm_vs_cd_test.npz`
+  - `data/ptbxl_500hz_norm_vs_hyp_test.npz`
+  - `data/ptbxl_500hz_norm_vs_mi_test.npz`
+  - `data/ptbxl_500hz_norm_vs_sttc_test.npz`
+
+Each `.npz` should contain:
+
+- `signals` – float32 array `[N, 12, L]` (z-scoring is done in the app).
+- `labels` – array `[N]` with 0 (NORM) / 1 (POS_CLASS) for the chosen task.
+- `reports` – object array `[N]` with clinical notes (strings).
+- `age` – array `[N]` (e.g. int or float).
+- `sex` – object array `[N]` (e.g. `'M'`, `'F'`, or empty).
+- `ecg_id` – array `[N]` with integer ECG identifiers.
+
+You can prepare these from PTB-XL using the same task definition and
+window length / sampling rate as in the training scripts, then upload
+them into this Space under the `data/` directory.
+
+### Run locally
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+On Hugging Face Spaces, `app.py` is loaded automatically.
+

app.py

@@ -0,0 +1,433 @@
+from __future__ import annotations
+
+"""
+MesomorphicECG XAI Gradio app for Hugging Face Spaces.
+
+This version focuses on:
+- Selecting sampling rate (100 / 500 Hz), model type (categorical vs single-linear),
+  and task (norm_vs_cd / norm_vs_hyp / norm_vs_mi / norm_vs_sttc).
+- Loading pre-packaged ECG examples from local binary .npz files in this Space.
+- Downloading the corresponding IMN checkpoint from
+  `SEARCH-IHI/mesomorphicECG` on the Hugging Face Hub.
+- Running inference and visualizing intrinsic feature attributions
+  (Impact = w * x) as a lead × segment heatmap plus per-lead ECG traces.
+
+Data binaries
+-------------
+For each (sampling_rate, task) pair you should provide a `.npz` file as
+configured in DATA_FILES below, with keys:
+
+    signals : float32 array [N, 12, L]
+    labels  : float32/int array [N] with 0 (NORM) / 1 (POS_CLASS)
+    reports : object array [N] of clinical notes
+    age     : array [N]
+    sex     : object array [N]
+    ecg_id  : array [N]
+"""
+
+import os
+from functools import lru_cache
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+import gradio as gr  # noqa: E402
+from huggingface_hub import hf_hub_download, list_repo_files  # noqa: E402
+
+import single_linear_imn_core as sl_core  # noqa: E402
+import categorical_imn_core as cat_core  # noqa: E402
+
+
+HF_MODEL_REPO = "SEARCH-IHI/mesomorphicECG"
+
+TASK_TO_POS = {
+    "norm_vs_mi": "MI",
+    "norm_vs_sttc": "STTC",
+    "norm_vs_cd": "CD",
+    "norm_vs_hyp": "HYP",
+}
+
+LEAD_NAMES = sl_core.DEFAULT_LEAD_NAMES
+
+
+# Mapping from (sampling_rate, task) -> local data binary.
+DATA_FILES: Dict[Tuple[int, str], str] = {
+    # 100 Hz
+    (100, "norm_vs_cd"): "data/ptbxl_100hz_norm_vs_cd_test.npz",
+    (100, "norm_vs_hyp"): "data/ptbxl_100hz_norm_vs_hyp_test.npz",
+    (100, "norm_vs_mi"): "data/ptbxl_100hz_norm_vs_mi_test.npz",
+    (100, "norm_vs_sttc"): "data/ptbxl_100hz_norm_vs_sttc_test.npz",
+    # 500 Hz
+    (500, "norm_vs_cd"): "data/ptbxl_500hz_norm_vs_cd_test.npz",
+    (500, "norm_vs_hyp"): "data/ptbxl_500hz_norm_vs_hyp_test.npz",
+    (500, "norm_vs_mi"): "data/ptbxl_500hz_norm_vs_mi_test.npz",
+    (500, "norm_vs_sttc"): "data/ptbxl_500hz_norm_vs_sttc_test.npz",
+}
+
+
+DATA_CACHE: Dict[Tuple[int, str], Dict[str, Any]] = {}
+MODEL_CACHE: Dict[Tuple[str, int, str], Dict[str, Any]] = {}
+
+
+def zscore_per_lead(x: np.ndarray) -> np.ndarray:
+    """Per-lead z-score normalization."""
+    mean = x.mean(axis=1, keepdims=True)
+    std = x.std(axis=1, keepdims=True).clip(min=1e-6)
+    return ((x - mean) / std).astype(np.float32)
+
+
+@lru_cache(maxsize=None)
+def _list_model_repo_files() -> List[str]:
+    return list_repo_files(repo_id=HF_MODEL_REPO, repo_type="model")
+
+
+def _resolve_ckpt_filename(model_type: str, sampling_rate: int, task: str) -> str:
+    if model_type == "single_linear":
+        category = f"single_linear_imn_{sampling_rate}hz"
+    else:
+        category = f"categorical_imn_{sampling_rate}hz"
+
+    prefix = f"{category}/{task}/"
+    files = _list_model_repo_files()
+    candidates = [f for f in files if f.startswith(prefix) and f.endswith(".ckpt")]
+    if not candidates:
+        raise FileNotFoundError(
+            f"No checkpoint (.ckpt) found in repo {HF_MODEL_REPO} under {prefix}. "
+            "Ensure upload_best_checkpoints_to_hf.py has populated this path."
+        )
+    best_style = [f for f in candidates if "best-imn-epoch=" in f]
+    chosen = sorted(best_style or candidates)[-1]
+    return chosen
+
+
+def load_imn_model(
+    model_type: str,
+    sampling_rate: int,
+    task: str,
+) -> Tuple[torch.nn.Module, str]:
+    key = (model_type, sampling_rate, task)
+    cached = MODEL_CACHE.get(key)
+    if cached and cached["model"] is not None:
+        return cached["model"], cached["device"]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    filename = _resolve_ckpt_filename(model_type, sampling_rate, task)
+    ckpt_local = hf_hub_download(repo_id=HF_MODEL_REPO, filename=filename)
+
+    if model_type == "single_linear":
+        model = sl_core.IMNLightning.load_from_checkpoint(ckpt_local, map_location=device)
+    else:
+        model = cat_core.IMNLightning.load_from_checkpoint(ckpt_local, map_location=device)
+
+    model.eval()
+    model.to(device)
+    MODEL_CACHE[key] = {"path": ckpt_local, "model": model, "device": device}
+    return model, device
+
+
+def load_data_binary(sampling_rate: int, task: str) -> Dict[str, Any]:
+    key = (sampling_rate, task)
+    if key in DATA_CACHE:
+        return DATA_CACHE[key]
+
+    path = DATA_FILES.get(key)
+    if path is None:
+        raise FileNotFoundError(f"No data file configured for (fs={sampling_rate}, task={task}).")
+    if not os.path.isfile(path):
+        raise FileNotFoundError(
+            f"Data file not found at '{path}'. "
+            "Upload a .npz with signals, labels, reports, age, sex, ecg_id."
+        )
+
+    with np.load(path, allow_pickle=True) as npz:
+        required = ["signals", "labels", "reports", "age", "sex", "ecg_id"]
+        missing = [k for k in required if k not in npz]
+        if missing:
+            raise KeyError(f"Data file '{path}' missing keys: {missing}")
+        data = {k: npz[k] for k in required}
+
+    DATA_CACHE[key] = data
+    return data
+
+
+def on_load_records(
+    sampling_rate: int,
+    task: str,
+    state: Optional[dict],
+):
+    try:
+        data = load_data_binary(int(sampling_rate), task)
+    except Exception as e:
+        return (
+            f"Load error: {e}",
+            gr.update(choices=[], value=None),
+            state or {},
+            "—",
+            "—",
+        )
+
+    signals = data["signals"]
+    labels = data["labels"]
+    reports = data["reports"]
+    age = data["age"]
+    sex = data["sex"]
+    ecg_id = data["ecg_id"]
+
+    N, C, L = signals.shape
+    pos_class = TASK_TO_POS.get(task, "MI")
+
+    records: List[Dict[str, Any]] = []
+    for i in range(N):
+        gt = pos_class if float(labels[i]) >= 0.5 else "NORM"
+        records.append(
+            {
+                "index": int(i),
+                "ecg_id": int(ecg_id[i]),
+                "gt": gt,
+                "report": str(reports[i]) if reports is not None else "",
+                "age": age[i] if age is not None else "",
+                "sex": str(sex[i]) if sex is not None else "",
+            }
+        )
+
+    choices = [f"{r['index']} | {r['ecg_id']} | {r['gt']} | age {r['age']} {r['sex']}" for r in records]
+    value = choices[0] if choices else None
+    state = {
+        "records": records,
+        "fs": int(sampling_rate),
+        "task": task,
+        "pos_class": pos_class,
+    }
+    report = (records[0]["report"] or "(no clinical notes)") if records else "—"
+    gt = records[0]["gt"] if records else "—"
+    status = (
+        f"Loaded {N} examples (fs={sampling_rate}Hz, {pos_class} vs NORM, L={L})."
+        if N > 0
+        else "No examples found in data file."
+    )
+    return status, gr.update(choices=choices, value=value), state, report, gt
+
+
+def on_select_record(choice: str, state: Optional[dict]):
+    if not state or not state.get("records") or not choice:
+        return "—", "—"
+    try:
+        idx = int(choice.split("|")[0].strip())
+    except Exception:
+        return "—", "—"
+    for r in state["records"]:
+        if r["index"] == idx:
+            return r["report"] or "(no clinical notes)", r["gt"]
+    return "—", "—"
+
+
+def explain_record(
+    model_type: str,
+    sampling_rate: int,
+    task: str,
+    record_choice: str,
+    state: Optional[dict],
+):
+    err = "Select a record and Load records first.", None, "—", "—", "—"
+    if not state or not state.get("records") or not record_choice:
+        return err
+    try:
+        rec_idx = int(record_choice.split("|")[0].strip())
+    except Exception:
+        return err
+    rec = next((r for r in state["records"] if r["index"] == rec_idx), None)
+    if not rec:
+        return err
+
+    fs = state["fs"]
+    pos_class_name = state.get("pos_class", "MI")
+    report = rec["report"] or "(no clinical notes)"
+    gt = rec["gt"]
+
+    try:
+        data = load_data_binary(int(sampling_rate), task)
+    except Exception as e:
+        return f"Data error: {e}", None, report, gt, "—"
+    try:
+        model, device = load_imn_model(model_type, int(sampling_rate), task)
+    except Exception as e:
+        return f"Checkpoint error: {e}", None, report, gt, "—"
+
+    signals = data["signals"]
+    if rec_idx < 0 or rec_idx >= signals.shape[0]:
+        return f"Invalid record index {rec_idx}.", None, report, gt, "—"
+
+    x = signals[rec_idx]  # [12, L]
+    if x.shape[0] != 12:
+        return f"Expected 12 leads, got {x.shape[0]}.", None, report, gt, "—"
+
+    signal_len_model = int(model.hparams["signal_len"])
+    if x.shape[1] != signal_len_model:
+        return (
+            f"ECG length {x.shape[1]} != model {signal_len_model}. "
+            "Ensure data binaries match the training window length.",
+            None,
+            report,
+            gt,
+            "—",
+        )
+
+    x = zscore_per_lead(x)
+    x_t = torch.from_numpy(x).float().unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        logits, gen_w, gen_b = model.model(x_t)
+        if model_type == "single_linear":
+            logit = logits.squeeze()
+            prob_pos = float(torch.sigmoid(logit).item())
+            w_used = gen_w[0, 0, :, :].cpu().numpy()
+        else:
+            probs = torch.softmax(logits, dim=1)
+            prob_pos = float(probs[0, 1].item())
+            w_used = gen_w[0, 1, :, :].cpu().numpy()
+
+    x_np = x.astype(np.float64)
+    impact = w_used * x_np  # [12, L]
+
+    # Window/stride heuristic by sampling rate
+    window = 50 if int(sampling_rate) == 100 else 250
+    stride = window // 2
+    seg_hm = sl_core.imn_weights_to_segments(impact, window=window, stride=stride)  # [12, T]
+
+    # Build simple figure: heatmap + 12 ECG traces
+    L = x_np.shape[1]
+    T = seg_hm.shape[1]
+
+    fig = plt.figure(figsize=(11, 10))
+    gs = fig.add_gridspec(14, 1, height_ratios=[2] + [1] * 12 + [0.5])
+
+    ax0 = fig.add_subplot(gs[0, 0])
+    im = ax0.imshow(seg_hm, aspect="auto", vmin=0.0, vmax=1.0, cmap="Reds")
+    ax0.set_yticks(range(12))
+    ax0.set_yticklabels(LEAD_NAMES)
+    ax0.set_xlabel(f"Segments (window={window}, stride={stride}, fs={fs}Hz)")
+    prob_str = f"P({pos_class_name})={prob_pos:.3f}"
+    pred = pos_class_name if prob_pos >= 0.5 else "NORM"
+    ax0.set_title(f"IMN Intrinsic Explanation | {pred} | {prob_str}")
+    fig.colorbar(im, ax=ax0, fraction=0.02, pad=0.01)
+
+    for lead in range(12):
+        ax = fig.add_subplot(gs[lead + 1, 0])
+        ax.plot(x_np[lead], linewidth=0.8, color="black", alpha=0.7)
+        ax.set_xlim(0, L - 1)
+        ax.set_ylabel(LEAD_NAMES[lead], rotation=0, labelpad=15, va="center")
+        ax.set_xticks([])
+
+    axf = fig.add_subplot(gs[13, 0])
+    axf.axis("off")
+    axf.text(
+        0.5,
+        0.5,
+        "Heatmap: |w(x)·x| aggregated over segments (higher = more contribution towards POS_CLASS).",
+        fontsize=9,
+        ha="center",
+        va="center",
+        wrap=True,
+        transform=axf.transAxes,
+    )
+
+    summary = (
+        f"**{pred}** | P({pos_class_name}) = {prob_pos:.3f} | "
+        f"Ground truth: **{gt}** | fs={fs}Hz, window={window}, stride={stride}"
+    )
+    return summary, fig, report, gt, f"{rec['ecg_id']}"
+
+
+def main():
+    demo = gr.Blocks(
+        title="MesomorphicECG XAI (IMN categorical + single-linear)",
+        theme=gr.themes.Soft(),
+    )
+    with demo:
+        gr.Markdown(
+            "# MesomorphicECG XAI\n"
+            "Interactive XAI viewer for Interpretable Mesomorphic Networks (IMN) on PTB-XL ECGs.\n\n"
+            "- Models and checkpoints from "
+            "[SEARCH-IHI/mesomorphicECG](https://huggingface.co/SEARCH-IHI/mesomorphicECG).\n"
+            "- Data samples loaded from binary `.npz` files stored in this Space.\n"
+            "- Heatmaps show segment-wise IMN contribution per lead."
+        )
+
+        with gr.Row():
+            sampling_rate = gr.Radio(
+                label="Sampling rate",
+                choices=[100, 500],
+                value=500,
+            )
+            model_type = gr.Radio(
+                label="Model type",
+                choices=["single_linear", "categorical"],
+                value="single_linear",
+                info="single_linear: single linear head; categorical: 2-class head.",
+            )
+            task = gr.Radio(
+                label="Task (positive class vs NORM)",
+                choices=list(TASK_TO_POS.keys()),
+                value="norm_vs_mi",
+            )
+            load_btn = gr.Button("Load records", variant="secondary")
+
+        load_status = gr.Markdown()
+        records_state = gr.State(value=None)
+
+        with gr.Row():
+            record_dd = gr.Dropdown(
+                label="Record (index | ecg_id | GT | age sex)",
+                choices=[],
+                value=None,
+            )
+
+        with gr.Row():
+            clinical_notes = gr.Textbox(
+                label="Clinical notes (report)",
+                value="",
+                lines=4,
+                max_lines=8,
+                interactive=False,
+            )
+            ground_truth = gr.Textbox(
+                label="Ground truth",
+                value="—",
+                interactive=False,
+            )
+
+        load_btn.click(
+            fn=on_load_records,
+            inputs=[sampling_rate, task, records_state],
+            outputs=[load_status, record_dd, records_state, clinical_notes, ground_truth],
+        )
+        record_dd.change(
+            fn=on_select_record,
+            inputs=[record_dd, records_state],
+            outputs=[clinical_notes, ground_truth],
+        )
+
+        run_btn = gr.Button("Run IMN explanation", variant="primary")
+
+        out_summary = gr.Markdown()
+        out_plot = gr.Plot()
+        out_notes = gr.Textbox(label="Clinical notes", lines=3, interactive=False)
+        out_gt = gr.Textbox(label="Ground truth", interactive=False)
+        out_meta = gr.Textbox(label="ECG ID", interactive=False)
+
+        run_btn.click(
+            fn=explain_record,
+            inputs=[model_type, sampling_rate, task, record_dd, records_state],
+            outputs=[out_summary, out_plot, out_notes, out_gt, out_meta],
+        )
+
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()
+

categorical_imn_core.py

@@ -0,0 +1,307 @@
+"""
+Core categorical IMN model definition for mesomorphicECG.
+
+This is a trimmed-down subset of
+`script_02022026_v7_IMN_GM_2_with_transition_net.py` containing only
+the pieces needed for inference:
+
+- ECG_IMN (categorical / 2-class hypernetwork)
+- IMNLightning (PyTorch Lightning wrapper)
+- imn_weights_to_segments (segment-wise aggregation helper)
+
+These definitions are compatible with checkpoints uploaded to
+`SEARCH-IHI/mesomorphicECG` under:
+  categorical_imn_100hz/<task>/
+  categorical_imn_500hz/<task>/
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from sklearn.metrics import roc_auc_score
+
+
+DEFAULT_LEAD_NAMES = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+
+
+def imn_weights_to_segments(impact_12L: np.ndarray, window: int, stride: int) -> np.ndarray:
+    """
+    Aggregates point-wise feature attribution (Impact) into segments for cleaner visualization.
+
+    impact_12L: [12, L]
+    Returns: [12, T] heatmap normalized to [0, 1] per-record.
+    """
+    assert impact_12L.ndim == 2
+    L = impact_12L.shape[1]
+    T = (L - window) // stride + 1
+    seg = np.zeros((12, T), dtype=np.float32)
+
+    for t in range(T):
+        s = t * stride
+        e = min(s + window, L)
+        seg[:, t] = np.abs(impact_12L[:, s:e]).mean(axis=1)
+
+    mx = seg.max() + 1e-9
+    seg = seg / mx
+    return seg
+
+
+class ECG_IMN(nn.Module):
+    """
+    Interpretable Mesomorphic Network for ECG with Transition Network (categorical).
+
+    Generates weights W [B, num_classes, 12, L] and biases b [B, num_classes].
+    Final logits for each class k: logits_k = sum(W_k * x) + b_k.
+    """
+
+    def __init__(self, input_channels: int = 12, signal_len: int = 1000, num_classes: int = 2, dropout: float = 0.2):
+        super().__init__()
+        self.num_classes = num_classes
+        self.C = input_channels
+        self.L = signal_len
+
+        # Hypernetwork backbone (encoder), input: [B, 1, 12, L]
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(1, 16, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(16),
+            nn.GELU(),
+        )  # -> [B, 16, 12, L]
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(16, 32, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.MaxPool2d(kernel_size=(1, 2)),
+        )  # -> [B, 32, 12, L/2]
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(32, 64, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.MaxPool2d(kernel_size=(1, 2)),
+        )  # -> [B, 64, 12, L/4]
+
+        self.dropout = nn.Dropout(dropout)
+
+        # Transition network: upsample to generate W [B, num_classes, 12, L]
+        self.transition = nn.Sequential(
+            # L/4 -> L/2, 64 -> 32
+            nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.Upsample(scale_factor=(1, 2), mode="nearest"),
+            # L/2 -> L, 32 -> 16
+            nn.Conv2d(32, 16, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(16),
+            nn.GELU(),
+            nn.Upsample(scale_factor=(1, 2), mode="nearest"),
+            # Final projection to num_classes channels (weights)
+            nn.Conv2d(16, num_classes, kernel_size=3, padding=1, bias=True),
+        )
+
+        # Bias generator: class-wise bias from global pooled features
+        self.bias_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.bias_head = nn.Linear(64, num_classes)
+
+    def forward(self, x: torch.Tensor):
+        """
+        x: [B, 12, L]
+        Returns:
+            logits: [B, num_classes]
+            generated_w: [B, num_classes, 12, L]
+            generated_b: [B, num_classes, 1]
+        """
+        B, C, L = x.shape
+
+        feat = x.unsqueeze(1)  # [B, 1, 12, L]
+        feat = self.conv1(feat)
+        feat = self.conv2(feat)
+        feat = self.conv3(feat)  # [B, 64, 12, L/4]
+        feat = self.dropout(feat)
+
+        # Weights W: [B, num_classes, 12, L]
+        generated_w = self.transition(feat)
+
+        # Bias b: [B, num_classes]
+        b_feat = self.bias_pool(feat).view(B, -1)
+        generated_b = self.bias_head(b_feat)
+
+        x_expanded = x.unsqueeze(1)  # [B, 1, 12, L]
+        weighted_input = generated_w * x_expanded  # [B, num_classes, 12, L]
+        logits = weighted_input.sum(dim=(2, 3)) + generated_b  # [B, num_classes]
+
+        return logits, generated_w, generated_b.unsqueeze(-1)
+
+
+class IMNLightning(pl.LightningModule):
+    """
+    PyTorch Lightning wrapper for ECG_IMN (categorical, 2-class).
+
+    Matches the training-time definition used in
+    `script_02022026_v7_IMN_GM_2_with_transition_net.py`
+    so `IMNLightning.load_from_checkpoint(...)` works for inference.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        signal_len: int,
+        dropout: float = 0.2,
+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        lambda_l1: float = 1e-4,
+        class_weights: list[float] | None = None,
+        scheduler_type: str | None = "cosine",
+        scheduler_params: dict | None = None,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.model = ECG_IMN(
+            input_channels=input_channels,
+            signal_len=signal_len,
+            dropout=dropout,
+        )
+
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.lambda_l1 = lambda_l1
+        self._class_weights = class_weights
+
+        self.scheduler_type = scheduler_type
+        self.scheduler_params = scheduler_params or {}
+
+        self.val_probs: list[torch.Tensor] = []
+        self.val_y: list[torch.Tensor] = []
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
+
+        if self.scheduler_type is None or self.scheduler_type == "none":
+            return optimizer
+
+        if self.scheduler_type == "cosine":
+            max_epochs = getattr(self.trainer, "max_epochs", None) or 100
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer,
+                T_max=max_epochs,
+                **self.scheduler_params,
+            )
+        elif self.scheduler_type == "step":
+            scheduler = torch.optim.lr_scheduler.StepLR(
+                optimizer,
+                step_size=self.scheduler_params.get("step_size", 10),
+                gamma=self.scheduler_params.get("gamma", 0.1),
+                **{k: v for k, v in self.scheduler_params.items() if k not in ["step_size", "gamma"]},
+            )
+        elif self.scheduler_type == "reduce_on_plateau":
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer,
+                mode="max",
+                factor=self.scheduler_params.get("factor", 0.5),
+                patience=self.scheduler_params.get("patience", 5),
+                **{k: v for k, v in self.scheduler_params.items() if k not in ["factor", "patience"]},
+            )
+        elif self.scheduler_type == "cosine_restarts":
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+                optimizer,
+                T_0=self.scheduler_params.get("T_0", 10),
+                T_mult=self.scheduler_params.get("T_mult", 2),
+                **{k: v for k, v in self.scheduler_params.items() if k not in ["T_0", "T_mult"]},
+            )
+        else:
+            return optimizer
+
+        if self.scheduler_type == "reduce_on_plateau":
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "val_auc",
+                },
+            }
+        else:
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": scheduler,
+            }
+
+    def _ce_weight(self):
+        if self._class_weights is None:
+            return None
+        return torch.tensor(self._class_weights, dtype=torch.float32, device=self.device)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits, gen_w, gen_b = self.model(x)
+
+        ce_loss = F.cross_entropy(logits, y, weight=self._ce_weight())
+        l1_loss = gen_w.abs().mean()
+        total_loss = ce_loss + (self.lambda_l1 * l1_loss)
+
+        pred = logits.argmax(dim=1)
+        acc = (pred == y).float().mean()
+
+        self.log("train_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("train_ce", ce_loss, on_step=False, on_epoch=True)
+        self.log("train_l1", l1_loss, on_step=False, on_epoch=True)
+        self.log("train_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+        return total_loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits, gen_w, _ = self.model(x)
+
+        ce_loss = F.cross_entropy(logits, y, weight=self._ce_weight())
+        prob = torch.softmax(logits, dim=1)[:, 1]
+        pred = logits.argmax(dim=1)
+        acc = (pred == y).float().mean()
+
+        self.log("val_loss", ce_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("val_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+
+        self.val_probs.append(prob.detach().cpu())
+        self.val_y.append(y.detach().cpu())
+
+    def on_validation_epoch_end(self):
+        if not self.val_y:
+            return
+        y_true = torch.cat(self.val_y)
+        y_score = torch.cat(self.val_probs)
+        auc = simple_auc_roc(y_true.float(), y_score.float())
+        self.log("val_auc", auc, on_step=False, on_epoch=True, prog_bar=True)
+        self.val_probs.clear()
+        self.val_y.clear()
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        logits, _, _ = self.model(x)
+        ce_loss = F.cross_entropy(logits, y, weight=self._ce_weight())
+        prob = torch.softmax(logits, dim=1)[:, 1]
+        pred = logits.argmax(dim=1)
+        acc = (pred == y).float().mean()
+
+        self.log("test_loss", ce_loss, on_step=False, on_epoch=True)
+        self.log("test_acc", acc, on_step=False, on_epoch=True)
+        return {"y": y.detach().cpu(), "p": prob.detach().cpu()}
+
+
+@torch.no_grad()
+def simple_auc_roc(y_true: torch.Tensor, y_score: torch.Tensor) -> float:
+    """
+    Simple AUROC helper, matching the training script.
+    """
+    y_true = y_true.detach().cpu().float()
+    y_score = y_score.detach().cpu().float()
+    if y_true.min() == y_true.max():
+        return float("nan")
+    return float(roc_auc_score(y_true.numpy(), y_score.numpy()))
+

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+torch
+numpy
+matplotlib
+pytorch-lightning
+scikit-learn
+huggingface-hub
+

single_linear_imn_core.py

@@ -0,0 +1,307 @@
+"""
+Core single-linear IMN model definition for mesomorphicECG.
+
+This file is a lightweight subset of
+`script_02022026_v7_IMN_GM_2_with_transition_net_with_one_linear_eq.py`
+containing only the pieces needed for inference:
+
+- ECG_IMN (single-linear hypernetwork)
+- IMNLightning (PyTorch Lightning wrapper)
+- imn_weights_to_segments (segment-wise aggregation helper)
+
+These definitions are compatible with checkpoints uploaded to
+`SEARCH-IHI/mesomorphicECG` under:
+  single_linear_imn_100hz/<task>/
+  single_linear_imn_500hz/<task>/
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from sklearn.metrics import roc_auc_score
+
+
+DEFAULT_LEAD_NAMES = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+
+
+def imn_weights_to_segments(impact_12L: np.ndarray, window: int, stride: int) -> np.ndarray:
+    """
+    Aggregates point-wise feature attribution (Impact) into segments.
+
+    impact_12L: [12, L] array of signed contributions (e.g. w * x).
+    Returns: [12, T] heatmap normalized to [0, 1] per-record.
+    """
+    assert impact_12L.ndim == 2
+    L = impact_12L.shape[1]
+    T = (L - window) // stride + 1
+    seg = np.zeros((12, T), dtype=np.float32)
+
+    for t in range(T):
+        s = t * stride
+        e = min(s + window, L)
+        seg[:, t] = np.abs(impact_12L[:, s:e]).mean(axis=1)
+
+    mx = seg.max() + 1e-9
+    seg = seg / mx
+    return seg
+
+
+class ECG_IMN(nn.Module):
+    """
+    Interpretable Mesomorphic Network for ECG (single-linear output).
+
+    Generates ONE set of weights W [B, 1, 12, L] and ONE bias b [B, 1].
+    Prediction: logit = sum(W * x) + b, P(pos) = sigmoid(logit).
+    """
+
+    def __init__(self, input_channels: int = 12, signal_len: int = 1000, dropout: float = 0.2):
+        super().__init__()
+        self.C = input_channels
+        self.L = signal_len
+
+        output_dim = 1  # single linear output
+
+        # Hypernetwork backbone (encoder), input: [B, 1, 12, L]
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(1, 16, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(16),
+            nn.GELU(),
+        )  # -> [B, 16, 12, L]
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(16, 32, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.MaxPool2d(kernel_size=(1, 2)),
+        )  # -> [B, 32, 12, L/2]
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(32, 64, kernel_size=(3, 15), padding=(1, 7), bias=False),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.MaxPool2d(kernel_size=(1, 2)),
+        )  # -> [B, 64, 12, L/4]
+
+        self.dropout = nn.Dropout(dropout)
+
+        # Transition network: upsample to generate W [B, 1, 12, L]
+        self.transition = nn.Sequential(
+            # L/4 -> L/2
+            nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.Upsample(scale_factor=(1, 2), mode="nearest"),
+            # L/2 -> L
+            nn.Conv2d(32, 16, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(16),
+            nn.GELU(),
+            nn.Upsample(scale_factor=(1, 2), mode="nearest"),
+            # Final projection to 1 channel (weights)
+            nn.Conv2d(16, output_dim, kernel_size=3, padding=1, bias=True),
+        )
+
+        # Bias generator: scalar bias from global pooled features
+        self.bias_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.bias_head = nn.Linear(64, output_dim)
+
+    def forward(self, x: torch.Tensor):
+        """
+        x: [B, 12, L]
+        Returns:
+            logits: [B, 1]
+            generated_w: [B, 1, 12, L]
+            generated_b: [B, 1, 1]
+        """
+        B, C, L = x.shape
+
+        feat = x.unsqueeze(1)  # [B, 1, 12, L]
+        feat = self.conv1(feat)
+        feat = self.conv2(feat)
+        feat = self.conv3(feat)
+        feat = self.dropout(feat)
+
+        # Weights W: [B, 1, 12, L]
+        generated_w = self.transition(feat)
+
+        # Bias b: [B, 1]
+        b_feat = self.bias_pool(feat).view(B, -1)
+        generated_b = self.bias_head(b_feat)
+
+        # Single-linear logit
+        x_expanded = x.unsqueeze(1)  # [B, 1, 12, L]
+        weighted_input = generated_w * x_expanded
+        logits = weighted_input.sum(dim=(2, 3)) + generated_b  # [B, 1]
+
+        return logits, generated_w, generated_b.unsqueeze(-1)
+
+
+class IMNLightning(pl.LightningModule):
+    """
+    PyTorch Lightning wrapper for ECG_IMN (single-linear).
+
+    This class matches the training-time definition used for checkpoints
+    in `script_02022026_v7_IMN_GM_2_with_transition_net_with_one_linear_eq.py`,
+    so that `IMNLightning.load_from_checkpoint(...)` works for inference.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        signal_len: int,
+        dropout: float = 0.2,
+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        lambda_l1: float = 1e-4,
+        pos_weight: float | None = None,
+        scheduler_type: str | None = "cosine",
+        scheduler_params: dict | None = None,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.model = ECG_IMN(
+            input_channels=input_channels,
+            signal_len=signal_len,
+            dropout=dropout,
+        )
+
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.lambda_l1 = lambda_l1
+        self.pos_weight_val = pos_weight
+
+        self.scheduler_type = scheduler_type
+        self.scheduler_params = scheduler_params or {}
+
+        self.val_probs: list[torch.Tensor] = []
+        self.val_y: list[torch.Tensor] = []
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
+
+        if self.scheduler_type is None or self.scheduler_type == "none":
+            return optimizer
+
+        if self.scheduler_type == "cosine":
+            max_epochs = getattr(self.trainer, "max_epochs", None) or 100
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer,
+                T_max=max_epochs,
+                **self.scheduler_params,
+            )
+        elif self.scheduler_type == "step":
+            scheduler = torch.optim.lr_scheduler.StepLR(
+                optimizer,
+                step_size=self.scheduler_params.get("step_size", 10),
+                gamma=self.scheduler_params.get("gamma", 0.1),
+                **{k: v for k, v in self.scheduler_params.items() if k not in ["step_size", "gamma"]},
+            )
+        elif self.scheduler_type == "reduce_on_plateau":
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer,
+                mode="max",
+                factor=self.scheduler_params.get("factor", 0.5),
+                patience=self.scheduler_params.get("patience", 5),
+                **{k: v for k, v in self.scheduler_params.items() if k not in ["factor", "patience"]},
+            )
+        else:
+            return optimizer
+
+        if self.scheduler_type == "reduce_on_plateau":
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "val_auc",
+                },
+            }
+        else:
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": scheduler,
+            }
+
+    def _get_pos_weight(self):
+        if self.pos_weight_val is None:
+            return None
+        return torch.tensor([self.pos_weight_val], device=self.device)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch  # y: [B] float
+        logits, gen_w, gen_b = self.model(x)
+        logits = logits.squeeze(1)  # [B]
+
+        bce_loss = F.binary_cross_entropy_with_logits(logits, y, pos_weight=self._get_pos_weight())
+        l1_loss = gen_w.abs().mean()
+        total_loss = bce_loss + (self.lambda_l1 * l1_loss)
+
+        probs = torch.sigmoid(logits)
+        preds = (probs > 0.5).float()
+        acc = (preds == y).float().mean()
+
+        self.log("train_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("train_bce", bce_loss, on_step=False, on_epoch=True)
+        self.log("train_l1", l1_loss, on_step=False, on_epoch=True)
+        self.log("train_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+        return total_loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits, gen_w, _ = self.model(x)
+        logits = logits.squeeze(1)
+
+        bce_loss = F.binary_cross_entropy_with_logits(logits, y, pos_weight=self._get_pos_weight())
+        prob = torch.sigmoid(logits)
+        pred = (prob > 0.5).float()
+        acc = (pred == y).float().mean()
+
+        self.log("val_loss", bce_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("val_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+
+        self.val_probs.append(prob.detach().cpu())
+        self.val_y.append(y.detach().cpu())
+
+    def on_validation_epoch_end(self):
+        if not self.val_y:
+            return
+        y_true = torch.cat(self.val_y)
+        y_score = torch.cat(self.val_probs)
+        auc = simple_auc_roc(y_true, y_score)
+        self.log("val_auc", auc, on_step=False, on_epoch=True, prog_bar=True)
+        self.val_probs.clear()
+        self.val_y.clear()
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        logits, _, _ = self.model(x)
+        logits = logits.squeeze(1)
+
+        bce_loss = F.binary_cross_entropy_with_logits(logits, y, pos_weight=self._get_pos_weight())
+        prob = torch.sigmoid(logits)
+        pred = (prob > 0.5).float()
+        acc = (pred == y).float().mean()
+
+        self.log("test_loss", bce_loss, on_step=False, on_epoch=True)
+        self.log("test_acc", acc, on_step=False, on_epoch=True)
+        return {"y": y.detach().cpu(), "p": prob.detach().cpu()}
+
+
+@torch.no_grad()
+def simple_auc_roc(y_true: torch.Tensor, y_score: torch.Tensor) -> float:
+    """
+    Simple AUROC helper, matching the training script.
+    """
+    y_true = y_true.detach().cpu().float()
+    y_score = y_score.detach().cpu().float()
+    if y_true.min() == y_true.max():
+        return float("nan")
+    return float(roc_auc_score(y_true.numpy(), y_score.numpy()))
+