PyTorch implementation of MUS-CDB: Mixed Uncertainty Sampling with Class Distribution Balancing for Active Annotation in Aerial Object Detection.
Recent aerial object detection models rely on a large amount of labeled training data, which requires unaffordable manual labeling costs in large aerial scenes with dense objects. Active learning is effective in reducing the data labeling cost by selectively querying the informative and representative unlabelled samples. However, existing active learning methods are mainly with class-balanced setting and image-based querying for generic object detection tasks, which are less applicable to aerial object detection scenario due to the long-tailed class distribution and dense small objects in aerial scenes. In this paper, we propose a novel active learning method for cost-effective aerial object detection. Specifically, both object-level and image-level informativeness are considered in the object selection to refrain from redundant and myopic querying. Besides, an easy-to-use class-balancing criterion is incorporated to favor the minority objects to alleviate the long-tailed class distribution problem in model training. To fully utilize the queried information, we further devise a training loss to mine the latent knowledge in the undiscovered image regions. Extensive experiments are conducted on the DOTA-v1.0 and DOTA-v2.0 benchmarks to validate the effectiveness of the proposed method. The results show that it can save more than 75% of the labeling cost to reach the same performance compared to the baselines and state-of-the-art active object detection methods.
-
Create a conda virtual environment and activate it.
conda create -n mus-cdb python=3.7 -y conda activate mus-cdb
-
Please install PyTorch 1.8.0 and torchvision 0.9.0 for CUDA 11.1 following the official instructions, e.g.,
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge
-
Please install mmcv-full, we recommend you to install the pre-build package as below.
pip install mmcv-full==1.4.5 -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.8.0/index.html
-
Please clone the MUS-CDB repository.
git clone https://github.com/zjw700/MUS-CDB.git cd MUS-CDB -
Please download the required package.
pip install -r requirements/build.txt
Please refer to data_preparation.md for dataset installation and segmentation.
Please modify the corresponding pretrained model path, they are located in:
Line 16 of configs/redet/redet_re50_refpn_1x_dota_le90.py: pretrained='$PRETRAINED_MODEL_PATH/'
Line 16 of configs/redet/redet_re50_refpn_1x_dota2_le90.py: pretrained='$PRETRAINED_MODEL_PATH/'We recommend you to use a GPU but not a CPU to train and test, because it will greatly shorten the time.
If you use only a single GPU, you can use the script.sh file directly as below:
chmod 700 ./script.sh
./script.sh The script uses the train.py and test.py functions to train and test where the value of config_path, sample, c_0_path, s_0_path and work_dir shound be changed according to your actual situation (e.g., config_path should be replaced by the path of the config file in the configs/_base_ folder).
Please note more parameters can be found and mofidy in the file of train.py and test.py .
The trained model for each cycle in active learning are Baidu Drive (Extraction code:ASLK).
This project is released under the Apache 2.0 license.
If you find this repository useful for your publications, please consider citing our paper.
@misc{2212.02804,
Author = {Dong Liang and Jing-Wei Zhang and Ying-Peng Tang and Sheng-Jun Hang},
Title = {MUS-CDB: Mixed Uncertainty Sampling with Class Distribution Balancing for Active Annotation in Aerial Object Detection},
Year = {2022},
Eprint = {arXiv:2212.02804},
}