# OrganLens [![arXiv](https://img.shields.io/badge/arXiv-2607.25164-b31b1b.svg)](https://arxiv.org/abs/2607.25164) [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97-Hugging%20Face-yellow)](https://huggingface.co/gezx1004/OrganLens) This repository is the official implementation of [**OrganLens: Organ-Specific Representation Learning for CT Foundation Models**](https://arxiv.org/abs/2607.25164). OrganLens learns organ-conditioned CT representations and provides one public workflow from raw NIfTI volumes to downstream disease metrics: ```text NIfTI + optional segmentation ↓ header-aware preprocessing and slice/organ pairs ↓ OrganLens backbone training ↓ teacher embedding extraction ↓ MLP training or bundled-MLP evaluation ↓ per-disease AUROC / AUPRC / accuracy ``` The released inference path uses the all-slice, soft-mask-area-weighted pooling method. Ablation-only methods and internal experiment utilities are intentionally excluded. ## Installation Python 3.9 or newer is required. An editable installation is recommended for training: ```bash cd OrganLens pip install -e . ``` For an inference-only environment: ```bash pip install . ``` Tested dependency versions are pinned in `requirement.txt`. `xformers` is optional and excluded because its binary must match PyTorch, Python, and CUDA. ## 1. Preprocess CT volumes For CT-RATE, OrganLens follows the preprocessing convention from the [official CT-CLIP repository](https://github.com/ibrahimethemhamamci/CT-CLIP), while making volume orientation explicit from the NIfTI header. Preprocessing starts directly from the released `.nii.gz` volumes. ### One NIfTI ```bash organlens-preprocess single \ --image /path/to/ct.nii.gz \ --mask /path/to/totalsegmentator.nii.gz \ --output-root /path/to/preprocessed \ --volume-id case_001 ``` `--mask` is optional. The default public representation is a `256³` float32 CT at approximately 2 mm, normalized to `[-1, 1]`, with a valid LPS affine. Orientation comes from the NIfTI header; organ position is never used to guess a flip. ### A dataset Prepare a manifest: ```csv volume_id,image_path,mask_path case_001,images/case_001.nii.gz,masks/case_001.nii.gz case_002,images/case_002.nii.gz,masks/case_002.nii.gz ``` Then run: ```bash organlens-preprocess batch \ --manifest /path/to/cases.csv \ --output-root /path/to/preprocessed \ --workers 8 ``` The output contains volumes, three slice views, per-volume audit JSON, and a training-ready `slices.csv`: ```text preprocessed/ volumes/images/.nii.gz volumes/masks/.nii.gz slices/images/{axial,coronal,sagittal}//*.npy slices/masks/{axial,coronal,sagittal}//*.npy metadata/.json slices.csv ``` See [PREPROCESSING.md](PREPROCESSING.md) for the exact interpolation, affine, slice-direction, and legacy-compatibility rules. ### Internal data-preparation APIs These operations are deliberately library functions, not extra commands: - TotalSegmentator label merging is implemented by `merge_totalsegmentator_labels` in `organlens/preprocessing.py`. - Existing slice trees can be expanded into one row per visible `(slice, organ)` pair by `generate_slice_organ_csv` in `organlens/data.py`. - Normal preprocessing already applies the merge and writes `slices.csv`; most users do not need to call either helper manually. The merge collapses bilateral kidneys, five lung lobes, three bowel labels, and the historical multipart muscle group. For example: ```python from organlens import merge_totalsegmentator_labels from organlens.data import generate_slice_organ_csv merged_mask = merge_totalsegmentator_labels(raw_totalsegmentator_array) summary = generate_slice_organ_csv( "/path/to/slices/images", "/path/to/slices/masks", "/path/to/train.csv", ) ``` ## 2. Train the OrganLens backbone The default, readable training recipe is `scripts/train_local_maskpred.sh`. Training is initialized from `pytorch_model.bin` in [gigaheart/gigaheart-pretrain](https://huggingface.co/gigaheart/gigaheart-pretrain). After downloading that checkpoint, provide three dataset/model paths: ```bash TRAIN_DATA=/path/to/preprocessed/slices.csv \ DATA_ROOT=/path/to/preprocessed \ PRETRAINED_WEIGHTS=/path/to/pytorch_model.bin \ bash scripts/train_local_maskpred.sh ``` Four GPUs: ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 \ N_GPUS=4 \ TRAIN_DATA=/path/to/preprocessed/slices.csv \ DATA_ROOT=/path/to/preprocessed \ PRETRAINED_WEIGHTS=/path/to/pytorch_model.bin \ OUTPUT_DIR=/path/to/output \ bash scripts/train_local_maskpred.sh ``` The defaults reproduce the released objective: organ-aware crops, 11 organ conditions, DINO/iBOT learning, and BCE + Dice mask prediction. Useful optional overrides are `BATCH_SIZE_PER_GPU`, `BASE_LR`, `EPOCHS`, `NUM_WORKERS`, `WARMUP_EPOCHS`, `OUTPUT_DIR`, `REPORT_TO`, and `MAX_ITER`. `MAX_ITER=1` is useful for a smoke test. Reporting defaults to TensorBoard; `REPORT_TO=wandb` requires a separate `wandb` installation. ## Released checkpoints The OrganLens release uses three files: [Download the released checkpoints from Hugging Face](https://huggingface.co/gezx1004/OrganLens): ```bash hf download gezx1004/OrganLens --local-dir ./organlens_checkpoints ``` | File | Size | Purpose | | --- | ---: | --- | | `teacher_checkpoint.pth` | 1.3 GB | Organ-conditioned backbone and mask decoder used to extract volume embeddings | | `organlens_ctrate_heads_524999.pt` | 397 MB | Bundled MLPs for 11 organ views and 18 CT-RATE diseases | | `organlens_radchest_heads_524999.pt` | 507 MB | Bundled MLPs for 11 organ views and 23 RAD-ChestCT diseases | The two head bundles do not duplicate the teacher weights. Raw-volume prediction requires `teacher_checkpoint.pth` plus the bundle for the target dataset. Embedding extraction alone requires only the teacher checkpoint. These are distinct from GigaHeart's `pytorch_model.bin`: that file initializes new backbone training, whereas `teacher_checkpoint.pth` is the resulting OrganLens model used for inference. ## 3. Extract embeddings with a teacher checkpoint The released teacher checkpoint is named `teacher_checkpoint.pth`. Extract all 11 organ embeddings from one CT-RATE volume: ```bash organlens-infer \ --dataset ctrate \ --volume /path/to/volume.nii.gz \ --teacher-checkpoint /path/to/teacher_checkpoint.pth \ --embedding-output /path/to/embeddings/case_001.pt ``` Extract any subset, for example heart, lung, and aorta: ```bash organlens-infer \ --dataset ctrate \ --volume /path/to/volume.nii.gz \ --teacher-checkpoint /path/to/teacher_checkpoint.pth \ --organs heart lung aorta \ --embedding-output /path/to/embeddings/case_001.pt ``` Valid organ names are: ```text spleen kidneys liver stomach pancreas lung esophagus trachea intestine heart aorta ``` `--organs` accepts one or more names. Omitting it extracts all 11. For RAD-ChestCT, use `--dataset radchest`. Keep the default 64 sampled frames: the CLI automatically uses frame batches of 16 for CT-RATE and 64 for RAD-ChestCT to reproduce the released embedding caches. The equivalent Python API is: ```python from organlens.inference import OrganLensEmbedder teacher = OrganLensEmbedder( "/path/to/teacher_checkpoint.pth", device="cuda", amp="bf16", ) payload = teacher.embed_volume( "/path/to/volume.nii.gz", organs=["heart", "lung", "aorta"], dataset="ctrate", ) heart = payload["embeddings"]["heart"] # float32 [1024] ``` Name each output `.pt`; the downstream commands then match labels and embeddings automatically. ## 4. Train MLP heads and compute metrics Label CSVs must contain `VolumeName` followed by binary disease columns: ```csv VolumeName,Cardiomegaly,Pericardial effusion,... case_001.nii.gz,1,0,... case_002.nii.gz,0,0,... ``` ### CT-RATE ```bash organlens-train-ctrate \ --embedding-root /path/to/embeddings \ --train-label-csv /path/to/train_labels.csv \ --valid-label-csv /path/to/valid_labels.csv \ --organ all \ --output-root /path/to/ctrate_heads ``` Defaults: independent `1024 → 512 → 1` MLPs, AdamW, learning rate `2e-3`, weight decay `0.05`, five epochs, and one warmup epoch. Every organ directory contains: ```text multitask_classifier.pt valid_metrics_by_disease.csv final_metrics.json ``` With `--organ all`, training also writes the portable `organlens_ctrate_heads.pt` bundle. ### RAD-ChestCT ```bash organlens-train-radchest \ --embedding-root /path/to/embeddings \ --train-label-csv /path/to/train_labels.csv \ --valid-label-csv /path/to/valid_labels.csv \ --test-label-csv /path/to/test_labels.csv \ --organ all \ --output-root /path/to/radchest_heads ``` The released RAD-ChestCT protocol runs 100 epochs. Each disease independently selects its best validation-AUROC epoch; the assembled classifier is evaluated on the test split once. The test split is never used for selection. ## 5. Evaluate an existing bundled MLP on embeddings This path does not load the 1.3 GB teacher or reprocess CT. It uses saved embeddings plus a released or locally trained bundle: ```bash organlens-evaluate \ --label-csv /path/to/test_labels.csv \ --embedding-root /path/to/embeddings \ --head-bundle /path/to/organlens_ctrate_heads_524999.pt \ --view matched \ --output-dir /path/to/evaluation ``` Outputs: ```text evaluation/ metrics_by_disease.csv final_metrics.json ``` `--view matched` uses the paper's disease-to-organ mapping, such as heart for Cardiomegaly, lung for Pleural effusion, aorta for arterial-wall calcification, and esophagus for hiatal hernia. Pass `--view heart`, for example, to evaluate every disease from one organ embedding. Empty, `-1`, and non-binary labels are skipped disease-by-disease. `--max-cases N` is available for smoke tests. ## 6. Benchmark raw CT-RATE volumes on four GPUs This entry point measures all 18 CT-RATE diseases separately for any requested subset of the 11 organs: ```bash organlens-benchmark-ctrate run \ --label-csv /path/to/valid_predicted_labels.csv \ --volume-root /path/to/CTRATE/dataset/valid_fixed \ --teacher-checkpoint /path/to/teacher_checkpoint.pth \ --head-bundle /path/to/organlens_ctrate_heads_524999.pt \ --organs heart lung pancreas \ --gpus 0 1 2 3 \ --output-dir /path/to/benchmark ``` Cases are sharded across GPUs, rather than organs. Each worker reads and preprocesses a volume once, constructs its 64 frames once, and then extracts all requested organ embeddings. Omitting `--organs` selects all 11. Completed cases are written atomically and skipped on rerun, so the same command resumes an interrupted benchmark. For `N` selected organs, outputs include the full `N × 18` table in `metrics_by_organ_disease.csv`, per-organ means in `metrics_by_organ.csv`, `auroc_matrix.csv`, per-case probabilities, failure records, and `final_metrics.json`. Use `--max-cases 4` for a four-GPU smoke test. ## Direct prediction for one volume Embedding extraction and bundled-head prediction can be combined: ```bash organlens-infer \ --dataset ctrate \ --volume /path/to/volume.nii.gz \ --teacher-checkpoint /path/to/teacher_checkpoint.pth \ --embedding-output /path/to/case_001.pt \ --head-bundle /path/to/organlens_ctrate_heads_524999.pt \ --prediction-output /path/to/case_001_predictions.json ``` For full checkpoint formats, exact all-slice pooling, and CT-RATE/RAD-ChestCT compatibility, see [INFERENCE.md](INFERENCE.md). ## Commands ```text organlens-preprocess organlens-infer organlens-train-ctrate organlens-train-radchest organlens-evaluate organlens-benchmark-ctrate ``` Equivalent scripts are retained under `scripts/`. ## Validation Run unit/integration tests: ```bash python -m unittest discover -s tests -v ``` Build a wheel: ```bash python -m pip wheel --no-deps . --wheel-dir dist ``` The wheel includes the modified DINOv2 implementation and the OrganLens ViT-L configuration required for inference. ## Citation If you find OrganLens useful, please cite: ```bibtex @misc{ge2026organlensorganspecificrepresentationlearning, title = {OrganLens: Organ-Specific Representation Learning for CT Foundation Models}, author = {Zhixuan Ge and Anqi Li and Sadeer Al-Kindi and Hanwen Xu and Wei Qiu}, year = {2026}, eprint = {2607.25164}, archivePrefix = {arXiv}, primaryClass = {cs.CV}, url = {https://arxiv.org/abs/2607.25164} } ``` ## Acknowledgment OrganLens builds on [DINOv2](https://github.com/facebookresearch/dinov2) from Meta AI Research. Upstream-derived files retain their original copyright headers. This repository is distributed under the Apache License 2.0; see [LICENSE](LICENSE).