doem97 / gen_sar_plora

• 🧭 Introduction
• 📝 Install
  • 1. Build CLSAR
  • 2. Build GENSAR
  • 3. Build Inference API
  • 4. Build ControlNet
• 🌌 Datasets
  • Download and Install
• 📊 Reproduce Our Results
  • 1. SAR Example
  • 2. From Scratch
  • 3. Fine-tuning
  • 4. LoRA
  • 5. 2LoRA
  • 6. pLoRA

Code and datasets for the paper "Non-Visible Light Data Synthesis: a Case Study for Synthetic Aperture Radar Imagery".

Please note: more updates are pending.

🧭 Introduction

This code space contains three parts:

• ./clsar: Image CLassification pipeline for SAR (Document of CLSAR).
• ./gensar: Image GENeration pipeline for SAR (Document of GENSAR).
• ./datasets: SAR ship classification dataset FUSRS and ORS ship dataset DOSRS (direct link to Datasets).

Please follow the instructions below to build the environments for GENSAR and CLSAR respectively, and download the datasets. After that, you can reproduce our results with our provided scripts.

📝 Install

Due to package conflicts, we build two seperate conda environments for GENSAR and CLSAR, respectively. Please follow the instructions below to build the environments.

1. Build CLSAR

Please refer Document of CLSAR for details.

1. Create a conda environment and activate it:

   # Install proper version of torch
   conda create -n clsar python=3.8 pytorch=1.10 cudatoolkit=11.3 torchvision==0.11.0 -c pytorch -y
   conda activate clsar

2. Then install MMCV and compile CLSAR:

   # Install mmcv package management tools
   pip3 install openmim
   mim install mmcv-full
   # Compile CLSAR folder as edit-able package
   cd ./clsar
   pip3 install -e .

Please refer to MMClassification Install for more detailed installation.

2. Build GENSAR

GenSAR depends on ControlNet, the Huggingface Diffusers and the SDWebUI as inference API. Torch version is torch2.1.0+cu118.

Please refer Document of GENSAR for details.

1. Conda Env with Torch 2.1

   # Create a conda environment and activate it
   conda env create --name gensar python=3.9
   conda activate gensar
   # Install torch2.0.1+cu118
   conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

2. Install Huggingface Dependencies. We recommand an editable install.

   # Compile diffusers from source
   cd gensar
   git clone https://github.com/huggingface/diffusers.git
   cd diffusers
   pip install -e ".[torch]"
   # (optional) xFormers for saving VRAM
   pip install xformers

3. Build Inference API

We employ SDWebUI API for quick batch inference.

1. Install and Start API backend. We recommand an editable install.

   conda activate gensar
   # Install SDWebUI
   cd ./inference
   bash ./webui.sh

2. Call API. By default the API will listen on addr localhost:7860. You can call inference API the same in template gensar/inference/scripts/inference/ors_sar_dosrscontour.py.

   # gensar/inference/scripts/inference/ors_sar_dosrscontour.py
   class ControlnetRequest:
       def __init__(self, url, prompt, path):
           self.url = url
           ...
   
       def build_body(self):
           self.body = {
               "prompt": self.prompt,
               "negative_prompt": "",
               ...
               "alwayson_scripts": {
                   "controlnet": {
                       "args": [
                           {
                               "enabled": True,
                               ...
                           }
                       ]
                   }
               },
           }

4. Build ControlNet

1. Build ControlNet. The gensar/controlnet branch is the same as in original build. Due to torch version conflict, we set a standalone environment for ControlNet. Note in our implementation, fine-tuning controlnet is not necessary. Directly employ pre-trained contour/segmentation ControlNet achieves same performance.

   cd ./controlnet
   conda env create -f environment.yaml
   conda activate ctrlnet

2. (Optional) Train. We provide training script in gensar/controlnet/scripts/train.sh:

   # Dataset is defined in gensar/controlnet/dataset_xxx.py
   python gensar/controlnet/scripts/train.py --batch_size 3 --gpus 4 --dataset fusrs_cam_v2 --ckpt_path ./checkpoints/xxx/

3. Convert ControlNet for Inference. To convert ControlNet resulting .ckpt (pyTorch Lightning checkpoint format) to .pth and deploy in Huggingface, we provide converting script in gensar/controlnet/convert_ckpt_pth.py. Please symbolink or save converted ControlNet weights in gensar/inference/extensions/sd-webui-controlnet/models.

🌌 Datasets

Download and Install

We constructed two datasets, the FUSRS for SAR ship classification and the DOSRS for training ORS LoRA. Download links are provided in Google Drive. The datasets should be put in ./datasets folder. The final folder structure should be like:

|-- clsar # image classification framework for SAR.
|-- datasets # our proposed datasets
|   |-- dosrs_v1 # ORS ship dataset
|   |-- fusrs_v2 # SAR ship dataset
|   `-- README.md
`-- gensar # image generation framework for SAR
    |-- controlnet # ControlNet for SAR
    |-- evaluation # KL and FID
    |-- finetune # train from Scratch or fine-tuning SD
    |-- gen_prompt # prompt construction
    |-- inference # inference pipeline
    `-- lora # ORS/SAR LoRA modules

Please refer Document of Datasets for more details.

📊 Reproduce Our Results

1. SAR Example

This is the raw SAR images from ./datasets/fusrs_v2/vgg_format/tanker.

2. From Scratch

First please link the dataset into ./finetune/data/fusrs_v2:

# Put the dataset
cd ./gensar/finetune
mkdir -p ./data/fusrs_v2
ln -s ../../datasets/fusrs_v2/hg_format ./data/fusrs_v2

We provide training scripts:

# Training from scratch
conda activate gensar
bash ./scripts/from_scratch.sh
# Inference
bash ./scripts/inf.sh

The resulting checkpoints are saved in finetune/output/fusrs_v2/. The generated images are saved in finetune/output/gen/.

3. Fine-tuning

First please follow the instruction in "From Scratch" to put the data. Then we provide fine-tuning scripts:

# Training from scratch
conda activate gensar
bash ./scripts/ft.sh
# Inference: please modify the checkpoint path
bash ./scripts/inf.sh

The resulting checkpoints are saved in finetune/output/fusrs_v2/. The generated images are saved in finetune/output/gen/.

4. LoRA

First please link the SAR dataset into lora/data/fusrs_v2:

# Put the SAR dataset
cd ./gensar/lora
mkdir -p ./data/fusrs_v2
ln -s ../../datasets/fusrs_v2/hg_format ./data/fusrs_v2

Train SAR LoRA:

# Training SAR LoRA
conda activate gensar
bash ./script_dgx/rank/256/256_fp32_s20000+100e_wp0_bs32_lr1e-03.sh

5. 2LoRA

First link the ORS dataset into lora/data/fusrs_v2:

# Put the ORS dataset
cd ./gensar/lora
mkdir -p ./data/dosrs_v1
ln -s ../../datasets/dosrs_v1/hg_format ./data/dosrs_v1

Train ORS LoRA:

# Training from scratch
conda activate gensar
bash ./script_dgx/train_dosrs.sh

The resulting LoRA module is in lora/output/DOSRS_v1

6. pLoRA

Please follow instructions above to put ORS and SAR data intolora/data. After that, the pLoRA training script is:

cd ./gensar/lora
conda activate gensar
bash ./script_dgx/cluster_lora/256/fusrs_clus_rank4_c4.sh