CerebralJEPA

CerebralJEPA is a representation model for structural brain MRI. It converts a 3D sMRI volume into compact embeddings that can be used for downstream research tasks such as cohort analysis, disease classification, age prediction, clustering, retrieval, or as features for task-specific models.

The model uses a 3D SwinUNETR encoder and reads NIfTI images (.nii or .nii.gz). During inference, the image processor standardizes the volume with RAS orientation, isotropic spacing, foreground estimation, robust intensity clipping, masked normalization, and foreground-centered crop/pad.

The default checkpoint is m4, and the default embedding is a 768-dimensional stage-4 brain representation.

Model Variants

Four pretrained checkpoints are provided:

Checkpoint	Representation setting	Output stages	Recommended use
m1	Global stage-4 representation	stage4	Basic whole-brain embedding extraction
m2	Multi-scale global stage-4 representation	stage4	Whole-brain embedding with global scale variation
m3	Global and local view training with stage-4 export	stage4	Stage-4 embedding with stronger view diversity
m4	Multi-stage local/global consistency: local stages 1-2, global stages 3-4	stage1, stage2, stage3, stage4	Default choice, especially when multi-level features are useful

Stage embedding dimensions:

stage1: 48
stage2: 96
stage3: 192
stage4: 768

Installation

Clone the repository and install it in a Python environment with the required dependencies. The examples below use the BrainST conda environment:

git clone https://huggingface.co/YOUR_NAME_OR_ORG/CerebralJEPA
cd CerebralJEPA
conda run -n BrainST python -m pip install -e .

The main dependencies are PyTorch, MONAI, Transformers, Nibabel, SimpleITK, Safetensors, and NumPy. They are listed in requirements.txt and environment.yml; users can install them in any compatible Python 3.10+ environment.

Quick Start

Generate an embedding for one structural MRI:

conda run -n BrainST python scripts/embed_smri.py \
  /path/to/smri.nii.gz \
  --model m4 \
  --output /path/to/embedding.npy

The saved array has shape [768] when using the default stage4 output.

To change the input crop/pad size:

conda run -n BrainST python scripts/embed_smri.py \
  /path/to/smri.nii.gz \
  --model m4 \
  --crop-shape 160 160 160 \
  --output /path/to/embedding.npy

Python Usage

from cerebraljepa_inference import CerebralJEPAPipeline

pipe = CerebralJEPAPipeline.from_pretrained("./model/m4", device="cuda")

result = pipe("/path/to/smri.nii.gz", return_tensors="pt")
embedding = result["embedding"]  # [768]

Use another checkpoint by changing the path:

pipe = CerebralJEPAPipeline.from_pretrained("./model/m1", device="cuda")

For lower-level control:

from cerebraljepa_inference import CerebralJEPAForImageEmbedding, CerebralJEPAImageProcessor

checkpoint = "./model/m4"
processor = CerebralJEPAImageProcessor.from_pretrained(checkpoint)
model = CerebralJEPAForImageEmbedding.from_pretrained(checkpoint).eval()

inputs = processor("/path/to/smri.nii.gz", return_tensors="pt")
outputs = model(**inputs)
embedding = outputs.embedding  # [1, 768]

Multi-Stage Features

m4 can return all available stage embeddings:

result = pipe(
    "/path/to/smri.nii.gz",
    stages="all",
    return_tensors="pt",
)

stage1 = result["embeddings"]["stage1"]  # [48]
stage2 = result["embeddings"]["stage2"]  # [96]
stage3 = result["embeddings"]["stage3"]  # [192]
stage4 = result["embeddings"]["stage4"]  # [768]

m1, m2, and m3 expose stage4 only.

Feature Maps

To return raw SwinTransformer feature maps before the 3D convolution refinement blocks:

result = pipe(
    "/path/to/smri.nii.gz",
    swin_features_only=True,
    return_tensors="pt",
)

swin0 = result["swin_features"]["swintransformer0"]
swin1 = result["swin_features"]["swintransformer1"]
swin2 = result["swin_features"]["swintransformer2"]
swin4 = result["swin_features"]["swintransformer4"]

From the command line:

conda run -n BrainST python scripts/embed_smri.py \
  /path/to/smri.nii.gz \
  --model m4 \
  --swin-features-only \
  --output /path/to/swin_features.npz

To save refined encoder feature maps:

conda run -n BrainST python scripts/embed_smri.py \
  /path/to/smri.nii.gz \
  --model m4 \
  --all-stages \
  --refined-features-only \
  --output /path/to/refined_features.npz

Projected Embeddings

The default output is the unprojected encoder embedding. The checkpoint also contains a 256-dimensional MLP projection head for compatibility with the original training objective.

outputs = model(**inputs, output_projected_embeddings=True)
embedding_768 = outputs.embedding
embedding_256 = outputs.projected_embedding

From the command line, projected embeddings are written with JSON output:

conda run -n BrainST python scripts/embed_smri.py \
  /path/to/smri.nii.gz \
  --model m4 \
  --include-projected \
  --output /path/to/embedding.json

Segmentation Preprocessing

For downstream segmentation workflows, CerebralJEPA also provides an intensity-only P2 preprocessor that preserves image geometry. It performs foreground estimation, foreground 0.5/99.5 percentile clipping, foreground z-score normalization, and background fill without orientation, spacing, crop/pad, or resize operations.

from cerebraljepa_inference import CerebralJEPASegmentationPreprocessor

preprocessor = CerebralJEPASegmentationPreprocessor()
sample = preprocessor(
    "/path/to/image.nii.gz",
    mask="/path/to/mask.nii.gz",
    return_tensors="pt",
)

image = sample["image"]  # [1, D, H, W]
mask = sample["mask"]    # [1, D, H, W]

Array-level use:

from cerebraljepa_inference import segmentation_p2_intensity_only

processed_image, foreground_mask = segmentation_p2_intensity_only(image_array)

Intended Use

CerebralJEPA is intended for research use in structural brain MRI representation learning and downstream model development. It is not a medical device and is not intended for clinical diagnosis, treatment selection, prognosis, triage, or any automated decision that affects patient care.

License

The code and model artifacts are released under the Apache License 2.0.