metadata

title: LoGoSAM_demo
app_file: app.py
sdk: gradio
sdk_version: 5.29.0

ProtoSAM - One shot segmentation with foundational models

Link to our paper here.
This work is the successor of DINOv2-based-Self-Supervised-Learning (Link to Paper).

Demo Application

A Gradio-based demo application is now available for interactive inference with ProtoSAM. You can upload your own images and masks to test the model. See README_DEMO.md for instructions on running the demo.

Abstract

This work introduces a new framework, ProtoSAM, for one-shot image segmentation. It combines DINOv2, a vision transformer that extracts features from images, with an Adaptive Local Prototype Pooling (ALP) layer, which generates prototypes from a support image and its mask. These prototypes are used to create an initial coarse segmentation mask by comparing the query image's features with the prototypes. Following the extraction of an initial mask, we use numerical methods to generate prompts, such as points and bounding boxes, which are then input into the Segment Anything Model (SAM), a prompt-based segmentation model trained on natural images. This allows segmenting new classes automatically and effectively without the need for additional training.

How To Run

1. Data preprocessing

1.1 CT and MRI Dataset

Please see the notebook data/data_processing.ipynb for instructions. For convenience i've compiled the data processing instructions from https://github.com/cheng-01037/Self-supervised-Fewshot-Medical-Image-Segmentation to a single notebook.
The CT dataset is available here: https://www.synapse.org/Synapse:syn3553734
The MRI dataset is availabel here: https://chaos.grand-challenge.org

run ./data/CHAOST2/dcm_img_to_nii.sh to convert dicom images to nifti files.

1.2 Polyp Dataset

Data is available here: https://www.kaggle.com/datasets/hngphmv/polypdataset?select=train.csv

Put the dataset data/PolypDataset/

2. Running

2.1 (Optional) Training and Validation of the coarse segmentation networks

./backbone.sh [MODE] [MODALITY] [LABEL_SET]

MODE - validation or training
MODALITY - ct or mri
LABEL_SET - 0 (kidneys), 1 (liver spleen)

for example:

./backbone.sh training mri 1

Please refer to backbone.sh for further configurations.

2.1 Running ProtoSAM

Put all SAM checkpoint like sam_vit_b.pth, sam_vit_h.pth, medsam_vit_b.pth into the pretrained_model directory.
Checkpoints are available at SAM and MedSAM.

./run_protosam.sh [MODALITY] [LABEL_SET]

MODALITY - ct, mri or polyp
LABEL_SET (only relevant if doing ct or mri) - 0 (kidneys), 1 (liver spleen) Please refer to the run_protosam.sh script for further configurations.

Acknowledgements

This work is largely based on ALPNet, DINOv2, SAM and is a continuation of DINOv2-based-Self-Supervised-Learning.

Cite

If you found this repo useful, please consider giving us a citation and a star!

@article{ayzenberg2024protosam,
  title={ProtoSAM-One Shot Medical Image Segmentation With Foundational Models},
  author={Ayzenberg, Lev and Giryes, Raja and Greenspan, Hayit},
  journal={arXiv preprint arXiv:2407.07042},
  year={2024}
}

@misc{ayzenberg2024dinov2,
      title={DINOv2 based Self Supervised Learning For Few Shot Medical Image Segmentation}, 
      author={Lev Ayzenberg and Raja Giryes and Hayit Greenspan},
      year={2024},
      eprint={2403.03273},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

ProtoSAM Segmentation Demo

This Streamlit application demonstrates the capabilities of the ProtoSAM model for few-shot segmentation. Users can upload a query image, support image, and support mask to generate a segmentation prediction.

Requirements

Python 3.8 or higher
CUDA-compatible GPU
Required Python packages (see requirements.txt)

Setup Instructions

Clone this repository:

git clone <your-repository-url>
cd <repository-name>

Create and activate a virtual environment (optional but recommended):

python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

Install the required dependencies:

pip install -r requirements.txt

Download the pretrained models:

mkdir -p pretrained_model
# Download SAM ViT-H model
wget -P pretrained_model https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
mv pretrained_model/sam_vit_h_4b8939.pth pretrained_model/sam_vit_h.pth

Update the model path in app.py:
- Set the reload_model_path in the config dictionary to the path of your trained ProtoSAM model.

Running the App

Start the Streamlit app with:

streamlit run app.py

This will open a browser window with the interface for the segmentation demo.

Usage

Upload a query image (the image you want to segment)
Upload a support image (an example image with a similar object)
Upload a support mask (the segmentation mask for the support image)
Use the sidebar to configure the model parameters if needed
Click "Run Inference" to generate the segmentation result

Model Configuration

The app allows you to configure several model parameters via the sidebar:

Use Bounding Box: Enable/disable bounding box input
Use Points: Enable/disable point input
Use Mask: Enable/disable mask input
Use CCA: Enable/disable Connected Component Analysis
Coarse Prediction Only: Use only the coarse segmentation model without SAM refinement

Notes

This demo requires a GPU with CUDA support
Large images may require more GPU memory
For optimal results, use high-quality support images and masks