Official implementation of Diffusion for Object-detection Domain Adaptation in Agriculture
DODA is a data synthesizer that can generate high-quality object detection data for new domains in agriculture, and help the detectors adapt to the new domains.
Model | Dataset | Resolution | Training Iters | Downlad Link |
---|---|---|---|---|
DODA-L2I | COCO | 512x512 | 30K | Google drive |
DODA-L2I | COCO | 256x256 | 100K | Google drive |
VAE | GWHD2021 | 256x256 | 170K | Google drive |
DODA | GWHD2021 | 256x256 | 80K | Google drive |
DODA-ldm | GWHD2021 | 256x256 | 315K | Google drive |
conda create -y -n DODA python=3.8.5
conda activate DODA
pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
pip install -r requirements.txt
bash Download_dataset.sh
python prepare_coco.py
python prepare_wheat_trainset.py # If you only want to test the model`s performance on GWHD, there is no need to run this line
python prepare_Terraref_testset.py
Generate images according to the bounding boxes of the COCO 2017 validation set:
First download the pretrained DODA-L2I to /models
folder, then run:
python generate_coco_testimg.py
Generate images according to the bounding boxes and reference images of the Terraref domain:
First download the pretrained DODA to /models
folder, then run:
python prepare_Terraref_testset.py
If you want to generate data to train the detector, first generate layout images using random_generate_layout_images.py
, then use generate_data_for_target_domain.py
to generate the data.
If you want to generate data for your own domain, please refer to generate_data_for_target_domain.py
You can try our method to generate images for wheat through the GUI:
python wheat_gradio_box2image.py
Please upload BOTH the reference image and layout image image respectively as shown:
PS: The demo reference image and layout image can be found in the
/figures
folder. More images can be found in/dataset
folder after runprepare_wheat_trainset.py
Or you can simply draw it yourself through drawing software. Each item should have a distinguishable color (with maximized values of the R, G, B channels), for example, (0, 0, 255)
, (255, 0, 255)
, etc. Below are some examples of possible layout images:
DODA training is divided into three parts, from first to last: VAE, LDM and controlnet. This repository reads the data set through a txt file, so first, please write the file names of all the images in your own dataset into a txt file.
Modify the config
in train_wheat.py
:
config = 'configs/autoencoder/DODA_wheat_autoencoder_kl_64x64x3.yaml'
Modify the txt_file
and data_root
in the config file to the path of the filenames txt file and the path to your own dataset.
then train the VAE by running:
python train_wheat.py
VAE is very robust, so we recommend skipping VAE training and using the pre-trained weight kl-f4-wheat.ckpt
we provide.
Modify the config
in train_wheat.py
:
config = 'configs/latent-diffusion/DODA_wheat_ldm_kl_4.yaml'
Modify the ckpt_path
in the config file DODA_wheat_ldm_kl_4.yaml
to the weight path of your VAE or the VAE provided by us.
Modify the txt_file
and data_root
in the config file to the path of the filenames txt file and the path to your own dataset.
then train the ldm by running:
python train_wheat.py
Modify the input_path
in tool_add_control.py
to the weight path of your ldm or the ldm provided by us, and modify output_path
to specify the name of the output weight.
Run tool_add_control.py
to add the ControlNet to the ldm:
python tool_add_wheat_control.py
Modify the resume_path
in train_wheat.py
to the path of the output weight.
Modify the config
in train_wheat.py
:
config = 'configs/controlnet/DODA_wheat_cldm_kl_4.yaml'
Modify the txt_file
and data_root
in the config file to the path of the filenames txt file and the path to your own dataset.
then train the cldm by running:
python train_wheat.py
Diffusion model is data hungry, and using more data always gives better results, so we strongly recommend mixing your data with GWHD for training. Mixing data can be achieved by putting all the images in your own dataset and the GWHD into one folder and writing the filenames of all images to one txt file.