Target Class Transformation for Segmentation Task Using U-GAN

Abstract

The paper presents a review of modern generative adversarial models for data augmentation, focusing on research, aimed at creating images and their corresponding segmentation masks. This task is particularly useful in cases where data is insufficient, hard to access, has confidential nature, or where labeling requires significant resources. The paper is aimed at the task of augmenting the minority class by transforming an image from another class and creating a segmentation mask. New approach is proposed for the simultaneous generation of the image and segmentation mask, using a generative adversarial network with U-Net generator. This generator takes an image of one class and noise, which is fed as an additional image channel. The generator tries to create an image of another class, minimizing changes in the original image and adding features of another along with the segmentation mask of the new class. The discriminator then determines whether the picture-mask pair is real or generated. The algorithm that applies only those changes of the generated image that are indicated by the created segmentation mask used to preserve the original appearance of the input image with minimal changes. This technique allows to obtain an image with features of the new class with minimal changes. The practical implementation of the proposed approach was conducted on a dataset of panoramic dental X-rays, based on which a set of individual teeth was created, some with fillings and some without. The experimental data set included 128 teeth without fillings and 128 with fillings. The GAN is trained to transform images without fillings into similar ones with fillings using all input images. Two experiments of 50 simulations each with different random states were conducted for training the segmentation model U-Net with ResNet-34 backbone to check the effectiveness of this augmentation. The first experiment used only real data for training, while the second included 64 additional images and masks created by the generator based on existing zero-class images. The average Jaccard score among all simulations for the first and second experiments were respectively 94.2 and 96.1. This result indicates that data generated using the proposed augmentation helps improve the quality of segmentation models and this approach can be combined with other augmentation techniques.