This repo contains our research code used in the paper "Unsupervised Generation of Labeled Training Images for Crop-Weed Segmentation in New Fields and on Different Robotic Platforms". The code in this repo has been frozen at the time of submission of the paper and will not be extended, but feel free to email me if you have any issues.
Datasets used in the paper are available online:
Weights of the following are also available online:
If you have a different setup, match the packages to your specific setup.
Clone this repo:
git clone --recurse-submodules https://github.com/PRBonn/StyleGenForLabels.git
cd StyleGenForLabels
Then, use the requirements.txt to setup the python environment of your choosing via pip/pip3:
pip install -r requirements.txt
To install PyTorch, we needed to use specific wheels to support our GPU sm archi:
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
Or, you can install PyTorch directly if this is not a problem for you:
pip install torch==1.10.2 torchaudio==0.10.1
You can use StyleGAN2 to pretrain your model.
OR if you want to train your own model:
Follow the instructions from stylegan2-pytorch to train the StyleGAN2 on the source and target domain.
Remember to use the pretrained weights to train StyleGAN2 for your source and target domain (both domains should be mixed into a single lmdb)!
Train the encoder:
cd pixel2style2pixel
Change the path in configs/paths_config.py to point to your dir with all the rgb images from source and test domain. The directories should contain both source and target rgb images combined together into a single directory.
'plant_train': '../dataset/train_images/',
'plant_test': '../dataset/test_images/',
Run the training script
python scripts/train.py \
--dataset_type=plants_seg_to_real_uav \
--exp_dir=<output_dir_path> \
--workers=<number_of_workers> \
--batch_size=<training_batch_size> \
--test_batch_size=<test_batch_size> \
--test_workers=<number_of_test_workers> \
--val_interval=<steps_per_validation> \
--save_interval=<steps_per_checkpoint> \
--image_interval=<steps_per_sample_saving> \
--encoder_type=GradualStyleEncoder \
--label_nc=<number_of_input_channels> \
--input_nc=<number_of_output_channels> \
--output_size=<output_image_size> \
--max_steps=<max_number_of_steps_to_train_for> \
--stylegan_weights=<path_to_stylegan2_checkpoint> \
--learning_rate=<learning_rate>
For example:
```sh
python scripts/train.py \
--dataset_type=plants_seg_to_real_uav \
--exp_dir=../../F_EN_training \
--workers=0 \
--batch_size=16 \
--test_batch_size=16 \
--test_workers=0 \
--val_interval=5000 \
--save_interval=10000 \
--image_interval=1000 \
--encoder_type=GradualStyleEncoder \
--label_nc=3 \
--input_nc=3 \
--output_size=512 \
--max_steps=250000 \
--stylegan_weights=../../checkpts/stylegan2_weights.pt \
--learning_rate=0.0001 \
--lpips_lambda=0 \
--id_lambda=0 \
--l2_lambda=0
cd stylegan2-pytorch
python closed_form_factorization.py --out <output_bin_name.pt> <stylegan2_checkpoint>
cd scripts
python patch_img.py --img_dirs <dir of input rgb> --out_dir <dir of output>
python latents_fr_real.py \
--sefa_path=<path_to_sefa_factors>
--data_path=<path_of_dir_with_source_images> \
--f_en_checkpoint_path=<path_of_F_EN_ckpt> \
--gmm_path=<path_to_where_gmm_will_be_saved> \
--out_dir_path=<path_to_where_the_generated_images_will_go> \
--test_workers=<number_of_cpu_threads_for_data_loading>
You can download the semantic segmentation network weights here
OR train your own network:
cd semantic_segmentation
vi ./config/config_train.yaml # change the path_to_dataset to point ot your source dataset
python train.py --config ./config/config_train.yaml --export_dir <output_dir>
create dataset from generated source images (multi/images --> images/rgb; leave annotations dir empty)
change the config_pred.yaml to point to the generated source dataset
train with the generated source dataset
cd semantic_segmentation
python predict.py --export_dir <output_dir> --config config/config_pred.yaml --ckpt_path erfnet_source.ckpt
copy the labels to the dataset dir (lightning_logs/version_XX/visualize/semantic_labels --> annotations)
cd scripts
python patch_img.py --img_dirs <dir of input rgb> --out_dir <dir of output>
python latents_fr_real.py \
--sefa_path=<path_to_sefa_factors>
--data_path=<path_of_dir_with_source_images> \
--f_en_checkpoint_path=<path_of_F_EN_ckpt> \
--gmm_path=<path_to_where_gmm_will_be_saved> \
--out_dir_path=<path_to_where_the_generated_images_will_go> \
--test_workers=<number_of_cpu_threads_for_data_loading>
python pkls_only.py \
--stylegan_ckpt=<path to the stylegan2 checkpoint> \
--f_en_checkpoint_path=<path to F_EN checkpoint> \
--source_latent_dir=<path to dir containing latent codes of the generated source dataset> \
--source_img_dir=<path to dir containing rgb images of the generated source dataset> \
--target_latent_dir=<path to dir containing latent codes of the generated target dataset> \
--target_img_dir=<path to dir containing rgb images of the generated target dataset> \
--out_dir=<path to the output_dir. Should not exist beforehand.>
cd scripts
python mask_labels.py \
--target_img_dir=<path_to_dir_with_generated_target_images> \
--labels=<path_to_dir_with_labels_to_refine> \
--output_dir=<output_dir_path>
cd semantic_segmentation
python train.py --config ... --export_dir ..
cd semantic_segmentation
python test.py --config ... --export_dir ... --ckpt_path ...
Open an issue or email Linn Chong at linn.chong@uni-bonn.de if you need help with something.
If you use this code for academic purposes, cite the paper:
@article{chong2023ral,
author = {Y.L. Chong and J. Weyler and P. Lottes and J. Behley and C. Stachniss},
title = {{Unsupervised Generation of Labeled Training Images for Crop-Weed Segmentation in New Fields and on Different Robotic Platforms}},
journal = ral,
volume = {8},
number = {8},
pages = {5259--5266},
year = 2023,
issn = {2377-3766},
doi = {10.1109/LRA.2023.3293356},
note = {accepted},
codeurl = {https://github.com/PRBonn/StyleGenForLabels}
}