/medficientsam

Efficient Segment Anything in Medical Images

Primary LanguagePythonApache License 2.0Apache-2.0

MedficientSAM: A Robust Medical Segmentation Model with Optimized Inference Pipeline for Limited Clinical Settings

screen screen

Getting Started

git clone --recursive https://github.com/hieplpvip/medficientsam.git
cd medficientsam
conda env create -f environment.yaml -n medficientsam
conda activate medficientsam

Environment and Requirements

System Ubuntu 22.04.5 LTS
CPU AMD EPYC 7742 64-Core Processor
RAM 256GB
GPU (number and type) One NVIDIA A100 40G
CUDA version 12.0
Programming language Python 3.10
Deep learning framework torch 2.2.2, torchvision 0.17.2

Results

Accuracy metrics are evaluated on the public validation set of CVPR 2024 Segment Anything In Medical Images On Laptop Challenge. The computational metrics are obtained on an Intel(R) Core(TM) i9-10900K.

Model Resolution Params FLOPs DSC NSD 2D Run Time 3D Run Time 2D Memory Usage 2D Memory Usage
MedficientSAM-L0 512x512 34.79M 36.80G 85.85% 87.05% 0.9s 7.4s 448MB 687MB
MedficientSAM-L1 512x512 47.65M 51.05G 86.42% 87.95% 1.0s 9.0s 553MB 793MB
MedficientSAM-L2 512x512 61.33M 70.71G 86.08% 87.53% 1.1s 11.1s 663MB 903MB

Reproducibility

The Docker images can be found here.

docker load -i seno.tar.gz
docker container run -m 8G --name seno --rm -v $PWD/test_input/:/workspace/inputs/ -v $PWD/test_output/:/workspace/outputs/ seno:latest /bin/bash -c "sh predict.sh"

To measure the running time (including Docker starting time), see https://github.com/bowang-lab/MedSAM/blob/LiteMedSAM/CVPR24_time_eval.py

Data Preparation

You need to participate in the challenge to access the dataset. After downloading it, copy .env.example to .env and modify CVPR2024_MEDSAM_DATA_DIR to correct path.

The directory structure should look like this:

CVPR24-MedSAMLaptopData
├── train_npz
│   ├── CT
│   ├── Dermoscopy
│   ├── Endoscopy
│   ├── Fundus
│   ├── Mammography
│   ├── Microscopy
│   ├── MR
│   ├── OCT
│   ├── PET
│   ├── US
│   └── XRay
├── validation-box
    └── imgs

Training

See train_scripts.

Inference

You can download the weights here.

To run MedficientSAM-L1 on the validation set:

python src/infer.py experiment=infer_finetuned_l1

See configs/experiment/infer_* for running other model variants.

Build Docker image

Export model

python src/export_onnx.py experiment=export_finetuned_l0_cpp output_dir=weights/finetuned-l1-augmented/e2_cpp

Build

docker build -f Dockerfile.cpp -t seno.fat .
slim build --target seno.fat --tag seno --http-probe=false --include-workdir --mount $PWD/test_input/:/workspace/inputs/ --mount $PWD/test_output/:/workspace/outputs/ --exec "sh predict.sh"
docker save seno | gzip -c > seno.tar.gz

References