DART-Seg

Text-prompted 3D CT segmentation. Name a structure, get its mask. 286 anatomical structures and lesions, 39M parameters, trained from scratch.

Usage

⚠️ ONLY PROMPT FOR THE STRUCTURES YOU NEED. Prompting for unnecessary or anatomically absent structures may cause hallucinations and false-positive segmentations.

pip install torch nibabel numpy safetensors huggingface_hub
hf download Koushik45048545309/DART-Seg --local-dir DART-Seg
cd DART-Seg
python dart_seg.py --image ct.nii.gz --prompts liver "left kidney" spleen --out preds/
python dart_seg.py --list        # supported prompts

After inference, DART-Seg may produce fragments. Use the below postprocessing to fix that.

# Optional: remove disconnected components of 3,000 voxels or fewer
pip install scipy

python - <<'PY'
from pathlib import Path
import nibabel as nib
import numpy as np
from scipy import ndimage

input_dir = Path("preds")
output_dir = Path("preds_cleaned")
output_dir.mkdir(parents=True, exist_ok=True)

MIN_VOXELS = 3000
structure = ndimage.generate_binary_structure(3, 3)  # 26-connectivity

files = list(input_dir.glob("*.nii.gz"))
if not files:
    raise SystemExit(f"No .nii.gz masks found in {input_dir.resolve()}")

for path in files:
    img = nib.load(str(path))
    data = np.asarray(img.dataobj)
    mask = data > 0

    components, count = ndimage.label(mask, structure=structure)
    sizes = np.bincount(components.ravel())
    keep = np.flatnonzero(sizes > MIN_VOXELS)
    keep = keep[keep != 0]
    cleaned = np.isin(components, keep).astype(np.uint8)

    header = img.header.copy()
    header.set_data_dtype(np.uint8)
    nib.save(
        nib.Nifti1Image(cleaned, img.affine, header),
        str(output_dir / path.name),
    )

    print(f"{path.name}: {int(mask.sum()):,} -> {int(cleaned.sum()):,} voxels")

print(f"Cleaned masks saved to: {output_dir.resolve()}")
PY

Note: Use hf download rather than git clone. The weights are stored with Xet, so a plain git clone without git-xet installed gives you placeholder files and fails with header too large. On older huggingface_hub versions, use huggingface-cli download instead of hf download.

Quote prompts and paths that contain spaces, e.g. --prompts "left kidney" or --image "my scan.nii.gz".

Python API (run from inside the downloaded folder so dart_seg.py is importable):

from dart_seg import DartSeg
seg = DartSeg.from_pretrained("Koushik45048545309/DART-Seg")
masks, affine = seg.predict_file("ct.nii.gz", ["liver", "pancreas"])

Model

• Input: CT (NIfTI), resampled to 2 mm; 64³ sliding window.
• Residual 3D U-Net; each prompt's text embedding attends to image features and produces its mask. All prompts share one image pass.
• Prompts are precomputed Qwen3-Embedding-4B vectors (text_bank.npz), so no language model is needed at inference. Only prompts listed in prompts.txt are supported.
• Runs on CPU or GPU. A full CT on CPU can take several minutes.

Limitations

Research use only, not for clinical decisions. Structures that are small at 2 mm (teeth, inner ear, small vessels) are less accurate.

License

CC BY-NC-SA 4.0.