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.
Due to package conflicts, we build two seperate conda
environments for GENSAR and CLSAR, respectively. Please follow the instructions below to build the environments.
Please refer Document of CLSAR for details.
-
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
-
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.
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.
-
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
-
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
We employ SDWebUI API for quick batch inference.
-
Install and Start API backend. We recommand an editable install.
conda activate gensar # Install SDWebUI cd ./inference bash ./webui.sh
-
Call API. By default the API will listen on addr
localhost:7860
. You can call inference API the same in templategensar/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, ... } ] } }, }
-
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
-
(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/
-
Convert ControlNet for Inference. To convert ControlNet resulting
.ckpt
(pyTorch Lightning checkpoint format) to.pth
and deploy in Huggingface, we provide converting script ingensar/controlnet/convert_ckpt_pth.py
. Please symbolink or save converted ControlNet weights ingensar/inference/extensions/sd-webui-controlnet/models
.
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.
This is the raw SAR images from ./datasets/fusrs_v2/vgg_format/tanker
.
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/
.
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/
.
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
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
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