--- pipeline_tag: image-segmentation license: apache-2.0 --- # Medal S: Spatio-Textual Prompt Model for Medical Segmentation [![Paper](https://img.shields.io/badge/Paper-Arxiv-b31b1b.svg)](https://arxiv.org/abs/2511.13001) [![OpenReview](https://img.shields.io/badge/OpenReview-Discussion-4CAF50.svg)](https://openreview.net/forum?id=9vCx66pnLn#discussion) [![HuggingFace](https://img.shields.io/badge/🤗%20HuggingFace-Model-yellow.svg)](https://huggingface.co/spc819/Medal-S-V1.0) [![GitHub](https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github)](https://github.com/yinghemedical/Medal-S) This repository provides guidance for training and inference of Medal S within the [CVPR 2025: Foundation Models for Text-Guided 3D biomedical image segmentation](https://www.codabench.org/competitions/5651/) Docker link for the 2025/05/30 testing submission: [Medal S](https://drive.google.com/file/d/1HRJqYUXajptGsKaXEhn-s3rGcnKIwGs7/view) ## Requirements **Python Version:** 3.10.16 ### Installation 1. **Create conda environment and install dependencies:** ```bash #!/bin/bash # Create environment conda create -n medals_local_test python=3.10 -y conda activate medals_local_test # Install packages pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu121 pip install transformers==4.51.3 monai==1.4.0 nibabel==5.3.2 tensorboard einops positional_encodings scipy pandas scikit-learn scikit-image batchgenerators acvl_utils # Install nnU-Net wget https://github.com/MIC-DKFZ/nnUNet/archive/refs/tags/v2.4.1.tar.gz tar -xvf v2.4.1.tar.gz pip install -e nnUNet-2.4.1 # Install dynamic network architectures cd model && pip install -e dynamic-network-architectures-main && cd .. ``` ## Training Guidance First, download the dataset from [Hugging Face: junma/CVPR-BiomedSegFM](https://huggingface.co/datasets/junma/CVPR-BiomedSegFM). * **Data Preparation**: Preprocess and organize all training data into a `train_all.jsonl` file using the provided script: `data/challenge_data/get_train_jsonl.py`. * **Knowledge Enhancement**: You can either use the pre-trained text encoder from SAT ([https://github.com/zhaoziheng/SAT/tree/cvpr2025challenge](https://github.com/zhaoziheng/SAT/tree/cvpr2025challenge)) available on [Hugging Face](https://huggingface.co/zzh99/SAT/tree/main/Pretrain), or pre-train it yourself following the guidance in this [repository](https://github.com/zhaoziheng/SAT-Pretrain/tree/master). As recommended by SAT, we **freeze** the text encoder when training the segmentation model. * **Segmentation**: The training script is located at `sh/cvpr2025_Blosc2_pretrain_1.0_1.0_1.0_UNET_ps192.sh`. Before training, NPZ files will be converted to the Blosc2 compressed format (from the nnU-Net framework). Training takes approximately 7 days with 2x H100-80GB GPUs for a 224x224x128 (1.5, 1.5, 3.0) spacing model, using a batch size of 2 per GPU. For a 192x192x192 (1.0, 1.0, 1.0) spacing model, it requires 4x H100-80GB GPUs with a batch size of 2 per GPU. You may modify the patch size and batch size to train on GPUs with less memory. ## Inference Guidance ### Quick Start The inference pipeline takes raw NIfTI images and free-form text prompts as input, and outputs segmentation predictions. All preprocessing (spacing conversion, normalization, resampling) is handled automatically. **Basic Usage:** ```bash # Using shell script (recommended) bash run_inference_medals_nifti.sh # Or using Python directly python inference_medals_nifti.py \ --input input.nii.gz \ --output output.nii.gz \ --config config_CT.json \ --mode stage1+stage2 \ --device cuda:0 \ --checkpoints ./checkpoints ``` ### Configuration Files The inference uses JSON configuration files to specify text prompts and modality-specific settings: **For CT images** (`config_CT.json`): - Supports multiple window types: `soft_tissue`, `bone`, `lung` - Each window type has its own text prompts and window settings - Multiple window types are processed separately and merged automatically ```json { "texts_soft_tissue": ["Aorta in whole body CT", "Liver in whole body CT"], "texts_bone": ["Vertebrae C1 in whole body CT"], "texts_lung": ["Left lung in whole body CT"], "window_settings": { "soft_tissue": {"window_level": 40, "window_width": 400}, "bone": {"window_level": 500, "window_width": 1500}, "lung": {"window_level": -600, "window_width": 1500} }, "modality": "CT" } ``` **For non-CT images** (`config_nonCT.json`): - Uses percentile-based normalization - Supports MRI, US, PET, microscopy ```json { "texts": ["Spleen in MRI"], "normalization_settings": { "percentile_lower": 0.5, "percentile_upper": 99.5, "preserve_zero": true }, "modality": "MRI" } ``` ### Input/Output **Input:** - Raw NIfTI image file (`.nii` or `.nii.gz`) with spacing metadata - Free-form text prompt(s) describing the anatomical structures or lesions to segment - Optional: Modality information (CT, MRI, US, PET, microscopy) **Output:** - Segmentation prediction as NIfTI file with same spacing/origin/direction as input - Label values correspond to the order of text prompts in the config **Example:** ```python # Input: raw CT image + text prompts Input: "chest_ct.nii.gz" Texts: ["Aorta in whole body CT", "Liver in whole body CT"] Labels: [1, 2] # Auto-generated if not specified # Output: segmentation mask Output: "chest_ct_stage1+stage2.nii.gz" # Voxels labeled as 1 = Aorta, 2 = Liver, 0 = Background ``` ### Core Inference Logic The inference pipeline automatically handles: 1. **Input Loading**: Reads raw NIfTI image with spacing metadata 2. **Preprocessing**: - CT: Window/level normalization (automatically selects appropriate window based on text prompts) - Non-CT: Percentile-based normalization - Spacing conversion and resampling to target resolution 3. **Text Encoding**: Encodes free-form text prompts using the knowledge encoder 4. **Segmentation**: Performs patch-based inference with optional two-stage refinement 5. **Post-processing**: Resamples prediction back to original image space ### Inference Modes - **`stage2_only`**: Single-stage high-resolution inference (faster) - **`stage1+stage2`**: Two-stage inference with ROI refinement (more accurate) ### Python API Example For programmatic usage, you can call the inference function directly: ```python from inference_medals_nifti import run_inference, load_config_from_json # Option 1: Using config file (recommended) config = load_config_from_json("config_CT.json") pred_array, inference_time = run_inference( image_path="input.nii.gz", output_path="output.nii.gz", modality=config["modality"], texts=config["texts"], label_values=[str(l) for l in config["labels"]], inference_mode="stage1+stage2", device="cuda:0", checkpoints_path="./checkpoints", window_settings=config.get("window_settings"), normalization_settings=config.get("normalization_settings"), window_type_mapping=config.get("window_type_mapping") ) # Option 2: Direct parameters pred_array, inference_time = run_inference( image_path="input.nii.gz", output_path="output.nii.gz", modality="CT", texts=["Aorta in whole body CT", "Liver in whole body CT"], label_values=["1", "2"], inference_mode="stage1+stage2", device="cuda:0", checkpoints_path="./checkpoints" ) ``` ### Important Notes - **Preprocessing**: All preprocessing steps (spacing, normalization, resampling) are handled internally. Do not preprocess images yourself to avoid discrepancies. - **Text Prompts**: Use descriptive, modality-specific prompts (e.g., "Aorta in whole body CT" rather than just "Aorta") - **Labels**: If not specified in config, labels are auto-generated as consecutive integers starting from 1 - **GPU Memory**: For large 3D images, adjust `category_batch_size` in the config to manage GPU memory usage ## Citation ``` @misc{shi2025medalsspatiotextualprompt, title={Medal S: Spatio-Textual Prompt Model for Medical Segmentation}, author={Pengcheng Shi and Jiawei Chen and Jiaqi Liu and Xinglin Zhang and Tao Chen and Lei Li}, year={2025}, eprint={2511.13001}, archivePrefix={arXiv}, primaryClass={cs.CV}, url={https://arxiv.org/abs/2511.13001}, } ``` ## Acknowledgements This project is significantly improved based on [nnU-Net](https://github.com/MIC-DKFZ/nnUNet/tree/master) and [SAT](https://github.com/zhaoziheng/SAT/tree/cvpr2025challenge). We extend our gratitude to both projects. Medal-S is developed and maintained by Medical Image Insights.