Add model details and instructions

README.md

@@ -1,3 +1,129 @@
----
-license: apache-2.0
----
+# ECG Image Classifier (MoE and MLP) on MedSigLIP Embeddings
+
+This repository provides two PyTorch ECG classifier checkpoints trained on top of frozen MedSigLIP image embeddings:
+
+- `moe_classifier_medsiglip.pt`: Mixture-of-Experts (MoE) classifier
+- `mlp_classifier_medsiglip.pt`: Dense feedforward (MLP) classifier
+
+These checkpoints expect embeddings produced by:
+
+- `google/medsiglip-448`
+
+The repository contains only the classifier heads. MedSigLIP weights are not included and must be obtained separately under Google’s license.
+
+---
+
+## Motivation
+
+This work was developed as part of the Google MedGemma Impact Challenge:  
+https://www.kaggle.com/competitions/med-gemma-impact-challenge/overview
+
+The goal is to build a lightweight, deployable ECG image classifier for chronic care screening, especially in low-resource clinical settings where ECG is often the most accessible diagnostic modality.
+
+---
+
+## Task and Data
+
+We formulate a supervised multi-label image classification task on 12-lead ECGs with five diagnostic categories:
+
+- NORM (normal)
+- MI (myocardial infarction)
+- STTC (ST-T changes)
+- CD (conduction disturbances)
+- HYP (hypertrophy)
+
+Training data combines:
+
+- PTB-XL, a large-scale dataset of raw 12-lead ECG waveforms in WFDB format with 16-bit precision
+- A supplementary ECG image dataset
+
+To enable image-based classification, raw PTB-XL waveforms are converted into realistic print-style ECG images using the open-source ECG image generator by Rahimi et al. This yields approximately 21,000 synthetic ECG images, which are combined with 713 real ECG images from the supplementary dataset.
+
+---
+
+## Model and Training
+
+ECG images are first encoded using MedSigLIP to obtain fixed-dimensional visual embeddings. Two lightweight classifiers are trained on top of these embeddings:
+
+- A dense feedforward network (MLP)
+- A Mixture-of-Experts (MoE) classifier
+
+The dataset is split into 60 percent training, 20 percent validation, and 20 percent testing. Both models are optimized with Adam using a learning rate of 1e-4 and weight decay of 1e-5. The MoE model additionally uses a load-balancing regularization term with lambda set to 0.1.
+
+For multi-label prediction, a uniform decision threshold of 0.3 is applied across all classes.
+
+---
+
+## Results
+
+On the held-out test set, the MoE classifier consistently outperforms the MLP baseline across all metrics. It achieves:
+
+- Lower Hamming loss: 0.167 vs 0.235
+- Higher ROC-AUC:
+  - Micro: 0.891 vs 0.827
+  - Macro: 0.879 vs 0.808
+- Higher F1 scores:
+  - Micro: 0.70 vs 0.61
+  - Macro: 0.67 vs 0.58
+
+Per-class F1 improves across all five diagnostic categories, with the largest gains observed for myocardial infarction and hypertrophy. Confusion matrix analysis indicates that the MLP baseline tends to trade precision for recall, producing more false positives and a lower overall F1. For this reason, the MoE classifier is used in the final application.
+
+---
+
+## Practical Implications
+
+Compared to using MedGemma alone, the MedSigLIP plus classifier pipeline provides more structured and reliable ECG predictions. In addition to discrete labels, the classifier outputs calibrated confidence scores. This supports threshold-based screening and triage, which is particularly useful in chronic care workflows and remote clinics where rapid ECG assessment can help prioritize referrals.
+
+---
+
+## How to Use
+
+### 1) Install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### 2) Run inference
+
+Single image with the MoE checkpoint:
+
+```bash
+python inference_loader.py \
+  --ckpt ./moe_classifier_medsiglip.pt \
+  --image ./sample_ecg.png \
+  --out ./preds_moe.json
+```
+
+Batch inference on a folder with the MLP checkpoint:
+
+```bash
+python inference_loader.py \
+  --ckpt ./mlp_classifier_medsiglip.pt \
+  --folder ./images \
+  --out ./preds_mlp.json
+```
+
+### 3) Optional arguments
+
+- `--model_id` (default: `google/medsiglip-448`)
+- `--device auto|cpu|cuda`
+- `--batch_size 16`
+- `--threshold 0.3` (overrides the checkpoint threshold)
+- `--hf_token <token>` (or set `HF_TOKEN` as an environment variable)
+
+### 4) Outputs
+
+The inference script returns:
+
+- `scores_by_class`: confidence scores for each diagnostic class
+- `predicted_labels`: labels above the decision threshold
+- `summary`: run metadata including checkpoint, model type, device, and embedding dimensions
+
+---
+
+## References
+
+[1] PTB-XL dataset  
+[2] Supplementary ECG image dataset used in this project  
+[3] Rahimi et al., open-source ECG image generator

inference_loader.py

@@ -0,0 +1,349 @@
+"""
+Standalone inference loader for ECG classifier checkpoints.
+
+Supports:
+- MoE checkpoints (experts.* + gate.* keys)
+- MLP checkpoints (fc1/fc2/out keys)
+
+Usage examples:
+  python inference_loader.py --ckpt ./moe_classifier_medsiglip.pt --image ./ecg.png
+  python inference_loader.py --ckpt ./mlp_classifier_medsiglip.pt --folder ./images --out ./preds.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+
+
+DEFAULT_MODEL_ID = "google/medsiglip-448"
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+
+
+class ExpertMLP(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden: tuple[int, ...] = (1028, 512, 256),
+        dropout: tuple[float, ...] = (0.15, 0.15, 0.10),
+    ):
+        super().__init__()
+        layers: list[nn.Module] = []
+        prev = in_dim
+
+        dropout_values = tuple(dropout)
+        if len(dropout_values) < len(hidden):
+            dropout_values = dropout_values + (0.0,) * (len(hidden) - len(dropout_values))
+        elif len(dropout_values) > len(hidden):
+            dropout_values = dropout_values[: len(hidden)]
+
+        for h, p in zip(hidden, dropout_values):
+            layers.append(nn.Linear(prev, h))
+            layers.append(nn.LayerNorm(h))
+            layers.append(nn.GELU())
+            layers.append(nn.Dropout(p))
+            prev = h
+
+        layers.append(nn.Linear(prev, out_dim))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class MoEClassifier(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        num_experts: int = 5,
+        gate_hidden: int = 512,
+        temperature: float = 1.0,
+        expert_hidden: tuple[int, ...] = (1028, 512, 256),
+        expert_dropout: tuple[float, ...] = (0.15, 0.15, 0.10),
+    ):
+        super().__init__()
+        self.temperature = temperature
+        self.experts = nn.ModuleList(
+            [
+                ExpertMLP(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    hidden=expert_hidden,
+                    dropout=expert_dropout,
+                )
+                for _ in range(num_experts)
+            ]
+        )
+        self.gate = nn.Sequential(
+            nn.Linear(in_dim, gate_hidden),
+            nn.ReLU(),
+            nn.Linear(gate_hidden, num_experts),
+        )
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        gate_logits = self.gate(x) / self.temperature
+        gate_w = torch.softmax(gate_logits, dim=-1)
+        expert_logits = torch.stack([expert(x) for expert in self.experts], dim=1)
+        mixed_logits = torch.sum(expert_logits * gate_w.unsqueeze(-1), dim=1)
+        return mixed_logits, gate_w, expert_logits
+
+
+class MLPClassifier(nn.Module):
+    def __init__(self, in_dim: int, hidden_1: int, hidden_2: int, out_dim: int):
+        super().__init__()
+        self.fc1 = nn.Linear(in_dim, hidden_1)
+        self.fc2 = nn.Linear(hidden_1, hidden_2)
+        self.out = nn.Linear(hidden_2, out_dim)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        return self.out(x)
+
+
+def get_device(device_arg: str) -> str:
+    if device_arg == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    if device_arg == "cuda" and not torch.cuda.is_available():
+        return "cpu"
+    return device_arg
+
+
+def collect_images(image: str | None, images: list[str] | None, folder: str | None) -> list[str]:
+    paths: list[str] = []
+    if image:
+        paths.append(image)
+    if images:
+        paths.extend(images)
+    if folder:
+        for name in sorted(os.listdir(folder)):
+            p = os.path.join(folder, name)
+            ext = os.path.splitext(p)[1].lower()
+            if os.path.isfile(p) and ext in IMAGE_EXTS:
+                paths.append(p)
+
+    out: list[str] = []
+    seen: set[str] = set()
+    for p in paths:
+        ap = str(Path(p).resolve())
+        if ap not in seen:
+            seen.add(ap)
+            out.append(ap)
+    return out
+
+
+def extract_features(output: Any) -> torch.Tensor:
+    if isinstance(output, torch.Tensor):
+        return output
+    if hasattr(output, "pooler_output") and output.pooler_output is not None:
+        return output.pooler_output
+    if hasattr(output, "last_hidden_state") and output.last_hidden_state is not None:
+        return output.last_hidden_state[:, 0, :]
+    raise TypeError(f"Unexpected image feature output type: {type(output)}")
+
+
+def build_classifier(ckpt: dict[str, Any]) -> tuple[nn.Module, str]:
+    state_dict = ckpt.get("state_dict")
+    if not isinstance(state_dict, dict) or not state_dict:
+        raise RuntimeError("Checkpoint missing state_dict.")
+
+    embed_dim = int(ckpt["embed_dim"])
+    num_classes = int(ckpt["num_classes"])
+
+    if any(key.startswith("experts.") for key in state_dict.keys()):
+        num_experts = int(ckpt.get("num_experts", 5))
+        expert_linear_layers: list[tuple[int, torch.Tensor]] = []
+        for key, value in state_dict.items():
+            if not (
+                key.startswith("experts.0.net.")
+                and key.endswith(".weight")
+                and isinstance(value, torch.Tensor)
+                and value.ndim == 2
+            ):
+                continue
+            layer_index = int(key.split(".")[3])
+            expert_linear_layers.append((layer_index, value))
+
+        if len(expert_linear_layers) < 2:
+            raise RuntimeError("Unable to infer expert architecture from checkpoint.")
+
+        expert_linear_layers.sort(key=lambda item: item[0])
+        expert_hidden = tuple(int(weight.shape[0]) for _, weight in expert_linear_layers[:-1])
+        gate_hidden = int(state_dict["gate.0.weight"].shape[0]) if "gate.0.weight" in state_dict else 256
+
+        model = MoEClassifier(
+            in_dim=embed_dim,
+            out_dim=num_classes,
+            num_experts=num_experts,
+            gate_hidden=gate_hidden,
+            temperature=1.0,
+            expert_hidden=expert_hidden,
+            expert_dropout=tuple(0.0 for _ in expert_hidden),
+        )
+        model_type = "moe"
+    elif {"fc1.weight", "fc2.weight", "out.weight"}.issubset(state_dict.keys()):
+        hidden_1 = int(state_dict["fc1.weight"].shape[0])
+        hidden_2 = int(state_dict["fc2.weight"].shape[0])
+        model = MLPClassifier(
+            in_dim=embed_dim,
+            hidden_1=hidden_1,
+            hidden_2=hidden_2,
+            out_dim=num_classes,
+        )
+        model_type = "mlp"
+    else:
+        raise RuntimeError("Unsupported checkpoint format.")
+
+    model.load_state_dict(state_dict, strict=True)
+    model.eval()
+    return model, model_type
+
+
+@torch.no_grad()
+def embed_images(
+    embedder: AutoModel,
+    processor: AutoImageProcessor,
+    image_paths: list[str],
+    batch_size: int,
+) -> tuple[np.ndarray, list[str]]:
+    embs: list[np.ndarray] = []
+    kept: list[str] = []
+    for i in range(0, len(image_paths), batch_size):
+        batch_paths = image_paths[i : i + batch_size]
+        batch_images = []
+        batch_kept: list[str] = []
+        for p in batch_paths:
+            try:
+                batch_images.append(Image.open(p).convert("RGB"))
+                batch_kept.append(p)
+            except Exception as exc:  # pragma: no cover
+                print(f"[skip] {p}: {exc}")
+
+        if not batch_images:
+            continue
+
+        inputs = processor(images=batch_images, return_tensors="pt").to(embedder.device)
+        out = embedder.get_image_features(**inputs)
+        feats = extract_features(out)
+
+        embs.append(feats.detach().cpu().numpy().astype(np.float32))
+        kept.extend(batch_kept)
+
+    if not embs:
+        return np.zeros((0, 0), dtype=np.float32), []
+    return np.concatenate(embs, axis=0), kept
+
+
+@torch.no_grad()
+def predict(
+    classifier: nn.Module,
+    X_emb: np.ndarray,
+    device: str,
+) -> np.ndarray:
+    xb = torch.from_numpy(X_emb).to(device)
+    logits_output = classifier(xb)
+    logits = logits_output[0] if isinstance(logits_output, tuple) else logits_output
+    probs = torch.sigmoid(logits).detach().cpu().numpy()
+    return probs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", type=str, required=True, help="Path to checkpoint .pt file")
+    parser.add_argument("--model_id", type=str, default=DEFAULT_MODEL_ID, help="HF model id for embedder")
+    parser.add_argument("--hf_token", type=str, default=None, help="HF token (or set HF_TOKEN env var)")
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cpu", "cuda"])
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--threshold", type=float, default=None, help="Override checkpoint threshold")
+
+    parser.add_argument("--image", type=str, default=None, help="Single image path")
+    parser.add_argument("--images", nargs="*", default=None, help="Multiple image paths")
+    parser.add_argument("--folder", type=str, default=None, help="Folder with images")
+    parser.add_argument("--out", type=str, default=None, help="Optional output JSON path")
+    args = parser.parse_args()
+
+    image_paths = collect_images(args.image, args.images, args.folder)
+    if not image_paths:
+        raise SystemExit("No images found. Use --image, --images, or --folder.")
+
+    device = get_device(args.device)
+    ckpt = torch.load(args.ckpt, map_location="cpu")
+    if not isinstance(ckpt, dict):
+        raise SystemExit("Checkpoint must be a dict.")
+
+    classifier, model_type = build_classifier(ckpt)
+    classifier.to(device)
+    classifier.eval()
+
+    embed_dim = int(ckpt["embed_dim"])
+    num_classes = int(ckpt["num_classes"])
+    threshold = float(args.threshold) if args.threshold is not None else float(ckpt.get("threshold", 0.5))
+    classes = ckpt.get("classes")
+    if not isinstance(classes, list) or len(classes) != num_classes:
+        classes = [f"class_{i}" for i in range(num_classes)]
+
+    token = args.hf_token or os.getenv("HF_TOKEN")
+    embedder = AutoModel.from_pretrained(args.model_id, token=token)
+    processor = AutoImageProcessor.from_pretrained(args.model_id, token=token)
+    embedder.to(device)
+    embedder.eval()
+
+    X_emb, kept_paths = embed_images(embedder, processor, image_paths, args.batch_size)
+    if X_emb.shape[0] == 0:
+        raise SystemExit("No images could be processed.")
+    if X_emb.shape[1] != embed_dim:
+        raise SystemExit(
+            f"Embedding dim mismatch: produced {X_emb.shape[1]} but checkpoint expects {embed_dim}."
+        )
+
+    probs = predict(classifier, X_emb, device=device)
+    preds = (probs >= threshold).astype(int)
+
+    summary = {
+        "checkpoint": str(Path(args.ckpt).resolve()),
+        "classifier_type": model_type,
+        "model_id": args.model_id,
+        "device": device,
+        "num_images": len(kept_paths),
+        "embed_dim": embed_dim,
+        "num_classes": num_classes,
+        "threshold": threshold,
+    }
+    if "num_experts" in ckpt:
+        summary["num_experts"] = int(ckpt["num_experts"])
+
+    results = []
+    for image_path, row_prob, row_pred in zip(kept_paths, probs, preds):
+        results.append(
+            {
+                "image_path": image_path,
+                "scores_by_class": {label: float(score) for label, score in zip(classes, row_prob)},
+                "predicted_labels": [label for label, y in zip(classes, row_pred) if int(y) == 1],
+            }
+        )
+
+    payload = {"summary": summary, "results": results}
+    print(json.dumps(summary, indent=2))
+    print(json.dumps(results[:3], indent=2))
+
+    if args.out:
+        out_path = Path(args.out).resolve()
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+        print(f"Saved output to: {out_path}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+pillow
+numpy
+tqdm