Automatic data augmentation for medical image segmentation using Adaptive Sequence-length based Deep Reinforcement Learning

Authors:
raw_author_name: Zhenghua Xu
raw_affiliation_string: State Key Laboratory of Reliability and Intelligence of Electrical Equipment, School of Health Sciences and Biomedical Engineering, Hebei University of Technology, Tianjin, China
raw_author_name: Gang Xu
raw_affiliation_string: School of Artificial Intelligence, Hebei University of Technology, Tianjin, China
raw_author_name: Yunxin Liu
raw_affiliation_string: State Key Laboratory of Reliability and Intelligence of Electrical Equipment, School of Health Sciences and Biomedical Engineering, Hebei University of Technology, Tianjin, China
raw_author_name: Miao Yu
raw_affiliation_string: State Key Laboratory of Reliability and Intelligence of Electrical Equipment, School of Health Sciences and Biomedical Engineering, Hebei University of Technology, Tianjin, China
raw_author_name: Hongwei Zhang
raw_affiliation_string: School of Computer Science and Engineering, Tianjin University of Technology, Tianjin, China
raw_author_name: Thomas Lukasiewicz
raw_affiliation_string: Department of Computer Science, University of Oxford, Oxford, United Kingdom; Institute of Logic and Computation, Vienna University of Technology, Vienna, Austria
raw_author_name: Junhua Gu
raw_affiliation_string: School of Artificial Intelligence, Hebei University of Technology, Tianjin, China

Abstract: Although existing deep reinforcement learning-based approaches have achieved some success in image augmentation tasks, their effectiveness and adequacy for data augmentation in intelligent medical image analysis are still unsatisfactory. Therefore, we propose a novel Adaptive Sequence-length based Deep Reinforcement Learning (ASDRL) model for Automatic Data Augmentation (AutoAug) in intelligent medical image analysis. The improvements of ASDRL-AutoAug are two-fold: (i) To remedy the problem of some augmented images being invalid, we construct a more accurate reward function based on different variations of the augmentation trajectories. This reward function assesses the validity of each augmentation transformation more accurately by introducing different information about the validity of the augmented images. (ii) Then, to alleviate the problem of insufficient augmentation, we further propose a more intelligent automatic stopping mechanism (ASM). ASM feeds a stop signal to the agent automatically by judging the adequacy of image augmentation. This ensures that each transformation before stopping the augmentation can smoothly improve the model performance. Extensive experimental results on three medical image segmentation datasets show that (i) ASDRL-AutoAug greatly outperforms the state-of-the-art data augmentation methods in medical image segmentation tasks, (ii) the proposed improvements are both effective and essential for ASDRL-AutoAug to achieve superior performance, and the new reward evaluates the transformations more accurately than existing reward functions, and (iii) we also demonstrate that ASDRL-AutoAug is adaptive for different images in terms of sequence length, as well as generalizable across different segmentation models.