A Data Augmentation Method Based on Image Generation for Rock Structural Surface Trace Detection With Single-Sample Class Imbalance

Fast and accurate detection of rock structural surface traces is of great importance in geology and engineering. Photogrammetry collects rock surface images from industrial cameras and uses edge detection algorithms to extract rock structural surface traces. With the advancement of deep learning technology, the detection performance of photogrammetry has been significantly improved, making it a low-cost and efficient detection method. However, the area of trace class in rock surface images is small, which suffers from a serious single-sample class imbalance problem. Traditional data augmentation methods make it difficult to solve this problem well. The single-sample class imbalance affects the prediction effect of the deep learning model, making the sensor data not effectively utilized. To address this challenge, we propose an image generation model based on generative adversarial networks (GANs) that aims to change the sample distribution in the rock surface image dataset to make it more suitable for prediction model training. The proposed model first generates trace masks from noise and then converts the trace masks to rock surface images. In order to be able to solve the problem of single-sample class imbalance, we use the trace proportion as a control condition to control the trace proportion in the generated sample. In addition, to improve the performance of the model, we designed a three-step training method and stabilized the training process using spectral normalization. The experiments show that our proposed method demonstrates superiority in solving the single-sample class imbalance problem, sample generation performance, and time consumption compared to other data enhancement methods. Finally, we conducted rock structural surface trace detection experiments using the generated rock surface samples with semantic segmentation models such as U-Net, Deeplabv3+, SegNeXt, SegFormer, and Swin-UNet. The results show that adding a certain proportion of class-balanced generated samples in the training dataset can improve the model’s prediction performance for rock structural surface traces.