Josafat-Mattias Burmeister, David Kuska, Stefan Reder, Rico Richter, Jan-Peter Mund, Jürgen Döllner
This repository implements a deep learning pipeline for semantic segmentation of tree species in high-resolution aerial imagery.
- Install a compatible version of PyTorch for your operating system.
- Clone this repository and install the Python package in editable mode:
git clone https://github.com/ai4trees/tree_species_seg.gitcd tree_species_segpython -m pip install -e .Alternatively, you can use our pre-configured Docker image to avoid setting up the dependencies manually (see the Docker section below).
To run the code, download and unzip the following datasets from the TreeAI Zenodo archive:
Place the unzipped folders 12_RGB_SemSegm_640_fL and 34_RGB_SemSegm_640_pL in a common base directory.
For running inference on the test set, you must also create a folder named SemSeg_test-images in the same base directory containing the test images.
Models can be trained using the train.py script in the scripts directory of our repository. The script has the following arguments:
- Path to a configuration file (YAML format): The configuration file specifies the dataset and model configuration. The configuration used for our competition submission is provided at
./configs/submission.yaml. --wandb-project(optional): Name of a Weights & Biases project for logging training metrics. If not provided, metrics will only be logged to the console. If provided, you will be asked to provide an API key for your Weights and Biases account at the beginning of the training.--run-name(optional): Name for the training run (used in Weights & Biases and checkpoint naming).--tags(optional): Tags for the Weights & Biases run.--checkpoint(optional): Path to an existing model checkpoint from which to resume training.
The training of our competition model can be reproduced using the following command:
python ./scripts/train.py --config ./configs/submission.yaml --wandb-project tree-ai --run-name submissionSet the dataset:base_dir option in the configuration file to the path of the base directory containing the dataset folders (12_RGB_SemSegm_640_fL and 34_RGB_SemSegm_640_pL). The training script automatically preprocesses the dataset and stores the preprocessed data in the folder specified by dataset:output_dir.
Model checkpoints will be saved in the directory defined by training:checkpoint_dir in the configuration file.
To obtain predictions from a trained model, you can use the inference.py script in the scripts directory. Two modes are available:
Use the tree_ai mode to run inference on the train/val/test subsets of the TreeAI dataset. This requires that you placed the corresponding TreeAI data folders in a common base directory, as described before.
The inference script has the following arguments:
- Path to a configuration file (YAML format). This should be the same configuration file as used for model training. The configuration file used for our submission is provided at
./configs/submission.yaml. - Path to the trained model checkpoint.
- Path to the output directory for predictions.
--output-format(optional): Output format for the model predictions. Must be either.npyor.png(default:.npy).--visualization-dir(optional): Path to a folder in which to save visualizations of the predictions. If not provided, no visualizations are created.--split: Subset of the TreeAI dataset to run inference on. Must be eithertrain,valortest(default:test).--batch_size(optional): Batch size. If not provided, the corresponding value from the configuration file is used.--num_workers(optional): Number of worker processes for data loading. If not provided, the corresponding value from the configuration file is used.
To reproduce the inference of our model for the TreeAI4Species competition, run:
python ./scripts/inference.py tree_ai --config ./configs/submission.yaml --checkpoint <path to model checkpoint file> --output-dir ./predictions --visualization-dir ./visualizationsUse the folder mode to run inference on your own images. The script has the following arguments:
- Path to the input folder containing the images on which to predict.
- Path to a configuration file (YAML format). This should be the same configuration file as used for model training. The configuration file used for our submission is provided at
./configs/submission.yaml. - Path to the trained model checkpoint.
- Path to the output directory for predictions.
--output-format(optional): Output format for the model predictions. Must be either.npyor.png(default:.npy).--visualization-dir(optional): Path to a folder in which to save visualizations of the predictions. If not provided, no visualizations are created.--batch_size(optional): Batch size (default: 8).--num_workers(optional): Number of worker processes for data loading (default: 8).--class_remapping_file(optional): Path to a JSON file containing a dictionary defining a mapping between consecutive class indices and class IDs. If not provided, the class indices are mapped to the class IDs of the TreeAI dataset per default.
An example for prediction on a custom folder:
python ./scripts/inference.py folder <path to input folder> --config ./configs/submission.yaml --checkpoint <path to model checkpoint file> --output-dir ./predictions --visualization-dir ./visualizationsWe provide a Docker image The Docker image that contains our package and its dependencies. You can start a Docker container from the Docker image with the following command (replace <path to your data> by a path to your data folder on the host system):
docker run --rm -it -v <path to your data>:/workspace/data josafatburmeister:josafatburmeister/tree_species_seg
Inside the Docker container you can run our training or inference scripts as usual, e.g.:
python ./tree_species_seg/scripts/train.py <training arguments>
This project is licensed under the MIT License - see the LICENSE file for details.
For questions, please contact: Josafat-Mattias Burmeister (burmeister@uni-potsdam.de)