Report-grounded segmentation of chest-CT findings

A 3D segmentation network that takes a chest CT and outputs, for each of 14 radiological finding categories, a voxel mask.

The method is a training-time one. Every finding in the ReXGroundingCT training set comes with the report sentence that describes it; during training each finding's pooled image region is aligned to the embedding of its own sentence. At inference there is no text: the alignment projector is not part of the deployed network, there is no prompt input, and one forward pass produces all 14 channels at once.

total loss  =  L_seg  +  lambda * mean_f ( 1 - cos( W p_f , e_f ) )
                                    ^ pooled region feature   ^ frozen sentence embedding

lambda = 0 reduces the code exactly to the vision-only baseline, which is how every comparison in this repo was run.

Results

ReXGroundingCT validation split, 200 cases / 341 findings, micro-averaged over ground-truth-present classes, 128x224x288 window, threshold 0.5, no TTA:

Dice	mIoU	NSD@2mm	Hit@0.1	inst. Prec.	inst. Rec.	inst. F1	HD95 (mm)
30.93	20.76	38.24	74.49	17.35	36.75	23.57	121.64

Layout

Path	What it is
train/	The method. Complete training-time code — the alignment loss, the instance-masked pooling, the dataloader that keeps L_seg byte-identical to the baseline, and the two offline asset builders. Start at train/README.md.
inference/	The released model. Self-contained preprocessing + inference + evaluation, verified end-to-end to reproduce the published per-case Dice from raw NIfTI. Start at inference/README.md and inference/PREPROCESSING.md.
figures/	Method figures (training route, label granularity, inference route) and the script that draws them.

The two halves are independent: you can run inference/ without ever touching train/.

Weights

The checkpoint (epoch 486, 107.5 M parameters, ~410 MiB) is not in git — it is attached to the latest release as ours_rexgroundingct.pth.

mkdir -p inference/weights
curl -L -o inference/weights/ours_rexgroundingct.pth \
  https://github.com/zhenyubu/26summer/releases/latest/download/ours_rexgroundingct.pth

Quick start (inference)

conda create -n ours python=3.11 && conda activate ours
pip install torch --index-url https://download.pytorch.org/whl/cu121   # match your CUDA
pip install -r inference/requirements.txt

cd inference
python smoke_test.py                                  # ~1 min, CPU, no data needed
python predict.py --in-dir /data/cts --out-dir preds  # NIfTI in, 14 masks out

Inputs are NIfTI in Hounsfield units. Measured on one A100-40GB with a 382x502x502 volume: forward pass 8.9 s at 4.5 GiB peak — an 8 GB card is enough. Preprocessing (60-80 s on CPU) is the bottleneck, and host RAM matters more than VRAM: the probability accumulator is one float32 volume per requested class, so all 14 classes on a scan that size costs ~5 GB. inference/PREPROCESSING.md covers the input contract and the failure modes that are silent rather than loud.

Notes for anyone re-running the training route

• Paths default to the development tree; override with CT_ROOT=/your/path, or edit ROOT at the top of train.py / build_inst.py / build_emb_tables.py.
• Use --no-lr-mirror: report sentences state laterality, so left-right mirroring contradicts the text target.
• Data is ReXGroundingCT; the preprocessed Dataset508 volumes and the _inst.b2nd sidecars are built locally by build_inst.py and are not redistributed here.