Tree CycleGAN with maximum diversity loss for image augmentation and its application into gear pitting detection

Visual detection is an available approach for measuring gear pitting. Unfortunately, the number of gear pitting images is limited, resulting in that the detection accuracy of gear pitting is unsatisfactory. In order to augment gear pitting samples with different styles, a novel Cycle Generative Adversarial Network based on a symmetric tree structure (Tree-CycleGAN) is proposed. In Tree-CycleGAN, a new type of generator with tree structure named tree generator is designed to produce various types of high quality target samples from the source-domain samples, and a maximum diversity loss is constructed to enlarge the difference between two arbitrary branches; then a similar tree reconstructor is designed for translating target samples into source samples. Two discriminators are designed for making the generated images approximate to the target images in two cyclic processes. Via inception score, structural similarity indexes and peak-signal-to-noise ratio, the quality and diversity of images obtained by Tree-CycleGAN are evaluated. Comparative results show the superiority of Tree-CycleGAN over other domain adaptation GANs. The proposed Tree-CycleGAN combined with U-net has been successfully applied to gear pitting detection. Experimental results prove that the proposed methodology precedes the basic U-net method without sample augmentation and the method based on CycleGAN and U-net. •A new image augmentation model named Tree-CycleGAN is built.•The symmetric tree generator and constructor are designed to achieve the diversity.•A maximum diversity loss is proposed to increase the diversity among generators.•An U-Net based method is proposed to accurately detect the gear pitting area.