Lung CT Nodule Classification — I-JEPA ViT-Small

Self-supervised I-JEPA ViT-Small/16 encoder pre-trained on lung CT images, with a linear probe for binary nodule classification (Nodule vs Healthy). Visual explanations (Grad-CAM, LIME, SHAP) are produced by the project code, not by the weights themselves.

Code, notebooks and web demo: https://github.com/Zainabfarih/XAI-MedVision

Files

File	Size	Description
`ijepa_best.pth`	~353 MB	I-JEPA encoder (ViT-Small/16), state dict key `context_encoder`
`probe_best.pth`	~15 KB	Linear probe head, state dict key `probe`

Results (test set, 3,065 images)

Metric	Value
Accuracy	0.8914
AUC-ROC	0.9582
Precision	0.9195
Recall	0.8560
F1-score	0.8866
Specificity	0.9262

Architecture

Encoder: ViT-Small/16 (timm vit_small_patch16_224), 384-dim embeddings, 196 patches.
Pre-training: I-JEPA, self-supervised, 100 epochs, no labels.
Probe: frozen encoder + linear head on the CLS token, 30 epochs.

Usage

import timm, torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

REPO = "zainabFarih/lung-ct-nodule-ijepa-vit-small"
device = "cuda" if torch.cuda.is_available() else "cpu"

encoder = timm.create_model("vit_small_patch16_224", pretrained=False,
                            num_classes=0, global_pool="")
enc_ckpt = torch.load(hf_hub_download(REPO, "ijepa_best.pth"), map_location=device)
encoder.load_state_dict(enc_ckpt["context_encoder"])

class LinearProbe(nn.Module):
    def __init__(self, dim=384, n=2, p=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p)
        self.fc = nn.Linear(dim, n)
    def forward(self, cls):
        return self.fc(self.dropout(cls))

probe = LinearProbe()
probe_ckpt = torch.load(hf_hub_download(REPO, "probe_best.pth"), map_location=device)
probe.load_state_dict(probe_ckpt["probe"])

encoder.eval(); probe.eval()
# logits = probe(encoder(x)[:, 0, :])   # x: (B, 3, 224, 224) normalised

Data & intended use

Trained on lung CT images derived from public datasets (LIDC-IDRI / Kaggle lung CT classification). Research and educational use only — not a medical device and not for clinical diagnosis. Verify the license of each source dataset before use.

Downloads last month: -