File size: 2,171 Bytes
427d150
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
# 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

1. Clone this repository:
```bash
git clone <your-repository-url>
cd <repository-name>
```

2. Create and activate a virtual environment (optional but recommended):
```bash
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate
```

3. Install the required dependencies:
```bash
pip install -r requirements.txt
```

4. Download the pretrained models:
```bash
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
```

5. 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:
```bash
./run_demo.sh
```

Or run it directly with:
```bash
python app.py
```

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

## Usage

1. Upload a query image (the image you want to segment)
2. Upload a support image (an example image with a similar object)
3. Upload a support mask (the segmentation mask for the support image)
4. Configure the model parameters using the checkboxes
5. 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