https://github.com/SteveImmanuel/rs-paint
[ICLR ML4RS 2025] Official implementation for the paper "Tackling Few-Shot Segmentation in Remote Sensing via Inpainting Diffusion Model"
https://github.com/SteveImmanuel/rs-paint
Last synced: 2 months ago
JSON representation
[ICLR ML4RS 2025] Official implementation for the paper "Tackling Few-Shot Segmentation in Remote Sensing via Inpainting Diffusion Model"
- Host: GitHub
- URL: https://github.com/SteveImmanuel/rs-paint
- Owner: SteveImmanuel
- License: apache-2.0
- Created: 2024-12-23T00:49:05.000Z (5 months ago)
- Default Branch: main
- Last Pushed: 2025-03-05T07:09:56.000Z (2 months ago)
- Last Synced: 2025-03-05T07:30:58.163Z (2 months ago)
- Language: Jupyter Notebook
- Homepage: https://rspaint.steveimm.id
- Size: 11.6 MB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
- Awesome-Segment-Anything - [code
README
# Tackling Few-Shot Segmentation in Remote Sensing via Inpainting Diffusion Model
**ICLR Machine Learning for Remote Sensing Workshop, 2025 (Oral)**
This repo contains the official code for training and generation for the paper "Tackling Few-Shot Segmentation in Remote Sensing via Inpainting Diffusion Model".
[](https://steveimmanuel.github.io/rs-paint)
[](https://www.apache.org/licenses/LICENSE-2.0)
[Steve Andreas Immanuel](https://steveimm.id) | [Woojin Cho](https://woojin-cho.github.io/) | [Junhyuk Heo](https://rokmc1250.github.io/) | [Darongsae Kwon](https://www.linkedin.com/in/darongsaekwon).
We introduce a image-conditioned diffusion-based approach for to create diverse set of novel-classes samples for semantic segmentation in few-shot settings in remote sensing domain. By ensuring semantic consistency using cosine similarity between the generated samples and the conditioning image, and using the Segment Anything Model (SAM) to obtain the precise segmentation, our method can train off-the-shelf segmentation models with high-quality synthetic data, significantly improving performance in low-data scenarios.
## Setup
Create a new conda environment and activate it with the following commands:
```
conda env create -f environment.yaml
conda activate rspaint
```
## Checkpoint Weights
All required checkpoints are available in [huggingface](https://huggingface.co/SteveImmanuel/RSPaint).
Download `sd_inpaint_samrs_ep74.ckpt` and `remoteclip.pt` and save them to the `checkpoints` directory. Optionally, if you also want to perform the mask refinement, download SAM checkpoint and also save it to the `checkpoints` directory.
## Generating Samples
>Note that the current version is not optimized. Generating a single sample with a resolution of 512x512 takes about 9GB of GPU memory. Using cpu is possible, but takes significantly longer.
Follow along the steps in the provided notebook [`notebooks/generate_samples.ipynb`](notebooks/generate_samples.ipynb) to generate samples using the trained model.
## Training
### Data preparing
To train the model, you can use any dataset with bounding box annotations.
The data structure must be like this:
```
.
├── bbox
│ ├── train
│ │ ├── 00001.txt
│ │ ├── 00003.txt
│ │ ├── 00004.txt
│ │ ├── ...
│ └── validation
│ │ ├── 00023.txt
│ │ ├── 00024.txt
│ │ ├── 00025.txt
│ │ ├── ...
└── images
├── train
│ │ ├── 00001.png
│ │ ├── 00003.png
│ │ ├── 00004.png
│ │ ├── ...
│ └── validation
│ │ ├── 00023.png
│ │ ├── 00024.png
│ │ ├── 00025.png
│ │ ├── ...
```
Update the `yaml` config file in the following lines:
```yaml
data:
target: main.DataModuleFromConfig
params:
batch_size: 8
wrap: False
train:
target: ldm.data.remote_sensing.RemoteSensingDataset
params:
state: train
dataset_dir:
arbitrary_mask_percent: 0.5
image_size: 512
version: openai/clip-vit-large-patch14
bbox_ratio_range: [0.1, 0.25]
validation:
target: ldm.data.remote_sensing.RemoteSensingDataset
params:
state: validation
dataset_dir:
arbitrary_mask_percent: 0.5
image_size: 512
version: openai/clip-vit-large-patch14
bbox_ratio_range: [0.1, 0.25]
```
Optionally, you can write custom dataloader and placing it in `ldm/data` directory. Then, change the `target` key to the custom dataloader.
To start training, use the following command:
```
python -u main.py \
--logdir --pretrained_model checkpoints/checkpoints/sd_inpaint_samrs_ep74.ckpt \
--base --scale_lr False --seed 20250110
```
## Citation
TBA
## Acknowledgements
This code is mainly based on [Paint by Example](https://github.com/Fantasy-Studio/Paint-by-Example). We thank the authors for their great work.