Guided-Chest-CT-LeJEPA

Source code on GitHub.

README

Guided Chest CT LeJEPA Pretraining

This repository contains the training and evaluation pipeline for LeJEPA (Latent-Euclidean Joint-Embedding Predictive Architecture) vision transformers pre-trained on the CT-RATE dataset. Three model variants are provided: a pure self-supervised baseline (LeJEPA-0), a single-source auxiliary variant with TotalSegmentator supervision (LeJEPA-1S), and a dual-source variant with both TotalSegmentator and ReXGroundingCT supervision (LeJEPA-2S).

---

Repository Structure

`` . ├── train_lejepa.py # LeJEPA pre-training entry point ├── train_gridsearch.py # Linear probing grid search ├── train_e2e_lora.py # LoRA fine-tuning ├── eval_rad.py # RAD-ChestCT transfer evaluation ├── test_ct_dataloader.py # Dataloader debugging ├── requirements.txt ├── Dockerfile │ ├── lejepa_core/ # Core LeJEPA architecture │ ├── lejepa_ssl_arch.py # SSLMetaArch training module │ ├── main_lejepa_trainer.py # Distributed training orchestration │ └── SIGReg.py # Sketched Isotropic Gaussian Regularization │ ├── models/ # Model definitions │ ├── vision_transformer.py # ViT backbone builders │ ├── lejepa_projector.py # LeJEPA projection head │ ├── colipri_pooling.py # 5 pooling schemes for linear probing │ └── e2e_colipri.py # End-to-end model wrapper │ ├── data/ # Data processing │ ├── guided_data_augmentation_CT_RATE.py # Multi-crop augmentation │ └── collate.py # Batch collation │ ├── dataloaders/ # Dataset loaders │ ├── datasetloader_web_ctrate.py # WebDataset loader │ ├── datasetloader_ctrate_multiscale.py # Multi-scale .npy loader │ ├── dataloader_embeddings.py # Pre-computed embedding loader │ ├── dataloader_linear.py # Linear probing loader │ └── dataloader_rad_embeddings.py # RAD-ChestCT loader │ ├── supervised_heads/ # Auxiliary supervision heads │ ├── organ_supervision.py # TotalSegmentator classification │ ├── soft_label_supervision.py │ └── example_supervised_head.py │ ├── scripts/ # Evaluation scripts │ ├── ctrate_generate_embeddings.py │ ├── knn.py │ ├── supervised_gap.py │ ├── supervised_3D.py │ ├── e2e_inference.py │ ├── evaluate_gap.py │ └── precompute_aggregations.py │ ├── utils/ # Shared utilities │ ├── config.py # OmegaConf configuration │ ├── dino_utils.py # Model init, LoRA, checkpointing │ ├── lejepa_schedule...