Official implementation of "UniverSeg: Universal Medical Image Segmentation" accepted at ICCV 2023.
Victor Ion Butoi*,
Jose Javier Gonzalez Ortiz*,
Tianyu Ma,
Mert R. Sabuncu,
John Guttag,
Adrian V. Dalca,
*denotes equal contribution
Given a new segmentation task (e.g. new biomedical domain, new image type, new region of interest, etc), most existing strategies involve training or fine-tuning a segmentation model that takes an image input and outputs the segmentation map.
This process works well in machine-learning labs, but is challenging in many applied settings, such as for scientists or clinical researchers who drive important scientific questions, but often lack the machine-learning expertiese and computational resources necessary.
UniverSeg enables users to tackle a new segmentation task without the need to train or fine-tune a model, removing the requirement for ML experience and computational burden. The key idea is to have a single global model which adapts to a new segmentation task at inference based on an input example set.
The universeg architecture is described in the model.py file.
We provide model weights a part of our release.
To instantiate the UniverSeg model (and optionally use provided weights):
from universeg import universeg
model = universeg(pretrained=True)
# To perform a prediction (where B=batch, S=support, H=height, W=width)
prediction = model(
target_image, # (B, 1, H, W)
support_images, # (B, S, 1, H, W)
support_labels, # (B, S, 1, H, W)
) # -> (B, 1, H, W)For all inputs ensure that pixel values are min-max normalized to the
We provide a jupyter notebook with a tutorial and examples of how to do inference using UniverSeg: Google colab | Nbviewer.
You can install universeg in two ways:
- With pip:
pip install git+https://github.com/JJGO/UniverSeg.git- Manually: Cloning it and installing dependencies
git clone https://github.com/JJGO/UniverSeg
python -m pip install -r ./UniverSeg/requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./UniverSeg)"If you find our work or any of our materials useful, please cite our paper:
@article{butoi2023universeg,
title={UniverSeg: Universal Medical Image Segmentation},
author={Victor Ion Butoi* and Jose Javier Gonzalez Ortiz* and Tianyu Ma and Mert R. Sabuncu and John Guttag and Adrian V. Dalca},
journal={International Conference on Computer Vision},
year={2023}
}
- Code is released under Apache 2.0 license
- Model Weights are released under OpenRAIL++-M license. According to the usage restriction the model must be only used for research purposes.