ProtoSAM Segmentation Demo

This Gradio 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 app with:

./run_demo.sh

Or run it directly with:

python app.py

This will start the server and provide a link to access the demo in your browser.

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)
Configure the model parameters using the checkboxes
Click "Run Inference" to generate the segmentation result

Model Configuration

The app allows you to configure several model parameters:

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